Variational Inference Failures Under Model Symmetries: Permutation Invariant Posteriors for Bayesian Neural Networks

TL;DR

This paper investigates how permutation symmetries in neural network weights cause multimodal posteriors that challenge variational inference, and proposes a symmetrization method to improve posterior approximation and predictive performance.

Contribution

The authors introduce a symmetrization mechanism for permutation invariant variational posteriors, improving the fit to true posteriors in Bayesian neural networks with weight symmetries.

Findings

Symmetrized posteriors better fit true posteriors.

Improved predictive accuracy with symmetrization.

Higher ELBO achieved using the proposed method.

Abstract

Weight space symmetries in neural network architectures, such as permutation symmetries in MLPs, give rise to Bayesian neural network (BNN) posteriors with many equivalent modes. This multimodality poses a challenge for variational inference (VI) techniques, which typically rely on approximating the posterior with a unimodal distribution. In this work, we investigate the impact of weight space permutation symmetries on VI. We demonstrate, both theoretically and empirically, that these symmetries lead to biases in the approximate posterior, which degrade predictive performance and posterior fit if not explicitly accounted for. To mitigate this behavior, we leverage the symmetric structure of the posterior and devise a symmetrization mechanism for constructing permutation invariant variational posteriors. We show that the symmetrized distribution has a strictly better fit to the true posterior, and that it can be trained using the original ELBO objective with a modified KL regularization term. We demonstrate experimentally that our approach mitigates the aforementioned biases and results in improved predictions and a higher ELBO.