PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, Rajesh Ranganath

arXiv:2501.06848v5cs.LGcs.CLcs.CV

TL;DR

Diffusion models struggle to generate samples matching user preferences, while existing approaches involve expensive training or restrict applicable rewards and model types. Feynman-Kac steering uses interacting particles and intermediate reward-based resampling at inference time, improving diffusion-model quality and controllability across image and text tasks. The approach can outperform fine-tuning and other inference-time methods, including with smaller models and off-the-shelf rewards.

  • Problem

    Diffusion models struggle with user-preference alignment, while fine-tuning is expensive and existing inference-time methods have reward or state-space limitations.

  • Method

    Feynman-Kac steering samples interacting particle systems and resamples particles during inference using potentials derived from intermediate rewards.

  • Results

    FK steering improves sample quality and controllability across image and text diffusion models, outperforming fine-tuning and other inference-time approaches.

  • Takeaways & Limitations

    Inference-time scaling with generic rewards can improve diffusion models across downstream tasks without training, including when steering smaller models.

  • Takeaways & Limitations

    FK steering and other inference-scaling approaches rely on strong reward functions, making reward modeling and automated evaluation important scope constraints.

Abstract

from arXiv · show

Diffusion models produce impressive results in modalities ranging from images and video to protein design and text. However, generating samples with user-specified properties remains a challenge. Recent research proposes fine-tuning models to maximize rewards that capture desired properties, but these methods require expensive training and are prone to mode collapse. In this work, we present Feynman-Kac (FK) steering, an inference-time framework for steering diffusion models with reward functions. FK steering works by sampling a system of multiple interacting diffusion processes, called particles, and resampling particles at intermediate steps based on scores computed using functions called potentials. Potentials are defined using rewards for intermediate states and are selected such that a high value indicates that the particle will yield a high-reward sample. We explore various choices of potentials, intermediate rewards, and samplers. We evaluate FK steering on text-to-image and text diffusion models. For steering text-to-image models with a human preference reward, we find that FK steering a 0.8B parameter model outperforms a 2.6B parameter fine-tuned model on prompt fidelity, with faster sampling and no training. For steering text diffusion models with rewards for text quality and specific text attributes, we find that FK steering generates lower perplexity, more linguistically acceptable outputs and enables gradient-free control of attributes like toxicity. Our results demonstrate that inference-time scaling and steering of diffusion models - even with off-the-shelf rewards - can provide significant sample quality gains and controllability benefits. Code is available at https://github.com/zacharyhorvitz/Fk-Diffusion-Steering .

1 Introduction

Diffusion models struggle with prompt adherence and user-preference control, while existing reward-based approaches trade off training cost, reward flexibility, or computational efficiency. FK steering addresses this with inference-time interacting particles and intermediate-potential resampling, improving image and text generation across the reported settings.

  • Method: Unlike fine-tuning and gradient guidance, FK steering supports arbitrary rewards, including nondifferentiable rewards, across discrete and continuous-state diffusion models.Fine-tuning requires expensive training, while gradient guidance is restricted to differentiable rewards and continuous-state models.
  • Method: FK steering samples multiple interacting diffusion processes, scores them with intermediate-reward potentials, and resamples promising particles while terminating lower-reward ones.The potentials are chosen so outputs approximate the reward-tilted distribution pθ(x0)exp(λr(x0)).
  • Results: With k = 4 particles, FK steering outperforms fine-tuning on prompt fidelity and aesthetic quality without training, while improving sample quality for text diffusion models.The reported text results include more competitive linguistic acceptability and perplexity relative to autoregressive samples.
  • Results: 0.8B-parameter models outperform 2.6B-parameter models on prompt fidelity using fewer FLOPs, while FK steering with k = 2 particles outperforms fine-tuned models.The text-to-image comparison uses ImageReward with base models and reports gains without reward gradients.
  • Results: FK steering also increases toxicity from 0.3% to 64.7% with k = 8 particles without gradient guidance and outperforms gradient guidance and best-of-n.This demonstrates gradient-free control of a rare specified attribute.
  • Conclusion: Across the considered settings, FK steering improves performance, highlighting inference-time scaling and steering as alternatives to reward-based fine-tuning.The framework is presented as applicable to both image and text diffusion models.

2 Related Work

Prior approaches steer diffusion models through fine-tuning, gradients, or particle methods, but each has important computational or applicability constraints. FK STEERING generalizes particle-based steering to broader rewards and discrete diffusion models.

  • Existing methods sample reward-tilted diffusion distributions through fine-tuning or inference-time steering, including universal guidance, best-of-n sampling, and SMC.
  • Fine-tuning requires training resources and couples the model to a reward, while gradient guidance requires differentiable rewards and continuous-valued diffusion models.
  • With k = 3 particles, FK-IPS outperforms fine-tuning in several settings without training.
  • FK STEERING extends particle-based Feynman-Kac methods beyond differentiable likelihoods to arbitrary rewards and discrete state spaces.
  • SVDD selects the highest-reward state at each step and copies it, reducing particle diversity.

3 Feynman-Kac Steering of diffusion models

FK STEERING uses interacting diffusion particles, intermediate reward potentials, and resampling to approximate reward-tilted target distributions during inference. Its framework supports different proposal generators, potentials, intermediate rewards, and resampling schedules.

  • Overview: FK STEERING targets ptarget(x0) ∝ pθ(x0)exp(λr(x0)) using particle-based approximations rather than direct sampling.
  • Overview: The method samples particles, scores them with potentials, and resamples them at intermediate generation steps.
  • Guarantee: As k →∞, the empirical particle distribution converges to the target distribution ptarget(x0).
  • Proposal generators: Proposal generators may use the diffusion transition kernel or reward-directed transitions such as gradient-based guidance.
  • Potentials: Importance sampling uses only a final reward potential, whereas FK STEERING scores intermediate states to favor paths likely to produce high-reward samples.
  • Potentials: Difference potentials reward particles whose intermediate rewards increase, while alternative potentials such as MAX can help when particles reach maximum reward early.
  • Resampling: Interval resampling restricts nontrivial potentials to selected steps, encouraging exploration while reducing sampling time and computation.
  • Intermediate rewards: Intermediate rewards can be learned from noisy states, including classifiers trained to approximate attributes during the diffusion process.

4 Experiments

The experiments evaluate FK STEERING for sample quality, prompt fidelity, controllability, and alternative potential or reward choices across image and text diffusion models. Across these settings, inference-time steering improves quality and control relative to fine-tuning, gradient guidance, and best-of-n baselines while scaling with particle count.

  • FK STEERING for sample quality: FK STEERING uses off-the-shelf rewards to steer text-to-image and text diffusion models for sample quality.Image experiments use ImageReward with GenEval, ImageReward, and HPS; text experiments use perplexity and linguistic acceptability rewards.
  • Text-to-image diffusion models: 11.5s versus 9.1s: FK STEERING SDv2.1 outperforms SDXL and its DPO fine-tuned version on GenEval and aesthetic quality with less sampling time.The comparison concerns a smaller steered model versus a larger model and its fine-tuned version.
  • Text-to-image diffusion models: With k = 4 particles, FK STEERING can outperform fine-tuned models on prompt fidelity and human preference alignment.The base model with k = 2 particles also achieves a higher GenEval score than DPO and DDPO fine-tuned models.
  • Text-to-image diffusion models: FK STEERING improves fine-tuned models further, increasing SDXL-DPO GenEval from 0.58 to 0.65 and SDv1.5 GenEval from 0.46 to 0.56.The table caption states that all reported metrics improve when fine-tuned models are steered.
  • Text-to-image diffusion models: FK STEERING with the model as proposal generator outperforms gradient guidance with faster sampling, while gradient-guided proposals improve performance at higher compute cost.The same comparison is reported for FK STEERING SDv1.5 with k = 4.
  • Text-to-image diffusion models: Scaling particles improves prompt fidelity and human preference alignment for all particles across all evaluated models.The effect is examined in Figure 4 using average particle performance.
  • Text diffusion models: 37.2 vs 79.2: MDLM trigram steering substantially improves perplexity, while its CoLA improvement is smaller at 35.3 vs 30.0.Across reward functions, FK STEERING with k = 4 outperforms best-of-n with n = 4 on the corresponding target metric.

5 Conclusion

The conclusion presents FK steering as an efficient, extensible inference-time approach that improves diffusion-model quality and controllability through particle scaling and configurable steering choices, while depending on strong rewards.

  • FK steering provides an efficient, extensible approach for inference-time steering of diffusion models.
  • Experiments show FK steering improves sample quality and controllability across image and text diffusion models, outperforming fine-tuning and other inference-time approaches.
  • FK steering supports plug-and-play improvements across downstream tasks by varying potentials, rewards, and samplers.Difference or max potentials with denoised-state intermediate rewards and the base diffusion model as proposal generator significantly improve performance.
  • Scaling the number of particles improves diffusion-model performance, while intermediate-reward resampling improves on best-of-n sampling.The text-to-image experiments also identify trade-offs between sample diversity, high rewards, and computation for better intermediate rewards.
  • Future work should examine variable or adaptive particle budgets and the limits of inference-time particle scaling and compute-performance trade-offs.
  • FK steering and other inference-scaling approaches rely on strong reward functions.Advancing automated evaluation and reward modeling is identified as a critical area for further improvement.

Societal impact

The paper discusses beneficial uses of controllable generation, including alignment and automated red-teaming, while acknowledging potential misuse by malicious actors.

  • FK steering can align language models with human preferences for personalization or safety and support automated red-teaming for deployment decisions.
  • Controllable-generation methods can also be used by malicious actors to generate harmful samples.
  • The authors believe the methods’ benefits for understanding and improving generative-model safety will ultimately outweigh these harms.

A Text to Image Experiments

The text-to-image experiments examine how the resampling schedule and λ affect particle diversity and reward-related performance using Stable Diffusion and ImageReward.

  • Inference timing is reported for FK steering on a single NVIDIA-A100 GPU and a two-device parallel implementation.
  • The experiments measure particle diversity with the CLIP encoder fθ for k particles.
  • The setup uses Stable Diffusion text-to-image models, ImageReward as the human-preference reward, and the difference potential.
  • The evaluation varies λ across [0,20,40,60,80] and resampling schedules across [0,70,75,80,85,90].
  • For every tested λ and resampling schedule, FK steering outperforms the base model on GenEval.
  • Lower λ yields higher CLIP diversity, whereas higher λ yields higher human-preference scores.

B Text Experiments

The text experiments use publicly available SSD-LM, MDLM, and GPT2-Medium checkpoints to generate conditioned continuations for controllable text-generation evaluation.

  • The experiments use publicly available SSD-LM, MDLM, and GPT2-Medium checkpoints.
  • Each experiment generates length-50 sequences conditioned on Han et al.’s controllable-generation prompts.
  • The evaluation generates 20 continuations for each of 15 prompts.

B.1 Baselines

The baseline configuration uses blockwise continuation generation for SSD-LM, with most experiments run at T = 500 and top-p sampling.

  • SSD-LM continuations are generated iteratively in blocks of 25.

B.2 FK STEERING Details

The text experiments use fixed FK STEERING settings, intermediate-text reward estimation, and evaluations of diversity, toxicity, and perplexity.

  • All FK STEERING text experiments use λ = 10.0, the difference-of-rewards potential, and 50 resampling events.Resampling occurs every 10 steps for SSD-LM and every 20 steps for MDLM.
  • The diversity tables report ImageReward and HPS scores, while GenEval scores measure prompt fidelity across FK STEERING particles.
  • Improved FK STEERING with MDLM evaluates 16 intermediate texts instead of the default four.
  • The toxicity evaluation reports steering-classifier and holdout-classifier toxicity together with GPT2-XL perplexity, averaged over three seeds.
  • Improved FK STEERING with SSD-LM fine-tunes toxicity classifiers on noised OpenWebText states, using 26K non-toxic and 26K toxic texts.

C.1 Choice of proposal distribution

The framework considers gradient-based and discrete normalization proposals, relates existing particle methods to FK sampling, and uses adaptive resampling to preserve diversity.

  • Proposal choices: Gradient-based guidance twists continuous-state diffusion transitions using reward gradients, requiring differentiable rewards.
  • Proposal choices: Discrete normalization provides a proposal option for discrete diffusion models such as MDLM.
  • Trade-off: Transition-kernel twisting methods can increase sampling time relative to the base model.
  • Existing methods: TDS uses SMC with a reward-difference potential and classifier-guided proposal generation for conditional sampling.
  • Framework scope: FK STEERING generalizes potentials and proposal generators to conditionally trained models, discrete diffusions, and non-differentiable rewards.
  • Existing methods: SVDD uses nested importance sampling with the base model as proposal, selecting one of k samples and copying it forward.
  • Adaptive resampling: Adaptive resampling skips resampling when ESS_t < k/2, encouraging particle diversity.

D FK STEERING samples

The image experiments examine how temperature, potentials, sampling steps, and interval resampling trade off particle diversity against prompt fidelity and sample quality.

  • Experimental setup: The SDXL experiments compare FK STEERING particles with k independent samples using ImageReward as the human-preference reward.
  • Effect of λ: Higher λ improves prompt fidelity but reduces particle diversity.
  • Effect of potential: The max potential reduces diversity compared with the difference potential.
  • Sampling steps: Increasing sampling steps from 100 to 200 increases sample diversity.
  • Interval resampling: Interval resampling encourages diversity even though overall diversity is reduced.
  • Prompt fidelity: FK STEERING selects samples that follow the prompt more closely than independent sampling.
Loading 2501.06848v5…