A General Framework for Inference-time Scaling and Steering of Diffusion Models

Thank you for highlighting the strength of our results. We also deeply appreciate your comments about the high-quality execution of our paper, the rigor/strength of our evaluations, and the clarity of our exposition.

Can FK steering make a Stable Diffusion 3 medium (2B) [A] model outperform the large (8B) [B] model?

SD3 2B + FK steering indeed outperforms SD3 8B. Stable Diffusion 3 2B with FK steering using just $k=3$ particles achieves a higher prompt fidelity score than the 8B parameter model.

Can FK steering make a timestep-distilled model like Flux-Schnell outperform Flux-dev?

Distillation + FK steering. Improving timestep distillation is an interesting future direction and a natural extension of FK steering.

Given that classifier guidance by Nichol & Dhariwal (2021) significantly improves these metrics, it is unclear if FK steering can achieve similar improvements.

Combining gradients with FK steering. Our ImageNet experiment show that FK steering can improve on classifier-guidance. We have also added ImageReward gradient-guidance results:

Here, FK steering both outperforms and improves gradient guidance.

I am curious as to why the GenEval performance does not match between Table 1 and Table 3. Additionally, for text-to-image tasks, it would be helpful to include more qualitative comparisons using multiple seeds.

• Average vs best particle performance. Table 1 has results from the best particle out of $k$ while Table 3 has results averaged over all $k$ particles. We do this to highlight that FK steering improves sample quality and prompt alignment metrics for all the particles. We will make this distinction more clear in the text and captions for the tables in the revised draft.

I'm interested in knowing whether the proposed method offers an orthogonal contribution to classifier-free guidance, or if it only proves effective at a specific classifier-free guidance scale.

• Thanks for this question. We ran FK steering with $k=4$ with a lower guidance scale in the table below. However, we note that higher guidance scales lead to better performance.

The evaluation could be improved by incorporating experiments on more up-to-date models and enhancing the metrics in the ImageNet evaluations.

• ImageNet. We will add these numbers in the final paper revision.
• Up-to-date models. We have included additional SD3 results above.

It would be beneficial to include a table or plot comparing the increase in inference time and memory cost

• Memory and sampling time increase. In table 6, we show the increase in sampling time from FK steering with $k=4$ particles. We include sample times for a single GPU run as well as parallel generation with 2 GPUs. We will update our draft to contain FLOP/memory details.

Showing whether performance continues to improve when increasing the number of particles beyond k=8 would help demonstrate the scalability of the proposed method.

• Beyond k=8. In figure 4, we include the GenEval and ImageReward scores for increasing values of $k$, from 2 to 16. The figure indicates that increasing particle count improves both scores. For our text experiments, we will include additional results in the revised version of the paper.

In the early denoising steps, the expected x_0 will be blurry as the diffusion models predict the “expectation” at each denoising step...How does the reward model play a meaningful role in the earlier denoising steps?

Thank you for this question. In figure 5 in the appendix, we plot the correlation of the reward $r_\phi(x_t) = r(x_0 = E_\theta[x_0 | x_t])$ with the reward at the terminal step $r(x_0)$. For ImageReward, we see a high correlation early in the generation. We demonstrate that learning the rewards will improve this correlation even further and can lead to improved performance. For instance, for toxicity we see a low correlation, prompting us to learn the intermediate rewards, which improves the generation's attribute accuracy as shown in table 4, see SSD-LM with learned $r_\phi$.

Thank you for highlighting the quality of our writing and experiments, and the practical applications of our approach.

Concerns regarding novelty claims in this paper.

Various fields such as statistics, posterior inference [Naesseth et al 2019], have used Feynman-Kac interacting particle systems for rare-event sampling and posterior inference with state-space models. We evaluate using FK-IPS to sample from target distributions defined by arbitrary rewards and diffusion models.

FK-IPS requires two things, (1) potential functions used to up-weight sample paths that yield high rewards, see eq 2, and (2) a method for sampling from the FK distributions. In this work, we provide:

• New potentials.
  1. We identify a simple condition in equation 3 that potentials must satisfy to consistently estimate the target distribution, and show that various new choices of potentials are possible.
  2. The traditionally used DIFFERENCE potential minimizes the variance of the potentials (see thm 10.1 in [Chopin et al., 2020]), not necessarily generating high-reward samples.
     • For bounded rewards (e.g. ImageReward), using the DIFFERENCE potential can paradoxically eliminate particles that achieve the maximum reward early in generation.
     • In tables 3 & 5, we show that a potential we develop that satisfies the product conditions, yields higher reward samples than the difference potential.
  3. We show that diffusion models offer many choices of intermediate rewards with different trade-offs between compute and knowledge of the terminal sample's reward. We propose two novel choices:
     • Learned from data. Using data and the noising process, we train intermediate rewards using regression. This objective generalizes learning a classifier trained on noisy states in Dhariwal et al., 2023 to arbitrary reward functions.
     • Many-sample intermediate reward. Using samples from $p_\theta(x_0 \mid x_t)$ we define intermediate rewards of the form: $\log \frac{1}{N} \sum_{i=1}^N \exp(r(x_0^i))$.
• Sampling.
  1. We show the effectiveness of the base model as the proposal generator, which is faster than techniques like gradient guidance and expands the use of SMC to discrete-state diffusion models and non-differentiable rewards.

With this framework, we show that inference-time steering of diffusion models is remarkably effective. FK steering:

• Outperforms fine-tuning. With only $2$ particles, FK steering outperforms DPO, DDPO fine-tuning approaches on prompt fidelity metrics for all the text-to-image models considered.
• Outperforms gradient-guidance. FK steering outperforms gradient guidance and is significantly faster (8.1s vs 20s for SDv1.5).
• Gradient-free control of discrete models: FK steering can be gradient-free and enables plug-and-play control of discrete-space diffusion models.
• Overcomes parameter counts. FK steering enables smaller models (0.8B) to outperform much larger models (2.6B) with less total compute.

This paper says that SVDD only selects a single sample, but that is only the case for the choice alpha=0, whereas for alpha > 0 SVDD seems very similar (if not identical?) to FK

Comparison with SVDD. SVDD selects a single proposal by sampling from a categorical distribution over the proposals $x_t^i$, with probability proportional to $\exp(r(x_t^i)/ \alpha)$ (see line 4 in algorithm 1). SVDD is a specific instantiation of nested importance sampling (algorithm 5 in Naesseth et al. 2019). FK steering uses SMC, which selects $k$ proposals instead of one, However, as noted in lines 198-204, it can use other particle-based samplers.

Prior works.

1. Nested diffusion sampling. Elata et al 2023 uses an inner denoising loop to sample a single $x_0$ for each state $x_t$, rather than using $E_{\theta}[x_0 | x_t]$, in the transition kernel. This approach samples from the diffusion many times, but does not retain multiple samples for each step $x_t$.

2. Conditional Sampling within Generative Diffusion Models We discuss the relevant method from this review paper, Wu et al 2023, which proposes a particle-based sampler for conditional sampling with continuous diffusion models with $\log p(y | x)$ as reward. In contrast, we steer with arbitrary rewards and present various choices of potentials.

Thank you for highlighting these references, we will add include them in the revised draft.

The other related works cited were posted after the ICML submission deadline. Due to limited space, we will add a discussion on these works in either the final draft or in the discussion period.

References

[Wu et al 2023] Wu, Luhuan, et al. "Practical and asymptotically exact conditional sampling in diffusion models." (2023)

[Chopin et al., 2020] Chopin, Nicloas et al. "An Introduction to Sequential Monte Carlo" (2020).

[Naesseth et al. 2019] Naesseth, Christian et al. "Elements of Sequential Monte Carlo." (2019).

Thank you for your feedback on our paper, and for noting the clarity of our writing. We discuss your concerns below.

The main claim—unifying inference-time steering methods into a single FK-based framework— might lack sufficient novelty

• We show that Feynman-Kac interacting particle systems, a well-known tool for rare-event sampling, are an effective and flexible framework for sampling from reward tilted diffusion models, where the reward can be any arbitrary scalar-valued function and the diffusion model can be continuous or discrete state.

  • In section F of the paper, we do show that prior works such as TDS and SVDD are specific instances of FK-IPS. However, we generalize beyond these choices and show the benefit of this generalization. We provide a detailed list of contributions below [also see response to r3 below].
  • Notably, these new choices for inference-time steering yield significant improvements.

• New choices yield higher-reward samples. The product of potential condition we identify in equation 3 unlocks novel choices (of potentials, samplers, and rewards).

  • Potential choices. For example, tables 3 & 5 show that new potential choices outperform traditional ones, such as the difference potential.
  • Intermediate rewards. Furthermore, we show that diffusion models enable many choices for estimating intermediate rewards. In tables 4 & 5, we show that these new choices improve performance.

Why did you choose simplistic methods Best-of-N, DPO, DDPO?

• Choice of DPO, DDPO. We select DPO and DDPO as benchmarks as they provide public checkpoints and code, as opposed to DRAFT or ELEGANT. DPO is used in prior work, such as [Esser et al. 2024]. We are happy to add DRAFT or ELEGANT as benchmarks in the revised draft.

• FK steering can boost fine-tuning. In Table 1, DPO does improve performance over the base model, but combining it with FK steering is even more performant. This indicates that fine-tuning and inference scaling are two complementary axes, and any improvements provided by DRAFT or ELEGANT could be further increased by using FK steering.

• Comparison to classifier-guidance. We have also added results comparing FK steering to gradient guidance. FK steering outperforms both DDPO, DPO, best-of-n, and gradient guidance.

• Best-of-N is simple yet effective. One exciting result in our paper is that best-of-n can even beat fine-tuning approaches (that require significant training compute) despite requiring no training.

Could you elaborate on the novelty provided by the MAX and SUM potentials?

• Theoretical justification. Based on eq. 3, many choices of potential provide a consistent approximation of the target distribution (see equation 3). However, when steering with arbitrary rewards, the DIFFERENCE potential does not necessarily yield high-reward samples.

  • Bounded rewards. Several choices of rewards are bounded. If a particle achieves maximum reward during sampling, the difference score will have to be negative and paradoxically will downweight this particle even though it has achieved a high reward. ImageReward is an example of such a bounded reward function.

• Empirical benefits. In tables 3 and 5, the MAX potential outperforms the DIFFERENCE and SUM potentials across model classes as well as different numbers of particles.

What distinguishes your intermediate rewards?

Prior works such as SVDD and TDS use the intermediate rewards $\log E_{p_\theta(x_0 \mid x_t)} \exp(r(x_0))$, either by approximating with the denoised expectation or by learning from model samples. In this work, we show that many choices of intermediate rewards can be used to consistently approximate the target distribution.

1. Unlike prior work that use intermediate rewards defined using model expectations, we show that the noise process and real data can be used to learn intermediate distribution via a regression objective.
2. Many-sample intermediate rewards: To the best of our knowledge, using samples from $p_\theta(x_0 | x_t)$ to define intermediate rewards $\log \sum_{i=1}^N \frac{1}{N} \exp(r(x_0^i))$ is a novel contribution of our work and was not proposed in prior works. The many-sample intermediate reward is a consistent estimator of $\log E_{p_\theta(x_0 | x_t)}\exp(r(x_0))$.

References

[Esser et al. 2024] Esser, Patrick, et al. "Scaling rectified flow transformers for high-resolution image synthesis." Forty-first international conference on machine learning. 2024.

Thank you for your thoughtful comments, and for highlighting that FK steering is a practical, effective approach for efficient, controllable generation with diffusion models. We also appreciate your assessment of our paper's clarity and strong theoretical grounding.

We address each of your concerns below. Let us know if there is any additional information we can include that would be helpful for your review.

Missing gradient guidance baseline in image experiments.

1. Requested gradient guidance results.
   • FK steering (with 4 particles, gradient-free) outperforms gradient guidance on all metrics and is significantly faster (8.1 vs 20 seconds for SDv1.5)!
   • FK steering combined with ImageReward gradient guidance performs even better but requires a significant increase in sampling time.

Gradient-guidance Results

Thank you for encouraging us to run these experiments. These results strengthen our paper and we will include them in our final draft.

2. Additional advantages over gradient guidance. Gradient-free steering has the following additional benefits:
   • Enables steering with non-differentiable rewards. Gradient-free steering methods enable the use of non-differentiable rewards such as perplexity from an autoregressive model or a trigram model, see table 2 in the paper for examples. Additionally, gradient-free steering can make use of closed-source reward models hosted via APIs, or non-differentiable constraints.
   • Enables steering discrete models. Gradient-free steering of masked diffusion models can enable attribute control (tables 2 and 4). More recently, LLaDA [Nie et al 2025] showed that discrete-state text diffusion models are competitive with auto-regressive models, making the development of gradient-free steering even more pertinent.

FK steering involves additional computation per particle due to resampling steps...this raises questions about its computational efficiency and whether its performance gains justify the added cost.

1. Computational efficiency. This is an excellent question. In the text-to-image generation experiments, using interval re-sampling (calling the reward model every 20 of 100 total steps) results in minimal compute overhead. We include timing results in the appendix (table 6) which show that FK steering results in a minimal increase in sampling time compared to best-of-n. For our most performant model, SDXL, the increase in sampling time is only 3%, see table below. We also note that as model size increases, the gap between FK and best-of-n reduces further.

2. Improved results justify computational overhead. In our controllable text generation experiments (see table 4) with the masked diffusion language model, FK steering with k=4 achieves attribute accuracy of 22.0%, significantly outperforming the best-of-8 attribute accuracy of 3.7%. Best-of-8 takes 19.3 seconds versus 20.7 seconds for FK steering(k=4) on a single A100 GPU. We have added the timing results to the revised draft.

Diversity is an important factor...incorporating it as an additional evaluation metric would provide a more comprehensive assessment of FK steering’s effectiveness and potential trade-offs.

• Diversity Results. Currently, we include CLIP-diversity scores in table 7, and provide 4 and 8 particle run samples, and highlight different components of the algorithm that can affect diversity. In the revised draft, we will include global diversity metrics as well.

...a direct comparison to recent adaptive guidance techniques and a more thorough analysis of computational efficiency trade-offs would further contextualize its impact within the broader literature.

• Adaptive guidance techniques. Other forms of conditioning, beyond the classifier and classifier-free guidance approaches we evaluate, could be used for the proposal generator in FK steering. We are very interested in exploring building on top of other recent methods in future work.

• Computational efficiency trade-offs. In table 6, we include the sampling time of FK steering compared to best-of-n. We discuss additional timing results above.

For gradient guidance, the paper should cite "Diffusion Posterior Sampling for General Noisy Inverse Problems"

Thank you. We will update our related works section to include this work in our final draft.

References

Nie et al. "Large Language Diffusion Models." arXiv preprint arXiv:2502.09992 (2025).