# Scalable Stochastic Gradient Riemannian Langevin Dynamics in Non-Diagonal Metrics

arXiv: https://arxiv.org/abs/2303.05101 · HTML: https://mlsanity.com/paper/2303.05101/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 57 references that are in the map.
- First posted on arXiv 9 Mar 2023.

## Builds on (most cited first, 30 of 34)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning](https://mlsanity.com/paper/1902.03932/) (2019)
- [A Complete Recipe for Stochastic Gradient MCMC](https://mlsanity.com/paper/1506.04696/) (2015)
- [Uncertainty Quantification and Deep Ensembles](https://mlsanity.com/paper/2007.08792/) (2020)
- [Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks](https://mlsanity.com/paper/1512.07666/) (2015)
- [Bayesian Neural Network Priors Revisited](https://mlsanity.com/paper/2102.06571/) (2021)
- [All You Need is a Good Functional Prior for Bayesian Deep Learning](https://mlsanity.com/paper/2011.12829/) (2020)
- [Model Selection in Bayesian Neural Networks via Horseshoe Priors](https://mlsanity.com/paper/1705.10388/) (2017)
- [On Uncertainty, Tempering, and Data Augmentation in Bayesian Classification](https://mlsanity.com/paper/2203.16481/) (2022)
- [Bayesian Dark Knowledge](https://mlsanity.com/paper/1506.04416/) (2015)
- [Challenges in Markov chain Monte Carlo for Bayesian neural networks](https://mlsanity.com/paper/1910.06539/) (2019)
- [On the Convergence of Stochastic Gradient MCMC Algorithms with High-Order Integrators](https://mlsanity.com/paper/1610.06665/) (2016)
- [Escaping Saddle Points with Adaptive Gradient Methods](https://mlsanity.com/paper/1901.09149/) (2019)
- [Benchmarking the Neural Linear Model for Regression](https://mlsanity.com/paper/1912.08416/) (2019)
- [QLSD: Quantised Langevin stochastic dynamics for Bayesian federated learning](https://mlsanity.com/paper/2106.00797/) (2021)
- [Scalable Second Order Optimization for Deep Learning](https://mlsanity.com/paper/2002.09018/) (2020)
- [BNNpriors: A library for Bayesian neural network inference with different prior distributions](https://mlsanity.com/paper/2105.06964/) (2021)
- [Efficiently Training Low-Curvature Neural Networks](https://mlsanity.com/paper/2206.07144/) (2022)
- [Stochastic Gradient Langevin Dynamics Algorithms with Adaptive Drifts](https://mlsanity.com/paper/2009.09535/) (2020)
- [Lagrangian Manifold Monte Carlo on Monge Patches](https://mlsanity.com/paper/2202.00755/) (2022)

## Built on by (most cited first, 8 of 8)

- [Riemannian Laplace Approximation with the Fisher Metric](https://mlsanity.com/paper/2311.02766/) (2023)
- [What's Inside Your Diffusion Model? A Score-Based Riemannian Metric to Explore the Data Manifold](https://mlsanity.com/paper/2505.11128/) (2025)
- [Devil is in the Details: Density Guidance for Detail-Aware Generation with Flow Models](https://mlsanity.com/paper/2502.05807/) (2025)
- [Warped geometric information on the optimisation of Euclidean functions](https://mlsanity.com/paper/2308.08305/) (2023)
- [On the Convergence of Locally Adaptive and Scalable Diffusion-Based Sampling Methods for Deep Bayesian Neural Network Posteriors](https://mlsanity.com/paper/2403.08609/) (2024)
- [Adaptive Stepsizing for Stochastic Gradient Langevin Dynamics in Bayesian Neural Networks](https://mlsanity.com/paper/2511.11666/) (2025)
- [Subspace Langevin Monte Carlo](https://mlsanity.com/paper/2412.13928/) (2024)
- [Feature Preserving Shrinkage on Bayesian Neural Networks via the R2D2 Prior](https://mlsanity.com/paper/2505.18280/) (2025)

## Abstract (arXiv)

Stochastic-gradient sampling methods are often used to perform Bayesian inference on neural networks. It has been observed that the methods in which notions of differential geometry are included tend to have better performances, with the Riemannian metric improving posterior exploration by accounting for the local curvature. However, the existing methods often resort to simple diagonal metrics to remain computationally efficient. This loses some of the gains. We propose two non-diagonal metrics that can be used in stochastic-gradient samplers to improve convergence and exploration but have only a minor computational overhead over diagonal metrics. We show that for fully connected neural networks (NNs) with sparsity-inducing priors and convolutional NNs with correlated priors, using these metrics can provide improvements. For some other choices the posterior is sufficiently easy also for the simpler metrics.