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

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.Source: §1 · p. 1
• 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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2
• 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.Source: §1 · p. 2

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.Source: §1 · pp. 1–2

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.Source: §2.1 · p. 2

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.Source: §2.2 · p. 2

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.Source: §3.1 · p. 3

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.Source: §3.2 · p. 3
• 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.Source: §3.2 · p. 4, §3.3.2 · p. 5

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.Source: §3.3 · pp. 4–5, §3.3.1 · p. 5

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.Source: §3.3.2 · p. 5

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.Source: §3.3.4 · p. 6

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.Source: §4.1.2 · p. 7, Figure 2 · p. 7
• 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.Source: §4.1.2 · pp. 7–8
• 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.Source: §4.1.2 · p. 8, Figure 4 · p. 8

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.Source: §4.1.2 · p. 8, Figure 5 · p. 9

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.Source: §4.2 · p. 9

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.Source: §4.3 · p. 10

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

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

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

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

• 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.Source: §C.2.5 · p. 27
• The upper-bound comparison follows because tempering makes both cM and cG no larger than their untempered values.Source: §C.2.5 · p. 27
• Setwise convergence of the empirical measure follows by substituting indicator functions for the test function.Source: §C.2.5 · p. 27

• Fine-tuning baselines use official implementations with minimal hyperparameter changes, but AlignProp checkpoints before reward collapse are used for comparison.Source: D IMPLEMENTATION DETAILS · p. 27
• Guidance baselines are adapted from official implementations, with guidance scaling and additional constants retained from those implementations.Source: D IMPLEMENTATION DETAILS · p. 28
• 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.Source: D IMPLEMENTATION DETAILS · p. 28
• 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.Source: D IMPLEMENTATION DETAILS · p. 28

• 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.Source: E TEST-TIME SCALING · p. 28
• DAS achieves superior performance with fewer particles and significantly less inference-time compute than the compared baselines.Source: E TEST-TIME SCALING · p. 28, Figure 8 · p. 29
• Best-of-N shows limited improvement beyond a compute threshold, while SMC scales better but remains below DAS.Source: E TEST-TIME SCALING · p. 28
• The proposal and tempering scheme improve sample efficiency, enabling better performance at lower cost and continued scaling with additional resources.Source: E TEST-TIME SCALING · p. 28

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.Source: F COMPARISON WITH DIFFUSION-BASED… · p. 30
• Grad parameterization incorporates the target score function, and omitting it can prevent diffusion-based samplers from fitting multimodal target distributions.Source: F COMPARISON WITH DIFFUSION-BASED… · p. 30, Figure 9 · p. 30

Experiment Setup

• The comparison includes fine-tuning, training-from-scratch, and parameterization variants evaluated on multimodal target distributions.Source: F COMPARISON WITH DIFFUSION-BASED… · p. 30

Empirical Validation

• DAS outperforms DDS in EMD, indicating that its samples are closer to the target distribution.Source: F COMPARISON WITH DIFFUSION-BASED… · p. 31

Conclusion

• DAS is a training-free solution for sampling from unnormalized target distributions under reward-alignment constraints.Source: F COMPARISON WITH DIFFUSION-BASED… · p. 31

Empirical Validation

• DAS demonstrates superior sampling from the reward-aligned target distribution on the 3d Swiss roll experiment.Source: §G.1 · p. 31, Figure 10 · p. 31

G.2 DAS WITH SDXL

• With SDXL, DAS achieves superior PickScore target optimization and HPSv2 cross-reward generalization compared with the listed baselines.Source: §G.2 · p. 31, Figure 12
• Across SD1.5 and SDXL, DAS improves PickScore, HPSv2, and ImageReward while maintaining TCE and LPIPS MPD diversity metrics.Source: §G.2 · p. 32, Table 3 · p. 32

• 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.Source: §G.3 · p. 32
• 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.Source: §G.3 · p. 32, Figure 14 · p. 33
• 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.Source: §G.3 · p. 32