Annealed Langevin Monte Carlo for Flow ODE Sampling

TL;DR

ALMC-ODE introduces an annealed Langevin Monte Carlo method based on a probability-flow ODE for efficient sampling from complex, multimodal distributions, outperforming traditional methods in challenging benchmarks.

Contribution

The paper develops ALMC-ODE, a novel sampling technique combining annealed Langevin chains with importance reweighting, providing theoretical guarantees and improved performance on multimodal targets.

Findings

ALMC-ODE outperforms HMC and direct ODE methods on multimodal benchmarks.

Theoretical analysis establishes a Jarzynski-type reweighting identity and variance bounds.

Numerical experiments demonstrate significant improvements in high-dimensional, multimodal sampling.

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