Adaptive Momentum Langevin Dynamics for Efficient Posterior Sampling

Abstract

Sampling from high-dimensional posterior distributions is a central challenge in Bayesian inference and noisy inverse problems. Standard first-order Langevin-based methods often suffer from slow convergence and sensitivity to step-size hyperparameters, particularly in annealed score-based inverse imaging pipelines. We propose Adaptive Momentum Langevin Dynamics (AMLD), a practical stochastic correction kernel that introduces a momentum variable into the annealed posterior sampling framework and equips it with an annealing-aware momentum retention schedule. The method is fully compatible with the SNIPS framework and retains its coordinate-wise adaptive step structure, acting as a lightweight drop-in replacement for the conventional first-order Langevin correction step. Extensive experiments on three representative image inverse problems—Gaussian deblurring, inpainting, and 4× super-resolution—demonstrate that AMLD consistently achieves strong PSNR and LPIPS performance, with competitive FID in most settings, compared to three state-of-the-art baselines (DDRM, DPS, SNIPS) under both nearly noiseless and noisy measurement conditions, while reaching target reconstruction quality using fewer sampling iterations. The proposed momentum-based sampler provides empirically improved exploration and robustness across evolving posterior landscapes, offering a practical and computationally efficient alternative to first-order annealed Langevin samplers in high-dimensional Bayesian inverse problems.