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

2026-01-10

AI Summary

This work addresses the challenge of aligning the generative distribution of a pre-trained diffusion model to a target distribution during inference—without retraining—while preserving its original generative capabilities. The authors propose Doob's matching framework, which introduces Doob’s h-transform into the guidance mechanism of diffusion models for the first time. By employing gradient-penalized regression, the method jointly estimates the h-function and its gradient, enabling efficient distribution alignment at inference time. The approach is theoretically grounded, establishing non-asymptotic convergence rates and proving approximation guarantees in terms of the 2-Wasserstein distance between the generated and target distributions.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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 $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.

Problem

• inference-time alignment
• diffusion models
• guidance
• target distribution
• Doob's h-transform

Innovation

• Doob's h-transform
• inference-time alignment
• diffusion guidance
• gradient-penalized regression
• non-asymptotic convergence