English

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

Machine Learning 2025-04-18 v3 Artificial Intelligence Computer Vision and Pattern Recognition Statistics Theory Statistics Theory

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.

Keywords

Cite

@article{arxiv.2501.05803,
  title  = {Test-time Alignment of Diffusion Models without Reward Over-optimization},
  author = {Sunwoo Kim and Minkyu Kim and Dongmin Park},
  journal= {arXiv preprint arXiv:2501.05803},
  year   = {2025}
}

Comments

ICLR 2025 (Spotlight). The Thirteenth International Conference on Learning Representations. 2025

R2 v1 2026-06-28T21:02:22.242Z