Inference-Time Scaling of Diffusion Models via Progressive Seed Pruning

AI summary

Overview

Research area: Computer vision / generative modeling — specifically inference-time scaling (test-time compute) for diffusion and flow-matching image generators using black-box reward models.

Technical level: Intermediate. The method itself is simple, but the framing assumes familiarity with diffusion samplers, DDIM/Euler solvers, best-of-N selection, importance sampling, tree search, and particle filtering.

Scope in one sentence: The paper proposes Progressive Seed Pruning (PSP), a training-free and gradient-free inference-time scaling strategy that abandons the usual constant-memory constraint, starts from a large pool of noise seeds, and progressively prunes candidates so that only promising trajectories receive the remaining denoising budget at a fixed total number of model evaluations.

What This Paper Is About

For text-to-image diffusion and flow-matching models, the random initial noise seed strongly influences final quality, so a natural way to spend extra test-time compute is to search over seeds guided by a reward model. Existing methods such as Best-of-N, importance-sampling approaches (e.g., FK-Steering), and tree search (e.g., DSearch) keep the number of parallel samples roughly constant throughout inference. The authors ask what happens if that constant-memory assumption is relaxed, and show that front-loading exploration — evaluating many seeds early, pruning aggressively, and spending the rest of the budget on a shrinking survivor set — uses a fixed compute budget more effectively.

Key Contributions

1. Identifying constant-memory inference as an unnecessary restriction. The paper argues that time-varying particle counts are a strong, underexplored inference-time scaling principle for diffusion and flow-matching models under black-box rewards, letting the method consider a larger pool of initial noise seeds.

2. Demonstrating consistent gains of PSP over baselines at matched compute. PSP is compared against Best-of-N, importance-sampling (FK-Steering), and tree-search (DSearch) baselines, plus additional methods (NTS, RBF, BFS, SVDD), across diffusion and flow-matching backbones on automated metrics and human evaluation, and its performance is shown to scale with compute.

3. Showing that determinism enables offline schedule search. Because PSP's outcome is fully determined by the schedule and the precomputed trajectories, pruning schedules can be simulated from cached intermediate rewards without rerunning the generator, enabling rapid per-task adaptation.

4. Demonstrating a training-free, gradient-free, determinism-friendly design. PSP works with deterministic solvers (DDIM with eta = 0 for diffusion, Euler Discrete for SD 3.5), unlike resampling-based methods that need stochasticity to generate diverse children.

Main Findings

• PSP leads on prompt alignment across all three backbones. Under matched compute (N̄ = 4), PSP achieves the highest GenEval scores in the main comparison: SD v1.5 = 0.574, SDXL = 0.645, SD 3.5 = 0.747, compared with BoN at 0.542, 0.629, and 0.747 respectively.

• Human evaluation favors PSP on all three backbones. SD v1.5 = 0.624, SDXL = 0.713, SD 3.5 = 0.841, against BoN at 0.594, 0.682, and 0.831. The study used 249 annotators from Prolific, 8,295 images (553 per method/backbone pair), 3 evaluations per image (24,885 ratings total), majority vote, and 80.3% annotator agreement.

• Reward maximization results are more mixed than alignment results. BFS reaches a higher ImageReward than PSP on SDXL (1.247 vs. 1.224), while PSP reaches higher ImageReward on SD v1.5 (0.827) and SD 3.5 (1.380). The authors attribute PSP's alignment advantage to being less susceptible to reward hacking, since it selects seeds rather than interfering with the sampling process through reward-seeking resampling.

• PSP remains effective where resampling methods struggle. On SD 3.5, BoN (ImageReward 1.336) outperforms FK-Steering (1.294), because resampling requires stochasticity that can degrade rectified-flow generation. PSP works with deterministic solvers.

• Scaling with compute. Scaling N̄ up to 16 by doubling initial seeds and survivor counts while keeping the same fractional prune points, PSP continues to improve and consistently outperforms BoN at matched compute. Regret relative to the 2N̄-seed upper bound decreases as N̄ grows.

• Intermediate reward rankings are informative early. ImageReward shows stronger correlation with final reward than HPS at most progress points. For the default schedule, the survival probability of the final best sample is 80% when pruning from 8 candidates at 25% inference progress and 90% when pruning from 4 at 50% progress, with regret close to 0.01 at both prune times except for the weaker SD v1.5 model.

• Overhead is modest relative to double-compute BoN. On a single H200, PSP (N̄ = 4) adds little runtime over BoN (N = 4) — SD v1.5: 2.99 s vs. 2.65 s; SDXL: 13.74 s vs. 12.48 s; SD 3.5: 37.59 s vs. 35.97 s — and is close to 2x faster than BoN (N = 8) at 4.70 s, 23.90 s, and 71.71 s. Peak VRAM stays close to BoN (N = 4) on SD v1.5 and SD 3.5 with VAE slicing (which adds less than 0.15 s), but remains noticeably higher on SDXL because its VAE is proportionally larger.

• Default schedule is near-optimal. A grid search over schedules preserving N̄ = 4, tuned on the IR Benchmark and tested on GenEval, showed little to no improvement over the untuned default schedule. Tuned strategies tend to give up one sample near the end in order to develop 3 samples beyond 50% of the inference process.

• PSP complements reward-based finetuning. PSP on a non-finetuned SD v1.5 (ImageReward 0.827, GenEval 0.574) exceeds standard sampling on a DPO-finetuned model (-0.022, 0.454). Applying PSP to the DPO-finetuned model yields further gains (0.907, 0.593 on SD v1.5; 1.365, 0.667 on SDXL) and outperforms BoN under similar compute (0.751, 0.562 and 1.300, 0.657).

• PSP generalizes to non-seed candidate pools. Applied to prompt selection over rewrites generated once per prompt with ChatGPT 5.2 and a fixed noise seed, PSP outperforms BoN at equal compute (SD v1.5: 0.782 vs. 0.641 ImageReward, 0.533 vs. 0.512 GenEval; SDXL: 1.319 vs. 1.256, 0.604 vs. 0.598).

• Reward choice matters for generalization. Supplemental experiments using HPS as the guidance signal are reported in Table S1; the authors state that IR generally leads to better generalization to GenEval, making it the best reward guidance to use in practice.

Methodology in Plain English

PSP treats image generation as a search over initial noise seeds rather than as a process to be steered mid-flight.

The method starts by drawing a large pool of independent noise seeds (k_T candidates). Instead of running every candidate all the way to a finished image, it advances them together and, at every step, asks the generator for its best current guess of the clean image — a denoised estimate that diffusion samplers and flow-matching solvers already compute internally. That guess is scored by a black-box reward model such as ImageReward. The lowest-scoring candidates are dropped, and only the survivors are advanced further. Pruning happens at a small number of predetermined steps with fixed survivor counts, so memory and runtime per interval are predictable.

Crucially, scoring uses no extra generator forward passes, because the denoised estimate is a byproduct of the denoising update. Compute is measured in denoising steps: if k_t trajectories are propagated at step t, total cost is C = sum of k_t over t, and the effective compute multiplier is N̄ = C/T, meaning the procedure uses the same compute as sampling N̄ times from the regular sampler.

The default schedule starts with 2N̄ seeds, prunes to N̄ at 25% progress, and prunes again to N̄/2 at 50% progress. This is motivated by the geometric halving identity: a halving schedule can evaluate an arbitrarily large initial pool at cost multiplier N̄ = 2, and scaling the schedule by a factor m gives N̄ = 2m. The authors use m = 2, hence N̄ = 4, matching the compute setting used by prior work.

Evaluation used Stable Diffusion v1.5 and Stable Diffusion XL (both diffusion, T = 64 steps) and Stable Diffusion 3.5 Large (flow matching, T = 32 solver steps), on GenEval prompts, with ImageReward as the guidance signal and GenEval, HPS v2, and human judgments as evaluation metrics, averaged over 3 random seeds.

Why This Matters

Impact on research. The paper reframes a design assumption — that inference-time search must hold parallelism constant — as an optional constraint, and shows that relaxing it yields gains without changing the number of model evaluations. It connects modern reward-guided diffusion to classical adaptive-particle and successive-halving ideas, and it argues that determinism is a first-class property because it makes schedule tuning a cheap, offline, simulation-only procedure.

Real-world applications (from the paper's framing):

• Deployment-time quality scaling: choosing compute multipliers at serving time, where the paper shows smaller backbones with larger N̄ can match or surpass larger backbones at similar total FLOPs.
• Prompt-alignment-sensitive generation: text-to-image systems where coarse features are decided early and seed choice dominates outcomes.
• Prompt rewriting pipelines: selecting among multiple rewritten prompts produced by an LLM, using the same pruning machinery with a fixed seed.
• Combined with finetuned models: applying PSP on top of DPO-finetuned backbones for further gains when extra inference compute is available.

Industry relevance. The motivation is explicitly systems-oriented: fixed prune points and fixed survivor counts make each interval's memory footprint and runtime predictable, which the authors frame as matching multi-GPU servers with elastic memory allocation. Overhead measurements on a single H200 and the near-2x speed advantage over double-compute BoN support the deployment argument.

Future Directions

1. Extending past scalar rewards. The authors state PSP presumes a scalar reward, making multi-objective or hard-to-quantify settings less natural — for example ControlNet-style constrained generation that must jointly balance text alignment with strict spatial consistency. They note many strict-spatial tasks admit scalar perceptual losses such as IoU or OKS.

2. Understanding when front-loading pays off. PSP's advantage depends on when the target features are decided. For objectives dominated by fine-grained detail generated late in the trajectory, such as aesthetic quality, early seed selection matters less and gains shrink — visible in the comparatively lower HPS for SDXL.

3. Richest schedule design space. The paper's default is a two-prune halving schedule; the number of initial candidates, prune timing, and survivor counts can all be varied at fixed budget. Offline simulation makes broader per-task schedule search cheap, and the authors note the default may be far from optimal in other domains.

4. Better reward models. The authors observe diminishing returns of ImageReward on GenEval beyond a high-IR regime, and state that PSP optimizes whatever guidance signal it is given and can directly benefit from stronger reward models as they become available.

Target Audience

Researchers and engineers working on diffusion and flow-matching generative models, inference-time compute scaling, and reward-guided generation. It is most useful for readers already comfortable with diffusion sampling internals (denoised estimates, DDIM, Euler solvers) and with inference-time search baselines, and for practitioners deploying text-to-image systems on multi-GPU infrastructure who need predictable per-interval memory and runtime. Readers interested in classical particle filtering and successive halving will also find the connection to those literatures relevant.

Authors’ abstract

Diffusion and flow-matching models dominate conditional image generation, yet inference-time scaling for these models is far less developed than for autoregressive language models. Because final quality is highly sensitive to the initial noise seed, many approaches spend extra compute on seed search or resampling under a black-box reward, but typically maintaining a constant memory footprint throughout inference. We show that relaxing this constraint enables an underexplored inference-time scaling axis: by front-loading exploration, evaluating many seeds early, and pruning aggressively, we can use a fixed compute budget more effectively. \emph{Progressive Seed Pruning} (\PSP) scores intermediate denoised estimates and progressively narrows the candidate set so that only promising trajectories are fully denoised, while keeping the total number of model evaluations fixed. Across diffusion and flow-matching backbones, \PSP \ consistently improves reward-guided selection and achieves higher GenEval scores (automated) and better human evaluation on prompt-alignment than best-of-$N$, importance-sampling, and tree-search baselines at matched compute. Project page: https://www.vision.caltech.edu/psp. Code: https://github.com/rogerioagjr/psp.