On permutation symmetries in Bayesian neural network posteriors: a
  variational perspective
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.
Simone Rossi, Ankit Singh, Thomas Hannagan
# On permutation symmetries in Bayesian neural network posteriors: a variational perspective

###### Abstract

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.

## 1 Introduction

Throughout the last decade, deep neural networks (dnn) have achieved significant success in a wide range of practical applications, becoming the fundamental ingredient for e.g., computer vision [e.g., 57, 25, 41, 64], language models [e.g., 24, 82, 17] and generative models [e.g., 54, 96, 97, 98, 102, 34]. Despite recent important advancements, understanding the loss landscape of dnn is still challenging. The characterization of its highly non-convex nature, its relation with architectural choices like depth and width and the connection with optimization and generalization are just some of the problems which have been the focus of extensive research in the last few years [e.g., 76, 26, 36, 32, 2, 30, 38, 79]. It is well known, for example, that one of the fundamental characteristics of deep neural networks is their ability to learn hierarchical features, and in this regards deeper networks seem to be exponentially more expressive than shallower models [e.g., 3, 5, 7, 8, 18, 109], leading the loss landscape to have many optima due to symmetries and over-parameterization [76, 26, 111, 94]. At the same time, the role of the depth of a model in relation with its width is far less understood [80], despite wide neural networks exhibiting important theoretical properties in their infinite limit behavior [e.g., 73, 23, 29, 53, 48, 37, 20].

Two notions that have been useful to shed light on the geometry of loss landscapes are that of _loss barriers_ and _mode connectivity_[36, 26]. The mode connectivity hypothesis states that given two points in the landscape, there exists a path connecting them such that the loss is constant or near constant (or, said differently, the loss barrier is null). We refer to _linear mode connectivity_ when the path connecting the two solutions is linear [32]. Recently, evidence has surfaced that stochastic gradient descent (sgd) solutions to the loss minimization problem can be linearly connected. Indeed, Entezari et al. [30] discuss the role of permutation symmetries from a loss connectivity viewpoint, conjecturing the possibility that mode connectivity is actually linear once accounting for all permutation invariances.
Additionally, [2] gathers compelling empirical evidence across several network architectures and tasks, that under such a permutation symmetry the loss landscape often contains a single, nearly convex basin.

In this work, we are taking a different perspective on this analysis. We are interested in the Bayesian treatment of neural networks, which results in a natural form of regularization and allows to reason about uncertainty in the predictions [101, 73, 66]. Bayesian inference for deep neural networks is notoriously challenging, as we wish to marginalize over multi-modal distributions with high dimensionality [47]. For this reason, there are various ways to approximate the posterior, involving techniques like variational inference [39, 14, 35, 62, 77], Markov chain Monte Carlo (mcmc) methods [75, 72, 27], possibly with stochastic gradients [19, 112, 33, 105, 70] and the Laplace approximation [68, 68, 84]. Indeed, fundamentally the Bayesian posterior and the loss landscapes are tightly interconnected: (i) solutions to the loss minimization problem are equivalent to maximum-a-posteriori (map) solutions, (ii) the loss landscape is equivalent to the un-normalized negative log-posterior. While in theory, given a dataset, the posterior is unique and the solution is global, many approximations will only explore local properties of the true posterior1. It's worth noting that a posterior over the parameters of the neural network induces a posterior on the functions generated by the model. Permutation symmetries play an important role in the geometry of the weight-space posterior, which are generally not reflected in function-space. While it is possible to carry out inference directly in function space, this poses a number of challenges [99, 71, 89, 61]. Fig. 1 illustrates this situation for a regression task on the Snelson dataset using a 3-layer dnn: on the left we compare two (approximate) solutions which have different weight-space posterior but similar function-space behavior. Notably, when we interpolate these two solutions (Fig. 1 on the right), we completely lose all capability of modeling the data. However, when we account for permutation symmetries in the posterior, we end up with solutions that once interpolated are still good approximations. This suggests that for any weight-space distribution, there exists a class of solutions which are functionally equivalent and linearly connected. This example motivates an informal generic statement:

Footnote 1: By local properties, we mean that despite theoretical convergence guarantees of many methods like variational inference and mcmc, in practice the true posterior for deep neural networks is still highly elusive; see for instance the empirical convergence analysis of Hamiltonian Monte Carlo (hmc) in [47].

_Solutions of approximate Bayesian inference for neural networks are linearly connected after accounting for functionally equivalent permutations._

While being similar to the one in [30, 2], if this was to hold true for approximate Bayesian neural networks (bnns) it would represent an important step in further characterizing the properties of the Bayesian posterior and the effect of various approximations. We purposely leave the previous statement broadly open regarding the choice of the approximation method to allow for a more general discussion. More specifically, in this paper we will analyze and focus our discussion on the variational inference framework, making a more specific conjecture:

**Conjecture 1**.: _Solutions of variational inference solutions for Bayesian inference for neural networks are linearly connected after accounting for functionally equivalent permutations._

Figure 1: **Permutation symmetries for regression on the Snelson dataset**[95]. _(Left)_ Two different solutions, with similar function-space behavior (showing \(\mu\pm 2\sigma\)). _(Right)_ Functions obtained when using two different strategies to interpolate between solutions. When we align the solutions, by taking into account the permutation symmetries, we retain the capability to model the data, indicating that there are solutions that once linearly interpolated exhibit no loss in performance. Note that, in weight space, the solution found with alignment is neither equal to solution 1 nor 2 (see the black curve which is a single function sample with fixed randomness).
Contributions.With this work, we aim at studying the linear connectivity properties of approximate solutions to the Bayesian inference problem and we make several contributions. (i) We extend the formalism of loss barrier and solution interpolation to bnns. (ii) For the variational inference setting, propose a matching algorithm to search for linearly connected solutions by aligning the distributions of two independent solutions with respect to permutation matrices. Inspired by [2], we frame the problem as a combinatorial optimization problem using approximation to the linear sum assignment problem. (iii) We then experiment on a variety of architectures and datasets, finding nearly zero-loss barriers for linearly connected solutions. In Fig. 2 we present a sneak-peek and a visualization of our findings, where we show that after weight distribution alignment we can find a permutation map \(P\) of the solution \(q_{1}\) such that it can be linearly connected through high density regions to \(q_{0}\).

## 2 Preliminaries on Bayesian deep learning

In this section, we review some basic notations on bnns and we review stochastic variational inference (svi), which is the main approximation method that we analyze in this paper. Let's consider a generic multilayer perceptron (mlp) with \(L\) layers, where the output of the \(l\)-th layer \(\mathbf{f}_{l}(\mathbf{\theta}_{l},\mathbf{x})\) is a vector-valued function of the previous layer output \(\mathbf{f}_{l-1}\) as follows,

\[\mathbf{f}_{l}(\mathbf{\theta}_{l},\mathbf{x})=\mathbf{W}_{l}a(\mathbf{f}_{l-1}(\mathbf{\theta}_{l-1}, \mathbf{x}))+\mathbf{b}_{l} \tag{1}\]

where \(a(\cdot)\) is a non-linearity, \(\mathbf{W}_{l}\) is a \(D_{l}\times D_{l-1}\) weight matrix and \(\mathbf{b}_{l}\) the corresponding bias vector. We shall refer to the parameters of the layer \(l\) as \(\mathbf{\theta}_{l}=\{\mathbf{b}_{l},\mathbf{W}_{l}\}\), and the union of all trainable parameters as \(\mathbf{\theta}=\{\mathbf{\theta}_{l}\}_{l=1}^{L}\).

The objective of using Bayesian inference on deep neural networks [67; 65] involves inferring a posterior distribution over the parameters of the neural network given the available dataset \(\{\mathbf{X},\mathbf{Y}\}=\{(\mathbf{x}_{i},\mathbf{y}_{i})\}_{i=1}^{N}\). This requires choosing a likelihood and a prior [73; 74; 103]:

\[p(\mathbf{\theta}\,|\,\mathbf{Y},\mathbf{X})=Z^{-1}p(\mathbf{Y}\,|\,\mathbf{\theta},\mathbf{X})p(\mathbf{ \theta}) \tag{2}\]

where the normalization constant \(Z\) is the marginal likelihood \(p(\mathbf{Y}\,|\,\mathbf{X})\). As usually done, we assume that the likelihood factorizes over observations, i.e. \(p(\mathbf{Y}\,|\,\mathbf{\theta},\mathbf{X})=\prod_{i=1}^{N}p(\mathbf{y}_{i}\,|\,\mathbf{\theta}, \mathbf{x}_{i})\).

Bayesian deep learning is intractable due to the non-conjugacy likelihood-prior and thus we don't have access to closed form solutions. Variational inference (vi) is a common technique to handle intractable Bayesian neural networks [12; 43; 39; 50]. Let \(\mathcal{P}(\mathbb{R}^{d})\) be the space of probability measures on \(\mathbb{R}^{d}\); vi reframes the inference problem into an optimization one, commonly by introducing a parameterized distribution \(q(\mathbf{\theta})\in\mathcal{P}(\mathbb{R}^{d})\) which is optimized to minimize the Kullback-Leibler (kl) divergence with respect to the true posterior \(p(\mathbf{\theta}\,|\,\mathbf{Y},\mathbf{X})\). In practice, this involves the maximization of the evidence lower bound (elbo) defined as

\[\mathcal{L}_{\textsc{elbo}}(q)\stackrel{{\mathrm{def}}}{{=}} \int\log p(\mathbf{Y}\,|\,\mathbf{\theta},\mathbf{X})\mathrm{d}q(\mathbf{\theta})-\textsc{kl} \left[q(\mathbf{\theta})\parallel p(\mathbf{\theta})\right] \tag{3}\]

whose gradients can be unbiasedly estimated with mini-batches of data [43] and the reparameterization trick [54; 55]. Despite its simple formulation, the optimization of the elbo hides several challenges, like the initialization of the variational parameters [87], the effects of over-parameterization on the quality of the approximation [88; 44; 58]. Here we are interested in how different solutions to Eq. (3) relate to each other in terms of loss barrier, which we will define formally in the following section.

## 3 Loss barriers

In the context of Bayesian inference, we are interested in the loss computed by marginalization of the model parameters with respect to (an approximation of) the posterior. As such, we use the predictive

Figure 2: **Permutations in multi-modal posterior. Log-posterior for MLP/CIFAR10, showing the two solutions (\(q_{0}\) and \(q_{1}\)) for which we can find a permutation map such that \(P_{\#}q_{1}\) can be linearly connected to \(q_{0}\) with low barrier (brighter regions).**
likelihood, a proper scoring method for probabilistic models [83], defined as

\[\log p(\mathbf{y}^{\star}\,|\,\mathbf{x}^{\star})=\log\int p(\mathbf{y}^{\star}\,|\,\mathbf{ \theta},\mathbf{x}^{\star})p(\mathbf{\theta}\,|\,\mathbf{Y},\mathbf{X})\mathrm{d}\mathbf{\theta} \approx\log\int p(\mathbf{y}^{\star}\,|\,\mathbf{\theta},\mathbf{x}^{\star})\mathrm{d}q( \mathbf{\theta}) \tag{4}\]

where \(q(\mathbf{\theta})\) is an approximation of the true posterior (parametric or otherwise), \(\{\mathbf{x}^{\star},\mathbf{y}^{\star}\}\) are respectively the input and its corresponding label a data point under evaluation. To keep the notation uncluttered for the remaining of the paper, we write the predictive likelihood computed for a set of points \(\{\mathbf{x}^{\star}_{i},\mathbf{y}^{\star}_{i}\}_{i=1}^{N}\) as a functional \(\mathcal{L}:\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}\), defined as

\[\mathcal{L}(q)\stackrel{{\mathrm{def}}}{{=}}\sum_{i=1}^{N}\log \int p(\mathbf{y}^{\star}_{i}\,|\,\mathbf{\theta},\mathbf{x}^{\star}_{i})\mathrm{d}q(\mathbf{ \theta}) \tag{5}\]

Let's assume two models trained with vi with two different initializations, random seeds, and batch ordering. Variational inference in the classic inverse sense \(\textsc{kl}\,[q\parallel p]\) is mode seeking, thus we expect the two runs to converge to different solutions, say \(q_{0}\) and \(q_{1}\). To test the loss barrier as we interpolate between the two solutions we need to decide on the interpolation rule. We decide to interpolate the solutions following the Wasserstein geodesics between \(q_{0}\) and \(q_{1}\). First, let's start with a few definitions. Let \(q\in\mathcal{P}(\mathbb{R}^{d})\) be a probability measure on \(\mathbb{R}^{d}\) and \(T:\mathbb{R}^{d}\to\mathbb{R}^{d}\) a measurable map; we denote \(T_{\#}q\) the _push-forward measure_ of \(q\) through \(T\). Now we can introduce the _Wasserstein geodesics_, as follows.

**Definition 1**.: _The Wasserstein geodesics between \(q_{0}\) and \(q_{1}\) is defined as the path_

\[q_{\tau}=\left((1-\tau)\mathrm{Id}+\tau T_{q_{0}}^{q_{1}}\right)_{\#}q_{0}, \qquad\tau\in[0,1] \tag{6}\]

_where \(\mathrm{Id}\) is the identity map and \(T_{q_{0}}^{q_{1}}\) is the optimal transport map between \(q_{0}\) and \(q_{1}\), which for Brenier's theorem [16], is unique._

While we could interpolate using a mixture of the two solutions, we argue that this choice is trivial and does not fully give us a picture of the underlying loss landscape. Indeed, Eq. (6) is fundamentally different from a naive mixture path \(\tilde{q}_{\tau}=(1-\tau)q_{0}+\tau q_{1}\). In case of Gaussian distributions, when \(q_{0}=\mathcal{N}(\mathbf{m}_{0},\mathbf{S}_{0})\) and \(q_{1}=\mathcal{N}(\mathbf{m}_{1},\mathbf{S}_{1})\), \(q_{\tau}\) is Gaussian as well [100] with mean and covariance computed as follows:

\[\mathbf{m}_{\tau} =(1-\tau)\mathbf{m}_{1}+\tau\mathbf{m}_{2}\] \[\mathbf{S}_{\tau} =\mathbf{S}_{1}^{-1/2}\left((1-\tau)\mathbf{S}_{1}+\tau\left(\mathbf{S}_{1}^ {1/2}\mathbf{S}_{2}\mathbf{S}_{1}^{1/2}\right)^{1/2}\right)^{2}\mathbf{S}_{1}^{-1/2} \tag{7}\]

which simplifies even further when the covariances are diagonal.

Now, we can define convexity along Wasserstein geodesics [4] as follows.

**Definition 2**.: _Let \(\mathcal{L}:\mathcal{P}(\mathbb{R}^{d})\to\mathbb{R}\), \(\mathcal{L}\) is \(\lambda\) geodesics convex with \(\lambda>0\) if for any \(q_{0},q_{1}\in\mathcal{P}(\mathbb{R}^{d})\) it holds that_

\[\mathcal{L}(q_{\tau})\leq(1-\tau)\mathcal{L}(q_{0})+\tau\mathcal{L}(q_{1})- \frac{\lambda\tau(1-\tau)}{2}\mathcal{W}_{2}^{2}(q_{0},q_{1}) \tag{8}\]

_where \(\mathcal{W}_{2}^{2}(q_{0},q_{1})\) is the Wasserstein distance defined as [51, 52, 104]_

\[\mathcal{W}_{2}^{2}(q_{1},q_{0})=\inf_{\gamma\in\Pi(q_{1},q_{0})}\int\|\mathbf{ \theta}_{1}-\mathbf{\theta}_{0}\|_{2}^{2}\mathrm{d}\gamma(\mathbf{\theta}_{1},\mathbf{ \theta}_{0}) \tag{9}\]

_with \(\Pi(\cdot,\cdot)\) being the space of measure with \(q_{0}\) and \(q_{1}\) as marginals._

While mathematically proving the geodesics convexity of the predictive likelihood for arbitrary architectures and densities is currently beyond the scope of this work, we can empirically define a proxy using the _functional loss barrier_, defined as follows.

**Definition 3**.: _The functional loss barrier along the Wasserstein geodesics from \(q_{0}\) and \(q_{1}\) is defined as the highest difference between the marginal loss computed when interpolating two solutions \(q_{0}\) and \(q_{1}\) and the linear interpolation of the loss at \(q_{0}\) and \(q_{1}\):_

\[\mathcal{B}(q_{0},q_{1})=\max_{\tau}\mathcal{L}(q_{\tau})-\left((1-\tau) \mathcal{L}(q_{0})+\tau\mathcal{L}(q_{1})\right) \tag{10}\]

_where \(q_{\tau}\) follows the definition in Eq. (6)._
This definition is a more general than the ones in [2; 30; 32] but we can recover [30] by assuming delta posteriors \(q_{i}=\delta(\mathbf{\theta}-\mathbf{\theta}_{i})\) and we can further recover [2; 32] by also assuming \(\mathcal{L}(q_{0})=\mathcal{L}(q_{1})\).

A comment on mixtures.In previous paragraphs, we argued that the mixture of distributions is not sufficient to capture the underlying complex geometry of the posterior. Now, we want to better illustrate this choice with a simple example. In Fig. 3 we plot the test likelihood with two interpolation strategies between two solutions (MLP on CIFAR10): the Wasserstein geodesics and the mixture. With mixtures, we see that the likelihood is pretty much constant during the interpolation, but this is very miss-leading: we don't see barriers not because they don't exist, but because the mixture simply re-weights the distributions, without continuously transporting mass in the parameter space.

## 4 Aligning distributions by looking for permutation symmetries

In this section, we formalize the algorithm that aligns the solutions of Bayesian inference through permutation symmetries of weight matrices and biases. Let \(\mathbb{S}(d)\) be the set of valid \(d\times d\) permutation matrices. Given a generic distribution \(q(\mathbf{\theta})\), we can apply a permutation matrix \(\mathbf{P}\in\mathbb{S}(D_{l})\) to a hidden layer output at layer \(l\), and if we define \(\mathbf{\theta}^{\prime}\) to be equivalent to \(\mathbf{\theta}\) with the exception of

\[\mathbf{W}^{\prime}_{l}=\mathbf{P}\mathbf{W}_{l},\quad\mathbf{b}^{\prime}_{l}=\mathbf{P}\mathbf{b}_{ l},\quad\mathbf{W}^{\prime}_{l+1}=\mathbf{W}_{l+1}\mathbf{P}^{\top}\,, \tag{11}\]

then \(P_{\#}q\) is the equivalent push-forward distribution for \(\mathbf{\theta}^{\prime}\), where \(P\) is the associated permutation map. Let us define the distribution over the functional output of the model as

\[q(\mathbf{f}(\mathbf{\theta},\cdot))=\int\delta\left(\mathbf{f}(\mathbf{\theta},\cdot)-\mathbf{f}( \widehat{\mathbf{\theta}},\cdot)\right)\mathrm{d}q(\widehat{\mathbf{\theta}})\,, \tag{12}\]

and, equivalently, the distribution on the function using the permuted parameters as

\[q(\mathbf{f}(\mathbf{\theta}^{\prime},\cdot))=\int\delta\left(\mathbf{f}(\mathbf{\theta}^{ \prime},\cdot)-\mathbf{f}(\widehat{\mathbf{\theta}}^{\prime},\cdot)\right)\mathrm{d}P _{\#}q(\widehat{\mathbf{\theta}}^{\prime})\,, \tag{13}\]

where in both cases \(\delta(\cdot)\) is the Dirac function. Then, it is simple to verify that the two models are functionally equivalent for any inputs,

\[q(\mathbf{f}(\mathbf{\theta},\cdot))=q(\mathbf{f}(\mathbf{\theta}^{\prime},\cdot))\,. \tag{14}\]

This implies that for any weight-space distribution \(q\), there exists a class of functionally equivalent solutions \(P_{\#}q\), in the sense of Eq. (14). These same considerations can be easily extended to other layers, by considering multiple permutation matrices \(\mathbf{P}_{l}\). For our analysis, given two solutions \(q_{0}\) and \(q_{1}\) we are interested in finding the permuted distribution \(P_{\#}q_{1}\), functionally equivalent to \(q_{1}\), in such a way that once interpolating using Eq. (6) we observe similar performance to \(q_{0}\) and \(q_{1}\). Formally, we can write

\[\operatorname*{arg\,min}_{P}\mathcal{D}(q_{1}(\mathbf{f}(\mathbf{\theta}^{\prime}, \cdot)),q_{0}(\mathbf{f}(\mathbf{\theta},\cdot)))=\operatorname*{arg\,min}_{P} \mathcal{D}(P_{\#}q_{1}(\mathbf{\theta}),q_{0}(\mathbf{\theta}))\,, \tag{15}\]

where \(\mathcal{D}\) is a generic measure of discrepancy.2

Footnote 2: To be formally correct, the l.h.s. is a discrepancy defined on stochastic processes while the r.h.s. is defined on random vectors. Under mild assumptions on the distribution on the parameters and the architectures, \(\mathcal{D}\) is well defined in both cases [see e.g., 81].

### Problem setup for permutation of vectors

We start from a single vector of parameters, disregarding for the moment the functional equivalence constrain. We will extend these results to matrices and multiple layers later. In practice considering

Figure 3: **Wasserstein geodesics and mixtures.** Test likelihood for mixture and the Wasserstein geodesics interpolation. Solutions are MLPs trained on CIFAR10.
Gaussian distributions, we know that if \(q=\mathcal{N}(\mathbf{m},\mathbf{S})\), then \(P_{\#}q=\mathcal{N}(\mathbf{P}\mathbf{m},\mathbf{PS}\mathbf{P}^{\top})\) and if \(q=\mathcal{N}(\mathbf{m},\text{diag}(\mathbf{s}^{2}))\) then \(P_{\#}q=\mathcal{N}(\mathbf{P}\mathbf{m},\text{diag}(\mathbf{Ps}^{2}))\)[11]. With the kl divergence kl\([P_{\#}q_{1}\parallel q_{0}]\) it's easy to verify that it leads to just a distance between means, disregarding any covariance information. While certainly this represents a valid choice, we argue that we can find a better solution by using the Wasserstein distance. For Gaussian measures, the Wasserstein distance has analytic solution:

\[\mathcal{W}_{2}^{2}(q_{1},q_{0}) =\left\lVert\mathbf{m}_{0}-\mathbf{m}_{1}\right\rVert_{2}^{2}+\mathrm{Tr }\bigg{(}\mathbf{S}_{1}+\mathbf{S}_{0}-2\left(\mathbf{S}_{1}^{1/2}\mathbf{S}_{0}\mathbf{S}_{1}^{1/ 2}\right)^{1/2}\bigg{)}=\] \[=\left\lVert\mathbf{m}_{0}-\mathbf{m}_{1}\right\rVert_{2}^{2}+\left\lVert \mathbf{S}_{0}-\mathbf{S}_{1}\right\rVert_{F}^{2}, \tag{16}\]

where \(\left\lVert\cdot\right\rVert_{F}\) denotes the Frobenius norm, \(\left\lVert\mathbf{A}\right\rVert=\sum_{ij}{a_{ij}}^{2}\), and where the second line is valid only if the covariances commute (\(\mathbf{S}_{1}\mathbf{S}_{0}=\mathbf{S}_{0}\mathbf{S}_{1}\)). In our case, then, we can simplify as follows:

\[\mathcal{W}_{2}^{2}(P_{\#}q_{1},q_{0})=\left\lVert\mathbf{m}_{0}-\mathbf{P}\mathbf{m}_{1} \right\rVert_{2}^{2}+\left\lVert\mathbf{s}_{0}-\mathbf{Ps}_{1}\right\rVert_{2}^{2}. \tag{17}\]

To summarize, the problem now can be written as:

\[\operatorname*{arg\,min}_{\mathbf{P}\in\mathbb{S}(d)}\left\lVert\mathbf{m}_{0}-\mathbf{ P}\mathbf{m}_{1}\right\rVert_{2}^{2}+\left\lVert\mathbf{s}_{0}-\mathbf{Ps}_{1}\right\rVert_{2}^{2}= \operatorname*{arg\,max}_{\mathbf{P}\in\mathbb{S}(d)}\left\langle\mathbf{P}\big{|}\bm {m}_{0}\mathbf{m}_{1}^{\top}+\mathbf{s}_{0}\mathbf{s}_{1}^{\top}\right\rangle_{F}\,, \tag{18}\]

where the expression \(\left\langle\mathbf{A}|\mathbf{B}\right\rangle_{F}\) is the Frobenius inner product, \(\left\langle\mathbf{A}|\mathbf{B}\right\rangle_{F}=\sum_{ij}A_{ij}B_{ij}\). Note that the r.h.s. of Eq. (18) is a valid instantiation of the linear assignment problem (lap) [10], which can be solved in polynomial time.

### From vectors to neural network parameters

Finally, we need to take into account that we have multiple layers and weight matrices, and that we are trying to find functionally equivalent solutions. For this, we decide to explicitly change our main objective by enforcing the functional equivalence constraint as follows:

\[\operatorname*{arg\,min}_{\{P_{i}\}}\mathcal{W}_{2}^{2}\left(P_{1\#}q_{1}^{(1 )},q_{0}^{(1)}\right)+\mathcal{W}_{2}^{2}\left(\left(P_{2}\circ P_{1}^{\top} \right)_{\#}q_{1}^{(2)},q_{0}^{(2)}\right)+\cdots+\mathcal{W}_{2}^{2}\left( \left(P_{L-1}^{\top}\right)_{\#}q_{1}^{(L)},q_{0}^{(L)}\right)\,,\]

where the notation \(\left(P_{l}\circ P_{l-1}^{\top}\right)\) represents the composition of the two permutation maps applied to rows and columns of the random weight matrices. More conveniently, this can be rewritten in terms of means and standard deviations. To leave the notation uncluttered, let's collect the means and the standard deviations for the layer \(l\) in \(\mathbf{M}^{(l)}\) and \(\mathbf{S}^{(l)}\), which are now both \(D_{l}\times D_{l-1}\) matrices, so that \(q^{(l)}=\prod_{ij}\mathcal{N}(M_{ij}^{(l)},S_{ij}^{(l)})\). Now we can write,

\[\operatorname*{arg\,max}_{\{\mathbf{P}_{i}\}_{i=1}^{L}} \left\langle\mathbf{M}_{0}^{(1)}\Big{|}\mathbf{P}_{1}\mathbf{M}_{1}^{(1)} \right\rangle_{F}+\left\langle\mathbf{S}_{0}^{(1)}\Big{|}\mathbf{P}_{1}\mathbf{S}_{1}^{(1 )}\right\rangle_{F}+\left\langle\mathbf{M}_{0}^{(2)}\Big{|}\mathbf{P}_{2}\mathbf{M}_{1}^{(2 )}\mathbf{P}_{1}^{\top}\right\rangle_{F}+\left\langle\mathbf{S}_{0}^{(2)}\Big{|}\mathbf{P }_{2}\mathbf{S}_{1}^{(2)}\mathbf{P}_{1}^{\top}\right\rangle_{F}+\] \[+\cdots+\left\langle\mathbf{M}_{0}^{(L)}\Big{|}\mathbf{M}_{1}^{(L)}\mathbf{P }_{L-1}^{\top}\right\rangle_{F}+\left\langle\mathbf{S}_{0}^{(L)}\Big{|}\mathbf{S}_{1}^ {(L)}\mathbf{P}_{L-1}^{\top}\right\rangle_{F}\,.\]

This optimization problem is more challenging than the one presented in Eq. (18): we are interested in finding permutation matrices to be applied concurrently to rows and columns of both means and standard deviations. This class of problems, also known as sum of bilinear assignment problems (solap), is NP-hard and no polynomial-time solutions exist. For this reason, we propose to use the setup in Ainsworth et al. [2] by extending it to our problem. In particular, by fixing all matrices with the exception of \(\mathbf{P}_{l}\), we observe that also in our case the problem can be reduced to a classic lap.

\[\operatorname*{arg\,max}_{\mathbf{P}_{l}} \left\langle\mathbf{M}_{0}^{(l)}\Big{|}\mathbf{P}_{l}\mathbf{M}_{1}^{(l)}\mathbf{P }_{l-1}^{\top}\right\rangle_{F}+\left\langle\mathbf{M}_{0}^{(l+1)}\Big{|}\mathbf{P}_{(l +1)}\mathbf{M}_{1}^{(l+1)}\mathbf{P}_{l}^{\top}\right\rangle_{F}+\] \[\left\langle\mathbf{S}_{0}^{(l)}\Big{|}\mathbf{P}_{l}\mathbf{S}_{1}^{(l)}\mathbf{P }_{l-1}^{\top}\right\rangle_{F}+\left\langle\mathbf{S}_{0}^{(l+1)}\Big{|}\mathbf{P}_{(l +1)}\mathbf{S}_{1}^{(l+1)}\mathbf{P}_{l}^{\top}\right\rangle_{F}=\] \[=\operatorname*{arg\,max}_{\mathbf{P}_{l}} \left\langle\mathbf{P}_{l}\right|\mathbf{M}_{0}^{(l)}\mathbf{P}_{l-1}\left(\mathbf{M}_{1} ^{(l)}\right)^{\top}+\left(\mathbf{M}_{0}^{(l+1)}\right)^{\top}\mathbf{P}_{l+1}\mathbf{M}_ {1}^{(l+1)}+\] \[\mathbf{S}_{0}^{(l)}\mathbf{P}_{l-1}\left(\mathbf{S}_{1}^{(l)}\right)^{\top}+ \left(\mathbf{S}_{0}^{(l+1)}\right)^{\top}\mathbf{P}_{l+1}\mathbf{S}_{1}^{(l+1)}\right\rangle _{F}. \tag{19}\]

As discussed in [2], going through each layer, and greedily selecting its best \(\mathbf{P}_{l}\), leads to a coordinate descent algorithm which guarantees to end in finite time. We present a pseudo-code in Algorithm 1.
## 5 Experiments

Now, we present some supporting evidence to Conjecture 1. We start by training two replicas of bnns with variational inference (we refer to the Appendix for additional details on the experimental setup). We then compute the marginalized barrier as \(\mathcal{B}(q_{0},q_{1})=\max_{\tau}\mathcal{L}(q_{\tau})-((1-\tau)\mathcal{L}(q _{0})+\tau\mathcal{L}(q_{1}))\) where \(\mathcal{L}(\cdot)\) is the predictive likelihood and \(\tau\in[0,1]\), from which we take 25 evenly distributed points. In particular, we seek to understand what happens to the vi solutions first for the naive interpolation from \(q_{0}\) and \(q_{1}\), and then for the interpolation after aligning \(q_{0}\) and \(P_{\#}q_{1}\). We experiment with mlps with three layers and ResNet20 [41] with various widths on MNIST [60], Fashion-MNIST [108] and CIFAR10 [56]. All models are trained without data augmentation [106] and with filter response normalization (frn) layers instead of BatchNorm. Finally, we set the prior to be Gaussian \(\mathcal{N}(\mathbf{0},\alpha^{2}\boldsymbol{I})\), with the flexibility of choosing the variance.

### Low-barrier interpolations

Fig. 4 shows the results with and without alignment. We see that regardless of the dataset and the model used, the performance degrades significantly when we move between the two solutions with the naive interpolation, showing the existence of barriers in the predictive likelihood for Gaussian vi solutions. However, with the alignment proposed in SS 4 and Algorithm 1, we recover zero barrier solutions for mlps on both MNIST and CIFAR10, and nearly-zero barrier for ResNet20 on CIFAR10. This holds both for the train and test splits, with quantifiably smaller barriers in the test set.

In Fig. 5 we study the effect of the width of a neural network in relation to the loss barrier by taking an mlp and a ResNet20 with an increasing number of hidden features. We see that wider models generally provide lower barriers: for mlps this holds true with and without alignment, while for the ResNet20 this is happening only after alignment. This extends some previous analysis done on loss-optimized networks. Specifically, Enetzari et al. [30] show that barriers seem to have a double descent trend, while Ainsworth et al. [2] discuss that low barrier solutions after accounting for symmetries are easier to find in wider networks.

Figure 4: **Zero barrier solutions. Comparison of loss barriers for standard vi (gray) and vi with alignment (orange). While loss barriers always appear between two solutions in the standard vi approach, in the case of vi with alignment there is no noticeable loss barrier for mlps and a nearly-zero loss barrier for ResNet20.**

Figure 5: **Effect of width. After distribution alignment, wider models exhibit lower likelihood barrier.**
We speculate that this might be due to the limiting behavior of Bayesian neural networks, which makes the posterior landscape Gaussian-like [48; 40]. While this does not fully explain the phenomenon observed, the existing connections between bnn and non-parametric models, like Gaussian Processes (gps) [83] and deep Gaussian processes (dgps) [22; 20; 93; 28], can provide additional insights on the role of symmetries in weight space [80].

Finally, as an additional check, we analyze the log-posterior with and without alignment by projecting the density into two dimensional slices, following the setup in [47; 36]. We study the two dimensional subspace of the parameter space supported by the hyperplane \(H=\{\mathbf{\theta}\in\mathbb{R}^{d}\,|\,\mathbf{\theta}=a\mathbf{\theta}_{a}+b\mathbf{\theta} _{b}+(1-a-b)\mathbf{\theta}_{c}\}\),, where \(a,b\in\mathbb{R}\) and \(\mathbf{\theta}_{a}\), \(\mathbf{\theta}_{b}\) and \(\mathbf{\theta}_{c}\) are the samples either from \(q_{0}\), \(q_{1}\) and \(q_{\tau}\) without alignment or from \(q_{0}\), \(P_{\#}q_{1}\) and \(P_{\#}q_{\tau}\) with alignment. With this configuration, all three samples always lie on this hyper-plane. In Fig. 6, we present the visualization of ResNet20 trained on CIFAR10. We see that the samples from \(q_{0}\) and \(P_{\#}q_{1}\) are connected by higher density regions than the ones between \(q_{0}\) and \(q_{1}\). This is in line with the results in Fig. 4, where we see that the loss barrier is lower after alignment.

### Analyzing the effect of the prior and testing the cold posterior effect

In all previous experiments we used a Gaussian prior \(\mathcal{N}(\mathbf{0},\alpha^{2}\mathbf{I})\) with fixed \(\alpha^{2}\); now we study the effect of a varying prior variance \(\alpha^{2}\) on the behavior of the loss barriers. We experiment this on a mlp trained on MNIST and Fashion-MNIST and on a ResNet20 (width x8) on CIFAR10. We report the results in Fig. 7. We can appreciate two behaviors: with alignment, there is no measurable effect of using different variances in finding zero-barrier solutions; on the contrary, without alignment we see that naive vi solutions are easier to interpolate with lower barrier when the prior is more diffused. At the same time, we see that higher variances produce bigger gaps between train barriers and test barriers. We speculate that this is due to overfitting happening with more relaxed priors, which makes low-barrier (but low-likelihood) solutions easier to find.

Additionally, several previous works have analyzed the effect of tempering the posterior in bnns[106; 112; 110; 47]. Specifically, we are interested in the distribution \(p_{T}(\mathbf{\theta}\,|\,\mathbf{Y})\propto(p(\mathbf{Y}\,|\,\mathbf{\theta},\mathbf{X})p(\mathbf{ \theta}))^{1/T}\), where \(T\) is known as the temperature. Note that starting from the above definition, we can write an equivalent elbo for vi which takes into account \(T\). For \(T<1\), we have cold posteriors, which are

Figure 6: **Posterior density visualization.** Analysis of the log-posterior computed for ResNet20. Samples from \(q_{0}\) and \(q_{1}\) are connected by lower density regions, while \(q_{0}\) and \(P_{\#}q_{1}\) are not.

Figure 7: **Effect of prior variance.** After distribution alignment, prior variance has low effect in finding zero-loss barriers, while with naive interpolation we see a decreasing trend the higher the variance is.
sharper than the true Bayesian posteriors, while for \(T>1\) we have warm posterior, which are more diffused. In Fig. 8 we see that barriers for cold posteriors with alignment are marginally closer to zero than for warm posteriors. Note that cold temperatures concentrate the distribution around the map, which motivates a further comparison with a non-Bayesian approach.

### Comparison with SGD solutions

Motivated by the results with cold posteriors, we also compare the behavior of barriers for vi versus map solutions obtained via sgd. Note that using the map solution we end up with the same setup of weight matching than in [2]. In Fig. 9 we report this comparison, which is carried out on ResNet20 for various width multipliers. For narrow models, we see that vi with alignment and map with weight matching both behave in a similar manner. Interestingly, as the model becomes wider, map solutions achieve marginally lower barriers than vi. We speculate that this might be due to the simple Gaussian parameterization for the approximate posterior, which doesn't completely capture the local geometry of the true posterior.

### Effect of normalization layers and data augmentation

We conclude this section with a discussion on the effects of normalization layers and data augmentation on the loss barriers for Bayesian neural networks.

Different normalization strategies can affect the overall geometry of the problem [30; 2]. As discussed in [49], interpolating with BatchNorm layers [45] is pathological due to _variance collapse_ of the feature representation in hidden layers. Additionally, batch-dependent normalization layers don't have a clear Bayesian interpretation, since the likelihood cannot factorize. For our experiments we choose to use the frn layer, as done in previous works [e.g., 47]. Note that frn layers are invariant to permutation units and therefore can be aligned without problems. To test the _variance collapse_ behavior, we analyze the variance of activations following the instructions in [49, SS3.1], with the sole difference that the activations are marginalized w.r.t. samples from the posterior. In Fig. 10 we can see that there isn't a pathological variance collapse after alignment. Finally, in Fig. 11 we compare another normalization layer, the LayerNorm (ln) [6]. Note that ln is also batch-independent, it has a clear Bayesian interpretation and it is invariant to permutation units. Indeed, we see that both normalization layers can be aligned, and ln exhibits lower barriers than frn.

Finally, in all previous experiments we skipped data augmentation, because the random augmentations introduce stochasticity which lacks a proper Bayesian interpretation in the formulation of the likelihood function (e.g. re-weighting of the likelihood due to the increase of the effective sample size [77]). Additionally, data augmentation can contribute to spurious effects difficult to disentangle (e.g., cold posterior effect [106]). Nonetheless, during the development of the method we didn't make an

Figure 8: **Effect of temperature.** After alignment, cold posteriors makes barriers marginally closer to zero

Figure 9: **VI versus sgd. vi and sgd behave equivalently after permutation alignment in narrow models. For wider models, sgd solutions reach lower barriers than vi.**
assumption on data augmentation and in Fig. 12 we experiment both with and without augmentation, showing that we are still able to recover similar low barrier solutions in both cases. Having said that, we advocate caution when using data augmentation in Bayesian neural networks, as it changes the shape of the posterior.

## 6 Related work

In earlier sections of the paper, we already briefly discussed and reviewed relevant works on mode connectivity, symmetries in the loss landscape and connection with gradient-based optimization methods. Here we discuss some relevant works on the connection to Bayesian deep learning. The work of Garipov et al. [36] sparked several contributions on exploiting mode connectivity for ensembling models, which is akin to Bayesian model averaging. For example, in [46] the authors propose to ensemble models using curve subspaces to construct low-dimensional subspaces of parameter space. These curve subspaces are, among others, the non-linear paths connecting low-loss modes (and consequently high-posterior density) in weight space. In [31], the authors attempt to explain the effectiveness of deep ensembles [59], concluding that it is partially due to the diversity of the sgd solutions in parameter space induced by random initialization. More recently, in [9] the authors reason about mode connecting volumes, which are multi-dimensional manifolds of low loss that connect many independently trained models. These mode connecting volumes form the basis for an efficient method for building simplicial complexes for ensembling. Here, we want to highlight that these works have not taken into account the permutation symmetries. Finally, in a concurrent submission, [107] proposes an algorithm to remove symmetries in mcmc chains for tanh networks.

## 7 Conclusions

By studying the effect of permutation symmetries, which are ubiquitous in neural networks, it is possible to analyze the fundamental geometric properties of loss landscapes like (linear) mode connectivity and loss barriers. While previously this was done on loss-optimized networks [2, 30, 32], in this work we have extended the analysis to Bayesian neural networks. We have studied the linear connectivity properties of approximate Bayesian solutions and we have proposed a matching algorithm (Algorithm 1) to search for linearly connected solutions, by aligning the distributions of two independent vi solutions with respect to permutation matrices. We have empirically validated our framework on a variety of experiments, showing that we can find zero barrier linearly-connected solutions for bnn trained with vi, on shallow models as well as on deep convolutional networks. This brings evidence for Conjecture 1 regarding the linear connectivity of approximate Bayesian solutions. Furthermore, we have studied the effect of various design hyper-parameters, like width, prior and temperature, and observed complex patterns of behavior, which would require additional research. In particular, the experiments raise questions regarding the relation between linear mode connectivity and the generalization of bnn, as well as the role of width with respect to limiting behaviors of non-parametric models like gps and dgp.
## Acknowledgments and Disclosure of Funding

SR wants to thank David Bertrand and the AIAO team at Stellantis for their work on the software and hardware infrastructure used for this work.
