Flat Posterior Does Matter For Bayesian Model Averaging

TL;DR

This paper investigates the importance of posterior flatness in Bayesian neural networks for model averaging, revealing that encouraging flat posteriors enhances generalization and proposing a new method, FP-BMA, to achieve this.

Contribution

The paper introduces FP-BMA, a novel training objective that explicitly promotes flat posteriors in Bayesian neural networks, improving generalization performance.

Findings

Most approximate Bayesian inference methods do not produce flat posteriors.

Flat posteriors lead to better generalization in Bayesian model averaging.

FP-BMA effectively captures flat posteriors and enhances downstream task performance.

Abstract

Bayesian neural networks (BNNs) estimate the posterior distribution of model parameters and utilize posterior samples for Bayesian Model Averaging (BMA) in prediction. However, despite the crucial role of flatness in the loss landscape in improving the generalization of neural networks, its impact on BMA has been largely overlooked. In this work, we explore how posterior flatness influences BMA generalization and empirically demonstrate that (1) most approximate Bayesian inference methods fail to yield a flat posterior and (2) BMA predictions, without considering posterior flatness, are less effective at improving generalization. To address this, we propose Flat Posterior-aware Bayesian Model Averaging (FP-BMA), a novel training objective that explicitly encourages flat posteriors in a principled Bayesian manner. We also introduce a Flat Posterior-aware Bayesian Transfer Learning scheme that enhances generalization in downstream tasks. Empirically, we show that FP-BMA successfully captures flat posteriors, improving generalization performance.