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

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

Problem Statement
Modern neural networks are typically overparameterized, resulting in loss landscapes where the global minima are not isolated points but form connected, continuous submanifolds of reparameterization-invariant solutions. While these solutions perform equally well on training data, they can induce diverse functional behaviors on unseen inputs. Existing approaches to uncertainty quantification face a dichotomy: methods that sample low-loss regions often fail to enforce exact interpolation of the training data, while methods restricted to the reparameterization-invariant set are often inherently local, limiting their ability to explore distinct minima valleys. The paper addresses the challenge of navigating the space of "good solutions" to sample diverse models that strictly minimize training loss while maintaining high predictive uncertainty on out-of-distribution (OOD) data.

Methodology: Dissipative Minima Sampler (DIMS)
The authors propose the Dissipative Minima Sampler (DIMS), a geometry-driven dynamical system designed to sample exactly from the minimum level sets of the loss function. The approach is grounded in Riemannian mechanics and classical dissipative systems.

1. Riemannian Formulation: The loss surface is treated as a Riemannian manifold $(M, g)$ embedded in $\mathbb{R}^{K+1}$. The metric tensor $G(\theta)$ is defined as a pullback metric, $G(\theta) = I_K + \nabla_\theta L(\theta)\nabla_\theta L(\theta)^\top$, which incorporates the curvature of the loss landscape into the geometry of the parameter space.
2. Dynamical System: The sampling process models a particle moving on this manifold. The dynamics are governed by a second-order ordinary differential equation (ODE) derived from the Euler-Lagrange equations. The system includes:
   • Kinetic Energy: $T(\alpha, \dot{\alpha}) = \frac{1}{2}\dot{\alpha}^\top G(\alpha)\dot{\alpha}$.
   • Potential Energy: $V(\alpha) = \kappa(L(\alpha) - L(\theta^*_{glob}))$, acting as a gravitational pull toward minima.
   • Dissipation: A friction term $-\eta(t)\dot{\alpha}$ is introduced to break energy conservation.
3. Speed-Dependent Friction: A key innovation is the definition of the dissipation function $\eta(t)$, which scales with the particle's speed (the square root of kinetic energy) and the angle between the velocity and the gradient. This ensures that energy is dissipated more rapidly when moving against the gradient, while allowing the particle to traverse the minimum level set (where the gradient is orthogonal to velocity) with minimal friction.
4. Sampling Procedure: The sampler initializes a particle at a trained minimum $\theta^*$ with a random velocity sampled from a distribution (typically a Laplace approximation respecting local geometry). The ODE is integrated numerically until the particle converges to a new point on the minimum submanifold $S^*(D)$.

Key Contributions

1. Deterministic Sampler for Minimum Submanifolds: DIMS provides a method to continuously visit distinct points on the minimum submanifolds of the training loss. It is formulated as a geometry-driven dynamical system with physically interpretable hyperparameters (friction coefficient $\eta_0$) that control exploration.
2. Theoretical Guarantees: The paper establishes that the proposed dynamics explore the loss surface before converging to the minimum manifold in finite time. Specifically, Theorem 1 guarantees convergence to the set of equilibria (local minima) under the strict saddle property, while Theorem 2 and 3 provide bounds on the convergence time and the arc-length traveled, ensuring non-trivial exploration.
3. Empirical Validation: The authors demonstrate that posterior approximations concentrated on the minimum manifold generalize well. DIMS improves predictive performance, calibration, and specifically OOD uncertainty quantification compared to existing Laplace-based approximations (Standard LA, Linearized LA, and Riemannian LA).

Experimental Results
The method was evaluated on regression (Snelson dataset) and classification tasks (Banana, UCI datasets, MNIST, Fashion MNIST).

• Regression: On the Snelson dataset, DIMS successfully converged to solutions that fit training data perfectly while exhibiting high uncertainty OOD, whereas the Riemannian Laplace Approximation (RLA) failed to guarantee interpolation when initialized from suboptimal positions.
• Classification: On UCI datasets, DIMS achieved competitive or superior Negative Log-Likelihood (NLL) and calibration metrics compared to RLA and other baselines. Notably, DIMS was robust to the friction hyperparameter $\eta_0$, often performing well with a fixed default value across diverse datasets.
• OOD Performance: On MNIST and Fashion MNIST, DIMS consistently outperformed other methods in AUROC scores for OOD detection, indicating superior ability to distinguish in-distribution from out-of-distribution samples.

Significance and Claims
The paper claims that DIMS offers a principled approach to uncertainty quantification in deep learning by explicitly targeting the geometry of the minimum loss submanifolds. Unlike traditional approximate posteriors that may spread over low-loss regions without enforcing exact data interpolation, DIMS guarantees convergence to the minimum manifold. The authors posit that this approach allows for the generation of diverse function-space samples that agree on training data but vary elsewhere, a property crucial for Bayesian inference. The work suggests that physics-inspired, geometry-aware dynamics can effectively navigate the complex loss landscapes of overparameterized networks, offering a robust alternative to existing sampling and approximation techniques. The authors acknowledge that the method is computationally demanding due to the need to solve ODEs in high-dimensional spaces but argue that the trade-off yields superior uncertainty estimates.