Annealed Langevin Monte Carlo for Flow ODE Sampling

2026-04-21

AI Summary

This work addresses the challenge of sampling from unnormalized multimodal distributions by proposing a novel framework that integrates annealed Langevin dynamics with a Jarzynski-type importance reweighting scheme. The approach leverages a probability flow ordinary differential equation (ODE) derived via stochastic interpolation, employing a time-varying transition kernel to bridge a Gaussian reference distribution and the target distribution. Within this formulation, the authors establish a reweighting identity tailored to the setting and derive an optimal backward kernel that minimizes the variance of importance weights. Empirical evaluations on high-dimensional Gaussian mixture models and a 64-dimensional Allen–Cahn field system demonstrate that the method substantially outperforms both Hamiltonian Monte Carlo and direct Monte Carlo ODE-based samplers, achieving a mean squared error in velocity field estimation that scales as 𝒪(1/n).

Technology Category

Search and Optimization: Sampling/Simulation-based SearchMachine Learning: Calibration & Uncertainty QuantificationReasoning under Uncertainty: Relational Probabilistic Models

Application Category

Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphsSearch and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for rankingUser Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

We propose Annealed Langevin Monte Carlo for Flow ODE Sampling (ALMC-ODE), a method for generating samples from unnormalized target distributions, with a particular emphasis on multimodal densities that are challenging for standard Markov chain Monte Carlo methods. ALMC-ODE is based on a probability-flow ordinary differential equation (ODE) derived from stochastic interpolants, which continuously transports a standard Gaussian reference distribution at $t = 0$ to the target distribution $ρ$ at $t = 1$. The key innovation lies in an annealed Langevin Markov chain that evolves through a sequence of intermediate distributions bridging the reference and the target. The resulting importance-weighted particles, reweighted via a Jarzynski-based scheme, yield a low-variance estimator of the velocity field governing the ODE. On the theoretical side, we establish a Jarzynski-type reweighting identity for general time-inhomogeneous transition kernels, characterize the optimal backward kernel that minimizes the variance of the importance weights, and prove an $\mathcal{O}(1/n)$ mean squared error bound for the resulting velocity-field estimator. Numerical experiments on challenging benchmarks, including Gaussian mixture models and a 64-dimensional Allen--Cahn field system, demonstrate that ALMC-ODE significantly outperforms both direct Monte Carlo ODE approaches and Hamiltonian Monte Carlo when applied to highly multimodal target distributions.

Problem

• multimodal distributions
• unnormalized target distributions
• sampling
• Markov chain Monte Carlo
• probability-flow ODE

Innovation

• Annealed Langevin Monte Carlo
• Flow ODE
• Stochastic Interpolants
• Jarzynski Reweighting
• Multimodal Sampling