Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Time-Reversed BSDEs for Accurate Gradient Estimation in Diffusion Models

Yuhang Mei, Amirhossein Taghvaei

Abstract

There is a growing literature adopting a stochastic optimal control (SOC) perspective to fine-tune diffusion models and related generative policies. A prominent class of methods, known as iterative diffusion optimization, solves the SOC problem by simulating the diffusion process, evaluating a loss function, and applying stochastic optimization algorithms, with adjoint matching emerging as a state-of-the-art approach. However, the adjoint process used in these methods is not adapted to the forward diffusion filtration, which can lead to unstable or high-variance gradient estimates. In this paper, we revisit gradient estimation in diffusion models through the lens of backward stochastic differential equations (BSDEs). We propose an alternative estimator based on a time-reversed BSDE formulation introduced in our prior work, which produces an adjoint process adapted to the underlying filtration. This adapted structure leads to more stable gradient estimates with potentially lower variance. We analyze the accuracy of the proposed estimator and compare it with adjoint matching. Numerical experiments on fine-tuning toy diffusion models demonstrate improved gradient stability and competitive performance.

Time-Reversed BSDEs for Accurate Gradient Estimation in Diffusion Models

Abstract

There is a growing literature adopting a stochastic optimal control (SOC) perspective to fine-tune diffusion models and related generative policies. A prominent class of methods, known as iterative diffusion optimization, solves the SOC problem by simulating the diffusion process, evaluating a loss function, and applying stochastic optimization algorithms, with adjoint matching emerging as a state-of-the-art approach. However, the adjoint process used in these methods is not adapted to the forward diffusion filtration, which can lead to unstable or high-variance gradient estimates. In this paper, we revisit gradient estimation in diffusion models through the lens of backward stochastic differential equations (BSDEs). We propose an alternative estimator based on a time-reversed BSDE formulation introduced in our prior work, which produces an adjoint process adapted to the underlying filtration. This adapted structure leads to more stable gradient estimates with potentially lower variance. We analyze the accuracy of the proposed estimator and compare it with adjoint matching. Numerical experiments on fine-tuning toy diffusion models demonstrate improved gradient stability and competitive performance.
Paper Structure (16 sections, 48 equations, 2 figures, 4 algorithms)

This paper contains 16 sections, 48 equations, 2 figures, 4 algorithms.

Table of Contents

  1. Introduction
  2. Problem Formulation
  3. Adjoint processes
  4. Notations
  5. Non-adapted adjoint process
  6. Adapted adjoint process
  7. Relationship between adapted and non-adapted adjoint processes
  8. Proposed methodology
  9. PDE solution of the BSDE
  10. Time-reversed diffusion
  11. Time-reversed BSDE
  12. Numerical procedure
  13. Numerical Experiments
  14. Verification in linear quadratic setting
  15. Optimal initial distribution for inverted pendulum
  16. ...and 1 more sections

Figures (2)

  • Figure 1: Numerical results for comparing the accuracy of gradient estimation methods. (a) influence of noise strength $\epsilon$ on MSE for the linear system in Sec. \ref{['sec:num-lin']}. The solid line represents the mean, and the shaded region represents the range from the minimum to the maximum across $10$ experiments. (b-c) Optimal initial distribution for the nonlinear system in Sec. \ref{['sec:num-ip']}. The heat map represents the optimization objective function and the points represent the support of the empirical distribution.
  • Figure 2: Numerical result for comparing TR-BSDE and adjoint matching on fine-tuning the diffusion model (Sec. \ref{['sec:num-diff']}): The black solid line denotes the target tilted distribution, the red dashed line denotes the fine-tuned result from the adjoint matching method, and the blue histogram denotes the fine-tuned result from TR-BSDE method.

Theorems & Definitions (1)

  • Remark 1