Entropy-MCMC: Sampling from Flat Basins with Ease
This addresses the challenge of suboptimal performance in Bayesian deep learning due to sampling from 'bad' modes, offering a solution for practitioners with limited computational budgets.
The paper tackles the problem of sampling from multi-modal posterior distributions in Bayesian deep learning by biasing sampling towards flat basins, which are associated with better generalization. The method outperforms existing baselines on classification, calibration, and out-of-distribution detection benchmarks.
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.