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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kim, Sunwoo, Kim, Minkyu, Park, Dongmin
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908324172136448
author Kim, Sunwoo
Kim, Minkyu
Park, Dongmin
author_facet Kim, Sunwoo
Kim, Minkyu
Park, Dongmin
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2501_05803
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Test-time Alignment of Diffusion Models without Reward Over-optimization
Kim, Sunwoo
Kim, Minkyu
Park, Dongmin
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
Statistics Theory
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.
title Test-time Alignment of Diffusion Models without Reward Over-optimization
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
Statistics Theory
url https://arxiv.org/abs/2501.05803