Entropy-MCMC: Sampling from Flat Basins with Ease

TL;DR

Entropy-MCMC introduces a novel sampling method that biases towards flat, well-generalizing modes in the posterior of deep neural networks, improving performance in Bayesian deep learning tasks.

Contribution

The paper proposes a new flat-basin biased MCMC sampling technique using an auxiliary guiding variable, with proven convergence and faster mixing than existing methods.

Findings

Successfully samples from flat posterior basins

Outperforms baselines on classification, calibration, OOD detection

Converges faster than existing flatness-aware methods

Abstract

Bayesian deep learning counts on the quality of posterior distribution estimation. However, the posterior of deep neural networks is highly multi-modal in nature, with local modes exhibiting varying generalization performance. Given a practical budget, targeting at the original posterior can lead to suboptimal performance, as some samples may become trapped in "bad" modes and suffer from overfitting. Leveraging the observation that "good" modes with low generalization error often reside in flat basins of the energy landscape, we propose to bias sampling on the posterior toward these flat regions. Specifically, we introduce an auxiliary guiding variable, the stationary distribution of which resembles a smoothed posterior free from sharp modes, to lead the MCMC sampler to flat basins. By integrating this guiding variable with the model parameter, we create a simple joint distribution that enables efficient sampling with minimal computational overhead. We prove the convergence of our method and further show that it converges faster than several existing flatness-aware methods in the strongly convex setting. Empirical results demonstrate that our method can successfully sample from flat basins of the posterior, and outperforms all compared baselines on multiple benchmarks including classification, calibration, and out-of-distribution detection.