🤖 machine learning

From Global to Local: A Scalable Benchmark for Local Posterior Sampling

This paper addresses the limitations of existing global convergence guarantees for stochastic gradient MCMC algorithms in degenerate neural network landscapes by introducing a scalable benchmark for local posterior sampling, which empirically demonstrates that RMSProp-preconditioned SGLD effectively captures local posterior geometry in models with up to 100 million parameters.

Original authors: Rohan Hitchcock, Jesse Hoogland

Published 2026-08-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Rohan Hitchcock, Jesse Hoogland

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the best spot to set up a campsite in a massive, foggy mountain range. In the world of artificial intelligence, this "mountain range" is called a loss landscape, and the "best spot" is the perfect configuration for a neural network to learn. For a long time, scientists thought these mountains had just a few distinct peaks and valleys, like a simple game board. But in reality, especially for the giant AI models we use today, the landscape is a bizarre, flat, and tangled mess. It's full of "degenerate" areas—vast, flat plains where you can move your tent in any direction without changing how good the view is.

To understand these complex models, scientists use a special tool called Stochastic Gradient MCMC (SGMCMC). Think of this as a very clever, slightly drunk hiker who wanders around the mountain, taking random steps to map out the terrain. The goal is to sample the "posterior," which is just a fancy way of saying "the map of all the good spots." The big question has always been: Can this hiker successfully explore the entire mountain range (global sampling)? The paper argues that for these messy, flat landscapes, the old rules say "no, it's impossible," but the hikers seem to be doing something useful anyway. The authors suggest we stop worrying about mapping the whole world and instead focus on whether the hiker can accurately describe the specific, weird valley they are currently standing in. This shift is crucial because understanding these local valleys helps us explain how AI models actually think and make decisions.

The paper introduces a new, scalable "test track" to see how well these hikers perform when they are stuck in these flat, degenerate valleys. The researchers built a specific type of model called a Deep Linear Network (DLN). Think of a DLN as a simplified, mathematical version of a neural network where we know exactly what the terrain should look like. In these networks, the "flatness" of the valley can be measured by a number called the Local Learning Coefficient (LLC). If you imagine the valley as a room, the LLC tells you how the volume of that room grows as you get closer to the center. In normal, non-flat valleys, this growth is predictable. But in the degenerate valleys of deep learning, the geometry is weird, and the LLC captures that weirdness.

The authors used these DLNs to create a benchmark where they knew the "true" LLC value. They then sent five different types of hikers (algorithms) into these valleys to see which one could measure the room's volume most accurately. The hikers included standard methods like SGLD and more advanced, adaptive versions like RMSProp-preconditioned SGLD and AdamSGLD. The results were clear: the standard hikers struggled, often getting confused by the flatness or taking steps that were too big, leading to chaotic results. However, the adaptive hikers, particularly the one using RMSProp, were much better at navigating the degenerate terrain. They could faithfully represent the local geometry, even in models with up to 100 million parameters.

Crucially, the paper points out that while we still don't have a mathematical proof that these hikers will eventually map the entire mountain (global convergence), they are empirically very good at describing the specific neighborhood they are in. The authors found that RMSProp-preconditioned SGLD was the most effective at capturing these local features. It was less sensitive to the size of the steps it took and provided more stable, accurate measurements of the local geometry compared to the others. Even when tested on a real-world language model (a transformer trained on a large dataset), the RMSProp method produced consistent results, whereas the standard method became unstable.

The paper concludes by highlighting a gap in our understanding: we know these algorithms work well locally in practice, but we lack the theoretical explanation for why they succeed when the old rules say they should fail. The authors suggest that the future of this field lies in shifting our focus from trying to prove global convergence (which might be impossible for these messy landscapes) to developing better local sampling guarantees. By using their new benchmark, researchers can now test and improve these algorithms to ensure they are accurately exploring the complex, degenerate valleys where modern AI lives, giving us better tools to understand and trust these powerful models.

Technical Summary: From Global to Local: A Scalable Benchmark for Local Posterior Sampling

Problem Statement

Neural networks possess highly complex, non-convex loss landscapes characterized by degenerate minima, where parameter changes do not necessarily alter the posterior density. While Stochastic Gradient MCMC (SGMCMC) algorithms, such as Stochastic Gradient Langevin Dynamics (SGLD), are widely used for Bayesian inference in these models, their theoretical understanding remains incomplete.

Existing global convergence guarantees for SGMCMC algorithms rely on strong assumptions—such as global Lipschitz continuity of the log-posterior gradient or the existence of specific Lyapunov functions—that are incompatible with the degenerate loss landscapes found in deep learning (e.g., deep linear networks). Theoretical results indicate that under these degenerate conditions, global convergence may not occur, and samplers may diverge. Despite this lack of theoretical justification, empirical evidence suggests that SGMCMC algorithms can successfully extract non-trivial local information from the posterior. This creates a gap: there is no rigorous framework to evaluate how well these algorithms explore the local geometry of degenerate minima, nor is there a theoretical explanation for their empirical success in this regime.

Methodology

To address this gap, the authors propose a shift in focus from global convergence to local posterior sampling. They introduce a scalable benchmark designed to evaluate a sampler's ability to capture known local geometric invariants, specifically the Local Learning Coefficient (LLC).

The Benchmark: Deep Linear Networks (DLNs)

The benchmark utilizes Deep Linear Networks (DLNs) as a testbed. DLNs are chosen because:

They exhibit degenerate loss landscapes similar to non-linear networks.
The ground-truth LLC values for specific points in their parameter space can be computed analytically using results from Singular Learning Theory (SLT) and algebraic geometry (Aoyagi, 2024; Lehalleur & Rimányi, 2024).

The authors generate two types of optimal parameters within the DLN solution space ( $W_0$ ):

Special Optimal Parameters: Points with minimal LLC, computable via closed-form formulas.
Generic Optimal Parameters: Points with higher LLC, computable via constrained integer optimization based on rank patterns.

This setup allows for the evaluation of samplers across varying degrees of local degeneracy within the same model architecture, scaling up to models with approximately 100 million parameters.

Evaluation Metric

The performance of an SGMCMC sampler is measured by its ability to estimate the LLC ( $\lambda$ ) using the estimator:
$\hat{\lambda}(w_0) = n\beta(\mathbb{E}_{w}^{\beta}[L_n(w)] - L_n(w_0))$
where the expectation is taken over the tempered posterior distribution. A sampler that faithfully represents the local posterior geometry should produce estimates $\hat{\lambda}$ that match the ground-truth $\lambda$ and exhibit low variance across different step sizes.

Key Contributions

Identification of Theoretical Gaps: The paper surveys existing global convergence guarantees for SGLD and related algorithms, demonstrating that their core assumptions (e.g., global Lipschitz conditions) are violated by degenerate loss landscapes. It highlights the lack of theoretical explanation for the empirical success of local sampling in these settings.
A Novel Scalable Benchmark: The authors introduce a benchmark that uses DLNs to provide ground-truth LLC values. Unlike previous work that compared estimates only against a lower bound, this method allows for the exact computation of ground-truth LLCs for both special and generic parameters, enabling precise evaluation of local geometry capture.
Empirical Evaluation of Common Algorithms: The benchmark is applied to five common SGMCMC algorithms: SGLD, RMSProp-preconditioned SGLD, Adam-preconditioned SGLD, SGHMC, and SGNHT.

Results

The evaluation across models ranging from 100K to 100M parameters yields the following findings:

Superiority of Preconditioned Samplers: RMSProp-preconditioned SGLD and Adam-preconditioned SGLD consistently outperform standard SGLD, SGHMC, and SGNHT. They achieve a superior mean-variance trade-off in LLC estimation and are less sensitive to the choice of step size.
Robustness and Stability:
- Standard SGLD often produces divergent estimates or NaNs at step sizes where the relative error appears low.
- RMSPropSGLD provides a clear failure signal (NaNs or massive loss spikes) when the step size is too large, whereas AdamSGLD degrades more gradually without obvious warning signs.
- In experiments with a real language model (a four-layer attention-only transformer), RMSPropSGLD produced consistent LLC-over-training curves, while standard SGLD failed to stabilize.
Impact of Degeneracy: Sampler performance varies significantly based on the degree of degeneracy. For instance, standard SGLD tends to underestimate the LLC at highly degenerate (special) parameters but performs accurately at generic ones. Conversely, preconditioned methods tend to overestimate LLC at generic parameters but maintain lower variance.

Significance and Claims

The paper argues that while global convergence guarantees for SGMCMC in degenerate landscapes are likely unattainable under current theoretical frameworks, local sampling remains a viable and practical objective.

Practical Guidance: The benchmark provides actionable guidance for practitioners, identifying RMSProp-preconditioned SGLD as the most effective method for faithfully representing local posterior geometry in large-scale models.
Theoretical Open Problem: The authors emphasize that the empirical success of these samplers in extracting local information, despite the lack of global convergence guarantees, remains an unexplained phenomenon. They posit that understanding the "effective support" of SGMCMC chains in practice is a central open problem for future theoretical research.
Future Directions: The work suggests that future benchmarks should be degeneracy-aware, moving beyond standard non-convex optimization test functions (like Rastrigin or Ackley) which often ignore the specific nature of degeneracy found in neural networks.

In conclusion, the paper establishes that shifting focus from global to local posterior sampling allows for the development of scalable, empirically verifiable benchmarks, revealing that preconditioned SGMCMC algorithms can successfully navigate the complex, degenerate geometries of modern neural network posteriors.