Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning

Hanyang Zhao, Haoxian Chen, Ji Zhang, David D. Yao, Wenpin Tang

TL;DR

The paper addresses aligning diffusion-based generative models with human preferences by formulating fine-tuning as a continuous-time reinforcement learning problem, where the backward diffusion dynamics are controlled via score functions treated as actions.It develops a rigorous theory for policy optimization in continuous time under stochastic (SDE) or deterministic (ODE) dynamics, including a policy gradient framework and regularization that can be expressed as a KL-bound over trajectories.The approach accommodates both stochastic samplers (e.g., VP-SDE-based backward processes) and deterministic samplers (e.g., DDIM/Rectified Flow), providing a unified view that connects score matching, control, and RL in a continuous-time setting.Although empirical results are deferred to an accompanying paper, the framework promises improved generation quality and controllability by aligning score-based priors with human feedback through principled continuous-time RL.

Abstract

Reinforcement Learning from human feedback (RLHF) has been shown a promising direction for aligning generative models with human intent and has also been explored in recent works for alignment of diffusion generative models. In this work, we provide a rigorous treatment by formulating the task of fine-tuning diffusion models, with reward functions learned from human feedback, as an exploratory continuous-time stochastic control problem. Our key idea lies in treating the score-matching functions as controls/actions, and upon this, we develop a unified framework from a continuous-time perspective, to employ reinforcement learning (RL) algorithms in terms of improving the generation quality of diffusion models. We also develop the corresponding continuous-time RL theory for policy optimization and regularization under assumptions of stochastic different equations driven environment. Experiments on the text-to-image (T2I) generation will be reported in the accompanied paper.

Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning

TL;DR

The paper addresses aligning diffusion-based generative models with human preferences by formulating fine-tuning as a continuous-time reinforcement learning problem, where the backward diffusion dynamics are controlled via score functions treated as actions.It develops a rigorous theory for policy optimization in continuous time under stochastic (SDE) or deterministic (ODE) dynamics, including a policy gradient framework and regularization that can be expressed as a KL-bound over trajectories.The approach accommodates both stochastic samplers (e.g., VP-SDE-based backward processes) and deterministic samplers (e.g., DDIM/Rectified Flow), providing a unified view that connects score matching, control, and RL in a continuous-time setting.Although empirical results are deferred to an accompanying paper, the framework promises improved generation quality and controllability by aligning score-based priors with human feedback through principled continuous-time RL.

Abstract

Reinforcement Learning from human feedback (RLHF) has been shown a promising direction for aligning generative models with human intent and has also been explored in recent works for alignment of diffusion generative models. In this work, we provide a rigorous treatment by formulating the task of fine-tuning diffusion models, with reward functions learned from human feedback, as an exploratory continuous-time stochastic control problem. Our key idea lies in treating the score-matching functions as controls/actions, and upon this, we develop a unified framework from a continuous-time perspective, to employ reinforcement learning (RL) algorithms in terms of improving the generation quality of diffusion models. We also develop the corresponding continuous-time RL theory for policy optimization and regularization under assumptions of stochastic different equations driven environment. Experiments on the text-to-image (T2I) generation will be reported in the accompanied paper.
Paper Structure (17 sections, 7 theorems, 71 equations)

This paper contains 17 sections, 7 theorems, 71 equations.

Table of Contents

  1. Introduction
  2. More Related works
  3. Formulation and Preliminaries
  4. Score Based Diffusion Models
  5. Continuous-time RL
  6. Main Results
  7. Fine-tuning Diffusion Models as Stochastic Control
  8. Directly aligning deterministic sampler
  9. Conclusions
  10. Connection between discrete-time and continuous-time sampler
  11. DDPM sampler is the discretization of VP-SDE
  12. DDIM sampler is the discretization of ODE
  13. ODE explanation by analyzing the derivative
  14. Exponential Integrator Explanation
  15. Theorem Proofs
  16. ...and 2 more sections

Key Result

Theorem 1

We have that, for any $c$, the discrepancy between the $p_{\theta}$ and $p_{\theta_{pre}}$ satisfies:

Theorems & Definitions (7)

  • Theorem 1
  • Theorem 2
  • Corollary 1
  • Theorem 3
  • Corollary 2
  • Theorem 4
  • Lemma 5: Theorem 5 of jia2022policy_gradient when $R\equiv 0$