WEIRDO: WEak resIdual Regularized DOob's h-transform diffusion alignment

2026-09-30

AI Summary

This study addresses the challenge of estimating guiding distributions to align target distributions during diffusion model inference, where existing methods suffer from regularization bias. Drawing upon stochastic optimal control and the Doob h-transform framework, this work proposes estimating the guidance term by penalizing the residual regularized least-squares risk via a dual Sobolev norm, thereby effectively eliminating such bias. Theoretically, it is demonstrated that the L2 squared error convergence rate surpasses the minimax rate for first-order derivatives of smooth regression functions, and high-probability bounds on the total variation distance are established. Numerical experiments further validate the effectiveness of the proposed approach.

Abstract

We study the problem of estimating the guidance that steers the distribution learned by a diffusion generative model toward a tilted target $q_0 \propto w\,p_0$ at inference time. Relying on the stochastic optimal control approach, we observe that the exact drift correction is the gradient of the logarithm of Doob's $h$-function, and we study the problem of estimating it from a sample. In the present paper, we assume that the score of the pretrained model is available, that the tilting weight is bounded and positive, and that the reference distribution has a bounded support, no smoothness of the weight is required. Introducing a penalized least-squares risk in which the penalty is the residual of the space-time harmonicity equation satisfied by the $h$-function, measured in a dual Sobolev norm, we derive high-probability bounds on the squared error of the resulting guidance estimate. Since the penalty vanishes at the target, the estimator is free of regularization bias, and in favourable scenarios its rate of convergence is faster than the minimax rate of estimating first-order derivatives of a smooth regression function. Assuming that $w$ is bounded and positive with $\mathbb{E}_{p_0}[w^{-\mathrm{s}}] < \infty$ for some $\mathrm{s} \in (0,\infty]$, and that the reference data are compactly supported, we prove that the guidance is estimable in squared $L^2$ at rate $\varepsilon_n^{\mathrm{s}/(\mathrm{s}+4)}$, where $\varepsilon_n = n^{-2(β-1)/(2(β-1)+d)}.$ We also transfer the obtained bounds to the total variation distance between the marginals of the estimated and the exactly guided samplers, and illustrate the performance of the suggested approach with numerical experiments.

Problem

• diffusion models
• distribution alignment
• Doob's h-transform
• stochastic optimal control
• guidance estimation

Innovation

• Diffusion alignment
• Doob's h-transform
• Stochastic optimal control
• Residual regularization
• Regularization bias