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

arXiv (Cornell University)

Abstract

Inference-time alignment for diffusion models aims to adapt a pre-trained reference diffusion model toward a target distribution without retraining the reference score network, thereby preserving the generative capacity of the reference model while enforcing desired properties at 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 variationally stable Doob's matching, a novel framework for provable guidance estimation grounded in Doob's $h$-transform. Our approach formulates guidance as the gradient of the logarithm of an underlying Doob's $h$-function and employs gradient-regularized 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 Wasserstein-2 distance. Finally, we show that variationally stable guidance estimators are adaptive to unknown low dimensionality, effectively mitigating the curse of dimensionality under low-dimensional linear subspace assumptions. Numerical experiments demonstrate the effectiveness of our proposed method and show the advantages compared to the baselines.

Authors 6

1. Jinyuan Chang

   Chinese Academy of Sciences · Southwestern University of Finance and Economics · Academy of Mathematics and Systems Science

   Affiliation as printed

   Joint Laboratory of Data Science and Business Intelligence , Southwestern University of Finance and Economics , Chengdu , Sichuan 611130 , China

   State Key Laboratory of Mathematical Sciences , Academy of Mathematics and Systems Science , Chinese Academy of Sciences , Beijing 100190 , China

2. Chenguang Duan Aachen Institut für Geometrie und Praktische Mathematik

   RWTH Aachen University

   Affiliation as printed

   Institut für Geometrie und Praktische Mathematik , RWTH Aachen University , Im Süsterfeld 2 , 52072 Aachen , Germany

3. Yuling Jiao

   Wuhan University

   Affiliation as printed

   School of Artificial Intelligence , Wuhan University , Wuhan , Hubei 430072 , China

4. Ruoxuan Li
5. Yi Xu

   Wuhan University

   Affiliation as printed

   School of Mathematics and Statistics , Wuhan University , Wuhan , Hubei 430072 , China

6. Jerry Zhijian Yang

   Wuhan University

   Affiliation as printed

   School of Mathematics and Statistics , Wuhan University , Wuhan , Hubei 430072 , China