PaperlayerView exact revision on arXiv

Source-linked AI summary

Best-of-$N$ Guidance for Test-time Diffusion Alignment

Richard Lee Kim, Yeongmin Kim, Gyuwon Sim, Taekyu Kim, Minsang Park, Il-chul Moon

arXiv:2610.05108v2cs.LGcs.AI

TL;DR

BoN uses reward only for post-hoc selection, limiting its use during generation, while prior reward-tilted approximations can require costly gradients or rollouts. BoNG performs online BoN selection and uses the selected particle to guide reverse diffusion. It improves single-output and average multi-output quality, ranking first in 29 of 36 comparisons against baselines.

  • Problem

    BoN uses reward only after generation to select one sample, while reward-tilted sampling is difficult to translate into denoising and its expected-future-reward term is generally intractable.

    Source: §3.2.1 · p. 3, §3.2.2 · p. 4

  • Method

    BoNG performs online selection of the highest-reward expected denoised particle and uses it to guide other particles during reverse diffusion without additional score-network feed-forward computations.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    Across 36 comparisons, BoNG achieved the highest score in 29 cases, ranking first in 80.56% of comparisons against baselines.

    Source: §5.1 · p. 7

  • Takeaways & Limitations

    BoNG improves both single-output performance and average multi-output quality with minimal additional inference cost.

    Source: §5.2 · p. 8, §6 · p. 10

    Existing approximations of expected future reward typically require costly gradient computation or rollout sampling, making them less efficient for test-time alignment.

    Source: §3.2.2 · p. 4

Abstract

from arXiv · show

Diffusion models achieve strong generative performance but often struggle to align generated samples with human preferences measured by a reward model. A simple yet effective algorithm for test-time alignment is Best-of-$N$ (BoN) sampling, which draws $N$ i.i.d. samples from a pre-trained diffusion model and outputs the single highest-reward sample. Despite its empirical success, BoN makes limited use of reward information, as it is incorporated only at the final selection stage without influencing the reverse diffusion trajectory during sampling. Consequently, BoN sampling does not improve the average alignment of generated samples and is primarily suited to single-output settings. We propose Best-of-$N$ Guidance (BoNG), a novel method that integrates the principle of BoN sampling directly into the reverse diffusion process. BoNG performs online BoN selection over denoising particles and adjusts the reverse diffusion process to steer the particle population toward higher-reward regions during generation. Specifically, by introducing an asymmetric guidance interaction among denoising particles, BoNG uses the current BoN particle as a guidance signal to the rest of the particle population. This particle-level interaction reshapes the sampling process toward higher-reward regions, enabling BoNG to improve not only the final best sample beyond Vanilla BoN sampling, but also the average quality of generated samples. Over 36 empirical comparisons, BoNG achieves the best performance in 29 cases, ranking first in 80.56% of the comparisons against SMC and Vanilla BoN sampling. BoNG also supports multi-output capability, achieving 1.3$\times$ ImageReward score of the latest sample-based guidance method with a 1.6$\times$ speedup. We release the code at https://github.com/aailab-kaist/BoNG.

1 Introduction

Diffusion outputs can miss user preferences, and existing test-time alignment methods face constraints in single- and multi-output settings. BoNG incorporates online Best-of-N selection into reverse diffusion to guide particles toward higher-reward regions.

1 Introduction

  • Diffusion outputs may not fully align with user preferences for input alignment and sample fidelity, motivating reward functions that evaluate condition alignment and output quality.Test-time alignment methods have drawn research interest relative to parameter scaling and fine-tuning approaches.Source: §1 · p. 1
  • BoN selects one sample after generation, while SMC is prone to particle collapse; gradient guidance can require costly neural backward passes and is prone to reward hacking.Guidance-based methods can adjust each sample's denoising direction to generate diverse outputs, while sample-based guidance avoids backward passes but retains rollout costs of 2–3× vanilla sampling.Source: §1 · pp. 1–2
  • BoNG consistently improves single-output performance over Vanilla BoN and SMC across various sample budgets.The introduction also reports empirical gains over prior guidance methods with minimal additional inference cost.Source: Figure 1 · p. 2, §1 · p. 2
  • BoNG selects the highest-reward expected denoised particle online and uses it to guide other particles during reverse diffusion without extra score-network feed-forward computations.This replaces iterative rollout sampling with selection at intermediate timesteps.Source: §1 · p. 2

2 Related Work

Prior test-time alignment methods include SMC, gradient guidance, and sample-based guidance. Their trade-offs include particle collapse, costly computation, and remaining rollout overhead.

2 Related Work

  • SMC suffers particle collapse, while gradient-guidance methods can require reward-model training or repeated neural backward passes, adding latency and memory overhead.SMC resamples using intermediate reward values; gradient methods approximate a time-dependent score to enable i.i.d. sampling.Source: §2 · p. 2
  • Sample-based guidance avoids particle collapse and neural backward propagation, but prior methods still require lookahead rollouts and can cost 2–3× vanilla diffusion sampling.LiDAR reduces per-step rollout sampling to one marginal rollout and is reported around 10× faster than gradient-guidance methods, while remaining 2–3× slower than vanilla sampling.Source: §2 · pp. 2–3

3 Preliminary

The preliminaries introduce the reverse diffusion process, formulate test-time alignment using a pretrained conditional model and reward, and describe BoN and reward-tilted sampling.

3.1 Diffusion Models

  • The reverse diffusion process generates data from terminal noise and uses a score model to estimate the score function in its stochastic differential equation.The forward process perturbs data with a Gaussian kernel, with terminal noise chosen so its distribution is close to a pure Gaussian prior.Source: §3.1 · p. 3

3.2 Test-time Human Alignment

  • Test-time alignment assumes a pretrained conditional diffusion model and a human preference reward, and aims to obtain samples with high reward by manipulating inference.The setup addresses generated samples that may not align well with the user input.Source: §3.2 · p. 3

3.2.1 Best-of-N Sampling

  • BoN draws N i.i.d. samples from the base distribution and returns the highest-reward sample, but uses no reward information during generation and does not directly address high-quality multi-output generation.Increasing N can improve the selected sample without changing the underlying generative model, but may require large N to reach a high-reward region.Source: §3.2.1 · p. 3

3.2.2 Sampling from a Reward-tilted Distribution

  • Reward-tilted sampling reweights the target distribution by reward, but translating this final-sample objective into iterative denoising is challenging and its expected-future-reward term is generally intractable.Prior approximations typically require costly gradient computation or rollout sampling; the score-based formulation replaces the score in the reverse process to yield reward-tilted samples.Source: §3.2.2 · p. 4

4 Method

BoNG selects the highest-reward one-step denoised particle and uses it to guide the rest during reverse diffusion. The method connects this selection to reward-tilted sample-based guidance and controls its strength through the guidance timestep window and scale.

4.1 Best-of-N Guidance (BoNG)

  • BoNG defines its guidance term from a Best-of-N target score at each timestep, steering particles toward the current highest-reward one-step denoised sample.The method uses one-step denoised samples as the candidate set and applies the BoN guidance term to adjust each particle.Source: §4.1 · p. 4
  • Non-best particles receive asymmetric guidance: they are attracted toward the current BoN particle while following its score-induced denoising direction.The interaction embeds Best-of-N selection behavior without additional score-network feed-forward computations.Source: §4.1 · p. 5

4.2 Connection to Prior Sample-based Test-time Alignment Methods

  • The analysis shows that hard Best-of-N selection is the large-λ limit of the SNIS estimator for reward-tilted guidance, using one-step denoised samples as a surrogate proposal set.A finite-λ relaxation uses soft SNIS weights, interpolating between the proposal-set average and the BoN target.Source: §4.2 · pp. 5–6

4.3 Instantiation of BoNG

  • Table 1 reports Best-of-N performance on GenEval prompts across sample counts and SD v1.5 and SDXL backbones, using ImageReward as the reward.Bold entries mark the top-performing results within each backbone group.Source: Table 1 · p. 7
  • BoNG’s guidance timestep window controls steering strength: applying guidance over more reverse steps yields stronger steering toward high-reward regions.A guidance scale separately controls the magnitude of guidance at each timestep.Source: §4.3 · pp. 6–7

5 Experiments

Across single-output comparisons, BoNG frequently outperforms Vanilla BoN and SMC, while multi-output results show gains across all three reported metrics. The analyses examine reward–diversity balance, computational efficiency, reward choice, matched-time comparisons, and scaling.

5.1 Single-Output Performance

  • BoNG ranks first in 29 of 36 comparisons (80.56%) against baselines and outperforms them across backbones and sample budgets.With N = 4, BoNG also outperforms Vanilla BoN with N = 8 in GenEval; it retains an advantage on FLUX.1-dev.Source: §5.1 · p. 7

5.2 Multi-Output Performance

  • BoNG achieves the best performance across all three metrics for four generated images, with very small test-time overhead.It has the strongest average performance among prior guidance methods while remaining close to vanilla diffusion sampling in inference time.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 Analysis

  • BoNG achieves a consistently better reward-diversity Pareto frontier than SMC by steering trajectories rather than duplicating and removing particles through resampling.This preserves diversity more effectively while maintaining strong reward improvement.Source: §5.3 · pp. 8–9, Figure 5 · p. 9
  • BoNG avoids the backward passes and lookahead rollouts required by gradient-guidance baselines, and achieves the steepest performance gain relative to computation.It also achieves better average performance than LiDAR, which relies on random lookahead samples rather than steering particles toward the current highest-reward denoising particle.Source: Figure 4 · p. 8, §5.3 · p. 9
  • Increasing CLIP weight raises CLIP scores, while increasing ImageReward weight raises ImageReward scores; mixing them yields frontier points balancing the metrics.With HPSv2 as guidance reward, BoNG improves both mean and BoN scores over Vanilla.Source: Figure 6 · p. 9, §5.3 · p. 10
  • Under matched end-to-end wall-clock time, BoNG achieves the best results on every metric with the shortest generation time compared with LiDAR.BoNG constructs guidance from the particles already being sampled, whereas LiDAR separately generates lookahead samples.Source: Table 4 · p. 9, §5.3 · p. 10
  • On SD v1.5 with DDPM 100 steps, BoNG consistently outperforms SMC and Vanilla BoN for N = 8, 16, and 32, with BoN performance improving as N grows.Its advantage also persists on the larger FLUX.1-dev backbone at every N.Source: Figure 7 · p. 9, Table 5 · p. 9, §5.3 · p. 10

6 Conclusion

The conclusion presents BoNG as an efficient test-time alignment framework, while the appendices derive a gradient identity and an importance-sampling approximation for reward-tilted posterior means.

6 Conclusion

  • BoNG integrates online Best-of-N selection into reverse diffusion, improving both single-output BoN performance and average multi-output quality with minimal additional inference cost.The conclusion also connects BoNG theoretically to sample-based guidance while avoiding its costly rollout procedure.Source: §6 · p. 10

A.1 Derivation of Lemma 4.1

  • Lemma 4.1 derives the expected-future-reward gradient using the forward Gaussian kernel and its Stein score.Source: §A.1 · pp. 14–15

A.2 Derivation of Lemma 4.2

  • Lemma 4.2 represents the reward-tilted posterior mean with self-normalized importance sampling, producing normalized softmax weights from proposal samples.Source: §A.2 · pp. 15–16

A.3 Derivation of Theorem 4.3

Theorem 4.3 establishes convergence of the reward-guidance approximation to the BoN guidance term as λ grows. The following subsections specify the experimental setup and baseline configurations.

A.3 Derivation of Theorem 4.3

  • As λ →∞, the self-normalized importance-sampling approximation of the reward-guidance gradient converges to the BoN guidance term for each particle.Source: §A.3 · pp. 16–17

B.1 Experiment Setup

  • Experiments use SD v1.5 and SDXL with DDPM 100-step and DDIM 50-step samplers, plus FLUX.1-dev with a 28-step Euler sampler.Results average three seeds except FLUX.1-dev, which uses one seed due to computational cost.Source: §B.1 · p. 17

B.2 BoNG Configuration

  • BoNG uses at most four guided transitions to match the SMC compute budget, with guidance windows and scales configured separately for each backbone and sampler.Source: §B.2 · p. 18

B.3 Baseline Configuration

  • The Vanilla baseline selects the highest-ImageReward sample from N standard diffusion samples; SMC uses FK-Steering's recommended configuration, and Best-k-of-N returns the top four of five unguided samples.Source: §B.3 · p. 18

B.4 Evaluation Metrics

The evaluation subsection defines image-quality, alignment, compositional-accuracy, and diversity measures, then gives a protocol for choosing BoNG's hyperparameters.

B.4 Evaluation Metrics

  • Evaluation uses ImageReward for preference and image quality, CLIP Score for prompt-image alignment, GenEval for compositional accuracy, and CLIP-image pairwise distance for diversity.Source: §B.4 · p. 18

C Hyperparameter Selection Protocol of BoNG

  • For new backbones or samplers, the protocol sets the window near reliable reward correlation, limits its size by the SMC resampling budget, and tunes guidance scale across the reward–diversity frontier.The guidance-scale sweep adds no extra compute.Source: C Hyperparameter Selection Protocol of… · pp. 18–19

C.1 Window Position

The section examines how BoNG's timestep window and guidance scale affect steering, reward, diversity, and computational cost. Its rules balance a reliable guidance signal with the ability to reshape particle content.

C.1 Window Position

  • The mid-range window [40, 44] is the only tested placement that improves all metrics simultaneously; early guidance has weak reward correlation, while late guidance cannot re-steer composition.At [80, 84], ImageReward improves but GenEval does not improve over Vanilla.Source: §C.1 · p. 19, Table 6 · p. 19

C.2 Window Size

  • Increasing window size K raises reward and gradually reduces diversity; at K = 4, reward exceeds Vanilla sampling with no loss of diversity (0.319 vs. 0.313).The additional cost scales linearly as NK(D + R).Source: §C.2 · p. 19, Table 7 · p. 20

C.3 Guidance Scale

  • At fixed window [40, 44] and K = 4, increasing guidance scale raises mean reward while gradually reducing diversity, selecting a frontier point at no extra compute.Source: §C.3 · p. 20, Table 8 · p. 20
  • Across s ∈ [1, 7.5], UG's ImageReward gains accompany falling CLIP score, while BoNG has a stronger ImageReward–CLIP trade-off than LiDAR; BoNG's CLIP score degrades at s = 10.Figure 8 compares the ImageReward–CLIP trade-off as guidance scale varies.Source: §C.3 · p. 20, Figure 8 · p. 21

D.1 Effect of Guidance Timestep Window

  • The timestep window trades off reward optimization against diversity and sets BoNG's added cost; the section reports the best performance–efficiency trade-off among compared methods.Source: §D.1 · pp. 21–22

D.2 Computational Cost

  • BoNG adds the smallest computational cost among the steering methods, requires no gradient computation, and matches Vanilla's memory footprint.Additional cost comes from scoring denoised estimates; the window size determines that overhead.Source: §D.2 · p. 22, Table 9 · p. 22

D.3 Anchor-Switching Ablation

The ablation tracks which particle supplies BoNG’s Best-of-N guidance across consecutive transitions. The anchor is dynamically re-elected rather than remaining fixed to one trajectory.

D.3 Anchor-Switching Ablation

  • Across 553 GenEval prompts, 1,659 transitions, and 3 runs, the experiment tracks the Best-of-N anchor over three consecutive transitions per prompt.The reported switch rate ranges from 0 (the anchor never changes) to 1.0 (it changes at every transition); mean, median, and standard deviation are reported across runs.Source: §D.3 · pp. 22–23, Table 11 · p. 23
  • The anchor is dynamically re-elected from the population, so no single trajectory dominates the guidance target.The authors conclude that the guidance signal emerges from particle interaction.Source: §D.3 · p. 23

E Finite-λ Soft Relaxation of BoNG

This section derives a rollout-free, finite-λ relaxation of BoNG by replacing exact proposal samples with one-step Tweedie estimates. It evaluates how λ shifts reward and diversity, while noting that the relaxation does not exactly target the finite-λ reward-tilted distribution.

Finite-λ construction

  • The rollout-free relaxation replaces exact proposal samples, which require full reverse-process rollouts, with one-step Tweedie estimates from the current particle population.These estimates are posterior means, ˆxj_0 = E_pθ(x0|xj_t,c)[x0], and yield the soft weights.Source: E Finite-λ Soft Relaxation of BoNG · p. 23
  • Smaller λ spreads guidance across candidates, while larger λ concentrates it increasingly on the highest-reward candidate; with a unique reward maximizer, λ → ∞ recovers hard-selection BoNG.The hard-selection algorithm is the version evaluated in the main experiments.Source: E Finite-λ Soft Relaxation of BoNG · p. 23

Approximation gaps

  • The relaxation is approximate: its estimates are not exact proposal samples, and guided particles need not retain the original marginal distribution.Thus, it does not exactly target the finite-λ reward-tilted distribution, and the exact importance-sampling interpretation does not directly transfer.Source: E Finite-λ Soft Relaxation of BoNG · p. 23

Empirical verification of the predicted scaling

F Additional Experiment Results

Additional experiments examine BoNG’s reproduction results, matched-budget comparison, transfer to HPSv3, composition with SMC, qualitative samples, and full sampling procedure. They report trade-offs and improvements across these settings, alongside the stated scope of each comparison.

F Additional Experiment Results

  • The additional experiments cover both multi-output and single-output settings, including reproduction, matched-time comparison, and a heavier guidance reward.The section also reports a BoNG-plus-SMC composition experiment.Source: F Additional Experiment Results · p. 24, §F.4 · p. 25

F.1 Unified-Codebase Reproduction of Table 2

  • In the unified-codebase reproduction, SMC has higher mean reward than BoNG but lower diversity (0.155 vs. 0.224).The authors attribute the lower diversity to resampling duplicating high-scoring ancestors; the reproduced relative ordering of Table 2 is preserved.Source: Table 13 · p. 24, §F.1 · p. 24

F.2 Comparison with LiDAR under Matched Wall-Clock Budgets

  • At N = 32, BoNG outperforms LiDAR on every metric with shorter end-to-end time under the reported matched wall-clock comparison.LiDAR uses n = 100 and δ = 5 to make its additional cost similar to BoNG’s NK(D + R) cost.Source: §F.2 · p. 24, Table 14 · p. 25

F.3 Performance with a Heavier Reward Model

  • With HPSv3 guidance, BoNG improves HPSv3 reward over Vanilla, while SMC attains a slightly higher HPSv3 score in this setting.ImageReward decreases under HPSv3 guidance for both methods; the authors attribute this to the reward models emphasizing different features.Source: Table 15 · p. 25, §F.3 · p. 25

F.4 Composing BoNG with SMC

  • Composing BoNG with SMC improves both mean and Best-of-N ImageReward at every particle count, without reward gradients or additional score-network evaluations.The authors describe guidance and resampling as orthogonal and composable, with BoNG modifying the score function and SMC resampling particles.Source: §F.4 · p. 25, Table 16 · p. 26

G Qualitative Comparison

  • Across backbones and samplers, BoNG produces samples described as better aligned with prompts than Vanilla and SMC samples.Figure 10 presents Best-of-N samples under the configurations of Table 1.Source: G Qualitative Comparison · p. 26, Figure 10 · p. 26

H Pseudocode of BoNG

  • Algorithm 3 combines Best-of-N selection with the guided update in one loop over reverse diffusion, propagating particles and applying guidance within the specified window.The procedure takes a prompt, particle count, diffusion steps, guidance scale and window, reward, and pretrained noise network as inputs.Source: §F.4 · p. 25, H Pseudocode of BoNG · pp. 26–27
Loading 2610.05108v2…