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

Authors: Sunwoo Kim, Minkyu Kim, Dongmin Park
Primary area: generative models
Keywords: diffusion models, alignment, reward over-optimization, sequential monte carlo samplers

Abstract:
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.

Decision: accept (spotlight)
Reviewer scores: 5, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1729232515247

thread_type:
official_review

<thread_object>
creator: reviewer_31k3

<content>
summary:
This paper propose a training-free sampling method based on Sequential Monte Carlo (SMC)  to align the diffusion models with specific objectives. Specially, Diffusion Alignment as Sampling (DAS) is designed to address the  limitations of the previous alignment approaches include fine-tuning and guidance methods. The author also provide theoretical analysis of DAS’s asymptotic properties and empirically validate DAS’s effectiveness across different tasks. Meanwhile, the authors conducte sufficient experiments to verify the validity of the DAS methodology.

strengths:
The paper is clearly written and there is a good discussion of the work involved. Based on the fact that the existing fine-tuning methods lead to the reward overoptimization problem while the guidance methods lead to the under-optimization problem, the authors propose the DAS method to alleviate these deficiencies. In addition, the authors provide a theoretical analysis of the method and give the relevant code, making the work very solid. Figure 1 illustrates the shortcomings of the existing methods as well as the advantages of the proposed method, and the experimental results are visualized by using an example of a mixed Gaussian distribution.

weaknesses:
+ The models underlying the experiments in this paper have some weaknesses, and the Stable Diffusion (SD) v1.5 model is somewhat outdated now. The Consistency model [1] and Flow model (SD3) [2]  are widely used nowadays, so I suggest the authors to conduct some experiments on the newer model so as to further illustrate the validity of the proposed method.

[1] Consistency models ICML-2024

[2] Scaling Rectified Flow Transformers for High-Resolution Image Synthesis ICML-2024

+ In addition to mixing Gaussian distributions, **Swiss rolls** are also commonly used to visualize whether a distribution has been learned or not, and due to their structural features, which can further reflect the model's ability to fit the distribution, the authors can give some visualizations that further illustrate the strengths of the proposed method.

questions:
Please refer to Weaknesses. I will also refer to other reviewers' comments
</content>
</thread_object>

</thread>
<thread>
creation date:
1729495603675

thread_type:
official_review

<thread_object>
creator: reviewer_qhss

<content>
summary:
The paper introduces a novel, training-free sampling method aimed at generating samples from a target distribution, with a specific focus on applications in Reinforcement Learning from Human Feedback (RLHF). The authors compare their approach against existing fine-tuning baselines and guidance techniques.

strengths:
- The introduction provides a clear overview of the problem.

- The proposed method appears promising and might be innovative (see Question 5.)

weaknesses:
- The choice of finetuning-based RLHF baselines may not be appropriate (see Question 1).

- The paper is sometimes hard to follow due to the delayed definition of new notations. For instance, the symbol $\gamma$ is used on line 208 but is not defined until line 250.

- The evaluation metrics used in the paper (line 355 and onward) are not explained, making it difficult to assess their relevance and meaning.

questions:
1. Baselines for Comparison: The authors correctly state that RLHF can be formulated as learning to sample from an unnormalized target distribution (Section 3.2). They show that current fine-tuning approaches in RLHF struggle to sample from multimodal target distributions, which highlights the limitations of these methods. However, this comparison may not be sufficient. There is significant research on using diffusion methods for sampling from multimodal distributions which would not fail at the examples presented in Figure 1 (e.g., [1], [2], [3] for continuous-time models, and [4] for discrete-time models). Including these approaches would provide a more convincing set of baselines. If this is not possible within the rebuttal phase's timeline, I believe it is necessary to at least mention this line of work in the paper. 

2. Clarification on Calculations: The calculation presented on line 153 and the following lines is unclear. Can you provide a detailed derivation?

3. Explanation of Evaluation Metrics: The evaluation metrics mentioned in line 355 and onward lack a clear explanation. Currently, understanding them requires consulting multiple references. Could you include a brief explanation in the paper for clarity?

4. Inference Time Comparison: How does the inference time of your method compare with fine-tuning techniques? It seems plausible that fine-tuning methods might produce samples more quickly. Is this the case?

5. Novelty of the Method: Is the proposed method entirely new, or is it simply novel in the context of RLHF? How does it compare to other Sequential Monte Carlo methods?

I will initially give a score of 3, but I am willing to update my score if my questions are properly addressed.

# References
[1] Zhang, Qinsheng, and Yongxin Chen. "Path Integral Sampler: A Stochastic Control Approach For Sampling." International Conference on Learning Representations.

[2] Berner, Julius, Lorenz Richter, and Karen Ullrich. "An Optimal Control Perspective on Diffusion-Based Generative Modeling." Transactions on Machine Learning Research.

[3] Vargas, Francisco, Will Sussman Grathwohl, and Arnaud Doucet. "Denoising Diffusion Samplers." Eleventh International Conference on Learning Representations.

[4] Sanokowski, Sebastian, Sepp Hochreiter, and Sebastian Lehner. "A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization." Forty-First International Conference on Machine Learning.

[5] Dongjun Kim, Yeongmin Kim, Se Jung Kwon, Wanmo Kang, Il-Chul Moon Proceedings of the 40th International Conference on Machine Learning, PMLR 202:16567-16598, 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730193921767

thread_type:
official_review

<thread_object>
creator: reviewer_j1vs

<content>
summary:
This paper proposes a training-free diffusion sampling method based on Sequential Monte Carlo (SMC) to sample from the reward-aligned target distribution. By incorporating tempering techniques, it offers a robust solution for aligning diffusion models with arbitrary rewards
while preserving general capabilities

strengths:
1. This paper is overall well-written and the motivation is clear. It aims to address the trade-off in diffusion models that align them with specific objectives while maintaining their versatility, which is a critical problem in generative modeling.

2. DAS’s effectiveness is comprehensively validated across diverse scenarios, including toy distribution simulation, single-reward, multi-objective, and online black-box optimization tasks.

weaknesses:
1. More intuitive explanations of SMC are suggested to add between the motivation and method to make it more consistent and intuitive since the introduction of SMC in supplementary material is a bit abstruse to understand, making the superiority of adopting SMC to address the training problem unclear.

2. How to choose hyperparameters such as $\gamma, \alpha$ and particles should be discussed across different scenarios.

questions:
Please see the weaknesses part above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698645436

thread_type:
official_review

<thread_object>
creator: reviewer_kkg4

<content>
summary:
The paper proposes DAS, a training-free approach for aligning diffusion models with specific objectives. It uses Sequential Monte Carlo (SMC) with tempering for reward alignment. This method is demonstrated across generative tasks, including single-reward and multi-objective cases, with performance comparable to fine-tuning methods in terms of target reward optimization and diversity.

strengths:
1. DAS does not require additional training, which reduces computational cost.
2.The use of SMC with tempering is justified through asymptotic properties.
3. DAS balances reward optimization and diversity, and is demonstrated across single-reward, multi-objective, and online settings.

weaknesses:
1. While DAS is compared with fine-tuning and guidance methods, comparisons to baselines like STEGANODE or controlled diffusion could have strengthened the evaluation.
2. DAS assumes differentiable reward functions, which may limit applicability in scenarios involving non-differentiable objectives.
3. Most experiments use Stable Diffusion v1.5, and additional models would have enhanced the generality of the findings.
4. The paper can do more image tasks. Currently it emphasizes findings on aesthetic score, which might not generalize well to other tasks.
5. Limitations: the setup of SMC with tempering, intermediate targets, and backward kernels can be technically demanding. And the effectiveness of DAS depends on the pre-trained model's quality, limiting performance on models with low initial diversity or reward alignment.

questions:
1. The method relies on specific tempering schemes and parameters, and the practical guidelines for selecting these could be more detailed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732216886565

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely appreciate the reviewers' positive feedback and valuable comments. Most reviewers agreed that (1) **the problem setting and motivation are clear** (Reviewer J1vs, QHss and 31k3), (2) **the methodology is sound and justified by theoretic analysis** (Reviewer Kkg4, J1vs and 31k3), (3) **the evaluation was performed extensively on diverse reward-alignment scenarios** (Reviewer Kkg4, J1vs, and 31k3), and (4) **the presentation is clear** (Reviewer J1vs and 31k3). During the rebuttal, we addressed the reviewers' remaining concerns by providing clarification with comprehensive additional experimental results (by adding the latest diffusion backbones, three baselines, one non-differentiable reward case, and a more challenging synthetic dataset; See the attached PDF file for visual results). We hope that the remaining concerns are successfully addressed by the rebuttal and are happy to answer more questions during the discussion period.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732804455861

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final Revision

comment:
We sincerely thank all reviewers for their constructive feedback that helped strengthen our work. The final revised paper incorporates several major improvements:

1. Enhanced Accessibility and Understanding
- Intuitive explanation of DAS methodology (Section 3.3 lines 196 ~ 207)
- Clear hyperparameter selection guidelines (Appendix D)
- Thorough comparison with diffusion-based samplers (Appendix E)
- Additional Swiss roll visualization (Appendix F.1)


2. Extended Applications
- Testing with SDXL architecture (Appendix F.2)
- Non-differentiable rewards via online black-box optimization (Section 4.3, Appendix F.3)
- Discussion of future works (Appendix G)


3. Improved Readability
- Optimized file size (<4MB)
- Improved notation timing and definitions in Section 3
- Enhanced evaluation metrics explanation (Section 4.1)
- Clarified optimality variable formulation (Section 3.1)

We believe these revisions address the reviewers' concerns while making the paper more accessible and practically useful. For easy reference, all changes in the revised manuscript are highlighted in green.

Again, we appreciate the thorough review process that helped improve the quality of our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper considers aligning (tilting) the diffusion model with reward. Formally, this corresponds to sampling from the following product of densities
$$p_{\text{target}} \propto p_{\text{pretrained}}(x)\exp(\beta r(x)),$$
where $p_{\text{pretrained}}(x)$ is the density of samples produced by the pre-trained diffusion model, $r(x)$ is a predefined reward model, and $\beta$ is the inverse temperature hyperparameter.

To sample from this product, the authors employ the Sequential Monte Carlo (SMC) framework analogously align (tilt) the pretrained marginals along the way by multiplying it with $\exp(\beta_t r(x))$. Furthermore, the authors introduce the transition kernel that minimizes the variance of the SMC weights (necessary for correct sampling), which results in a very practical scheme which is demonstrated empirically.

The reviewers unanimously recognize both the practical and theoretical contributions of the paper, which leaves no doubts about its acceptance.