PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-time Alignment of Diffusion Models without Reward Over-optimization

Sunwoo Kim, Minkyu Kim, Dongmin Park

arXiv:2501.05803v3cs.LGcs.AIcs.CVmath.ST

TL;DR

Diffusion models need alignment to target objectives without losing the diversity and generalization associated with their pretrained distributions. The paper introduces DAS, a training-free test-time Sequential Monte Carlo method with tempering, and reports comparable or superior rewards to fine-tuning while preserving diversity and cross-reward generalization across several optimization settings.

  • Problem

    Pretrained diffusion models may not represent desired target distributions, while fine-tuning can over-optimize rewards and approximate guidance may optimize target rewards inadequately.

  • Method

    DAS uses tempered Sequential Monte Carlo at test time to evaluate and resample multiple diffusion candidates according to reward and proximity to the pretrained model distribution.

  • Results

    DAS achieves comparable or superior target rewards to fine-tuning while excelling in diversity and cross-reward generalization across single-reward, multi-reward, and online optimization experiments.

  • Takeaways & Limitations

    DAS provides a versatile test-time approach for aligning diffusion models with diverse rewards without fine-tuning, while reducing alignment cost and complexity.

  • Takeaways & Limitations

    Existing methods struggle to sample multimodal reward-aligned target distributions: fine-tuning misses modes, while approximate guidance produces low rewards.

Abstract

from arXiv · show

Diffusion models excel in generative tasks, but aligning them with specific objectives while maintaining their versatility remains challenging. Existing fine-tuning methods often suffer from reward over-optimization, while approximate guidance approaches fail to optimize target rewards effectively. Addressing these limitations, we propose a training-free, test-time method based on Sequential Monte Carlo (SMC) to sample from the reward-aligned target distribution. Our approach, tailored for diffusion sampling and incorporating tempering techniques, achieves comparable or superior target rewards to fine-tuning methods while preserving diversity and cross-reward generalization. We demonstrate its effectiveness in single-reward optimization, multi-objective scenarios, and online black-box optimization. This work offers a robust solution for aligning diffusion models with diverse downstream objectives without compromising their general capabilities. Code is available at https://github.com/krafton-ai/DAS.

1 INTRODUCTION

DAS is introduced as a test-time Sequential Monte Carlo method for reward-aligning diffusion models while preserving generalization. Experiments span aesthetic and preference rewards, multiple backbones and objectives, and single-reward, multi-objective, and online settings.

  • Contribution: DAS uses Sequential Monte Carlo to guide diffusion sampling toward high-reward candidates while preserving model generalization.Tempering techniques are designed to improve sample efficiency in the number of required candidates.
  • Contribution: Unlike fine-tuning, DAS performs alignment at test time without training or extensive hyperparameter tuning.The method is evaluated on Stable Diffusion v1.5 and also shows consistency with SDXL.
  • Experiments: DAS outperforms all fine-tuning baselines on LAION aesthetic score and PickScore without sacrificing cross-reward generalization or output diversity.The reported pattern extends to SDXL, CLIPScore, and Compressibility.
  • Experiments: The method is validated across single-reward optimization, multi-objective optimization, and online black-box optimization.These scenarios are presented as tests of alignment across diverse downstream objectives.
  • Contribution: The paper provides theoretical analysis of DAS’s asymptotic properties and the benefits of tempering in SMC sampling for diffusion models.The contributions also include empirical validation across diverse optimization scenarios.

2 RELATED WORK

Prior diffusion-model alignment methods primarily use fine-tuning or guidance. Fine-tuning optimizes rewards through learned model updates, while guidance uses reward or classifier information to steer sampling, often through approximations.

  • Fine-tuning: Fine-tuning methods include supervised preference-based training, reinforcement learning, and direct backpropagation through differentiable objectives.These approaches adapt model parameters to increase alignment with reward signals.
  • Guidance: Guidance methods steer diffusion sampling using classifiers or reward models, with recent approaches designed to work with off-the-shelf models.Classifier guidance may require additional training, whereas recent methods approximate guidance directly.

3 DIFFUSION ALIGNMENT AS SAMPLING (DAS)

DAS formulates diffusion alignment as sampling from a reward-aligned distribution and uses tempered Sequential Monte Carlo to generate high-reward samples while retaining coverage of multimodal targets. The method addresses fine-tuning’s mode-seeking behavior and approximate guidance’s inexactness, with theoretical guarantees for asymptotic exactness and sample efficiency.

  • 3.1 PROBLEM SETUP: ALIGNING DIFFUSION MODELS WITH REWARDS: Reward alignment is framed as sampling from a target distribution that combines the pretrained distribution with exponentiated reward, but only an unnormalized density is available.Evaluating this density requires running a probability flow ODE, making direct sampling non-trivial.
  • 3.2 LIMITATIONS OF EXISTING METHODS: Fine-tuning can miss modes in multimodal targets because reverse-KL variational objectives encourage mode-seeking behavior and reduced diversity.The paper connects this behavior to reward over-optimization and low diversity in fine-tuned models.
  • 3.2 LIMITATIONS OF EXISTING METHODS: Approximate guidance replaces the true posterior score with a prediction based on noisy-data estimates, causing larger errors early in sampling when noise is high.The approximation uses Tweedie-based clean-data prediction and may therefore fail to sample exactly from the reward-aligned target.
  • 3.3 SAMPLING FROM REWARD-ALIGNED TARGET DISTRIBUTION VIA TEMPERED SMC: DAS uses SMC particles, intermediate target distributions, resampling, and diffusion-based backward kernels to take incremental guided steps toward the reward-aligned target.Tempering reduces the particle requirements of traditional SMC and makes reward-aligned diffusion sampling more practical.
  • 3.3.2 INTERMEDIATE TARGETS: APPROXIMATE POSTERIOR WITH TEMPERING: Tempering interpolates between the diffusion prior and target, and the schedule λ_t = (1 + γ)^(t − 1) works well in this setting.The paper presents this density-tempering adaptation as novel among SMC applications to diffusion sampling.
  • THEORETICAL ANALYSIS: Under regularity conditions, DAS is asymptotically exact, and its asymptotic-variance analysis provides theoretical justification for tempering’s sample-efficiency benefits.The convergence result concerns sample estimates of expectations under the reward-aligned target distribution.

4 EXPERIMENTS

Experiments evaluate DAS across single-reward alignment, tempering, multi-objective optimization, and online black-box optimization. Across these settings, DAS targets high rewards while preserving generalization, diversity, prompt alignment, or exploration without additional training.

  • Single-reward alignment: DAS consistently achieves high target rewards while maintaining cross-reward generalization and diversity, mitigating the over-optimization seen in fine-tuning methods.Fine-tuning methods generally have high target rewards but low diversity and poor generalization, whereas guidance methods preserve these properties but under-optimize target rewards.
  • Single-reward alignment: DAS preserves aesthetic appeal, animal features, and diversity on the “crocodile” prompt, unlike fine-tuning outputs that deviate from pretrained behavior.Fine-tuning samples show reduced diversity in colors, backgrounds, and appearances, indicating reward over-optimization.
  • Single-reward alignment: DAS substantially outperforms fine-tuning methods for PickScore across all reported metrics and produces aligned images across color, count, composition, location, style, and unusual-scene prompts.The qualitative comparison uses unseen prompts and samples generated with only 4 particles.
  • Tempering ablation: Tempering enables high aesthetic scores and ImageReward with 4 or 8 particles, while reducing latent-manifold deviation compared with untempered SMC.Without tempering, SMC suffers from over-optimization even with 32 particles; the authors recommend low γ with tuned α for quality and efficiency.
  • Multi-objective optimization: DAS achieves Pareto-optimal solutions for weighted combinations of aesthetic score and CLIPScore without fine-tuning or model interpolation.Across reward weights, DAS provides prompt alignment and aesthetic quality, while direct fine-tuning on weighted averages fails to improve the Pareto front.
  • Online black-box optimization: In online black-box optimization, DAS matches SEIKO on aesthetic scores while outperforming it on unseen rewards and diversity, supporting broader exploration without repeated diffusion-model fine-tuning.DAS directly samples from the distribution aligned with the updated surrogate model, whereas SEIKO uses a fine-tuned model.

5 CONCLUSIONS

DAS aligns diffusion models with rewards without fine-tuning, while preserving generalization and diversity across single- and multi-reward settings and supporting online exploration.

  • DAS uses Sequential Monte Carlo at test time to align diffusion models with rewards without fine-tuning.
  • DAS achieves comparable or superior target rewards to fine-tuning methods in single- and multi-reward experiments.
  • DAS preserves diversity and cross-reward generalization while optimizing target rewards.
  • Online optimization demonstrates efficient exploration of diverse, high-quality solutions.

REPRODUCIBILITY STATEMENT

The reproducibility materials provide proofs, pseudocode, and implementation details for DAS, including adaptive resampling, tempering, and online adaptation.

  • Complete proofs for the theoretical results are provided in Appendix C.
  • Pseudocode covers adaptive resampling, adaptive tempering, and adaptation to the online setting.
  • Adaptive resampling uses Effective Sample Size to decide when resampling is necessary, helping maintain particle diversity.
  • The implementation uses the Srinivasan Sampling Process resampling scheme to reduce variance relative to simple multinomial resampling.
  • The full DAS algorithm initializes particles from the prior, runs reverse-time sampling, and returns weighted particle approximations of the target distribution.

A.3 PSEUDOCODE FOR ONLINE BLACK-BOX OPTIMIZATION

The appendix describes online black-box optimization with DAS and situates it within the SMC and Feynman-Kac framework used for particle approximations and theoretical analysis.

  • Unlike SEIKO, the online DAS procedure uses DAS to generate new samples rather than fine-tuning the diffusion model at each update.
  • DAS online optimization generates samples, obtains noisy black-box feedback, updates a dataset, and trains reward and uncertainty surrogates iteratively.
  • The online procedure returns the sequence of target-aligned distributions produced across iterations.
  • For aesthetic and JPEG compressibility experiments, the surrogate model is an MLP operating on frozen CLIP embeddings.
  • Sequential Monte Carlo uses importance sampling and resampling to recursively approximate state-space models defined through Feynman-Kac structures.
  • Under bounded-potential assumptions, particle filtering and SMC samplers are asymptotically exact.

B.3 CENTRAL LIMIT THEOREMS AND STABILITY OF ASYMPTOTIC VARIANCES

The analysis characterizes convergence and variance stability for particle approximations, then connects these guarantees to DAS under boundedness and mixing assumptions.

  • Particle-filtering estimation error follows a central-limit-theorem scaling of O(N^-1/2), with asymptotic variance determining the convergence rate.
  • Cumulative asymptotic variance can grow as sampling errors accumulate, motivating strongly mixing Markov kernels.
  • Under bounded-density and bounded-potential assumptions, the Markov kernels and associated processes become strongly mixing.
  • The asymptotic variance is bounded uniformly in time under the stated assumptions.
  • Tempering can lower the uniform upper bound on asymptotic variance.
  • For DAS, the theoretical results require bounded rewards, bounded guidance gradients, and effectively bounded state supports.
  • Under these assumptions, DAS has asymptotically exact particle approximations of its final target distribution.

C.2.5 PROOF OF PROPOSITION 3

The proof shows that tempering does not increase DAS’s asymptotic variance bound and therefore supports improved sample efficiency, with corresponding setwise convergence results.

  • Tempering yields an asymptotic variance upper bound no larger than without tempering.The proof compares the bounds through cM and cG, whose relevant expressions are increasing in the tempering parameters.
  • The upper-bound comparison follows because tempering makes both cM and cG no larger than their untempered values.
  • Setwise convergence of the empirical measure follows by substituting indicator functions for the test function.

D IMPLEMENTATION DETAILS

The implementation uses fixed diffusion and particle settings for DAS, with task-dependent KL coefficients and a practical hyperparameter-selection recipe.

  • Fine-tuning baselines use official implementations with minimal hyperparameter changes, but AlignProp checkpoints before reward collapse are used for comparison.
  • Guidance baselines are adapted from official implementations, with guidance scaling and additional constants retained from those implementations.
  • DAS uses 100 diffusion time steps, γ = 0.008, and typically 16 particles across experiments.Online black-box optimization instead uses 4 particles, while ablation studies are excluded from the general setting.
  • The KL coefficient α varies by reward scale and task, using 0.01 for aesthetic score, 0.0001 for PickScore, and 0.005 for multi-objective and online optimization.

E TEST-TIME SCALING

DAS scales test-time computation more effectively than Best-of-N and SMC, achieving stronger target and unseen rewards while using less compute.

  • DAS consistently outperforms Best-of-N and SMC as inference compute increases on PickScore and HPSv2.PickScore is the target reward, whereas HPSv2 measures an unseen reward.
  • DAS achieves superior performance with fewer particles and significantly less inference-time compute than the compared baselines.
  • Best-of-N shows limited improvement beyond a compute threshold, while SMC scales better but remains below DAS.
  • The proposal and tempering scheme improve sample efficiency, enabling better performance at lower cost and continued scaling with additional resources.

F COMPARISON WITH DIFFUSION-BASED SAMPLERS

DAS is compared with training-based diffusion samplers and fine-tuning approaches for multimodal target sampling and broader diffusion-model alignment. Across these settings, DAS preserves modes, improves distribution matching, and generalizes across backbones and rewards.

  • Key Differences: Fine-tuning starts from pretrained models, whereas diffusion-based samplers typically train from scratch with random initialization.Fine-tuning incorporates pretrained prior knowledge but may increase susceptibility to mode collapse.
  • Key Differences: Grad parameterization incorporates the target score function, and omitting it can prevent diffusion-based samplers from fitting multimodal target distributions.
  • Experiment Setup: The comparison includes fine-tuning, training-from-scratch, and parameterization variants evaluated on multimodal target distributions.
  • Empirical Validation: DAS outperforms DDS in EMD, indicating that its samples are closer to the target distribution.
  • Conclusion: DAS is a training-free solution for sampling from unnormalized target distributions under reward-alignment constraints.
  • Empirical Validation: DAS demonstrates superior sampling from the reward-aligned target distribution on the 3d Swiss roll experiment.
  • G.2 DAS WITH SDXL: With SDXL, DAS achieves superior PickScore target optimization and HPSv2 cross-reward generalization compared with the listed baselines.
  • G.2 DAS WITH SDXL: Across SD1.5 and SDXL, DAS improves PickScore, HPSv2, and ImageReward while maintaining TCE and LPIPS MPD diversity metrics.

G.3 NON-DIFFERENTIABLE REWARDS

DAS handles non-differentiable rewards through online black-box optimization without modifying its core algorithm. On JPEG compressibility, it outperforms the pretrained model and DDPO while preserving CLIPScore and diversity.

  • DAS incorporates non-differentiable rewards as black-box objectives through its online black-box optimization framework.This enables optimization without modifying the core DAS algorithm.
  • DAS-UCB achieves the best JPEG compressibility score, outperforming both the pretrained model and DDPO.JPEG compressibility is measured as negative file size after quality-factor-95 compression.
  • DAS-UCB maintains CLIPScore and diversity relative to the pretrained model, mitigating reward over-optimization.Qualitative results show reduced background complexity while preserving key semantic features.
Loading 2501.05803v3…