PaperlayerView exact revision on arXiv

Source-linked AI summary

Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control

Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, Sergey Levine

arXiv:2402.15194v2cs.LGcs.AIstat.ML

TL;DR

Goal-directed fine-tuning can exploit imperfect rewards, causing reward collapse through low genuine rewards, reduced diversity, and deviation from the training distribution. ELEGANT frames fine-tuning as entropy-regularized control against a pretrained diffusion model. The authors report theoretical and empirical evidence of diverse, high-reward samples and mitigation of overoptimization across image and biological sequence generation.

  • Problem

    Goal-directed fine-tuning can exploit learned reward proxies beyond their training-data support, reducing diversity and producing samples with low genuine rewards.

    Source: §1 · p. 2

  • Method

    ELEGANT fine-tunes diffusion models as entropy-regularized stochastic control against the pretrained model, optimizing drift and initial distribution.

    Source: §1 · p. 2

  • Results

    The authors report that ELEGANT generates samples with high genuine rewards and diversity, with theoretical and empirical support across image and biological sequence generation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that their approach helps alleviate overoptimization while facilitating the generation of diverse samples.

    Source: §9 · p. 9

    The stated small-error guarantee for learned rewards is on pdata and does not extend beyond its support.

    Source: §4 · p. 3

Abstract

from arXiv · show

Diffusion models excel at capturing complex data distributions, such as those of natural images and proteins. While diffusion models are trained to represent the distribution in the training dataset, we often are more concerned with other properties, such as the aesthetic quality of the generated images or the functional properties of generated proteins. Diffusion models can be finetuned in a goal-directed way by maximizing the value of some reward function (e.g., the aesthetic quality of an image). However, these approaches may lead to reduced sample diversity, significant deviations from the training data distribution, and even poor sample quality due to the exploitation of an imperfect reward function. The last issue often occurs when the reward function is a learned model meant to approximate a ground-truth "genuine" reward, as is the case in many practical applications. These challenges, collectively termed "reward collapse," pose a substantial obstacle. To address this reward collapse, we frame the finetuning problem as entropy-regularized control against the pretrained diffusion model, i.e., directly optimizing entropy-enhanced rewards with neural SDEs. We present theoretical and empirical evidence that demonstrates our framework is capable of efficiently generating diverse samples with high genuine rewards, mitigating the overoptimization of imperfect reward models.

1. Introduction

Goal-directed fine-tuning can exploit imperfect rewards, reducing diversity and moving samples away from the training distribution. The paper introduces ELEGANT, an entropy-regularized control method, and reports theoretical and empirical evidence for mitigating reward collapse across image and biological sequence generation.

1. Introduction

  • Reward collapse arises when fine-tuning exploits learned nominal rewards beyond their training-data support, producing repeated samples with low genuine rewards.Diversity is also desirable in its own right, including for generating protein-sequence candidates for wet-lab experiments.Source: §1 · p. 2
  • ELEGANT frames fine-tuning as entropy-regularized stochastic optimal control, combining terminal reward with regularization against the pretrained diffusion model for both drift and initial distribution.The regularization is intended to preserve the pretrained model's trajectory bridges and limit deviation from that model.Source: §1 · p. 2
  • The authors show that jointly optimizing drift and initial distribution can yield diverse, high-nominal-reward samples within the training distribution's support, where nominal rewards approximate genuine rewards.They present this as a way to mitigate reward collapse.Source: §1 · p. 2
  • ELEGANT is reported to generate samples with high genuine rewards and diversity, with stronger theoretical support and superior empirical performance to KL-penalized RL across image and biological sequence generation.The authors apply the method to both domains.Source: §1 · p. 2

2. Related Works

The paper situates its fine-tuning method among diffusion, guidance, and reward-optimization approaches. It distinguishes ELEGANT through entropy regularization, learning an initial distribution, and theoretical and empirical support for mitigating reward collapse.

Diffusion models

  • Diffusion and bridge-matching models use SDEs, while the paper's fine-tuning methods apply to models with diffusion terms but not purely ODE-based models.The paper directs readers to Shi et al. (2023) for a broader review.Source: §2 · p. 2

Guidance

  • Guidance steers diffusion samples using an auxiliary differentiable objective; the experiments compare ELEGANT with a reward-gradient guidance baseline.Classifier-based guidance was introduced by Dhariwal and Nichol (2021).Source: §2 · p. 2

Fine-tuning as RL/control

  • Unlike related reward-backpropagation methods, ELEGANT adds entropy regularization and learns an initial distribution, which the authors say are important for targeting the desired distribution.The paper presents theoretical results and empirical evidence for more efficient reward-collapse mitigation.Source: §2 · p. 2

3. Preliminaries

The preliminaries describe continuous-time diffusion models as SDEs and establish notation for the pretrained process and its distributions. Proofs are deferred to Appendix C.

3. Preliminaries

  • A continuous-time diffusion model is described by an SDE with drift and diffusion coefficients and an initial distribution; training learns the drift so the terminal distribution corresponds to the data distribution.The paper notes that some works use the opposite time convention for the initial distribution.Source: §3 · p. 3
  • The pretrained SDE induces a terminal distribution pdata that represents the data distribution, and the paper defines notation for process measures, marginals, joint distributions, and conditional distributions.The paper uses these distributions to describe trajectories and endpoint-conditioned processes.Source: §3 · p. 3
  • The proofs are deferred to Appendix C.Source: §3 · p. 3

4. Desired Properties for Fine-Tuning

The paper frames fine-tuning around high reward, sample diversity, and proximity to the pretrained diffusion model, using a target distribution with entropy regularization. It explains how removing that regularization can lead to overoptimization beyond the pretrained distribution's support.

4. Desired Properties for Fine-Tuning

  • The stated small-error guarantee for the learned reward applies on pdata; it does not extend outside pdata's support.Source: §4 · p. 3
  • Fine-tuning should preserve high rewards, sample diversity, and sufficient proximity to the pretrained diffusion model; proximity matters because learned rewards tend to be accurate on pdata's support.The paper describes the genuine reward as usually unknown and the nominal reward as a data-learned proxy.Source: §4 · p. 3
  • The objective combines a reward term with entropy regularization against pdata, and its target distribution provides a tractable, principled fine-tuning target.The reward term is intended to preserve high rewards, while the entropy term preserves diversity and proximity to pdata.Source: §4 · p. 3

4.1. Importance of KL Regularization

  • As α approaches zero, the target distribution tends toward a Dirac delta at x⋆, the reward-optimal point within pdata's support.Source: §4.1 · p. 4
  • Without KL regularization, optimization may leave pdata's support, where nominal rewards may be inaccurate and high nominal reward can correspond to poor genuine reward.The paper calls this overoptimization and notes that generated images or biological sequences may no longer retain their domain characteristics.Source: §4.1 · p. 4

5. Entropy-Regularized Control with Pre-Trained Models

The paper formulates diffusion fine-tuning as entropy-regularized control that learns a drift and initial distribution to steer a pretrained SDE toward a reward-weighted target. Its optimal marginals preserve the pretrained process’s conditional bridges, connecting reward optimization with staying within the pretrained data support.

5.1. Stochastic Control Formulation

  • The learned drift u and initial distribution ν maximize terminal reward while penalizing trajectory KL to the pretrained process, yielding target p_tar at T.The target is reached by sampling the SDE with the optimal pair (u⋆, ν⋆).Source: §5.1 · p. 4
  • At intermediate t, the optimal marginal combines the value function with the pretrained marginal, using a normalizing constant independent of t.This gives an explicit form for the optimal marginal distribution.Source: §5.1 · p. 4

5.2. Feynman–Kac Formulation

  • The value function is higher for states that can reach high-reward regions by following the pretrained diffusion, and the optimal drift moves toward states where higher rewards are easier to attain.The optimal control equals the scaled reward gradient at T, but not at earlier times.Source: §5.2 · p. 5

5.3. Bridge Preserving Property

  • The optimal and pretrained SDEs have identical posterior distributions conditioned on terminal states, a consequence of entropy regularization that preserves bridges and helps keep samples within pretrained-data support.The method can be viewed as bridge-matching fine-tuning with the terminal distribution set to the reward-weighted target.Source: §5.3 · p. 5

6. Learning an Optimal Initial Distribution via Entropy-Regularized Control

The second control problem learns a stochastic initial distribution, which standard neural SDE solvers typically do not provide. The method uses an SDE to represent that distribution, avoiding the potential multimodal misspecification of a Gaussian model.

6. Learning an Optimal Initial Distribution via Entropy-Regularized Control

  • Standard neural SDE solvers typically fix the initial point or assume a Dirac distribution, whereas this control problem requires learning a stochastic initial distribution that can act as a sampler.A Gaussian model with a neural-network-parameterized mean may be significantly misspecified when the target distribution is multimodal.Source: §6 · p. 5
  • The method approximates the target initial distribution with an SDE, whose induced distributions can represent intricate multimodal distributions.The reference SDE starts from a typically simple initial distribution.Source: §6 · p. 5
  • Solving the second control problem and following its learned SDE from −T to 0 produces samples from ν⋆ at time 0.The theorem justifies the marginal distribution at time 0 induced by that SDE.Source: §6 · p. 6

7. Algorithm

ELEGANT estimates a value function, then solves two stochastic control problems with neural SDE solvers to learn drift terms and an initial distribution. Its implementation faces approximation error from neural-network expressiveness and discretization.

7. Algorithm

  • ELEGANT first learns the value function, then solves a control problem with a neural SDE solver, and finally solves a second control problem using the learned initial distribution.The learned drift terms are used during sampling by following the learned SDE.Source: §7 · p. 6

7.1. Value Function Estimation

  • Because regressing exp(r(x_T)/α) on x_0 may be numerically unstable for small α, the authors recommend directly regressing r(x_T) on x_0, assuming r(x_T)=k(x_0)+ε with noise under P_data.Source: §7.1 · p. 6

7.2. Sources of Approximation Errors

  • Approximation errors arise from neural SDE solver expressiveness and discretization, as well as from the expressiveness of the value-function network.Discretization contributes errors in both the solver and sampling phases.Source: §7.2 · pp. 6–7

8. Experiments

The experiments evaluate whether ELEGANT can generate high-reward, diverse samples while staying close to the training distribution. Across biological sequences and images, results indicate improved reward, diversity, or training efficiency versus baselines, with reduced overoptimization.

8. Experiments

  • Experiments compare ELEGANT, NO KL, PPO + KL, and Guidance using learned nominal rewards and measure Reward, KL-Div, and Div against a pretrained-data-support goal.The authors evaluate the mean metrics across trials; the nominal reward is trained to approximate genuine reward on the pretrained model's data support.Source: §8 · p. 7

8.1. Protein and DNA Sequences

  • Across GFP and TFBind, ELEGANT exceeds PPO + KL and Guidance in reward while maintaining lower KL divergence and higher diversity.The supplied tables cover GFP and TFBind; their captions give setup details and pretrained-model baselines, not the individual result values.Source: §8.1 · p. 7, Table 1 · p. 8, Table 2 · p. 8
  • Without entropy regularization, samples can have high nominal reward but near-zero diversity; increasing α trades a slight average-reward decline for more diverse samples.Figure 2 shows TFBind histograms for genuine reward r⋆(x) and nominal reward r(x).Source: §8.1 · pp. 7–8, Figure 2 · p. 8

8.2. Quantitative Evaluation of Overoptimization

  • On TFBind, NO KL achieves high nominal r but not genuine r⋆, whereas ELEGANT circumvents overoptimization.The comparison is presented in Table 3; its caption specifies α = 0.01 for ELEGANT and PPO and K = 0.8T for Truncation.Source: §8.2 · p. 8, Table 3 · p. 8

8.3. Images

  • For aesthetic fine-tuning, ELEGANT trains significantly faster than PPO + KL and keeps KL consistently below NO KL; Figure 4 shows generated examples.Figure 3 plots reward and KL divergence against reward queries.Source: Figure 3 · p. 8, Figure 4 · p. 8, §8.3 · p. 9

9. Conclusion

The conclusion presents ELEGANT as an efficient, theoretically and empirically grounded approach that helps alleviate overoptimization and supports diverse samples. The appendices situate the method among related work and connect its inference formulation to classifier guidance.

9. Conclusion

  • The paper concludes that its computationally efficient, theoretically and empirically grounded approach helps alleviate overoptimization and facilitates diverse samples.Future work will investigate newer diffusion models tailored to biological or chemical applications.Source: §9 · p. 9

A. Additional Related Works

  • For language-model fine-tuning, KL regularization is described as the most commonly used approach, and the paper compares against a KL-penalized RL baseline.The authors characterize that baseline as analogous to the dominant approach in language-model fine-tuning.Source: A. Additional Related Works · p. 13
  • Related path-integral sampling work differs in target distribution and fine-tuning focus; its common Dirac-delta initial-distribution assumption is unsuitable for this work.The paper also distinguishes its diffusion-model focus from standard-RL control work.Source: A. Additional Related Works · p. 13
  • The paper relates classifier-based guidance to MALA-style sampling and notes its practical instability from the need to train numerous classifiers.Its experiments compare ELEGANT with a reward-gradient guidance baseline.Source: A. Additional Related Works · p. 13

B. Connection with Classifier Guidance

  • Appendix B relates the optimal SDE to classifier guidance: under a reward approximation involving k(x_t) and noise, the inference-time approximation takes the same form as Equation (12), up to a constant.The Gaussian-noise condition is stated for the equivalence.Source: B. Connection with Classifier Guidance · pp. 13–14

C. Proofs

The proofs establish the optimal control and initial-distribution result through an intuitive bridge-preserving argument and formal derivations. The diffusion-model overview describes SDE-based generation, including finite-horizon and bridge-matching considerations.

C.1. Intuitive proof of Theorem 1

  • The intuitive proof presents a bridge-preserving argument and explains why optimization over the induced path distribution works with Markovian drift.The optimal drift is itself Markovian, as noted in the remark.Source: §C.1 · p. 15

C.2. Formal Proof of Theorem 1

  • The formal proof establishes the trajectory distribution induced by the optimal control and optimal initial distribution.It checks that the constructed distribution is valid and derives a closed-form optimal initial distribution.Source: §C.2 · pp. 15–16

C.3. Another Formal Proof of Theorem 1

  • The alternate formal proof verifies that the proposed marginal density satisfies the Kolmogorov forward equation.It also gives the optimal initial distribution and derives the optimal drift using Lemma 1.Source: §C.3 · pp. 16–17

C.4. Proof of Lemma 1

  • The proof of Lemma 1 uses the Hamilton–Jacobi–Bellman equation to characterize the optimal control and value function, then applies Feynman–Kac.The displayed steps derive the control and invoke the Feynman–Kac formula for the conclusion.Source: §C.4 · pp. 17–18

C.5–C.7. Proofs of Lemmas 2–4

  • The Lemma 3 proof obtains joint and conditional distributions by marginalizing at different points along the trajectory.It identifies forward and backward conditional distributions.Source: §C.6 · p. 18

C.8. Proof of Theorem 2

  • Theorem 2’s proof is omitted because it is the same as the proof of Theorem 1.Source: §C.8 · p. 18

D. Diffusion models

  • The diffusion-model overview describes SDEs trained so their time-T marginal follows the data distribution, with DDPMs learning a score for time-reversed generation.A finite-horizon forward SDE may fail to reach a predefined prior; the Schrödinger Bridge and bridge-matching formulations address the setup through reference processes connecting initial and data distributions.Source: D. Diffusion models · p. 19

E. Details of Implementation

The implementation appendix describes neural-SDE training, value-function estimation, baselines, datasets, network and hyperparameter references, and image-task procedures.

E.1. Implementation Details of Neural SDE

  • Neural-SDE parameters are updated with SGD using automatic differentiation, with Adam used to select the practical learning rate.The procedure simulates trajectories using an off-the-shelf discretization method.Source: §E.1 · p. 20

E.2. More Refined Methods to Learn Value Functions

  • The refined value-function method uses Monte Carlo estimates and empirical risk minimization on sampled trajectories conditioned on initial states.Source: §E.2 · p. 20

E.3. Baselines

  • The baselines include PPO + KL updates around the pretrained drift and guidance using a trained time-conditional reward model at sampling.The PPO + KL implementation modifies DDPO with an added KL penalty; guidance combines the pretrained score with a reward-conditioned score.Source: §E.3 · p. 21

F.1.1. DATASET.

  • TFBind uses 65,792 eight-base DNA sequences, while GFP experiments select 33,637 samples from a dataset of 56,086 amino-acid sequences.TFBind uses half its dataset to train an imperfect learned reward; GFP models sequence differences from a baseline sequence.Source: §F.1.1 · p. 21

F.1.2. STRUCTURE OF NEURAL NETWORKS.

F.1.3. HYPERPARAMETERS

F.1.4. ADDITIONAL RESULTS

F.2.1. DETAILS OF IMPLEMENTATION

  • Image experiments use DDIM with 50 steps and LoRA plus gradient checkpointing to back-propagate through the sampling and decoding process.Guidance is trained on AVA’s more than 250,000 evaluations; image-task training uses four A100 GPUs and animal prompts sampled uniformly from 50 classes.Source: §F.2.1 · pp. 22–23

F.2.2. ADDITIONAL RESULTS

ELEGANT outperforms the reported baselines in aesthetic-score fine-tuning, while the comparisons expose trade-offs in reward collapse, image fidelity, and optimization stability. Training curves and additional generated images provide further comparisons under the stated evaluation setup.

F.2.2. ADDITIONAL RESULTS

  • ELEGANT with α = 10 consistently generates images with higher scores than the baseline, and the improvement persists to a lesser degree at α = 1.The qualitative samples also suggest improved capacity for generating high-fidelity images.Source: §F.2.2 · p. 24
  • NO KL exhibits severe reward collapse by epoch 10 at reward ≈7, after which rising reward scores do not correspond to meaningful generations.The passage also reports that NO KL images fall short of ELEGANT in scores, notably for hippopotamus and crocodile depictions.Source: §F.2.2 · p. 24
  • Strong guidance produces more colorful images at the expense of image quality, while weak guidance or target levels yield weak reward-guided generation.The Figure 7 qualitative comparison uses target 10 and guidance level 200 to balance generation performance and fidelity.Source: §F.2.2 · p. 24
  • PPO + KL maintains stable generation quality and robust KL divergence, but limited training samples make it inefficient in training time and synthesis performance.Source: §F.2.2 · p. 24
  • PPO achieves higher reward than PPO + KL but still falls short of ELEGANT in the reported training curves.Figure 6 plots reward and KL divergence against the number of reward queries.Source: §F.2.2 · p. 24, Figure 6 · p. 25
  • Figure 7 shows additional samples from ELEGANT and baselines using 12,800 reward inquiries, with evaluation prompts unseen during training.The guidance model is trained on AVA.Source: Figure 7 · p. 26
Loading 2402.15194v2…