Don't Stop Me Yet: Sampling Loss Minima via Dissipative Riemannian Mechanics

2026-05-14

AI Summary

This work addresses the challenge of efficiently and accurately sampling from continuous solution sets of minima in modern neural network loss landscapes, which often arise due to reparameterization invariance. To this end, the authors propose DiMS, a novel sampler that, for the first time, enables exact sampling from such solution manifolds. DiMS is grounded in dissipative Riemannian mechanics and constructs a dynamical system driven by gravitational, frictional, and kinetic forces, achieving both global exploration and local convergence on the level set of loss minima. Leveraging a physics-inspired hyperparameter mechanism to modulate exploration strength, DiMS demonstrates substantially superior performance over existing sampling methods in uncertainty quantification tasks within Bayesian inference.

Abstract

The minima of modern neural network loss functions are typically not isolated, rather they form connected components of reparameterization invariant solutions on the training data. Analytically characterizing these solutions is a hard problem, but sampling approaches are feasible. By construction, existing methods either spread over low-loss regions, and thus do not sample reparameterization invariant solutions exactly, or are inherently local, which limits exploration of other minima valleys. We propose sampling such reparameterization invariant models using a dynamical system based on kinetic energy, subject to a gravitational pull and a friction term that dissipates energy from the system. Our proposed sampler, DiMS, is guaranteed to sample exactly from the minimum level sets and depends on physically motivated hyperparameters which allows control over the exploration capabilities of the sampler. We consider uncertainty quantification in Bayesian inference as the motivating problem and observe improved performance compared to previously proposed approaches.

Problem

• loss minima
• reparameterization invariance
• sampling
• neural networks
• uncertainty quantification

Innovation

• reparameterization invariance
• loss landscape sampling
• dissipative dynamics
• Riemannian mechanics
• Bayesian uncertainty quantification