On permutation symmetries in Bayesian neural network posteriors: a variational perspective
This work addresses the challenge of approximate inference in Bayesian neural networks by leveraging permutation symmetries, offering incremental improvements in understanding loss landscape geometry.
The paper tackles the problem of understanding loss landscape connectivity in Bayesian neural networks by extending the concept of marginalized loss barriers and solution interpolation, and proposes a matching algorithm to align distributions with respect to permutation symmetries, resulting in nearly zero marginalized loss barriers for linearly connected solutions across various architectures and datasets.
The elusive nature of gradient-based optimization in neural networks is tied to their loss landscape geometry, which is poorly understood. However recent work has brought solid evidence that there is essentially no loss barrier between the local solutions of gradient descent, once accounting for weight-permutations that leave the network's computation unchanged. This raises questions for approximate inference in Bayesian neural networks (BNNs), where we are interested in marginalizing over multiple points in the loss landscape. In this work, we first extend the formalism of marginalized loss barrier and solution interpolation to BNNs, before proposing a matching algorithm to search for linearly connected solutions. This is achieved by aligning the distributions of two independent approximate Bayesian solutions with respect to permutation matrices. We build on the results of Ainsworth et al. (2023), reframing the problem as a combinatorial optimization one, using an approximation to the sum of bilinear assignment problem. We then experiment on a variety of architectures and datasets, finding nearly zero marginalized loss barriers for linearly connected solutions.