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

AI summary

Overview

Research area: Machine learning — reinforcement learning and generative modeling, specifically the alignment/fine-tuning of diffusion models with reward signals.

Technical level: Advanced. The paper builds directly on KL-regularized reinforcement learning, soft Q-functions, soft Bellman equations, Tweedie's formula, and consistency models, and most of its argument lives in the equations rather than in prose.

Scope (one sentence): The paper proposes and evaluates SQDF, a KL-regularized RL fine-tuning method for diffusion models that uses a training-free, differentiable approximation of the soft Q-function to apply reward gradients directly to the denoising steps.

What This Paper Is About

Diffusion models generate high-likelihood samples but often need to be aligned with a downstream objective such as aesthetics, text-image alignment, or molecular bioactivity. Existing fine-tuning methods either optimize the reward too aggressively — producing high-reward but unnatural images and losing diversity — or require unstable value-function training or high-variance Monte Carlo gradient estimates to avoid that collapse. The goal of this paper is to fine-tune diffusion models using reward gradients while keeping the model close to its pre-trained distribution, so that reward improves without semantic or diversity collapse.

Key Contributions

1. SQDF, a KL-regularized reinforcement learning method for diffusion alignment that applies a reparameterized policy gradient through a training-free, differentiable one-step approximation of the soft Q-function, avoiding separate value-function training.
2. A discount factor γ applied over denoising steps, which down-weights early denoising steps in the credit assignment, motivated by bounds derived on the discounted soft Q-function.
3. Use of a consistency model as the reference model for predicting the clean sample, replacing the unreliable high-noise posterior mean estimate that Tweedie's formula provides.
4. Off-policy training with a replay buffer, which the authors use to manage the reward–diversity trade-off and improve mode coverage, plus an empirical study separating the effect of each component.

Main Findings

• SQDF mitigates over-optimization in text-to-image fine-tuning. When optimizing the LAION aesthetic score and the HPSv2 human preference score, gradient-based baselines ReFL and DRaFT reached high target rewards but showed sharp declines in alignment and diversity; DDPO failed to reach a comparable aesthetic score and exhibited rapid diversity collapse. SQDF achieved the highest alignment and diversity at equivalent reward levels.
• SQDF dominates KL-augmented baselines on trade-off curves. When the KL-regularization coefficient was swept for both SQDF and KL-augmented versions of other methods, SQDF achieved higher rewards while maintaining better alignment and diversity, and mostly occupied the Pareto frontier across metrics.
• SQDF leads online black-box optimization. In the experimental setting adopted from Uehara et al., SQDF-UCB reached a target aesthetic score of 7.87 (0.15) and SQDF-Bootstrap 7.87 (0.11), versus SEIKO-UCB at 7.49 (0.18), SEIKO-Bootstrap at 7.80 (0.11), and PPO+KL at 6.63 (0.45). SQDF also produced positive ImageReward alignment (1.14 and 1.10) where SEIKO variants and PPO+KL reported negative values (−1.69, −1.08, −1.35). Diversity was highest for SQDF (LPIPS-div 0.51 and 0.49; DreamSim-div 0.54 and 0.53).
• Removing the discount factor helps the raw reward but hurts quality. Setting γ = 1 reached a higher final aesthetic score, but optimized more slowly in early epochs and produced a significant drop in alignment and diversity.
• The consistency model drives efficiency. In the ablation, removing the consistency model lowered the target reward from 7.87 to 7.10.
• The replay buffer preserves diversity. Removing the buffer raised the aesthetic score to 8.06 but reduced diversity (DreamSim-div 0.56, LPIPS-div 0.55 versus 0.58 and 0.56 with the buffer).
• DDIM comparison. SQDF with 2-step DDIM sampling improved optimization substantially over Tweedie's formula, but 4-step DDIM produced unstable training, attributed to high variance in the gradients; the consistency model was adopted as the Pareto solution.
• Not reported: the paper does not report training set sizes, number of prompts used, or compute/hardware budgets. Experiments were run with three random seeds and, for all tasks, with Stable Diffusion v1.5; Stable Diffusion XL is mentioned in the introduction as a target of evaluation, while the experiment section states that Stable Diffusion v1.5 was used for all tasks.

Methodology in Plain English

The diffusion denoising process is treated as a finite-horizon Markov decision process, where each denoising step is an action and a reward is given only at the final clean sample. The researchers write down a KL-regularized objective: maximize reward while staying close to the pre-trained model, which acts as the reference policy that keeps outputs natural and diverse.

The key trick is that the optimal soft Q-function for this setup can be approximated by the reward of an estimate of the clean image, without training any Q-network. Because that estimate is differentiable, the reward gradient can be pushed back through a reparameterized sample (adding scaled noise to the predicted mean) into the policy parameters. This yields a low-variance, sample-efficient update that never backpropagates through the full denoising chain.

Three additions stabilize and improve this: a discount factor that reduces the credit given to early, noisy denoising steps; a consistency model that predicts the clean image more accurately at all noise levels than the single-step Tweedie estimate; and a replay buffer that lets the method reuse high-reward and diverse samples off-policy, controlling the reward-diversity trade-off.

Evaluation proceeds in two settings. First, Stable Diffusion v1.5 is fine-tuned on differentiable rewards — the LAION aesthetic predictor and HPSv2 — and measured on prompt-alignment scores (HPSv2, ImageReward, PickScore) to detect semantic collapse and on pairwise LPIPS distance and DreamSim cosine distance to detect diversity collapse, against DDPO, DRaFT, ReFL, and KL-regularized variants of DDPO and DRaFT. Second, an online black-box optimization task iteratively fine-tunes the model using a black-box aesthetic oracle under a limited query budget, against SEIKO-UCB, SEIKO-Bootstrap, and PPO+KL.

Why This Matters

The work targets a central failure mode of reward optimization in generative models: pushing a model toward a reward proxy eventually degrades what made the model good in the first place. By folding KL regularization into an RL formulation that still exploits reward gradients, SQDF suggests that a training-free soft Q-function approximation is enough to get the benefit of both worlds, and it does so without a separate value network, which prior work has described as notoriously hard to train in diffusion settings.

Real-world applications:

• Text-to-image synthesis, where aesthetic quality and prompt alignment must both be preserved during alignment.
• Video generation, named in the paper as a domain where diffusion models are used.
• Biological molecule and molecular bioactivity design, where a reward such as bioactivity must be optimized without leaving the space of plausible molecules.
• Online black-box optimization of expensive or hard-to-differentiate oracles under a limited query budget, such as an experimental assay or a human evaluation pipeline.

Industry relevance: Alignment of generative models to human preference is a standard post-training step, and a method that reaches higher reward with less collapse and lower sample cost is directly relevant to teams shipping image and video generation systems, as well as to scientific discovery pipelines that use generative models under query budgets.

Future Directions

• Better one-step estimators. The authors name leveraging more advanced one-step distillation models as a promising direction, since the quality of the soft Q estimate caps the effectiveness of the update.
• Smarter buffer management. Improved replay-buffer strategies are proposed as a route to better control of the reward-diversity trade-off and mode coverage.
• Robustness to imperfect reward proxies. The black-box experiments show models being driven out of distribution by an inaccurate proxy; how to make training robust in that regime remains open.
• Scaling and generalization. The paper notes that Stable Diffusion XL was among the models targeted, but all reported experiments use Stable Diffusion v1.5, so behavior at larger scale and on other domains such as video or molecules is not established here.

Target Audience

Researchers and practitioners working on reinforcement learning for generative models, diffusion model alignment, or preference-based fine-tuning who already understand policy gradients and the diffusion reverse process. The paper is not beginner-friendly: it assumes familiarity with MDPs, soft Bellman equations, and reparameterization. It is most directly useful to readers evaluating KL-regularized RL objectives, reward over-optimization, or reward-diversity trade-offs in generative model post-training.

Authors’ 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.