Sampling via Stochastic Interpolants by Langevin-based Velocity and Initialization Estimation in Flow ODEs

TL;DR

This work tackles the challenge of sampling from unnormalized Boltzmann densities, especially when targets are multimodal, by constructing a probability flow ODE via linear stochastic interpolants: $X_t=tX_1+(1-t)X_0$, transported from an easily sampleable Gaussian-convolved initialization $p_{X_{T_0}}$ to the target $p_{X_1}$. Sampling is decomposed into two Langevin-based tasks: (i) draw samples from $p_{X_{T_0}}$ and (ii) estimate the velocity field $u(t,x)$ through Monte Carlo/Langevin estimates of ${ m E}[X_1|X_t=x_t]$ using $p_{X_1|X_t}$, then integrate the PFODE $rac{d}{dt}oldsymbol{ abla}igl( ext{psi}igr)(t,x)=u(t, ext{psi}(t,x))$. The method incorporates RMSprop-style preconditioning to accelerate Langevin dynamics and stabilize flow integration, with convergence guarantees and extensive experiments on multimodal 2D and high-dimensional distributions as well as Bayesian mixtures. Empirically, SSI achieves superior multimodal recovery than ULA, MALA, pULA, and HMC, and scales to high dimensions and Bayesian inference, illustrating its practical impact for robust probabilistic sampling.

Abstract

We propose a novel method for sampling from unnormalized Boltzmann densities based on a probability-flow ordinary differential equation (ODE) derived from linear stochastic interpolants. The key innovation of our approach is the use of a sequence of Langevin samplers to enable efficient simulation of the flow. Specifically, these Langevin samplers are employed (i) to generate samples from the interpolant distribution at intermediate times and (ii) to construct, starting from these intermediate times, a robust estimator of the velocity field governing the flow ODE. For both applications of the Langevin diffusions, we establish convergence guarantees. Extensive numerical experiments demonstrate the efficiency of the proposed method on challenging multimodal distributions across a range of dimensions, as well as its effectiveness in Bayesian inference tasks.

Figures (8)

• Figure 1: Density $p_{X_t}$ in \ref{['eq:gmm_interpolant']} with $m=2$ at three time points, illustrating the pitchfork bifurcation.
• Figure 2: Comparison of SSI against baseline methods for the rings distribution in Section \ref{['section:experiments:two:dim']}. The blue points represent $10^{4}$ particles generated by each method.
• Figure 3: Comparison of SSI against baseline methods for the MoG7x7 distribution in Section \ref{['section:experiments:two:dim']}. The blue points represent $10^{4}$ particles generated by each method.
• Figure 4: Comparison of SSI against baseline methods for the MoG40 distribution in Section \ref{['section:experiments:two:dim']}. The blue points represent $10^{4}$ particles generated by each method.
• Figure 5: Pairwise marginal distributions on ${\mathbb{R}}^{2}$ for the Many Well distribution. The orange points represent samples from the ground truth distribution, while the blue points represent the particles generated by SSI.

Theorems & Definitions (23)

• Theorem 2.1
• Example 3.1: Gaussian mixture
• Lemma 3.2
• Proposition 4.1
• Remark 4.2: Early-stopping
• Proposition 5.1
• proof
• Remark 5.2: Regularizing effect of Gaussian convolution
• Lemma 5.3: Linear growth of velocity field
• Theorem 5.4: Lipschitz continuity of velocity field