Can Microcanonical Langevin Dynamics Leverage Mini-Batch Gradient Noise?

Scaling inference methods such as Markov chain Monte Carlo to high-dimensional models remains a central challenge in Bayesian deep learning. A promising recent proposal, microcanonical Langevin Monte Carlo, has shown state-of-the-art performance across a wide range of problems. However, its reliance on full-dataset gradients makes it prohibitively expensive for large-scale problems. This paper addresses a fundamental question: Can microcanonical dynamics effectively leverage mini-batch gradient noise? We provide the first systematic study of this problem, establishing a novel continuous-time theoretical analysis of stochastic-gradient microcanonical dynamics. We reveal two critical failure modes: a theoretically derived bias due to anisotropic gradient noise and numerical instabilities in complex high-dimensional posteriors. To tackle these issues, we propose a principled gradient noise preconditioning scheme shown to significantly reduce this bias and develop a novel, energy-variance-based adaptive tuner that automates step size selection and dynamically informs numerical guardrails. The resulting algorithm is a robust and scalable microcanonical Monte Carlo sampler that achieves state-of-the-art performance on challenging high-dimensional inference tasks like Bayesian neural networks. Combined with recent ensemble techniques, our work unlocks a new class of stochastic microcanonical Langevin ensemble (SMILE) samplers for large-scale Bayesian inference.

💡 Research Summary

This paper tackles the scalability bottleneck of Microcanonical Langevin Monte Carlo (MCLMC), a recent state‑of‑the‑art MCMC method that explores posterior distributions on a constant‑energy manifold. While MCLMC outperforms traditional Hamiltonian Monte Carlo (HMC) and its No‑U‑Turn Sampler (NUTS) variant on a variety of benchmarks, it requires full‑dataset gradients at every integration step, making it impractical for modern deep learning workloads. The authors ask a fundamental question: can the microcanonical dynamics exploit the stochastic gradient noise that naturally arises from mini‑batch training?

To answer this, they first formulate a naïve stochastic version of MCLMC (SMILE‑naive) by simply replacing the full‑batch gradient in the microcanonical SDE with a mini‑batch estimate. They show that, unlike isotropic Langevin noise, the gradient noise introduced by mini‑batching is typically anisotropic and position‑dependent, with covariance (V(\theta)). By extending the continuous‑time analysis of Robnik & Seljak (2024), they prove (informally, Theorem 3.1) that anisotropic noise induces a non‑vanishing drift term that shifts the stationary distribution away from the true posterior, creating a systematic bias. Empirical verification on three 10‑dimensional analytical posteriors (Gaussian, Rosenbrock, Funnel) confirms that SMILE‑naive suffers substantial second‑moment bias, often exceeding that of well‑tuned SGLD or SGHMC.

To mitigate this bias, the paper introduces a gradient‑noise preconditioning scheme. Assuming the local covariance (V(\theta)) can be estimated, they apply a linear transformation (\theta’ = L(\theta_0)^{\top}\theta) where (L L^{\top}=V). In the transformed coordinates the noise becomes isotropic, restoring the theoretical guarantees of microcanonical dynamics. Because full covariance estimation is infeasible for large neural networks, the authors adopt a diagonal approximation and maintain moving‑average estimates of per‑parameter gradient standard deviations (\sigma). The preconditioned variable is then (\theta’ = (\sqrt{d}/|\sigma|),(\theta\odot\sigma)), which reduces to the identity mapping under isotropic noise. This “pSMILE‑naive” method dramatically lowers the bias across all noise structures examined.

A second major contribution is an adaptive step‑size tuner based on the energy error (\Delta E) that naturally arises from the microcanonical integrator. Since the exact dynamics conserve total energy, any change in energy per step is a direct proxy for numerical discretisation error. The authors monitor the variance of (\Delta E) and compare it to the intrinsic variance from stochastic gradients and finite‑sample Monte Carlo error. When (\Delta E) exceeds a prescribed tolerance, the step size is reduced; when it is comfortably below the tolerance, the step size is increased. This energy‑variance‑based tuner acts as a “numerical guardrail,” preventing the integrator from destabilising in high‑dimensional, highly curved posteriors while still allowing aggressive exploration when the dynamics are well‑behaved.

Combining diagonal preconditioning with the adaptive tuner yields the final algorithm, called pSMILE. The authors also extend pSMILE to an ensemble setting (SMILE‑E) by running multiple short, warm‑started chains and aggregating their samples, a strategy known to improve exploration in multimodal Bayesian neural network (BNN) posteriors.

Extensive experiments on Bayesian neural networks for CIFAR‑10, Fashion‑MNIST, and SVHN demonstrate that pSMILE‑E achieves state‑of‑the‑art performance. Compared to the strong baseline of scale‑adapted SGHMC, pSMILE‑E improves test accuracy by 1–2 %, reduces negative log‑likelihood and expected calibration error by roughly 10–15 %, and increases effective sample size per second by a factor of ~1.8. Ablation studies show that preconditioning alone cuts bias by ~60 % and that the adaptive tuner alone substantially improves numerical stability, especially on the Funnel posterior where naïve SMILE diverges. Moreover, the method exhibits markedly reduced sensitivity to hyper‑parameters: a default mini‑batch size of 256 and an initial step size of (10^{-3}) work well across all tasks.

In summary, the paper establishes that microcanonical Langevin dynamics are theoretically robust to isotropic noise but vulnerable to the anisotropic, data‑dependent noise typical of stochastic gradients. By introducing a practical diagonal preconditioning scheme and an energy‑error‑driven step‑size controller, the authors convert MCLMC into a scalable stochastic sampler (SMILE) that retains its original sampling efficiency while operating on mini‑batches. This work opens a new avenue for large‑scale Bayesian inference, suggesting future research directions such as full Riemannian preconditioning, non‑Gaussian noise models, and extensions to streaming or federated data settings.