Test-time Alignment of Diffusion Models without Reward Over-optimization
Sunwoo Kim, Minkyu Kim, Dongmin Park
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext cc717519-30cc-4641-9cd6-a294c415893bCited by top-tier papers35
- Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-based DecodingXiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia et al.NeurIPS 2025 · 147 citations
- Reasoning with Sampling: Your Base Model is Smarter Than You ThinkAayush Karan, Yilun DuICLR 2026 · 87 citations
- Inference-time scaling of diffusion models through classical searchXiangcheng Zhang, Haowei Lin, Haotian Ye, James Y. Zou et al.ICLR 2026 · 57 citations
- Inference-Time Text-to-Video Alignment with Diffusion Latent Beam SearchYuta Oshima, Masahiro Suzuki, Yutaka Matsuo, Hiroki FurutaNeurIPS 2025 · 50 citations
- Inference-Time Scaling for Flow Models via Stochastic Generation and Rollover Budget ForcingJaihoon Kim, Taehoon Yoon, Jisung Hwang, Minhyuk SungNeurIPS 2025 · 43 citations
Builds on48
- Learning Transferable Visual Models From Natural Language SupervisionAlec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh et al.ICML 2021 · 47,906 citations
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
Related papers
- Diffusion Alignment as Variational Expectation-MaximizationJaewoo Lee, Minsu Kim, Sanghyeok Choi, Inhyuck Song et al.ICLR 2026 · 2 citations
- Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding OptimisationTaehoon Kim, Henry Gouk, Timothy M. HospedalesCVPR 2026 · 1 citation
- Test-Time Guidance for Flow-Based Generative Models via Parallel Tempering on Source DistributionsShih-Hsin Wang, Joel Keller, Taos Transue, Drake Brown et al.ICML 2026
- Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-FunctionHyeongyu Kang, Jaewoo Lee, Woocheol Shin, Kiyoung Om et al.ICLR 2026 · 5 citations
- Diffusion Tree Sampling: Scalable inference‑time alignment of diffusion modelsVineet Jain, Kusha Sareen, Mohammad Pedramfar, Siamak RavanbakhshNeurIPS 2025 · 41 citations
