AdamMCMC: Combining Metropolis Adjusted Langevin with Momentum-based Optimization

TL;DR

This paper introduces AdamMCMC, a novel algorithm that combines MALA with Adam optimization to efficiently estimate epistemic uncertainty in deep neural networks by sampling from a tempered posterior distribution.

Contribution

It presents a new MCMC method that integrates momentum-based optimization with Langevin dynamics, providing theoretical guarantees and practical efficiency for uncertainty quantification.

Findings

Proves the chain admits the Gibbs posterior as invariant distribution.

Demonstrates efficiency on high-energy particle physics classifier.

Shows improved uncertainty estimation in neural networks.

Abstract

Uncertainty estimation is a key issue when considering the application of deep neural network methods in science and engineering. In this work, we introduce a novel algorithm that quantifies epistemic uncertainty via Monte Carlo sampling from a tempered posterior distribution. It combines the well established Metropolis Adjusted Langevin Algorithm (MALA) with momentum-based optimization using Adam and leverages a prolate proposal distribution, to efficiently draw from the posterior. We prove that the constructed chain admits the Gibbs posterior as invariant distribution and approximates this posterior in total variation distance. Furthermore, we demonstrate the efficiency of the resulting algorithm and the merit of the proposed changes on a state-of-the-art classifier from high-energy particle physics.