A General Framework for Inference-time Scaling and Steering of Diffusion Models
Raghav SinghalZachary HorvitzRyan TeehanMengye RenZhou YuKathleen McKeownRajesh Ranganath
Presents Feynman-Kac steering, a training-free framework that resamples particle trajectories at intermediate generation steps using arbitrary reward functions, allowing smaller diffusion models to outperform fine-tuned models and larger architectures in sample quality and prompt fidelity.
Diffusion-based generative models have achieved strong performance across images and text, yet ensuring that outputs consistently adhere to user prompts or safety preferences remains a major challenge. Standard solutions—such as fine-tuning models with preference datasets, gradient-based guidance, or generating multiple outputs and picking the best (best-of-n sampling)—are computationally expensive, inflexible, or limited to differentiable settings. The article introduces and evaluates Feynman-Kac (FK) steering, an inference-time framework designed to steer diffusion models toward user-defined reward functions without requiring model retraining.
To achieve this, FK steering tracks multiple candidate generation trajectories (called particles) and iteratively scores and resamples them during intermediate diffusion steps. Paths with higher likelihood of achieving high final rewards are multiplied, while low-reward paths are pruned. The researchers evaluated this approach across standard image generation models (Stable Diffusion variants) and text diffusion models (continuous and discrete), testing alignment metrics, text fluency, and rare-attribute steering such as toxicity red-teaming.
Key findings demonstrate that FK steering consistently outperforms conventional fine-tuning and inference baselines with minimal compute overhead. With as few as two particles, FK steering of base text-to-image models surpassed models fine-tuned specifically on preference alignment benchmarks. In addition, steering a smaller 0.8-billion parameter model achieved higher prompt fidelity and human preference scores than a substantially larger 2.6-billion parameter model while requiring less wall-clock time (9.1 seconds versus 11.5 seconds). In text diffusion tasks, FK steering lowered perplexity, improved grammatical acceptability, and dramatically enhanced rare-attribute detection—boosting toxicity detection rates from under 1% up to 64.7% for automated safety testing.
These results show that scaling compute during inference offers a viable and often superior alternative to costly model retraining. Organizations can maintain a single, frozen base model and apply diverse, plug-and-play reward functions dynamically at deployment, reducing training costs and improving adaptability. For teams implementing generative systems, the article recommends adopting particle-based inference steering when fine-tuning is impractical and tuning steering hyperparameters—such as resampling schedules and temperature—to balance output quality against sample diversity. Because highly aggressive steering can reduce output variety, future work should focus on optimizing reward estimators and testing broader production workflows.
- Paper: Loss-Guided Diffusion Models for Plug-and-Play Controllable Generation, Jiaming Song et al. (2023). This paper establishes the principles of loss-guided Monte Carlo sampling for steerable diffusion, which directly motivates and provides the baseline context for Feynman-Kac interacting particle steering.
- Paper: Classifier-Free Diffusion Guidance, Jonathan Ho et al. (2022). It introduces classifier-free guidance, establishing standard inference-time conditioning mechanisms that the source seeks to generalize and improve upon for arbitrary reward functions.
- Paper: Diffusion Models Beat GANs on Image Synthesis, Prafulla Dhariwal et al. (2021). It establishes gradient-based classifier guidance for diffusion models, providing the core guidance paradigm that Feynman-Kac steering seeks to surpass without requiring gradient approximations.
- Paper: Score-Based Generative Modeling through Stochastic Differential Equations, Yang Song et al. (2021). It establishes the continuous-time stochastic differential equation framework underlying modern diffusion sampling trajectories used by interacting particle systems.
- Paper: Denoising Diffusion Probabilistic Models, Jonathan Ho et al. (2020). It provides the foundational denoising diffusion probabilistic modeling formulation upon which particle trajectories and reverse-time sampling schemes are built.
- Paper: Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm, Qiang Liu et al. (2016). It introduces particle-based transport methods for posterior approximation, laying theoretical foundations for interacting particle inference methods.
- Paper: Diffusion-LM Improves Controllable Text Generation, Xiang Lisa Li et al. (2022). It introduces controllable text generation via diffusion models, defining the discrete text steering problem addressed in the source paper.
- Paper: Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, Aaron Lou et al. (2024). It provides the discrete score entropy diffusion modeling formulation used when applying inference-time steering mechanisms to language models.
No sufficiently relevant recommendations were found.
