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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Chang, Jinyuan, Duan, Chenguang, Jiao, Yuling, Xu, Yi, Yang, Jerry Zhijian
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866918321867194368
author Chang, Jinyuan
Duan, Chenguang
Jiao, Yuling
Xu, Yi
Yang, Jerry Zhijian
author_facet Chang, Jinyuan
Duan, Chenguang
Jiao, Yuling
Xu, Yi
Yang, Jerry Zhijian
contents 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 $h$-transform. Our approach formulates guidance as the gradient of 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 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_06514
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching
Chang, Jinyuan
Duan, Chenguang
Jiao, Yuling
Xu, Yi
Yang, Jerry Zhijian
Machine Learning
Numerical Analysis
Optimization and Control
Statistics Theory
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 $h$-transform. Our approach formulates guidance as the gradient of 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 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.
title Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching
topic Machine Learning
Numerical Analysis
Optimization and Control
Statistics Theory
url https://arxiv.org/abs/2601.06514