Jinyuan Chang, Chenguang Duan, Yuling Jiao, Ruoxuan Li, Yi Xu, Jerry Zhijian Yang
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.
We introduce Flow-Based Distribution Matching (FBDM), which learns self-supervised representations through non-adversarial velocity regression. The encoder and auxiliary field are trained jointly by regressing analytic velocities along spherical paths from learned representations to reference targets. Both augmented views share a reference target, with a capacity limit on each reference center and an explicit view-alignment loss. A structured reference supports more components than the auxiliary flow dimension while retaining geometric separation. The auxiliary field is discarded after pretraining; only the backbone is retained downstream. On three ResNet-18 benchmarks with matched auxiliary parameter counts, FBDM provides 2.17 to 2.36 times higher per-epoch training throughput than critic-based Distribution Matching (DM), with about 0.21 GiB additional peak allocated GPU memory. Experiments show that FBDM is competitive with existing SSL methods from CIFAR to ImageNet and performs comparably to DM on their shared benchmarks. Under stated conditions, our population analysis bounds nearest-centroid classification error on the initial normalized pretraining representation by the FBDM loss.
Jinyuan Chang, Yuling Jiao, Lican Kang, Junjie Shi
In this work, we propose a novel deep bootstrap framework for nonparametric regression based on conditional diffusion models. Specifically, we construct a conditional diffusion model to learn the distribution of the response variable given the covariates. This model is then used to generate bootstrap samples by pairing the original covariates with newly synthesized responses. We reformulate nonparametric regression as conditional sample mean estimation, which is implemented directly via the learned conditional diffusion model. Unlike traditional bootstrap methods that decouple the estimation of the conditional distribution, sampling, and nonparametric regression, our approach integrates these components into a unified generative framework. With the expressive capacity of diffusion models, our method facilitates both efficient sampling from high-dimensional or multimodal distributions and accurate nonparametric estimation. We establish rigorous theoretical guarantees for the proposed method. In particular, we derive optimal end-to-end convergence rates in the Wasserstein distance between the learned and target conditional distributions. Building on this foundation, we further establish the convergence guarantees of the resulting bootstrap procedure. Numerical studies demonstrate the effectiveness and scalability of our approach for complex regression tasks.
Yuling Jiao, Lican Kang, Jerry Zhijian Yang, Jincheng Ying
Diffusion policies represent multimodal action distributions, but a $Q$-weighted update does not specify how to sample from the resulting target distribution. We propose Schrödinger--Föllmer Actor--Critic (SFAC), an offline-to-online reinforcement learning (RL) method for Kullback--Leibler (KL)-regularized policy improvement through conditional diffusion. A minimax Bellman critic estimates the action-value function, which defines an exponentially tilted target policy. A Doob $h$-transform expresses this update as a correction to the reference diffusion drift. The target drift is a weighted conditional mean of Brownian-bridge labels. We fit it by weighted regression using endpoints from the current policy and pointwise action-value evaluations. In the small-update regime, the KL-regularized update follows the natural policy-gradient direction, and the Doob correction represents the same local change in the space of diffusion drifts. Under suitable conditions, we derive finite-sample bounds that separate the errors from critic estimation, weighted neural drift regression, and diffusion discretization through a reverse-KL actor bound. The average-performance bound depends explicitly on the comparison coefficients over the policy iterations. Synthetic experiments assess approximation to prescribed advantage-tilted targets and sensitivity to the diffusion-step count. Across six offline-to-online continuous-control tasks, the $Q$-guided bridge-regression implementation achieves higher final-checkpoint returns than its BC initialization after $500K$ online interactions.