Sampling method improves stability beyond smooth potentials
Muon meets Tamed Langevin: Momentum Preconditioning beyond Convex and gradient-Lipschitz Potentials
Machine Learning
Summary
Sampling complex probability distributions on matrices is hard when the mathematical landscape is bumpy or lacks smooth slopes. This paper introduces a new approach that tweaks the way momentum is handled in a common physics-inspired process called Langevin dynamics. The authors prove their method still targets the correct distribution and converges quickly even under relaxed smoothness conditions. They also show the algorithm remains stable over time, making it practical for sampling tasks where previous methods struggled.
What this means in practice
- •For machine learning engineers: Improve stability and efficiency in sampling complex models used in machine learning with non-smooth potential functions.
- •For computational chemists: Enable more reliable sampling of molecular configurations in simulations where energy landscapes are non-convex and irregular.
Authors
Nikolaos Makras, Sotirios Sabanis
Abstract
We consider the problem of sampling from Gibbs distributions on matrix spaces whose potential energies are neither convex nor globally gradient-Lipschitz. We introduce a family of non-quadratic kinetic energies that lead to a new underdamped Langevin system with momentum preconditioning, in which the gradient of the kinetic energy acts as a smooth spectral taming of the momentum. We prove that, under these relaxed assumptions on the potential, the resulting dynamics leaves the target Gibbs measure invariant, and we establish exponential convergence to equilibrium in a weighted total variation distance. Finally, we show that the corresponding Euler-Maruyama discretization admits moment bounds that are uniform in time, without any modification of the potential gradient, which ensures the stability of the resulting sampling algorithm.