Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function

TL;DR

This paper introduces SQDF, a novel reinforcement learning-based fine-tuning method for diffusion models that improves alignment with downstream objectives while maintaining diversity and naturalness.

Contribution

We propose SQDF, a KL-regularized policy gradient method with innovations like a discount factor, consistency models, and replay buffer to enhance diffusion model fine-tuning.

Findings

SQDF achieves higher target rewards in text-to-image tasks.

SQDF maintains diversity and naturalness better than existing methods.

SQDF demonstrates high sample efficiency in black-box optimization.

Abstract

Diffusion models excel at generating high-likelihood samples but often require alignment with downstream objectives. Existing fine-tuning methods for diffusion models significantly suffer from reward over-optimization, resulting in high-reward but unnatural samples and degraded diversity. To mitigate over-optimization, we propose Soft Q-based Diffusion Finetuning (SQDF), a novel KL-regularized RL method for diffusion alignment that applies a reparameterized policy gradient of a training-free, differentiable estimation of the soft Q-function. SQDF is further enhanced with three innovations: a discount factor for proper credit assignment in the denoising process, the integration of consistency models to refine Q-function estimates, and the use of an off-policy replay buffer to improve mode coverage and manage the reward-diversity trade-off. Our experiments demonstrate that SQDF achieves superior target rewards while preserving diversity in text-to-image alignment. Furthermore, in online black-box optimization, SQDF attains high sample efficiency while maintaining naturalness and diversity. Our code is available at https://github.com/Shin-woocheol/SQDF.