Time-Reversed BSDEs for Accurate Gradient Estimation in Diffusion Models

TL;DR

This paper introduces a time-reversed BSDE approach for more stable and accurate gradient estimation in diffusion models, addressing limitations of existing adjoint matching methods.

Contribution

It proposes a novel time-reversed BSDE estimator that produces an adapted adjoint process, improving gradient stability and reducing variance in diffusion model optimization.

Findings

Enhanced gradient stability demonstrated in toy experiments

Lower variance in gradient estimates compared to adjoint matching

Competitive performance in fine-tuning 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.