Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Inference-Time Alignment for Diffusion Models via Doob's Matching

Jinyuan Chang, Chenguang Duan, Yuling Jiao, Yi Xu, Jerry Zhijian Yang

TL;DR

This work introduces Doob's matching, a principled inference-time alignment framework for diffusion models that tilts samples toward a target distribution without retraining the base score network. Guidance is recast as the gradient of the log-Doob's $h$-function and estimated via gradient-penalized regression to jointly recover $h$ and $\nabla h$, ensuring a consistent Doob's guidance estimator. The authors establish non-asymptotic convergence rates for the guidance estimator and for the induced controllable diffusion in the $2$-Wasserstein distance, with explicit guidance on hyperparameters and truncation/scaling for stable sampling. The methodology supports a broad range of applications, including Bayesian posterior sampling and reward-guided generation, while preserving the base model's generative capabilities. Overall, Doob's matching provides rigorous, end-to-end guarantees for inference-time alignment of diffusion models against target distributions.

Abstract

Inference-time alignment for diffusion models aims to adapt a pre-trained diffusion model toward a target distribution without retraining the base score network, thereby preserving the generative capacity of the base model while enforcing desired properties at the inference time. A central mechanism for achieving such alignment is guidance, which modifies the sampling dynamics through an additional drift term. In this work, we introduce Doob's matching, a novel framework for guidance estimation grounded in Doob's $h$-transform. Our approach formulates guidance as the gradient of logarithm of an underlying Doob's $h$-function and employs gradient-penalized regression to simultaneously estimate both the $h$-function and its gradient, resulting in a consistent estimator of the guidance. Theoretically, we establish non-asymptotic convergence rates for the estimated guidance. Moreover, we analyze the resulting controllable diffusion processes and prove non-asymptotic convergence guarantees for the generated distributions in the 2-Wasserstein distance.

Inference-Time Alignment for Diffusion Models via Doob's Matching

TL;DR

This work introduces Doob's matching, a principled inference-time alignment framework for diffusion models that tilts samples toward a target distribution without retraining the base score network. Guidance is recast as the gradient of the log-Doob's -function and estimated via gradient-penalized regression to jointly recover and , ensuring a consistent Doob's guidance estimator. The authors establish non-asymptotic convergence rates for the guidance estimator and for the induced controllable diffusion in the -Wasserstein distance, with explicit guidance on hyperparameters and truncation/scaling for stable sampling. The methodology supports a broad range of applications, including Bayesian posterior sampling and reward-guided generation, while preserving the base model's generative capabilities. Overall, Doob's matching provides rigorous, end-to-end guarantees for inference-time alignment of diffusion models against target distributions.

Abstract

Inference-time alignment for diffusion models aims to adapt a pre-trained diffusion model toward a target distribution without retraining the base score network, thereby preserving the generative capacity of the base model while enforcing desired properties at the inference time. A central mechanism for achieving such alignment is guidance, which modifies the sampling dynamics through an additional drift term. In this work, we introduce Doob's matching, a novel framework for guidance estimation grounded in Doob's -transform. Our approach formulates guidance as the gradient of logarithm of an underlying Doob's -function and employs gradient-penalized regression to simultaneously estimate both the -function and its gradient, resulting in a consistent estimator of the guidance. Theoretically, we establish non-asymptotic convergence rates for the estimated guidance. Moreover, we analyze the resulting controllable diffusion processes and prove non-asymptotic convergence guarantees for the generated distributions in the 2-Wasserstein distance.
Paper Structure (36 sections, 28 theorems, 185 equations, 1 algorithm)

This paper contains 36 sections, 28 theorems, 185 equations, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Organization
  4. Preliminaries on Diffusion Models
  5. Forward and time-reversal process
  6. Path measure and filtration
  7. Training phase: base score matching
  8. Inference phase: sampling
  9. Controllable Diffusion Models via Doob's Matching
  10. Problem setup
  11. Controllable diffusion models with Doob's $h$-transform
  12. Dynamics under the target path measure
  13. Controllable diffusion process under the base measure
  14. Doob's matching with gradient penalty
  15. Limitations of vanilla regression
  16. ...and 21 more sections

Key Result

Proposition 3.1

The Radon-Nikodym derivative process $L_t$, as defined in eq:time:dependent:likelihood, admits the following representation: where $h^{*}:[0,T]\times{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$, referred to as Doob's $h$-function, is defined as the conditional expectation of the terminal weight: Furthermore, the log-likelihood ratio satisfies the following SDE:

Theorems & Definitions (60)

  • Remark 2.1: Initialization
  • Example 1: Bayesian inverse problems
  • Example 2: Reward-guided generation
  • Example 3: Transfer learning for diffusion models
  • Proposition 3.1: Doob's $h$-function
  • Proposition 3.2
  • Remark 3.3: Novikov condition
  • Proposition 3.4
  • Remark 3.5
  • Proposition 4.1
  • ...and 50 more