Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching
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 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 variationally stable Doob's matching, a novel framework for provable 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-regularized 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. Finally, we show that variationally stable guidance estimators are adaptive to unknown low dimensionality, effectively mitigating the curse of dimensionality under low-dimensional subspace assumptions.
Cite
@article{arxiv.2601.06514,
title = {Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching},
author = {Jinyuan Chang and Chenguang Duan and Yuling Jiao and Yi Xu and Jerry Zhijian Yang},
journal= {arXiv preprint arXiv:2601.06514},
year = {2026}
}