Skip to content

flash-bon/flash-bon

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Flash-BoN: Instant Drafts for Inference-Time Scaling in Diffusion Models

arXiv Project Page

Ruchit Rawal, Reza Shirkavand, Sayak Paul, Yuxin Wen, Heng Huang, Yizheng Chen, Tom Goldstein, Gowthami Somepalli

Accepted to ECCV 2026.

Budget-constrained, verifier-guided best-of-N generation for diffusion models. Given a prompt and a wall-clock budget, Flash-BoN drafts cheap candidates (TaylorSeer-style caching + layer/timestep skipping), selects the best with a pairwise vision-language tournament, refines the top-k to full quality by resuming denoising from a captured state, and scores them with VQAScore.

Image
On the GenAI-Bench quality-vs-compute frontier, Flash-BoN dominates prior inference-time scaling methods (BoN, ZOS, BFS, DFS) across Wan2.1 1.3B / 14B and FLUX.1-dev.

How it works

Cheap drafts are generated by caching/skipping most of the diffusion transformer's compute for the early denoising steps; a pairwise VLM tournament selects the most promising ones across batches; the global top-k are then refined to full quality by resuming denoising from each draft's captured state. Because drafting is near-free, far more candidates fit in the same budget.

Image

Additional findings

  • NFEs are a misleading efficiency metric. Counting function evaluations ignores verifier overhead; under realistic wall-clock budgets, plain Best-of-N already matches or outperforms guided search methods like BFS/DFS/ZOS.
  • Gains grow with model scale. Flash-BoN's improvement over Best-of-N widens on larger models, with up to a +8% AUC gain on Wan2.1 14B and FLUX.1-dev, and leads across all ten GenAI-Bench categories on Wan2.1 1.3B.
  • Diversity drives quality. Draft-pool diversity (measured with the Vendi Score) correlates strongly with downstream performance (Pearson r = 0.75) — because drafting is near-free, Flash-BoN explores a much more diverse candidate pool per budget.
  • Composable with other techniques. Stacking Flash-BoN with Reflection-Tuning (prompt optimization) adds +16% AUC; stacking with BFS adds +6% AUC.
  • Ties are common. Over 80% of prompts produce ties on the top pointwise verifier score, which is why Flash-BoN's selection stage combines pointwise scoring with a multi-stage Elo tournament rather than relying on pointwise scores alone.
  • Speeds up RL post-training too. Used to build training pairs, Flash-BoN (as Flash-Flow-GRPO) reaches the same reward as the baseline in 10x fewer training steps.

See the project page for qualitative examples and full ablations.

Layout

run.py                 end-to-end: prompt + budget -> draft -> select -> refine -> VQAScore
run_examples.sh        copy-paste commands (per method / model / budget)
configs/               per-model caching configs (flux_dev, wan_1_3b, wan_14b)
envs/                  pinned pip requirements for the three conda envs
assets/                README figures
fbon/
  paths.py             single source of filesystem locations (cache dir, server interpreters)
  models.py            model registry + load_pipeline (+ block/top-level compile)
  flux.py, wan.py      pipeline subclasses with draft-capture / resume
  caching.py           TaylorSeer caching + layer/timestep skipping (diffusers hooks)
  cache_state.py       GenerationCache + batch slice/merge
  config.py            load a caching config JSON
  ranking.py           multi-stage pairwise tournament (ELO across batches)
  verifier.py          tournament wrapper used for in-loop selection
  vqascore.py          VQAScore final-eval client
  servers.py           auto-launch/reuse the vLLM verifier + VQAScore servers
  serve_vqascore.py    the VQAScore HTTP server

Usage

CUDA_VISIBLE_DEVICES=0 python run.py --prompt "a photo of four giraffes" --budget_seconds 150 --model_key wan-1.3b

Outputs are written to outputs/<method>/<model>/budget<N>s/<prompt>/: result.json (timings, rankings, VQAScore) and final_topk/rank<r>_seed<seed>.png. See run_examples.sh for copy-paste commands.

Environment setup

Three conda environments (Python 3.12), one per process. run.py runs in the main env and launches the other two automatically (each on --server_gpu), talking to them over HTTP. Pinned requirements are in envs/.

# 1. main -- diffusion + this repo (run.py runs from source; do NOT pip-install the project)
conda create -n flashbon_env python=3.12 -y && conda activate flashbon_env
pip install -r envs/requirements-main.txt
# optional FlashAttention-2 (needs nvcc + a C++ compiler to build); else run with --attn_backend sdpa
pip install flash-attn==2.8.3 --no-build-isolation

# 2. vLLM verifier -- serves Qwen/Qwen2.5-VL-7B-Instruct
conda create -n flashbon_vllm python=3.12 -y && conda activate flashbon_vllm
pip install -r envs/requirements-vllm.txt

# 3. VQAScore final metric -- t2v_metrics
conda create -n flashbon_vqa python=3.12 -y && conda activate flashbon_vqa
pip install -r envs/requirements-vqa.txt
conda install -c conda-forge ffmpeg -y                # t2v_metrics checks `ffmpeg -version` at import
python -c "import t2v_metrics; print('ok')"           # verify the env is good

Each server is launched by conda activate-ing its env (flashbon_vllm / flashbon_vqa under $FBON_CACHE_DIR by default; see fbon/paths.py). If your envs live elsewhere or have other names, set FBON_VLLM_ENV / FBON_VQA_ENV (a conda env name or prefix path); set FBON_CONDA_SH if conda can't be auto-located.

Citation

If you find this work useful, please cite:

@misc{rawal2026flashboninstantdraftsinferencetime,
      title={Flash-BoN: Instant Drafts for Inference-Time Scaling in Diffusion Models},
      author={Ruchit Rawal and Reza Shirkavand and Sayak Paul and Yuxin Wen and Heng Huang and Yizheng Chen and Tom Goldstein and Gowthami Somepalli},
      year={2026},
      eprint={2607.04461},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2607.04461},
}

About

(ECCV 2026): Official code for Flash-BoN: Instant Drafts for Inference-Time Scaling in Diffusion Models

Topics

Resources

Stars

16 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors