Technical Summary: Guiding Posterior Exploration with Optimizer-Derived Geometry

Problem Statement

Bayesian Deep Learning (BDL) aims to provide principled uncertainty quantification for neural networks, a critical requirement for high-stakes applications. While sampling-based methods like Markov Chain Monte Carlo (MCMC) offer a rigorous approach to approximating posterior distributions, their practical application is hindered by the immense computational cost of exploring high-dimensional, multimodal, and ill-conditioned posterior landscapes.

Recent scalable inference strategies, such as the two-stage Bayesian Deep Ensemble (BDE), attempt to mitigate this by first optimizing to find multiple distinct modes (akin to Deep Ensembles) and then initializing independent MCMC chains from these solutions to explore local basins of attraction. However, a significant inefficiency remains: the optimization and sampling stages are typically treated as disjoint processes. Consequently, the geometric information about the loss landscape (specifically local curvature) gathered during the optimization phase is discarded. The subsequent sampling phase is forced to rediscover this local curvature from scratch, often requiring a lengthy and costly "warmup" or scale-adaptation period to adjust the sampler's preconditioning matrix.

Methodology

This paper proposes a method to bridge the gap between optimization and sampling by reusing statistics generated by adaptive optimizers (specifically AdamW) to inform the initialization of Stochastic Gradient MCMC (SGMCMC) samplers.

Optimizer-Derived Geometry

Adaptive optimizers like AdamW maintain exponential moving averages of gradients: a first-moment vector ($\mu_t$) acting as momentum and a second-moment vector ($\nu_t$) acting as a proxy for gradient variance (and thus local curvature). The authors argue that these statistics, computed as a byproduct of the optimization phase, serve as a "free," high-quality preconditioner for the subsequent sampling phase.

Preconditioning Strategies

The authors formalize several strategies to transfer these optimizer states to the mass matrix $M$ of an SGHMC sampler:

1. Gradient-based: $M = \text{diag}(g_T^2 + \epsilon I)$, using the final gradient. This is purely local but susceptible to noise.
2. First-moment: $M = \text{diag}(\mu_T^2 + \epsilon I)$, using the momentum term. This offers temporal smoothing.
3. Second-moment: $M = \text{diag}(\sqrt{\nu_T} + \epsilon I)$, using the squared gradient accumulator. This approximates the diagonal Fisher Information Matrix (FIM) and captures global curvature information more robustly.

Layer-wise Normalization

To enhance numerical stability and prevent extreme step sizes for individual parameters, the authors evaluate a layer-wise strategy. Instead of assigning a unique preconditioning value to every parameter, they compute a single scalar per layer by averaging the optimizer statistics of that layer's parameters and normalizing across layers. This acts as a regularizer and upper-bounds the maximum parameter-wise step size.

Theoretical Justification

The paper provides a theoretical link (Proposition 3.1) showing that if the posterior potential is locally quadratic, setting the mass matrix $M$ to the Hessian $H^*$ at the mode induces an isotropic Gaussian approximation with a unit condition number. This achieves optimal local mixing. Since the full Hessian is intractable, the second-moment statistic $\nu_T$ serves as a computationally efficient diagonal proxy for $H^*$, effectively transferring the "effective metric tensor" that successfully navigated the optimization landscape to the sampling phase.

Key Contributions

1. Theoretical Framework: The authors derive theoretically grounded methods for using optimizer properties (specifically from AdamW) to precondition SGMCMC samplers, bridging the optimization and sampling phases.
2. Empirical Evaluation: They evaluate multiple optimizer-derived geometry transfer strategies across diverse modalities (tabular, image, language) and architectures (FCNs, CNNs, ResNets, Transformers).
3. Efficiency and Stability Analysis: The study assesses the ability of these strategies to accelerate sampler convergence, reduce or eliminate the need for warmup, and improve robustness to hyperparameters (specifically step size) without additional computational costs.
4. Geometric Analysis: The paper analyzes local geometry patterns in realistic BNNs, uncovering that static, optimizer-derived preconditioning preserves the distinct geometric information of each mode found in Phase I, whereas adaptive samplers (like AdaSGHMC) tend to "overwrite" this initial geometry, reducing ensemble diversity.

Results

The proposed approach was evaluated on datasets including bikesharing, protein, F-MNIST, CIFAR-10, and a modernized Shakespeare dataset.

• Predictive Performance: Optimizer-derived preconditioning consistently maintains or improves predictive performance (RMSE and Accuracy) compared to standard Deep Ensembles and Vanilla SGHMC.
• Uncertainty Quantification (UQ): The method yields significant improvements in UQ metrics (Log Pointwise Predictive Density and Perplexity), often outperforming the Deep Ensemble baseline by over 10% in LPPD.
• Warmup Elimination: Strategies based on gradients and first-moments achieve peak performance immediately with zero warmup, effectively eliminating the costly scale-adaptation phase. The second-moment strategy benefits from longer warmup but still outperforms baselines.
• Robustness: Preconditioned variants maintain stable performance across a broad range of step sizes, whereas Vanilla SGHMC degrades sharply with suboptimal step sizes.
• Comparison to Adaptive Methods: Surprisingly, the static optimizer-derived preconditioning outperforms adaptive methods like AdaSGHMC. The authors attribute this to the fact that AdaSGHMC's internal adaptation overwrites the initial geometric diversity of the ensemble chains, leading to self-similar mass matrices and reduced generalization.

Significance

The paper claims that transferring optimizer-derived geometry to SGMCMC samplers is an effective, zero-cost strategy for improving Bayesian Deep Ensembles. By utilizing statistics already computed during the optimization phase, the method:

• Reduces Computational Cost: It eliminates or significantly shortens the sampling warmup phase, offering a more efficient path to posterior exploration.
• Enhances Stability: It provides a numerically stable starting point that is robust to hyperparameter choices.
• Preserves Diversity: By using a static preconditioner derived from the specific mode found in Phase I, it preserves the geometric diversity of the ensemble, which the authors argue is crucial for superior generalization compared to adaptive samplers that homogenize the ensemble's geometry.

The authors conclude that this approach allows practitioners to achieve state-of-the-art uncertainty quantification and predictive performance with minimal additional computational overhead, making Bayesian inference more accessible for high-stakes domains.