Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond

Wenpin Tang, Fuzhong Zhou

TL;DR

The paper addresses reward-collapse and diversity issues in fine-tuning diffusion samplers by formulating entropy-regularized objectives and solving them via stochastic control.It derives a closed-form tilt of the pretrained distribution, develops a path-space control framework with Girsanov-based connections, and provides a Hamilton–Jacobi–Bellman solution to obtain optimal controls and initial distributions.Beyond entropy, the work extends to general $f$-divergence regularizers, deriving surrogate reward transformations and corresponding HJB equations to guide sampling and initialization.Numerical experiments on Stable Diffusion v1.5 show that Forward KL and $\gamma$-divergence outperform standard KL regularization, achieving higher aesthetics-reward with less drift and artifacts, especially at larger exploration levels.

Abstract

This paper aims to develop and provide a rigorous treatment to the problem of entropy regularized fine-tuning in the context of continuous-time diffusion models, which was recently proposed by Uehara et al. (arXiv:2402.15194, 2024). The idea is to use stochastic control for sample generation, where the entropy regularizer is introduced to mitigate reward collapse. We also show how the analysis can be extended to fine-tuning with a general $f$-divergence regularizer. Numerical experiments on large-scale text-to-image models--Stable Diffusion v1.5 are conducted to validate our approach.

Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond

TL;DR

The paper addresses reward-collapse and diversity issues in fine-tuning diffusion samplers by formulating entropy-regularized objectives and solving them via stochastic control.It derives a closed-form tilt of the pretrained distribution, develops a path-space control framework with Girsanov-based connections, and provides a Hamilton–Jacobi–Bellman solution to obtain optimal controls and initial distributions.Beyond entropy, the work extends to general $f$-divergence regularizers, deriving surrogate reward transformations and corresponding HJB equations to guide sampling and initialization.Numerical experiments on Stable Diffusion v1.5 show that Forward KL and $\gamma$-divergence outperform standard KL regularization, achieving higher aesthetics-reward with less drift and artifacts, especially at larger exploration levels.

Abstract

This paper aims to develop and provide a rigorous treatment to the problem of entropy regularized fine-tuning in the context of continuous-time diffusion models, which was recently proposed by Uehara et al. (arXiv:2402.15194, 2024). The idea is to use stochastic control for sample generation, where the entropy regularizer is introduced to mitigate reward collapse. We also show how the analysis can be extended to fine-tuning with a general -divergence regularizer. Numerical experiments on large-scale text-to-image models--Stable Diffusion v1.5 are conducted to validate our approach.
Paper Structure (9 sections, 10 theorems, 53 equations, 3 figures, 1 table)

This paper contains 9 sections, 10 theorems, 53 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Diffusion models
  3. Entropy-regularized fine-tuning and stochastic control
  4. Entropy-regularized fine-tuning
  5. Stochastic control
  6. Solve the stochastic control problem
  7. Extension to regularization by $f$-divergence
  8. Numerical experiments
  9. Conclusion

Key Result

Lemma 3.1

Assume that $\int \exp\left(\frac{r(y)}{\alpha} \right) p_{\tiny \hbox{pre}}(y) dy < \infty$. We have: where $C:=\int \exp\left(\frac{r(y)}{\alpha} \right) p_{\tiny \hbox{pre}}(y) dy$ is the normalizing constant.

Figures (3)

  • Figure 1: Generated images of the fine-tuned models with $\alpha=10$.
  • Figure 2: Generated images of the fine-tuned models with $\alpha=1$.
  • Figure 3: Generated images of the fine-tuned models with $\alpha=0.1$.

Theorems & Definitions (20)

  • Lemma 3.1
  • proof
  • Proposition 3.2
  • proof : Proof of Proposition \ref{['prop:TVd']}
  • Proposition 3.3
  • proof
  • Proposition 3.4
  • proof
  • Proposition 3.5
  • proof
  • ...and 10 more