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.

For artificial intelligence to be reliable, it must accurately estimate its own uncertainty. Bayesian neural networks trained with Microcanonical Langevin dynamics are a powerful method for this, but so far required processing the entire, potentially large dataset at once, making it too slow to scale for modern AI. This is the first work to analyze both on a theoretical and practical level what happens when we use smaller data batches instead of the full dataset. We show that using smaller data batches introduces uneven noise, which causes the algorithm to become biased and unstable. We resolve these problems by developing a preconditioning technique to balance the noise, alongside an adaptive guardrail system that constantly monitors the system's internal dynamics to prevent catastrophic failures and improve its efficiency. The result is SMILE: a fast, robust algorithm that successfully scales advanced uncertainty quantification to large-scale neural networks.