On permutation symmetries in Bayesian neural network posteriors: a variational perspective
Authors: Simone Rossi
Ankit Singh
Thomas Hannagan
Venue: NeurIPS 2023 poster (2023)
Keywords: Bayesian deep learning
approximate inference
permutation symmetries

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. Building on the work of Ainsworth et al. (2023), we frame 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.

On permutation symmetries in Bayesian neural
network posteriors: a variational perspective
Simone Rossi Stellantis, France
Ankit Singh Stellantis, India
Thomas Hannagan Stellantis, France
Abstract
The clusive nature of gradient-based optimization in neural networks is tied to their loss landscape geometry. which poorly understood. However recent work has brought solid evidence that there is essentially no loss barrier between the local so- lutions of gradient descent, once accounting for sright-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, first extend the formalism of marginalized loss barrier and solution interpolation to BNNs, before proposing matching algorithm search for linearly connected solutions. This is achieved by aligning the distributions of two independent approximate Bayesian solutions with respect permutation matrices. We build on the results of Ainsworth al. (2023), reframing the problem as a combinatorial optimization one, using an approximation 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.
Introduction
Throughout the last decade, deep neural networks (DNNs) have achieved significant success in a wide range of practical applications, becoming the fundamental ingredient for e.g., computer vision [c.g.. 57, 27,25,4 41, 64], language models [e.g.. 24, 82, 17] and generative models [e.g.. 54, 96, 97, 98, 102, 34]. Despite recent important ad advancements, understanding the loss landscape of DNNS 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 (.. 5, 7, 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 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 the landscape, there exists a path connecting them such that the loss constant or near constant (or, said differently, the loss barrier is null). We refer to lineur mode connectivity when the path connecting the two solutions 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.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Solution 1
Solution 2
Naive interpolation
Aligned interpolation
2-
m 0
2
2
0
2
0
2
0
2
Figure 1: Permutation symmetries for regression on the Snelson dataset [95]. (Left) Two different solutions, with similar space behavior (showing μl 2ar). (Right) Functions obtained when using different strategies to interpolate between solutions. When we align the solutions, by taking into account the permutation symmetrics, 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 nor (see the black curve which is single function samplc with fixed randomness).
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 natural form of regularization and allows to reason about uncertainty in the predictions [101, 73, 66]. Bayesian inference for deep neural networks 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- e-posteriori (MAP) solutions, (ii) the loss landscape 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 posterior'. 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. illustrates this situation for a regression task on the Snelson dataset using 3-layer DNN: on the left we compare two (approximate) solutions which have different weight-space posterior but similar function-sppce behavior. Notably, when interpolate these two solutions (Fig. 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 informal generic conjecture:
notoriously challenging,
Solutions of approximate Bayesian inference for neural networks are linearly connected after ac- counting for functionally equivalent permutations.
While being similar to the one in [30, 21, if this conjecture 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 conjecture broadly open regarding the choice of the approximation method allow for a more general discussion. More specifically, in this paper we will analyze and focus our discussion on the variational inference framework, making more specific conjecture:
Conjecture 1. Variational inference solutions for approximate Bayesian inference in neural networks are linearly connected after accounting for functionally equivalent permutations.
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].

---

Log- pog-posterior for CIFARIO
Contributions. With this work, we aim at studying the linear connectivity properties of approximate solutions to the Bayesian inference problem and we make several con- tributions. (i) We extend the formalism of loss barrier and solution interpolation to BNNs. (ii) For the variational in- ference setting, propose a matching algorithm search for linearly connected solutions by aligning the distributions of two independent solutions with respect to permutation matri- ces. Inspired by [2]. we frame the problem as 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 bar- riers 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 of the solution şi such that it can be linearly connected through high density regions to qu-
P#4
40
Figure 2: Permutations in multi-modal posterior. Lag-posterior for MLP/CI- FARI0, showing the two solutions (go and 41) for which we can find permutation map such that P= can be linearly con- nected to do with low barrier (brighter re- gions).
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 generic multilayer perceptron (MIP) with I. layers, where the output of the 1-th layer f( x)) is vector-valued function of the previous layer output fi-1 as follows,
f0 W (a((ft-100-1111
where a(-) is a non-linearity, Wi is a Dr x Dj_1 weight matrix and by the corresponding bias vector. We shall refer the parameters of the layer as 0; {b), Wi), and the union of all trainable parameters as 0 (0:)--1
The objective of using Bayesian inference on deep neural networks [67, 65] involves inferring posterior distribution over the parameters of the neural network given the available dataset ( [X, Y) {( y,) - This requires choosing likelihood and prior [73, 74, 103]:
p(0 Y X) = >Y 0, X )(())
(2)
where the normalization constant Z is the marginal likelihood p(Y XX. As usually done, we assume p(y, 0,z;).
that the likelihood factorizes over observations, i.e. p(Y 0,, X) =
Bayesian deep learning is intractable due the non-conjugacy likelihood pariihood--rior and thus don't 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 P(RR) be the space of probability measures on Rd; VI reframes the inference problem into optimization one, commonly by introducing parameterized distribution (()) )) E P(RR)) which is optimized to minimize the Kullback-Leibb (KL)
divergence with respect the true posterior p(0) Y. X). In practice, this involves the maximization of the evidence lower bound (ELBO) defined as
(g)
del
f logg(Y 0, X d))
KI. (a()) || [()]
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 verrparraeetrizzaioon on the quality of the approximation [88, 44, 58]. Here we are interested in how different solutions to Eq. (3) relate 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 the model parameters with respect to (an approximation of) the posterior. As such, use the predictive
3

---

Let's assume two models trained with with two different initializations, random seeds, and batch ordering. Variational inference in the classic inverse sense KL [g || p] is mode seeking, thus we expect interpolate likelihood, a proper scoring method for probabilistic models [83], defined as
log ply" 1lyy"I x=)) log / ply" 00,x") )p(@ Y. dd log / ply" 0, 0 )dq(()
(4)
where (0)) an approximation of the true posterior (parametric or otherwise), {æ* { are re- spectively the input and its corresponding label data point under evaluation. To keep the notation uncluttered for the remaining of the paper, we write the predictive likelihood computed for set of points {2), V Ji as functional L: P(RR") - IR, defined as
((g) def Σ log f ==1
/ plut 0, o))d)(()))
(5)
the runs to converge different solutions, say qo and Q1- To test the loss barrier between the two solutions we need to decide on the interpolation rule. We decide to interpolate the solutions following the Wasserstein geodesics between Yo and 91- First, let's start with a few definitions. Let E (P(Rd)) be probability measure on R and T R + IR measurable map; denote Tæq the push-forward measure of q through T. Now can introduce the Wasserstein geodesics, as follows.
Definition 1. The Wasserstein geodesics between qu and qı is defined as the path
4r = ((1 - r)Id + TT 40
T TE [0.1
(6)
where Id is the identity map and T is the optimal transport map between and qı which for Brenier's 's theorem [16]. unique.
While we could interpolate using a mixture of the two solutions, we argue that this choice is trivial and does fully give picture of the underlying loss landscape. Indeed, Eq. (6) fundamentally different from a naive mixture path or = (1 - T) + Tqi In case of Gaussian distributions, when
Qu = (m,,, Sy) and q (m,, S1), 1,), dr Gaussian as well [100] with mean and covariance computed as follows:
m, (1 -) + mm2
S, = 1-2 (1- + + (s
S
1/2) "0("" 1/2)
ST 5-1/2
which simplifies even further when the covariances are diagonal.
Now, we can define convexity along Wasserstein geodesics [4] as follows. Definition 2. Let L: P(RB) + R, L is À geodesics conver with À >> for any go 4n E P(Rd)
it
holds that
((g,) S (1 ((--7)4 () + (((1)) Ar r) w W2 91)
(8)
where W2 (go, 41) is the Wasserstein distance defined as [104,5,, 52]
inf yETT(a .90) / [100 (0
Wa
(9)
with II(..)) being the space of measure with go and qı 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 proxy using the functional loss barrier, defined as follows.
Definition 3. The functional loss barrier along the Wasserstein geodesics from qu and qı is defined as the highest difference between the marginal loss computed when interpolating two solutions go and qu and the linear interpolation of the loss go and q1:
B(go. 91) max L(g+) - ((1 - ((g)) + L(q1))
(10)
where 4u follows the definition in Eq. (6).

---

This definition is more general than the ones in [2, 30, 32] but we can recover [30] by assuming delta posteriors 9; 5(0 0,) and we can further recover [2, 32] by also assuming ( (())).
comment on mixtures. In previous para- graphs, argued that the mixture of distributions is not sufficient to capture the underlying complex geometry of the posterior. Now, want to better illustrate this choice with a simple example. In Fig. we plot the test likelihood with two inter- polation strategies between two solutions (MLP on CIFARIO): the Wasserstein geodesics and the mixture. With mixtures, we see that the likelihood pretty much constant during the interpolation, but this is very miss-leading: we don't see barriers because they don't exist, but because the mix- ture simply re-weights the distributions, without continuously transporting mass in the parameter space.
+
Predictive livelihood
CC Cldd 24
Mecred
2.00
Vicaas
Interpolation coefficient
Figure Wasserstein geodesics and mixtures. Test likelihood mixture Wasserstein geodesics interpolation. Solutions are MLPs trained on CIFAR1O.
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 S(d) be the set of valid d x d permutation matrices. Given generic distribution (()), we can apply a permutation matrix P E S(Di) to a hidden layer output layer and if define f' be equivalent to f with the exception of
W PW 1, Pbr,
W Wi+1
W W4+1 WT+PP
(11)
then P q the equivalent forward distribution for o', where P is the associated permutation map. Let us define the distribution over the functional output of the model as
(f(().)) = / (f f(0.-) dq(().
(12)
and, equivalently, the distribution on the function using the permuted parameters as
q(f(0' f ( f(@0 d
(13)
where both cases 5(-) is the Dirac function. Then, it is simple to verify that the two models are functionally equivalent for any inputs,
=
(14)
This implies that for any weight-space distribution g. there exists class of functionally equivalent solutions gæg. in the sense of Eq. (14). These same considerations can be easily extended to other layers, by considering multiple permutation matrices P,. For our analysis, given two solutions qu and qi we are interested in finding the permuted distribution PH functionally equivalent to Qı. in such a way that once interpolating using Eq. (6) we observe similar performance to go and yı. Formally, we can write
arg min D(q 9). = arg min D(P+q)
(15)
where D a generic measure of discrepancy.
4.1 Problem setup for permutation vectors
We start from single vector of parameters, disregarding for the moment the functional equivalence constraint. We will extend these results to matrices and multiple layers later. practice considering
"To be formally correct, the lh..s a discrepancy defined on stochastic processes while the rh.s. is defined on random vectors. Under mild assumptions on the distribution on the parameters and the architectures, D is well defined in both cases [see e.g.. 81].
5

---

Gaussian distributions, we know that if (m, S), then P=q=N(Pm. PSP and if q (mm, dia)(s")) then -= -o (Pm,, diag( P>²)) [11]. With the KI. divergence KI. KI. (. gol it's easy to verify that it leads to just distance between means, disregarding any covariance information. While certainly this represents a valid choice, we argue that we can find better solution by using the
Wasserstein distance. For Gaussian measures, the Wasserstein distance has analytic solution:
112 + Tr S + So 2 - SoS)(2) 1/2
(11, 20) = m m2
|mm mm 1|2 + ||50 -111
(16)
where ||-||, denotes the Frobenius norm, ||||| = Eij aij², and where the second line is valid only if the covariances commute SS = 1)) In our case, then, we can simplify as follows:
(P 2o) = mm Pm. 11 + |180 P
(17)
To summarize, the problem now can be written
arg min ||mo Pm1ll 11 |0 Ps 11 PE)(d)
arg max p| + +008 PES(d)
(18)
where the expression (A[B) F. is the Frobenius inner product, (A Eis Asy Note that the r...s. of Eq. (18) is valid instantiation of the linear assignment problem (LAP) [10], which can be solved in polynomial time.
4.2 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 explicitly change our main objective by enforcing the functional equivalence constraint as follows:
arg min W (Pi (P.}
(1) (1"" + W3 ((22 o PT)
(L) + W ((PT_)) Pi 9u
where the notation (P) P,-. 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 s collect the means and the standard deviations for the layer in M and (()), which are now both Dr x Di-1 matrices, so that Now we can write,
() Ilg N(M s(O)
(M6)) PIM ())
+ (
PiS + (M.') P2M 2 PTT) + ((( P2S P
arg max (P.)
+
(M) 1 ((PPPI-)) PI-1); + (s 1 (()PL_)) PL-1)
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-tin time solutions exist. For this reason, we propose to use the setup in Ainsworth al. [2] by extending to our problem. In particular, by fixing all matrices with the exception of Pt, we observe that also in our case the problem can be reduced to classic L.A..
arg max
(
P.M PT-1) F +
(M!
(+1)
M
(1+1) P.T)
(5
() PT11) + (s
Pat (41+1) P.T) =
P, M (0) Pi-1 (M + (M!ttl) T
Pu (+11)
arg max
P
5 P:-1 (s
(s PHtSH+1)
(19)
As discussed in [2], going through each layer, and greedily selecting its best Pt, leads to coordinate descent algorithm which guarantees to end in finite time. We present pseudo-code in Algorithm 1.

---

Algorithm 1: Algorithm to align variational inference solutions
Data: Variational inference solutions Qo and 91
Result: Permutation matrices P: P: f- I, Vi (L..... L.::
while not converged do
for E conndomPermt 1.............. 1) do
P. = arg max (P, (+), where is the r.h.s. of Eq. (19)
MLP on MNIST
MLP on CIFARIO
ResNet20 on CIFAR10
Predictive liseilhood
0.0
Predictio ljkellhood
1.0
1.--
Predictive likelihood
Googg0oo00000oooooo 0000000000000000 000000000000ooooo
0.5-
1.5
1.0
2.0
0.25
0.50
0.75
1.00
0.25
0.75
1.00
0.00
0.25
0.50
0.75
1000
Interpolation coefficient
Interpolation coefficient
Interpolation coefficient
VI (Train)
VI with distr. alignment (Train)
VI (Test)
VI with distr. alignment (Test)
Figure 4: Zero barrier solutions. Comparison of loss barriers for standard VI (gray) and with alignment (orange). While loss barriers always appear between solutions in the standard vI approach, in the case of VI with alignment there is no noticeable loss barrier for MIPS and a nearly-zero loss barrier for ResNet20.
5 Experiments
Now, we present some supporting evidence to Conjecture 1. We start by training two replicas of BNNS with variational inference (we refer the Appendix for additional details on the experimental setup).
We then compute the marginalized barrier as B Bigo, 41) max, ((gr) ((1 C)) + + T(g
where C(-) is the predictive likelihood and 7 E [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 do and Ql, and then for the interpolation after aligning Qo and P# We experiment with MI.PS with three layers and ResNet20 [41] with various widths on MNIST [60], Fashion-MNISTT [108] and CIFAR 10 [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 N I2), with the flexibility of choosing the variance.
5.1 Low-barrier interpolations
Fig. 4 shows the results with and without alignment. We 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 ğ 4 and Algorithm 1, we recover zero barrier solutions for MLPS on both MNIST and CIFARIO 10, and nearly-zero barrier for ResNet20 CIFARIO. 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 neural network in relation to the loss barrier by taking an MLP and a ResNet20 with an increas- ing number of hidden features. We see that wider models generally provide lower barriers: for MI.Ps this holds true with and without align- ment, while for the ResNet20 this is happening only after alignment. This extends some previ- ous analysis done on loss loss-optimized networks. Specifically, Entezari et al. [30] show that barri- ers seem to have a double descent trend, while Ainsworth al. [2] discuss that low barrier solutions after accounting for symmetries are casier to find in wider networks.
MLP MNIST
ResNet20 CIFAR10
Barrier in the predictive likllhooo
Barrier in the predictive hhlldooo
2t.
2?
Number of features
Number of features
(Train)
VI (Tesl)
with align. (Train)
Vi with align. (Test)
Figure 5: Efffec of width. After distribution alignment, wider models exhibit lower likelihood barrier.
7

---

Posterior without alignment
Posterior with alignment
P# 9+
Note:
0.5
33.96
33962
33955
33966
33.62
33.558
(a)
x10
(b)
x10'
Figure 6: Posterior density visualization. Analysis of the log-posterior computed for ResNet20. Samples from Qo and qi connected by lower density regions, while go and P# 91
not.
MLP MNIST
MLP FMNIST
ResNet20 CIFARLO
Barrier the gedictive likellhood
Barrier in the gredictive likellhoodd
Barrier in the predicti to
10-2
10
10P
10
10t
10
10-
10
10*
Prior
Prior a²
VI (Train)
with distr. alignment (Train)
(Test)
with distr. alignment (Test)
Figure 7: Effeet of prior variance. After distribution alignment, prior variance has low effect in finding zero-loss barriers, while with naive interpolation we see decreasing trend the higher the variance is.
We speculate that this might be due the limiting behavior of Bayesian neural networks, which makes the posterior landscape Gaussian-like [48, 40J. While this does not fully explain the phenomenon observed, the existing connections between BNNs and non-parametric models, like Gaussian Processes (GPS) [83] and deep Gaussian processes (DGPs) [22, 20, 28, 93], can provide additional insights on the role 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 (1 b)0=}. where b TR and 0., O, and f the samples either from To- 2 and dr without alignment or from Qo. P# and P#0+ with alignment. With this configuration, all three samples always lie on this hyper-plane. In Fig. we present the visualization of ResNet20 trained on CIFARIO. We see that the samples from qo and Рæqı are connected by higher density regions than the ones between go and gi. This is in line with the results in Fig. where we that the loss barrier is lower after alignment.
= {0 e TRd 0 = at, + 10,
5.2 Analyzing the effect of the prior and testing the cold posterior effect
In all previous experiments used Gaussian prior A(0,, a² I) with fixed a now study the effect of a varying prior variance a² on the behavior the loss barriers. We experiment this on a MLP trained on MNIST and FaShion-MNIST and ResNet20 (width x8) 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; the contrary, without alignment 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 likeliioohoo) solutions casier find.
Additionally, several previous works have analyzed the effect of tempering the posterior BNNs [106,
1 112, 110, 47]. Specifically, we are interested in the distribution pr(0 Y) x (p(Y O, X (0)) 1 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

---

MLP MNIST
MLP - FMNIST
ResNet20 CIFAR1O
Barrier the predictive likeliieed
Barrier in the predictior likellhoood
prodictive like ihood
Barrier in the
0-
10"
10
10-1
10
10a
10-
10"
Posterior temperature
Posterior temperature
Posterior temperature
VI (Train)
with distr. alignment (Train)
VI (Teat)
with distr. alignment (Test)
Figure 8: Effect of temperature. After alignment, cold posteriors makes barriers marginally closer zero
sharper than the true Bayesian posteriors, while for T > 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.
5.3 Comparison with SGD solutions
ResNet20 CIFAR10
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 carried out on ResNet20 for various width multipliers. For narrow mod-
Barrier in the predictive likellhood
2.0
1.0
0.5
els, we see that VT with alignment and MAP with weight matching both behave in similar manner. Interestingly, as the model becomes wider, MAP solutions achieve marginally lower barriers than VI. We speculate that this might due to the sim- ple Gaussian parameterization for the approximate posterior, which doesn't 't completely capture the local geometry of the true posterior.
Features multiplier
SGD
VI with alignment
SGD with weight matching
Figure 9: VI versus SGD, VI and SGD behave equiv- alently after permutation alignment in narrow mod- els. For wider models, SGD solutions reach lower barriers than VI.
5.4 Effect of normalization layers and data augmentation
We conclude this section with discussion the effects of normalization layers and data augmenta- tion 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 batch-dependent normalization layers don't '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. behavior, we analyze the variance of activations following the instructions in [49, $3.11 13. with the sole difference that the activations are marginalized w.r.t samples from the posterior. In Fig. 10 we can that there isn't 't pathological variance collapse after alignment. Finally, in Fig. 11 compare another normalization layer, the LayerNorm (LN) [6]. Note that LN is also batch-independent, has clear Bayesian interpretation and it is invariant permutation units. Indeed, we see that both normalization layers can be aligned, and LN exhibits lower barriers than FRN.
test the variance collapse
Finally, in all previous experiments we skipped data augmentation, because the random augmenta- tions 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 didn't 't make an

---

ResNet20 CIFAR10
ResNet20 -IIARRO
ResNet20 CIFARR0
Predictive likelihood
Predictive likelihood
Variance
0.5
2-
0.0
9
12
15
18
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Layer index
Interpolation coefficient
Interpolation coefficient
Interpolation after alignment interpolation
Norm. layer FRN
Allgnment
DA
Allgnment Fallse
LN
False
False
Figure 10: Variance collapse. Variance collapse is present for the naive interpolation, but it is reduced after alignment.
Figure 11: Effect of normaliza- tion layers. Both FRN and LN can be aligned, and LN exhibits lower barriers.
Figure 12: Effect of data aug- mentation. With without DA, we are able to recover simi- lar low barrier solutions.
assumption on data augmentation and Fig. 12 we experiment both with and without augmentation, showing that 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- praeent-based optimization methods. Here we discuss some relevant works 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 cow-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 I. the authors attempt to explain the effectiveness of deep ensembles [59], concluding that it is partially due to the diversity of the SGD solutions parameter space induced by random initialization. More recently. in [9] the authors reason about mode connecting volumes, which multi- ditttiidmmenionall 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, want 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.
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 done 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 (A (Algorithm I) to search for linearly connected solutions, by aligning the distributions of two independent solutions with respect permutation matrices. We have empirically validated our framework on a variety of experiments, showing that we can find zero barrier linearly-conneeted solutions for BNNS trained with V1, on shallow models as well as on deep convolutional networks. This brings evidence for Conjecture 1 regarding the linear connectivity of variational inference approximations for BNNS. Furthermore, we have studied the effect of various design typer-parameteer 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 BNNs, well as the role of width with respect to limiting behaviors of nor-marait models like GPs and DGPS.
10

---

Acknowledgments and Disclosure of Funding
The authors want to thank David Bertrand and the AIAO team at Stellantis for their work on the software and hardware infrastructure used for this work.
References
[1] M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, Goodfellow, A. Harp. G. Irving, M. Isard, Y. Jia, R. Joze- fowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, Shlens, B. Steiner, I. Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Va- sudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng. TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems, 2015. Software available from tensorflow.org.
[2] S. Ainsworth, J. Hayase, and S. Srinivasa. Git Re-Basin:
Merging Models modulo Permutation Symmetries. In The Eleventh International Conference on Learning Representations, 2023.
[3] 7. Allen-Zhu and Y. Li. What Can ResNet Learn Efficiently, Going Beyond Kernels? In Wallach, H. Larochelle, Beygelzimer, F. d' d'Alché-Bue, Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[4] L. Ambrosio, N. Gigli, and G. Savare. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics. ETH Zürich. Birkhäuser Basel, 2008. ISBN 9783764387228.
[5] S. Arora, S. S. Du, W. Hu, Z. Li, R. R. Salakhutdinov, and R. Wang. On Exact Computation with an Infinitely Wide Neural Net. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[6] J. L. Ba, R. Kiros, and G. E. Hinton. Layer Normalization, 2016.
[7] Bai and J. D. Lee. Beyond Linearization: On Quadratic and Higher- Order Approximation of Wide Neural Networks. In International Conference on Learning Representations, 2020.
[8] Y. Bengio, O. Delalleau, and N. Le Roux. The Curse of Dimensionality for Local Kernel Machines. Technical Report TR-1258, March 2005.
[9] G. Benton, W. Maddox, S. Lotfi, and A. G. G. Wilson. Loss Surface Simplexes for Mode Connecting Volumes and Fast Ensembling. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Muchine Learning Reseurch, pages 769-779. PMLR, 18-24 Jul 2021.
[10] D. Bertsekas. Network Optimization: Continuous and Discrete Methods. Athena scientific optimization and computation series. Athena Scientific, 1998. ISBN 9788865290299,
[11] C. M. Bishop. Pattern recognition and machine learning. Springer, 1st ed. 2006. Corr. 2nd printing 2011 edition, 2006. ISBN 0387310738.
[12] Blei, S. Mohamed, and Ranganath. Variational inference: Foundations and modern
methods. 2016.
[13] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational Inference: A Review for Statisti- cians. 112(518):859 877, 2017.
[14] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight Uncertainty in Neural Network. In F. Bach and Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1613-1622. PMLR, 2015.
[15] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milnn and Q. Zhang. JAX: Composable Transforma- tions of Python+NumPy Programs, 2018.
[16] Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on Pure and Applied Mathematics, 44(4):375-417, 1991.
11

---

[17] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan. P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, Agarwal, A. Perbbert-Vosss G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. GGay, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language Models are Few-Shot Learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[18] M. Chen, Y. Bai, J. D. Lee, T. Zhao, H. Wang. C. Xiong, and R. Socher. Towards Understanding Hierarchical Learning: Benefits of Neural Representations. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 22134-22145. Curran Associates, Inc., 2020.
[19] T. Chen, E. Fox, and C. Guestrin. Stochastic Gradient Hamiltonian Monte Carlo. In E. P. Xing and Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 1683-1691. PMLR, 2014.
[20] K. Cutajar, E. V. Bonilla, P. Michiardi, and M. Filippone. Random feature expansions for deep Gaussian processes. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 884-893. PMLR, 2017-08.
[21] M. Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. In C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
[22] A. C. Damianou and N. D. Lawrence. Deep Gaussian Processes. In Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2013, volume 31 of Proceedings of Machine Learning Research, pages 207-215. JMLR.org. 2013.
[23] A. G. de G. Matthews, J. Hron, M. Rowland, R. E. Turner, and Z. Ghahramani. Gaussian Process Behaviour in Wide Deep Neural Networks. In Proceedings of the 6th International Conference on Learning Representations, 2018.
[24] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[25] A. Dosovitskiy, L.. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. De- hghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An Image Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.
[26] F. Draxler, K. Veschgini, M. Salmhofer, and Hamprecht. Essentially No Barriers in Neural Network Energy Landscape. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1309-1318. PMLR, 10-15 Jul 2018
[27] S. Duane, A. Kennedy, B. J. Pendleton, and D. Roweth. Hybrid Monte Carlo. 195(2):216- 222, 1987.
[28] M. M. Dunlop, M. A. Girolami, A. M. Stuart, and A. L. Teckentrup. How Deep Are Deep Gaussian Processes? 19(1):2100-214 2018-01.
[29] V. Dutordoir, J. Hensman, M. van der Wilk, C. H. Ek, Z. Ghahramani, and N. Durrande.
Deep Neural Networks as Point Estimates for Deep Gaussian Processes. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 9443 9455. Curran Associates, Inc., 2021.
[30] R. Entezari, H. Sedghi, O. Saukh, and B. Neyshabur. The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks. In International Conference on Learning Representations, 2022.
[31] S. Fort, C. H. Hu, and B. Lakshminarayanan. Deep Ensembles: Loss Landscape Perspective, 2020.
12

---

[32] J. Frankle, G. K. Dziugaite, D. Roy, and M. Carbin. Linear Mode Connectivity and the Lottery Ticket Hypothesis. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning. volume 19 of Proceedings of Machine Learning Research, pages 3259-3269. PMIR, 13-18 Jul 2020.
[33] Franzese, D. Milios, M. Filippone, and P. Michiardi. Revisiting the Effects of Stochasticity for Hamiltonian Samplers. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 6744-6778. PMLR, 17-23 Jul 2022.
[34] G. Franzese, S. Rossi, L. Yang. A. Finamore, D. Rossi, M. Filippone, and P. Michiardi. How Much Is Enough? A Study on Diffusion Times in Score-Based Generative Models. Entropy, 25(4), 2023.
[35] Y. Gal and Z. Ghahramani. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning. volume 48 of Proceedings of Machine Learning Research, pages 1050-1059, PMLR, 2016.
[36] Т. Garipov, P. Izmailov, D. Podoprikhin, D. P. Vetrov, and A. G. Wilson. Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa Desa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages
Curran Associates, Inc., 2018.
[37] A. Morriga-Alonse C. E. Rasmussen, and L. Aitchison. Deep Convolutional Networks as shallow Gaussian Processes, In International Conference on Learning Representations, 2019.
[38] C. Godfrey, D. Brown, T. Emerson, and H. Kvinge. On the Symmetries Deep Learning Models and their Internal Representations. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[39] A. Graves. Practical Variational Inference for Neural Networks. In J. Shawe-Taylor, R. S. Zemel, P. I.. Bartlett, F. Pereira, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 24, pages 2348-2356. Curran Associates, Inc., 2011.
[40] B. He, B. Lakshminarayanan, and Y. W. Teh. Bayesian Deep Ensembles via the Neural Tangent Kernel. In H. Larochelle, M. Ranzato, Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33, 2020.
[41] K. He, X. Zhang, S. Ren, and J. Sun. Deep Residual Learning for Image Recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770-778, 2016.
[42] G. E. Hinton and D. van Camp. Keeping the Neural Networks Simple by Minimizing the Description Length of the Weights. In Proceedings of the Sixth Annual Conference on Computational Learning Theory, 1993. ISBN 0897916115.
[43] M. D. Hoffman, D. M. Blei, C. Wang, and J. W. Paisley. Stochastic Variational Inference. 14 (1):1303-134)) 10000-137,,2111
[44] T. Huix, S. Majewski, A. Durmus, E. Moulines, and A. Korba. Variational Inference of overparameterizze Bayesian Neural Networks: a theoretical and empirical study, 2022.
[45] S. Ioffe and C. Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In F. Bach and D. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 448-456, Lille, France, 07-09 Jul 2015. PMLR.
[46] P. Izmailov, W. Maddox, P. Kirichenko, T. Garipov, D. P. Vetrov, and A. G. Wilson. Subspace Inference for Bayesian Deep Learning. In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence, UAI 2019, Tel Aviv, Israel, July 22-25, 2019, volume 115 of Proceedings of Machine Learning Research, pages 1169-1179. AUAI Press, 2019.
[47] P. Izmailov, S. Vikram, M. D. Hoffman, and A. G. Wilson. What are bayesian neural network posteriors really like? In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 4629-4640. PMLR, 2021.
13

---

[48] A. Jacot, F. Gabriel, and C. Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems, volume 31, pages 8571 -571-. Curran Associates, Inc., 2018.
[49] K. Jordan, H. Sedghi, O. Saukh, R. Entezari, and B. Neyshabur. REPAIR: REnormalizing Permuted Activations for Interpolation Repair. In The Eleventh International Conference on Learning Representations, 2023.
[50] M. I. Jordan, Z. Ghahramani, T. S. Jaakkola, and L. K. Saul. An Introduction to Variational Methods for Graphical Models. 37(2):183-233, 1999-11-01.
[51] L. V. Kantorovich. On the transfer of masses. 37:227-229, 1942.
[52] L. V. Kantorovich. On problem of Monge. 3:225-226, 1948.
[53] M. E. E. Khan, Immer, E. Abedi, and M. Korzepa. Approximate Inference Turns Deep Networks into Gaussian Processes. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alchéé Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[54] D. P. Kingma and M. Welling, Auto-Encoding Variational Bayes. In International Conference on Learning Representations, 2014.
[55] D. P. Kingma, T. Salimans, and M. Welling. Variational Dropout and the Local Reparameteri- zation Trick. In Advances in Neural Information Processing Systems 28, pages 2583. Curran Associates, Inc., 2015.
[56] A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
[57] A. Krizhevsky, Sutskever, and G. E. Hinton. ImageNet Classification with Deep Convo- lutional Neural Networks. In F. Percira, C. J. C. Burges, L. Bottou, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., 2012.
[58] R. Kurle, R. Herbrich, T. Januschowski, B. Wang, and J. Gasthaus. On the detrimental effect of invariances the likelihood for variational inference. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[59] B. Lakshminarayanan. A. Pritzel, and C. Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fer- gus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems 30, pages 6402-6413. Curran Associates, Inc., 2017.
[60] Y. LeCun, C. Cortes, and C. Burges. MNIST handwritten digit database. ATT Labs [Online]. Available: http://yann. lew.:/annneeeec..cmo wwwpaa/n.o///ee////iiitt..oo 2, 2010.
[61] J. Lee, J. Sohl-dickstein, J. Pennington, R. Novak, S. Schoenholz, and Y. Bahri. Deep Neural Networks as Gaussian Processes. In International Conference on Learning Representations, 2018.
[62] Y. Li and Y. Gal. Dropout Inference in Bayesian Neural Networks with D. Precup and Y. W. Tch, editors, Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages 2052-2061. PMLR, 2017.
In
[63] Q. Liu and D. Wang. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. In D. D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems 29, pages 2378-2386 Curran Associates, Inc., 2016.
[64] Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. ConvNet for the 2020s. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[65] J. J. C. MacKay. Bayesian Model Comparison and Backprop Nets. In J. E. Moody, J. Hanson, and R. Lippmann, editors, Advances in Neural Information Processing Systems 4, pages 839-846. Morgan Kaufmann, 1991.
[66] D. J. C. MacKay. A Practical Bayesian Framework for Backpropagation Networks. 4(3): 1992-05.
14

---

[67] J. C. Mackay. Bayesian methods for backpropagation networks. In E. Domany, J. L. van Hemmen, and K. Schulten, editors, Models of Neural Networks III, chapter 6, pages 21 1-1245.. Springer, 1994.
[68] D. J. C. MacKay. Choice of Basis for Laplace Approximation. 33(1):77-86, 1998.
[69] C. Maddison, A. Mnih, and Y. W. Teh. The Concrete Distribution: Continuous Relaxation of Discrete Random Variables, In Proceedings the 5th International Conference on Learning Representations. OpenReview.net, 2017.
[70] S. Mandt, M. D. Hoffmann and D. M. Blei. Stochastic Gradient Descent as Approximate Bayesian Inference. Journal of Machine Learning Research, 18(134):1-35, 2017.
[71] A. G. d. G. Matthews, Hensman, R. Turner, and Z. Ghahramani. On Sparse Variational Methods and the Kullback-Leible Divergence between Stochastic Processes. In A. Gretton and C. C. Robert, editors, Proceedings of the 19th International Conference on Artificial Intelligence and Statistics, volume 51 of Proceedings of Machine Learning Research, pages 231-239, Cadiz, Spain, 09-11 May 2016. PMLR.
[72] N. Metropolis and S. Ulam. The monte carlo method. 44(247):335-34) 1949. PMID: 18139350.
[73] R. M. Neal. Bayesian Learning for Neural Networks. PhD thesis, 1994.
[74] M. Neal. Priors for infinite networks. Technical report crg-tr-94-1, University of Toronto,
1994.
[75] R. M. Neal. MCMC Using Hamiltonian Dynamics, chapter 5.
CRC Press, 2011.
[76] B. Neyshabur, Li, S. Bhojanapalli, Y. LeCun, and N. Srebro. The role of over overovorraaaa parametrization in generalization of neural networks. In International Conference on Learning Representations, 2019.
[77] K. Osawa, S. Swaroop, M. E. E. Khan, A. Jain, R. Eschenhagen, R. E. Turner, and R. Yokota. Practical Deep Learning with Bayesian Principles. In Advances in Neural Information Pro- cessing Systems, volume 32, pages 4287 4299. Curran Associates, Inc., 2019.
[78] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. Detito, M. Raison, Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An Imperative Style, High- Pegh-eerffrrance Deep Learning Library. In H. Wallach, H. Larochelle. A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024-803 8035. Curran Associates, Inc., 2019.
[79] F. Pittorino, A. Ferraro, G. Perugini, C. Feinauer, C. Baldassi, and Zecchina. Deep Networks Toroids: Removing Symmetries Reveals the Structure of Flat Regions in the Landscape Geometry. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 17759-17781 PMLR, 17-23 Jul 2022.
[80] G. Pleiss and J. P. Cunningham. The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective. In A. Beygelzimer, Dauphin, Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[81] Y. Polyanskiy and Y. Wu. Strong Data- Pata-Procsssing Inequalities for Channels and Bayesian Networks. In E. Xarlen, M. Madiman, and F. M. Werner, editors, Convexity and Concentration, pages 211-249 New York, NY, 2017. Springer New York.
[82] A. Radford and K. Narasimhan. Improving Language Understanding by Generative Pre- Training. 2018.
[83] C. E. Rasmussen and C. K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press, 2005. ISBN 0262 18251X..3X.
[84] H. Ritter, A. Botev, and D. Barber. A Scalable Laplace Approximation for Neural Networks. In Proceedings of the 6th International Conference on Learning Representations. OpenReview.net, 2018.
[85] H. Robbins and S. Monro. A Stochastic Approximation Method. 22(3):400 1951.
15

---

[86] G. Roeder, Y. Wu, and D. Duvenaud. Sticking the Landing: Simple, Lower- Variance Gradient Estimators for Variational Inference. In I. Guyon, U. von Luxburg. S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages 6925-6934, 2017.
[87] S. Rossi, P. Michiardi, and M. Filippone. Good Initializations of Variational Bayes for Deep Models. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5487-5497. PMLR, 2019.
[88] S. Rossi, S. Marmin, and M. Filippone. Deep Learning. In Advances in Neural Information Processing Systems 33, 2020.
Variational Inference for Bayesian
[89] T. G. J. Rudner, Z. Chen, Y. W. Teh, and Y. Gal. Tractable Function-Spaat -pacce Variational Inference in Bayesian Neural Networks. A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[90] F. R. Ruiz, M. Titsias, and D. Blei. The Generalized Reparameterization Gradient. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[91] D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning Representations by Back- propagating Errors. (23(688)) 1986-10.
[92] T. Salimans and D. A. Knowles. Fixed-Form Variational Posterior Approximation through Stochastic Linear Regression. 8(4):837-885 2013.
[93] H. Salimbeni and M. Deisenroth. Doubly Stochastic Variational Inference for Deep Gaussian Processes. In Advances in Neural Information Processing Systems 30, pages 4588-4599,
Curran Associates, Inc., 2017.
[94] B. Simsek, F. Ged, A. Jacot, F. Spadaro, C. Hongler, W. Gerstner, and J. Brea. Geometry of the Loss Landscape in Overparameterized Neural Networks: Symmetries and Invariances. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9722-9732.
PMLR, 18-24 Jul 2021.
[95] E. Snelson and Z. Ghahramani. Sparse Gaussian Processes using Pseudo-Inputs. In Advances in Neural Information Processing Systems 18, pages 1257-1264. MIT Press, 2006.
[96] J. Sohl-Dickstein, Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[97] J. Song, C. Meng, and S. Ermon. Denoising Diffusion Implicit Models. In International Conference on Learning Representations, 2021.
[98] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations, 2021.
[99] S. Sun, G. Zhang, J. Shi, and R. B. Grosse. Functional Variational Bayesian Neural Net- works. In Proceedings of the 7th International Conference on Learning Representations. OpenReview.net, 2019.
[100] A. Takatsu. Wasserstein geometry of Gaussian measures. Osaka Journal of Mathematics, 48 (4):1005 - 1026, 2011.
[101] Tishby, Levin, and Solla. Consistent inference of probabilities in layered networks: predictions and generalizations. In International 1989 Joint Conference Neural Networks, pages 403-409 vol.2, 1989.
[102] B.-H. Tran, S. Rossi, Milios, Michiardi, E. V. Bonilla, and M. Filippone. Model Selection for Bayesian Autoencoders. In Advances in Neural Information Processing Systems, volume 34, pages 19730-19742. Curran Associates, Inc., 2021.
[103) B.-H. Tran, S. Rossi, D. Milios, and M. Filippone. All You Need is Good Functional Prior for Bayesian Deep Learning. In The Journal of Machine Learning Research, 2022.
[104] Villani. Optimal Transport: Old and New, volume 338. Springer Science Business Media, 2008
16

---

[105] M. Welling and Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In Proceedings of the 28th International Conference on International Conference Machine Learning, pages 681-688, Madison, WI, USA, 2011. 1. Omnipress. ISBN 9781450306192
[106] F. Wenzel, K. Roth, B. S. Veeling, J. Swigtkowski, L. Tran, S. Mandt, J. Snoek, 'T. Salimans, R. Jenatton, and S. Nowozin. How Good is the Bayes Posterior in Deep Neural Networks Really? Proceeding of the 37th International Conference on Machine Learning. 2020.
[107] J. G. Wiese, L. Wimmer, T. Papamarkou, B. Bischl, S. Günnemann, and D. Rügamer. Towards Efficient MCMC Sampling in Bayesian Neural Networks by Exploiting Symmetry, 2023,
[108] Xiao, K. Rasul, and R. Vollgraf. Fashion- -NNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms. CoRR, abs/ 71211708.77777 07747, 2017.
[109] G. Yehudai and O. Shamir. On the Power and L Limitations of Random Features for Understand- ing Neural Networks. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[110] C. Zeno, Golan, A. Pakman, and D. Soudry. Why Cold Posteriors? On the Suboptimal Generalization of Optimal Bayes Estimates. In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
[111] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. International Conference on Learning Representations, 2017.
[112] R. Zhang, C. Li, J. Zhang, C. Chen, and A. G. Wilson. Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning. In Proceedings of the 8th International Conference on Learning Representations. OpenReview.net, 2020.
17

---

A Future work and open questions
Besides the specific open questions discussed above, anticipate some possible future work. The framework we proposed to analyze the loss barriers in BNNS is general and can be applied to ap- proximations other than VI. Future work could undertake to extend this analysis the Laplace approximation [68]. However, this raises methodological challenges since Algorithm is not directly applicable, due to the use of dense covariance matrix (the inverse of the Hessian). Additionally, would be worth extending these results to sample- based inference, like stochastic gradient Hamilto- nian Monte Carlo (SGHMC) [19], or particle-based inference, like Stein variational inference [63]. This would call for careful analysis, starting from the solution of Eq. (6), which would require approximations [21].
B Additional details on the alignment method
In the main paper, to align the distributions with respect to permutation matrices we argue to use the Wasserstein distance rather than the KL divergence. Indeed, by considering the KL divergence KL [1#01 qol between Gaussians we have
KL. (+#91 go] = log det diag(so) (s)) log det diag (Ps1) ) + Tr(diag(Ps1 1))+
(20)
(mo Pm.) diag(s ( - Pm))
(21)
It's easy to verify that the first three terms do not depend on P, leading to just a distance between means and disregarding any covariance information. In the figure below, we visualize the difference between doing LAP with the KL cost and LAP with the Wasserstein cost.
Original distributions
LAP with KL cost
LAP with Wasserstein cost
-2
Figure A13: Alignment using different objectives. Given distributions symmetrical w.r.t. the y plane, using the KL cost LAP results the identity permutation (which fails to recover the symmetry), while the Wasserstein cost better aligns the two distributions
C Experimental setup
If not otherwise stated, we start by training two replicas of BNNS with variational inference and we compute the marginalized barrier as B(go, Q1) max, L(q-)- ((1 (()) + +L(11) where L(-) is the predictive likelihood and r E [0, 1]. from which we take 25 evenly distributed points. In particular, we seek to understand what happens the VI solutions with and without alignment applied to one of the two distribution (71 in our case). We experiment with MI.PS with three layers and ResNet20 [41] with various widths on MNIST [60], Fashion-
[108] and CIFARIO [56]. All models are trained without data augmentation, because the random augmentations introduce stochasticity which lacks a proper Bayesian interpretation in the formulation of the likelihood function [106]. Finally, we set the prior to be Gaussian (0, a T), with the flexibility of choosing the variance. All models are trained using the classic SGD optimizer with momentum [85, 91] using the reparameterization trick [54] with one sample during training and 128 samples during testing We use the categorical distribution and the Gaussian distribution as classification and regression likelihood, respectively. Tables 2 and 3 show details on the MI.Ps and convolutional neural networks (CNNs) base architectures used in our experimental campaign, while Table 1 reports the hyper- parameters used in the experiments. Note that differently from Entezari et al. [30] and Ainsworth
18

---

Table 1: Hyperparameters used for the experiments
Dataset
CIFARIO ResNet20
MNIST
Model
MLP
MLP
Data Aug.
False
False
False
Batch size
500
500
500
Temperature
1.0
1.0
1.0
Test samples
128
128
128
Train samples
1
1
1
std. init
0.01
0.01
0.01
Base features
16
512
512
Prior
0.01
0.0025
0.01
Learning rate
0.000001
0.000001
0.000001
Train epochs
1000
1000
1000
al. [2], we don't use data augmentation. A possible protocol for handling data augmentation in BNNs is presented by Osawa et al. [77] and involves carefully tuning the likelihood temperature to correctly counting the number of data points.
Table 3: ResNet20
Layer
Dimensions
Conv2D
x x Din
Table 2: MLP
x 3, 16
Layer
Dimensions
Residual Block
x 3, 16
x
Linear-ReLUU
512 x Din
3
Linear-ReL.U
512 x 512
Residual Block
3 3,32
x
Linear-ReLU
512 512
Linear Softuax
Dour 512
3 64
Residual Block
X 3
3, 64
AvgPool
8
Linear-Softman
Don X 64
C.1 Computing platform
The experiments have been performed using JAX [15] and two AWS p4d.2444.a 24xlarge instances with NVIDIA A100 GPUs. Experiments were conducted using in the eu-west- I region, which has carbon efficiency of 0.62 kgCO2eg/kWh. cumulative of 6500 hours of computation was performed on GPUs and it includes interactive sessions as well as small experiments with very low GPU usage, providing pessimistic estimation the true utilization. Total emissions are estimated be 1007.5 kgCO,eq of which 100 percents were directly offset by AWS.
D Additional results
We present timings obtained by profiling the time needed to solve the SOLAP with the Wasserstein cost, as well as time for the deterministic case [2]. In Fig. we show results for MLP and ResNer20 architectures, varying the model width. It is evident that, in the majority of cases, the algorithm completes within a minute. Morcover, as anticipated, in case of V1 solving our distribution alignment problem for wide neural networks is more computationally demanding compared to merely matching weights from SGD solutions.
Finally, we also test our setup on the CIFAR100 dataset [56]. Surprisingly, we were not able to replicate the same level performance as the other cases. In Fig. A15, we see that, despite converging well, we fall short to find zero-barrier solutions. Similarly to the comments of Ainsworth al. [2], we also stress that the failure to align distributions does not rule out the existence of a proper permutation map that the algorithm couldn't find. Nonetheless, this raises number of questions: the Bayesian posterior is the product of two ingredients, the prior and the likelihood, conditioned to observing dataset.
19

---

ResNe120
Time [poc]
10'
10-1
10
1024 2048
16
512
Model width
Model width
Figure A14: Timings. Profile of the algorithms align the solutions and to match weights from SGD solutions.
ResNet20 on CIFAR100
Predictive likellhood
g00o
VI (Train)
VI with distr. aligement (Train)
0.4
VI (Test)
Interpolation coefficient
VI with distr. alignment (Test)
Figure A15: Alignment failure. The method proposed fails to recover zero-barrier solutions for CIFAR 100.
Log posterior
Finally, as an additional check, we analyze the log- posterior 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 the form
Model go
Model y)
Model Pé4g
= 10
=0 E E 0 a0, + 10, + -1-a b)oc)
5.54
3.60
3.66
where a, b R and O.. O, and Og are the means of go. In and Q1 With this configuration, all three solutions lie on this hyper-plane. In Fig. 6, we present the visual- ization of ResNet20 trained CIFAR10. We see that the distributions do and Pa 91 are connected by higher density regions than the ones between yo and @ı. Also, as expected the symmetries arise from the form of the likelihood, and the prior has comparable strength with respect to the three posteriors. Later, we also study in more details the effect of the prior's variance in finding low-barrier solutions.
Log-likelihood
Log: prior
-175000 -1 1000000
-25000
3.680
3.732
3.79
>12"
Figure A16: Posterior density visualization. All three solutions are local approximation posterior, but qo and Ps Qu are connected by lower density regions.
E A primer on variational inference for Bayesian neural networks
is classic tool tackle intractable Bayesian inference [50, 13]. casts the inference problem into an optimization-basee procedure to compute tractable approximation of the true posterior. Assume generic parametric model f parameterized by some unknown parameters 0 (i.e. f(- 0) and collection of data y E N corresponding to some input points X = { [z; E }= In our setting, we have probabilistic model p(y f(x; O)) with parameters 0, a prior distributions on them p( and set of observations {X, y/ In a nutshell, the general recipe of VI consists of (i) introducing a set Q of distributions: (ii) defining a tractable objective that "measure" the distance between any arbitrary distribution y(0) E Q and the true posterior p(0 y); and finally (iii) providing a programmatic way to find the distribution 7() that minimizes such distance. In practice, ((0) has some free parameters v (also known variational parameters), which are optimized such that
20

---

the approximating distribution g is as closer as possible to the true posterior p(0 yy. We can derive the variational objective starting from the definition of the KL.
KL (g(0; || p y)] = [() [log (0; v) - log p(0 yy] =
(22)
= T qe [log ( v) - log p 0) - [(g) + logp(y)
Rearranging we have that
log) KL () I| p(0 y) = = [lg (0;;)) - log (0) log (())
(23)
The r.h.s. of the equation defines our variational objective, also known as ELBO, that can be arranged as follows.
LELBO = -Bqeew)) log p () - KI. (1(02)) p())
(24)
Model fitting
Regularization teem
This formulation highlights the property of this objective, which is made of two components: the first one is the expected log-liielihood under the approximate posterior q and measures how the model fits the data. The second term, the other hand, has the regularization effect of penalizing posteriors that are far from the prior as measured by the KL. Before diving into the challenges of optimization of the ELBO, we shall spend a brief moment discussing the form of the approximating distribution q- One of the simplest and easier choice is the mean field approximation [42], where each variable 0; is taken be independent with respect to the remaining 0... Effectively, this imposes factorization of the posterior,
K
g(0 v) = II (0:;))
(25)
1=1
where v; is the set of variational parameters for the parameter 0,. On top of this approximation, g(0)) is often chosen to be Gaussian,
((,)) = =N(,,
(26)
Now, the collection of all means and variances {ui, of) defines the set of variational parameters to optimize.
For BNNS the analytic evaluation of the ELBO (and its gradients) is always untractable due the non- linear nature of the expectation of the log-likelihood under the variational distribution. Nonetheless, this can be easily estimated via Monte Carlo integration [72], by sampling NMC times from qu.
NMC
1
Ey((; logp(y =
log ply 0,), with θ, ~ ((:
(27)
NMC
In practice, this is as simple as re-sampling the weights and the biases for all the layers NMC times and computing the output for each new sample.
We now have tractable objective that needs to be optimized with respect to the variational parameters v. Very often the KI. term is known, making its differentation trivial. On the other hand the expectation of the likelihood is not available, making the computation of its gradients more challenging. This problem can be solved using the so-called reparameterizatioo trick [92, 54). The reparameterization trick aims at constructing 0 as an invertible function T of the variational parameters v and of another random variable e, so that 0 T(E; v). Generally, T that suits this constraint might not exists; Ruiz et al. [90] discuss how to build "weakly" dependent transformation T for distributions like Gamma, Beta and Log-normal. For discrete distributions, instead, one could use a continuous relaxation, like the Concrete [69]. E is chosen such that its marginal p(s) does not depend on the variational parameters. With this parameterization, T separates the deterministic components of q from the stochastic ones, making the computation of its gradient straightforward. For a Gaussian distribution with mean μ and variance a², T corresponds to as simple scale-location transformation of an isotropic Gaussian noise,
0 (u, +
μ + as with 8 NO
(28)
This simple transformation ensures that p(E) = (0, 1) does not depends on the variational parame- ters [u,a²) o2). The gradients of the E.I.O can be therefore computed
V CELBO = Eple)
Vo logp(y 0) (= V [(E - KL (q( || p(0)) (29)
21

---

The gradient V e logp(y 0) depends on the model and it can be derived with automatic differentation tools [1, 78). while V (e v) doesn't have any stochastic components and therefore can known deterministically. Note that the reparameterization trick can be also used when the KT. is not analitically available. In that case, we would end up with,
V CELBO Ep(e) logp(y 0) log (0;; v)
VuT(E:
(30)
Roeder et al. [86] argue that when we believe that ((0;v) se p(y 0), Eq. (30) should be prefered over Eq. (29) even computing analitically the KL possible. Note that this case very unlikely for BNN posteriors, and that the additional randomness introduced by the Monte Carlo estimation of the KL could be harmful.
In case of large datasets and complex models, the formulation summarized in Eq. (29) can be computationally challenging, due the evaluation of the likelihood and its gradients NMC times. Assuming factorization the likelihood,
N ply 0) plylf(X:0) = II p f(zx;)))
(31)
==1
this quantity can be approximated using mini-batching [39, 43]. Recalling y as the set of labels of dataset with N examples, by taking B C y as random subset of y. the likelihood term can be estimated in an unbiased way as
N Σ log p(y 9). MA
(y
(32)
where M is the number of points in the minibatch. At the cost of increase "randomness", we can use Eq. (29) to compute the gradients of the ELBO with the minibatch formulation in Eq. (32). Stochastic optimization, e.g. any version of SGD, will converge to local optimum provided with a decreasing learning rate and sufficient gradient updates [85].
22


Reviews:
[{"review_id": "LhCkrD2sTX", "reviewer": "Reviewer_Kqgj", "summary": "This paper paper extends linear mode connectivity modulo permutation to Bayesian neural networks posteriors. Authors do this by imposing a Wasserstein metric on the space of distributions and look at how log likelihood changes along the geodesic between two distributions obtained via approximate Bayesian inference. For experiments, authors study BNNs parametrized by multivariate Gaussians with diagonal covariance. Analogous to weight matching algorithm proposed in [1], authors propose to find permutation that minimizes Wasserstein distance between the two distributions. For Gaussian distribution with diagonal covariance, authors propose an algorithm similar to weight matching that also takes into account the variances of the Gaussian. In experiments authors show that this heuristic for finding permutation that leads to small change in  log likelihood along the geodesic.\n\n*[1] Git Re-Basin: Merging Models modulo Permutation Symmetries. Ainsworth et al.*", "questions": "Empirically, how important is the role of adding covariance matrix in the objective when computing the permutations? For instance, using standard weight-matching in this setting just assumes that the covariance matrix is identity. it would be useful to see a table with comparison of different approaches. ", "limitations": "This paper limits itself on studying connectivity modulo permutation for approximate Bayesian NNs posteriors found via gradient based algorithms. It’s not clear if this phenomenon holds for other approaches which should be interesting future work. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Overall a well-written paper that extends previously studied linear mode connectivity modulo permutation of neural networks to Bayesian neural networks.\n- Experiments demonstrate that the heuristic proposed to compute permutations works for network architectures and datasets studied in the paper.", "weaknesses": "- Limitations of the proposed heuristic are not extensively discussed in the paper. For instance model architectures with batch norm typically fails without additional fixes [2].\n- It’s not clear if the proposed algorithm out performs standard weight-matching / activation matching / STE estimator based approach to finding permutation discussed in [1].\n\n*[2] REPAIR: REnormalizing Permuted Activations for Interpolation Repair. Jordan et al.*", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his/her interesting comments:\n\n* **Model architecture**: we agree that the choice of model architecture can impact the ability to find good solutions with low/zero barrier. One of this choice is the width of the neural network, for which the wider the model, the lower the barrier is. We are aware that the choice of the normalization strategy (especially batch-dependent) can affect the overall geometry of the problem (exactly how demonstrated in Fig.3 in your reference). Additionally, batch-dependent normalization layers don't have a clear Bayesian interpretation (technically, the likelihood in line 84 cannot factorize). That's why we chose to use the Filter-Response-Normalization layer, as done in other previous work [e.g. 40 in paper]. Having said that, we think it would be possible to apply the REPAIR method to our case as well, but I would expect the FRN to behave similarly to LN. For quick sanity check, we analyze the variance of activation following the instructions in [2], with the sole difference that the activations are marginalized w.r.t. samples from the posterior; for both naïve interpolation and aligned interpolation we then compute the variance ratio as discussed in Section 3.1 of [2] (we take $\\tau=0.5$ as middle point between the two “endpoints”). Results for a couple of models are available in the rebuttal PDF. We can see that there seems to be a level of collapse for the various configurations, which despite being less pathological than the one shown in [2, Figure 2] and improving using the distribution aligned method, appears to be an effect present in the Bayesian settings as well. Finally, we realized that REPAIR is not cited in the current version of the paper, we will make sure to add it in the next version. Thanks for the interesting point!\n\n* **Performance w.r.t. to [1]**: Figure 8 suggests that weight-matching (which was overall performing the best among the three in [1]), can achieve slightly better performance w.r.t. distribution alignment especially for wide networks (albeit with overall lower likelihood).\n\n* **Effect of (co)variance**: please refer to the table below, where we report the barrier with different alignment objectives:  the naïve interpolation (without any alignment) and two types of alignment objective (full, which reflects the correct objective as derived in Eq. (21) from the Wasserstein distance, and mean, which disregards the information regarding the variance). For the moment, we report this for two ResNet architectures and one MLP, all trained with CIFAR10. As we can see, we consistently get better results when we use the proper full objective. For the next version of the manuscript, we will make sure to complete the table.\n\n| Method              | Model      | Barrier      |\n| ------------------- |:---------- | ------------ |\n| Aligned (full)      | ResNet20x2 | **1.486** |\n| Aligned (mean only) | ResNet20x2 | 1.492     |\n| No alignment        | ResNet20x2 | 1.911     |\n| Aligned (full)      | ResNet20   | **1.789** |\n| Aligned (mean only) | ResNet20   | 1.993     |\n| No alignment        | ResNet20   | 2.067     |\n| Aligned (full)      | MLP        | **0.028** |\n| Aligned (mean only) | MLP        | **0.028** |\n| No alignment        | MLP        | 0.880     |\n"}, {"review_id": "osnIUsUHDq", "reviewer": "Reviewer_2rWU", "summary": "The authors study the geometry of SGD-trained Gaussian mean-field variational approximations to the posteriors of Bayesian neural networks (BNN). In large part, the authors propose extensions of the method and analysis of Ainsworth et al. [1] from MAP-estimated neural networks to BNNs. Notably, the authors informally conjecture that a permutation symmetry exists in the solutions that approximate Bayesian methods can find, similar to the permutation symmetry of MAP solutions SGD finds. To test this hypothesis, the authors first use optimal transport theory to extend the linear interpolation and the loss barrier framework of [1] to the variational posterior approximation setting. Then, they show that they can align individual mean-field Gaussian weight posteriors by approximately solving a bilinear assignment problem using an algorithm analogous to the one proposed by [1] for MAP solutions. The authors demonstrate empirically that zero-barrier interpolations exist for permutation-aligned variational posteriors of non-trivial architectures (ResNets) trained on non-trivial datasets (MNIST, FashionMNIST and CIFAR-10).\n\n## References\n\n[1] Ainsworth, S. K., Hayase, J., & Srinivasa, S. (2022). Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836.", "questions": " ", "limitations": " ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Understanding the properties of approximate solutions to the posteriors of BNNs is one of the central challenges of Bayesian deep learning; hence extending the method and analysis of Ainsworth et al. to the variational BNN setting is an important step towards this goal. The proposed extension of the interpolation framework and the alignment procedure using optimal transport theory seems natural, and the experimental methodology is sound and reasonably thorough. Besides the experiments analogous to the ones of Ainsworth et al., the authors also investigate the effect of prior variance and the cold posterior effect, the limit of which is the MAP solution. The paper is mostly well-written and easy to follow.", "weaknesses": "While I did not find any significant weaknesses in the work, there are a couple of points that, if improved/clarified, could significantly strengthen the paper. \n\n### Conjecture 1\n\nFirst, I found Conjecture 1 too broad and imprecise, even compared to the conjecture given by Entezari et al. [1], mainly for two reasons. The conjecture states that:\n\n\"Solutions of approximate Bayesian inference for neural networks are linearly connected after accounting for functionally transparent permutations.\"\n\nFirst, I think including the class of approximate Bayesian methods is too broad because it covers variational inference (VI), Markov chain Monte Carlo (MCMC), and the Laplace approximation, which all yield significantly different solutions. Hence, since the authors only study mean-field VI in the paper, I think it would be better to restrict the conjecture to this case only and potentially extend the conjecture to the other methods if future work supports it with some empirical evidence.\n\nRelated to the first point, it is thus unclear if there is a universal notion of \"linear connectedness\" for the approximate solutions I mention above. Furthermore, as both Entezari et al. [1] and Ainsworth et al. [2] pointed out, the modes' linear connectedness appears to be a particular feature of SGD. Hence, I would suggest that the authors rephrase their conjecture to something like\n\n\"SGD-based solutions of mean-field variational Bayesian inference for neural networks are linearly connected after accounting for functionally transparent permutations.\"\n\nThus, this is also a more direct generalization of the conjecture given in [1]. What do the authors think?\n\n### Connection between the interpolation and the alignment process\nThe authors define the interpolation of two variational posteriors as well as the alignment procedure using optimal transport (OT) theory. However, as far as I can tell, these two things are not obviously connected. While I am not claiming that the authors' choice to use OT for these definitions is unnatural, it currently seems more of a choice of convenience than one dictated by theory. In particular, it seems that the authors chose this precisely because the alignment procedure reduces to a bilinear assignment problem, and the interpolating density is also Gaussian. Could the authors clarify whether there is some deeper theory that would justify the authors' choices? Does the alignment objective in eq (20) follow somehow from the interpolation method in eq (6)?\n\nRelatedly, the authors state the following on L116: \n\n\"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.\" \n\nWhat the authors mean by the word \"trivial\" here is unclear. They appear to mean that interpolating with arithmetic mixtures is an \"obvious\" choice, not \"trivial\". The properties of the mixture choice are not at all obvious to me, and it is not clear why it doesn't provide us with information about the loss landscape. Again, I am not saying that the authors' choice of using OT is wrong and the arithmetic mixtures are better or equally useful; but they need to give arguments (either theoretical or empirical) why they think it is uninteresting. \n\nIn a similar vein, I wonder if using geometric mixtures, i.e. using $q_\\tau \\propto q_0^{1 - \\beta} \\cdot q_1^{\\beta}$ for $\\beta \\in [0, 1]$, would have similar behaviour, or if it would lead to a different alignment procedure.\n\n### Figure 5\nGiven the interpolation and alignment process, is the hyperplane in Figure 5 meaningful? The hyperplane is defined using a simple linear interpolation of the weight posteriors' mean parameters, which seems to go against all the previous machinery the authors argued for earlier in the paper. Perhaps if the authors want to include such a hyperplane, perhaps there's a definition that can be made using OT, or they could draw samples $ \\theta_a \\sim q_0, \\theta_b \\sim q_1, \\theta_c \\sim P_\\sharp q_1$ and linearly interpolate those?\n\n### Miscellaneous\nEq (21) has a small mistake, the argmax ranges over $i \\in [0:L]$, but the objective only involves permutations with indices up to $L - 1$. Furthermore, Eq (21) could be compactified by defining P_0 = P_L = I and writing the sum using $\\sum$ notation.\nLine 4 in Algorithm 1 should be broken up into two or three lines, and the font size should be increased.\nThe font size in Figures 3-8 is too small and should be increased to match at least the font size of the captions.\nThe training procedure needs clarification. Did the authors use Bayes by Backprop [3] or the local reparameterization trick [4]?\n\n## References\n\n[1] Entezari, R., Sedghi, H., Saukh, O., & Neyshabur, B. (2021). The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprint arXiv:2110.06296.\n\n[2] Ainsworth, S. K., Hayase, J., & Srinivasa, S. (2022). Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836.\n\n[3] Blundell, C., Cornebise, J., Kavukcuoglu, K., & Wierstra, D. (2015, June). Weight uncertainty in neural network. In International conference on machine learning (pp. 1613-1622). PMLR.\n\n[4] Kingma, D. P., Salimans, T., & Welling, M. (2015). Variational dropout and the local reparameterization trick. Advances in neural information processing systems, 28.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his/her interesting comments. Below we reply inline to the reviewer's questions:\n\n* **Strength of conjecture 1**: We agree with you on this. We decided to have a broad conjecture to leave room for possible extensions to Laplace approximation (easier) and SG-MCMC methods (definitely more challenging). Since this point was raise by Reviewer T6Qt as well, we think it's appropriate to rephrase it to limit the context under analysis. Thanks for the suggestion.\n* **Connection between interpolation and alignment**: That is correct: technically Eq (20) and Eq (6) are independent and indeed other choices can be made (interpolation using the mixtures and alignment the KL divergence). For the interpolation, the choice of the Wasserstein geodesic was dictated by the geometry of the interpolation paths between the two distributions (akin to linear interpolation for the Euclidean metric, more on this in the next point). Additionally, the use of Wasserstein distance for the alignment conveniently allows us to re-interpret the problem as an assignment problem, for which efficient routines are available. In the appendix, we briefly show why we think the KL distribution is not the best objective we can use to align distributions. Indeed, other choices can be made (see new experiment run for Reviewer Kqgj).\n* **\"Mixtures are trivial\"**: we realized that \"trivial\" was a poor choice of word in this context. What we meant by that is that the mixture of distribution is not sufficient to capture the underlying complex geometry of the posterior. To aid this conversation, please have a look at the Figure 1 in the PDF. Here, we plot the test likelihood with two interpolation strategies: OT vs mixture (both without any alignment). 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. \nWe will make sure to reword this passage more appropriately.\n\n* **Figure 5**: Thanks for the suggestion. Indeed, this might not be the best way to visualize the underlying structure. Nonetheless, we believe that the visualization has still value, as it shows the different connectivity property between the aligned and the not-aligned solutions. Having said that, we took the Reviewer's opinion into account and we prepared another visualization (see Figure 2 and 3 in the rebuttal PDF). In this case we sampled from the two solutions, as well as from the interpolated one (for $\\tau=0.5$), with and without alignment. We then use these three sample as support to build the projecting hyperplane. For sake of visualization, these three samples once projected are not on a line, but in the original space they are. With alignment, we see that the three samples all lie in a relatively flat region on the posterior, in contrast to the case of no alignment where we once more see the raising of a barrier when connecting the three samples.\n\n* **Miscellaneous**: Thanks for spotting the mistake in Eq (21), you are right. Font size and figure placement were known issues, unfortunately we didn't find other arrangements that would fit in 9 pages. Eventually, if allowed to add an additional page for the camera-ready, we will make sure to fix this. For the training procedure, we refer back to the details in the Appendix (we are indeed using the reparameterization trick, not the local)"}, {"review_id": "7pggcR4qt5", "reviewer": "Reviewer_Kavz", "summary": "This work permutes together the distributions of bayesian neural network (BNN) parameters in the context of variational inference (VI) so that they are linearly connected. This is done by adapting recent work on permuting together SGD solutions to be linearly connected, by optimizing for similarity of means and variances from VI in the permutation assignment objective, instead of similarity of parameters (as in Ainsworth et al. 2022). The problem setup and permutation objective are described and derived in detail. Experiments show that this method is comparable in loss barrier to directly permuted SGD solutions.", "questions": "Given that barrier is higher when aligning VI distributions versus aligning the MAP parameters (akin to directly aligning SGD solutions), what are the potential benefits of optimizing alignment from a distributional standpoint? If the goal is to minimize barrier, why not simply align the MAP? Some motivation would be helpful here.", "limitations": "The main limitation (which is briefly discussed in section 7, but not strongly emphasized throughout) is that the assumptions on the posterior distribution are too rigid: namely that parameters are independent Gaussians. The authors conjecture (line 239) that this may be the cause of reduced performance relative to Ainsworth et al. (2023). Given the derivation in (16) already admits arbitrary covariances, it seems very interesting to consider the case where the covariance of the posterior is non-diagonal (lines 275-276). This could also lead to significant innovations over the existing alignment algorithm of Ainsworth et al. (2023), which is largely unchanged in this work.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Bringing permutation alignment into a probabilistic setting is novel, and the formalization of the methodology is satisfying. The paper is well-organized, and a number of variations (variance of prior, temperature) are considered in the experiments. Despite its focus on BNNs, this work also has relevant implications for the SGD solution setting, since it would be nice to be able to extend linear mode connectivity from individual solutions to families of SGD solutions whose parameters are stochastic (e.g. due to randomness in initialization/training).", "weaknesses": "In both the temperature and prior variance experiments, a comparison of barriers at the limit of 0 variance/temperature (corresponding to direct alignment of the MAP solution) for CIFAR-10 would be interesting, as it would put the observed difference between VI alignment and SGD alignment in figure 8 into context. Specifically, it would be nice to see figures 6 and 7 replicated on CIFAR-scale networks, and a reference line  included in each figure to indicate the barrier achieved by SGD alignment.\n\nFrom a presentation standpoint, figures 4-8 are arranged strangely (5 is out of order) and too small to read clearly. In particular, the text and lines should be larger, some of the margins smaller, and more distinct colors used. Figures 1-3 are less important from a readability standpoint, but may also benefit from the same adjustments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for insightful comments and discussion points.\n\n* **Temperature and prior for ResNet20/CIFAR**: Thanks for the suggestion, we will take this into consideration. During the limited time span of this rebuttal, we were able to run the ResNet comparison with different prior variances and different temperatures (see Figure 7 in the rebuttal PDF). Modulo the fact that finding zero-barrier is more difficult on ResNet (vs MLPs), the comments for prior variances in the paper generally applies for this experiment as well: \"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\". Similar for the temperature comparison, for which \"we see that barriers for cold posteriors with alignment are marginally closer to zero than for warm posteriors\". To conclude, in both cases the results are consistent with the comments in the paper, but we will add these additional experiments to strengthen the discussion. Thanks for the suggestion!\n* **Presentation**: Thanks, this is already taken into consideration for the next version of the manuscript (see comment to Reviewer 2rWU). We will make sure to improve the readability of the plots.\n* **Aligning using MAP**: Thanks for the question. Indeed, we could align the MAP solutions and take the permutation symmetry found. Still, we don't think this will work in practice for few reasons: (i) independent models do not share the same permutation matrices, (ii) if we look only at the means, the new results requested by Reviewer Kqgj show that this is indeed not the best objective for alignment.\n* **Limitation on the posterior**: Yes, indeed extensions to non-diagonal covariances are possible. We haven't investigated this possibility yet for a couple of reasons: (i) full covariances with variational inference are intractable for these large models, resorting in various possible approximations to computationally scale to these networks (low-rank, structured covariance, etc.); (ii) with non-diagonal covariances the steps to go from Eq. (16) to the formulation in Eq. (19) would not be possible. A way to approach this could be by learning the permutation $P$ using a new ELBO together with a straight-through estimator:\n\n$$\n\\max_{\\widetilde{q_1}} \\mathcal L_{elbo}\\left(\\frac 1 2 (Id+T_{q_0}^{\\widetilde{P}\\\\#\\widetilde q_1})\\\\# q_0 \\right) \\quad \\text{where} \\quad \\widetilde{P} = \\arg\\min_P \\mathcal W(P_\\\\# q_1, q_0)\n$$\n\nTo make it tractable, the r.h.s. needs to approximate the full covariance with its diagonal components, while the l.h.s. should be computed with its full covariance. We believe that such discussion requires some careful analysis, which was beyond the scope of this work."}, {"review_id": "WX23t6jS52", "reviewer": "Reviewer_bvHV", "summary": "The authors extend the recent idea of linear mode connectivity up to permutation symmetry to the setting of Bayesian neural networks. They demonstrate that two different variational approximations to the Bayes posterior enjoy mode connectivity along the Wasserstein geodesic of one distribution, and a suitably permuted version of the other. Such a permutation is discovered by replacing the L2 distance in previous work with the more suitable Wasserstein distance, and similarly,  the objective can then be relaxed to a layer-wise linear assignment problem, leading to a tractable coordinate descent algorithm. The authors verify their results numerically, showing that two approximate solutions can indeed be connected for modern networks such as ResNet20 on CIFAR10.", "questions": "1. Is data augmentation employed for the experiments involving the cold posterior effect? Data augmentation has been observed to be the main driver of the CPE [1, 2, 3] and hence using it would probably lead to a more visible effect. It has also been recently justified that tempering is a principled way to use data augmentation in Bayesian frameworks [4], so it would not affect the validity of the approximation. \n2. Are there gains in terms of test accuracy of the (tempered) approximate Bayesian posteriors versus a standard SGD baseline? \n\n\n[1] What are Bayesian Neural Network Posteriors Really Like?,\nIzmailov et al.\n\n[2] Data augmentation in Bayesian neural networks and the cold posterior effect,\nNabarro et al.\n\n[3] Disentangling the Roles of Curation, Data-Augmentation and the Prior in the Cold Posterior Effect,\nNoci et al.\n\n[4] How Tempering Fixes Data Augmentation in Bayesian Neural Networks,\nBachmann et al.", "limitations": "The authors have addressed the limitations of their work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Studying mode connectivity for approximate Bayesian inference is a natural follow-up question to previous work, while at the same time requiring non-trivial extensions such as the Wasserstein geodesic and Wasserstein distance. It is very nice and somewhat surprising that the resulting objective (which is seemingly involved) can be relaxed in a very similar spirit, leading to a tractable problem. \n2. The experimental setup is quite carefully created and a lot of ablations for different \tparameters such as the prior variance, the temperature and the width of the network are performed, giving a very complete picture.", "weaknesses": "1. The proposed algorithm and the setup seem to heavily rely on the specific approximation method, namely variational inference with a Gaussian distribution of diagonal covariance as the variational family. Even if the diagonal covariance assumption is relaxed, it is not obvious to me how one can guarantee tractability, let alone moving to multi-modal approximations such as MCMC where not even the Wasserstein geodesic or distance are known in closed form. This is somewhat unsatisfying, since the power of BNNs and the Bayesian posterior in general only really starts to unfold once multiple modes are leveraged. It is thus somewhat unclear how much a unimodal approximation to the Bayes posterior really captures and how much it is really different from a simple point estimate.\n2. In general, the problem seems to get less and less interesting, the more precise the approximation to the Bayes posterior becomes. This is simply because the Bayes posterior would incorporate all possible modes (given the prior gives them some mass) and hence there is no other posterior to align with. This is different from the SGD setup, where multiple modes always remain a problem, precisely due to the point-wise nature of the algorithm. The authors are very careful in the main text and always refer to approximate Bayesian inference, but I think it would be helpful to clarify this discrepancy.\n3. I understand why the authors use the log-likelihood as a metric to evaluate connectivity, given that it is a proper scoring rule, but it is also very difficult to interpret how meaningful a decrease in likelihood is for practice. This is in contrast to test accuracy, where we have a better understanding of the scale. It would be helpful if the authors could provide the same plots for test accuracy instead of log-likelihood. That way it would also be easier to assess how meaningful the barrier in Fig. 3 for ResNet20 actually is. It would also be helpful to reproduce the same plots as in Fig. 3 without adding the non-aligned connectivity score since this massively increases the scale. That way it is actually difficult to tell how connected the aligned solutions are. They are obviously way more connected than the baseline, which is nice but it would be better to see it in more detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback and insightful comments.\n\n* **Tractability of Approximation Methods**: The reviewer raises a valid concern regarding the reliance of our proposed algorithm on specific approximation methods, such as variational inference with a Gaussian distribution of diagonal covariance. We agree that the power of Bayesian neural networks lies in capturing multiple modes, and unimodal approximations may limit their expressiveness. While our method is tractable for the chosen approximation, we acknowledge that extending it to more complex multi-modal approximations like MCMC could be challenging but not intractable (e.g. the Wasserstein geodesics can be approximated with the Sinkhorn algorithm [17 in paper]).\n\n* **Precision of Approximations and Interpretability**: The reviewer correctly notes that as the approximation to the Bayes posterior becomes more precise, the problem of  connectivity becomes less interesting, as the Bayes posterior would naturally incorporate all possible modes. While we agree with the reviewer, we also need to acknowledge that despite the best efforts from the community, the BNN posterior is still elusive: even carefully tuned MCMC methods cannot fully characterize the full extent of its nature. And they still needs orders of magnitude more of gradient evaluations. We believe that this characterization of permutation symmetries will have value for SG-MCMC methods as well.\n\n* **Comparison with SGD**: Regarding the comparison with SGD, it's still up to debate whether SGD solutions are actual global minimizers (the same Ainsworth's paper shows that we can find better solutions than the ones found by pure SGD, supporting the hypothesis that SGD solutions are not global). In this regard, we can view SGD as an approximation method of the exact minimization problem $\\min_\\theta\\ell(\\theta)$, for some loss function $\\ell(\\cdot)$, in the same way in which we can view approximate Bayesian inference (including sampling) as an optimization problem over space of probability measures w.r.t. the true posterior (although it's clear that the nature of the symmetries are not equivalent for the two setting).\n\n* **Evaluation Metric**: We appreciate the reviewer's suggestion regarding using test accuracy as an additional metric for evaluating mode connectivity. Test accuracy is indeed more interpretable in practice, and we agree that including it in the evaluation would provide valuable insights. In response to this suggestion, we will incorporate plots of test accuracy alongside log-likelihood in the revised manuscript, offering a more comprehensive assessment of the model's performance and connectivity. In the meanwhile, please check in the rebuttal PDF the replication of Figure 3 with the accuracy, along side with the requested zoom.\n\n* **Data Augmentation and Cold Posterior Effect**: We appreciate the reviewer pointing out the potential impact of data augmentation on the cold posterior effect (CPE). No, data augmentation is never employed in our experiments (see general comments). Having said that, during this rebuttal we were able to run a comparison on the effect of the data augmentation alone (without tempered posteriors). In Figure 8 of the rebuttal PDF, we can see that the behavior for both with and without DA is very similar. We will conduct additional experiments with data augmentation and temperature scaling to provide a more comprehensive analysis of its influence on our results. Due to the limited time we were unable to have the comparison DA+temperature, but it will be added for the next version.\n\n* **Test Accuracy Gains**: While analyzing the effect of tempered posterior is not our primary scope in our work, we report the results in the table below for ResNet20.\n\n| Method      | Accuracy |\n| ----------- | -------- |\n| Tempered VI | 0.8281   |\n| VI          | 0.8580   |\n| SGD         | 0.8546   |\n\nBecause we are using a clean likelihood (no batch statistics, and no DA) the results with VI are slightly worse but in line with the literature (see Appendix K4 in [92 from paper])"}, {"review_id": "4CTFuvep72", "reviewer": "Reviewer_T6Qt", "summary": "The authors conjecture that after accounting for permutation symmetries in overparametrized neural networks that lead to same functional behavior, the low-loss solutions are linearly connected. In the context of Bayesian neural networks with variational inference, the authors use this conjecture to propose a simple matching algorithm based on Linear Assignment Problem, to find an equivalent variational distribution but attains linear connectivity with another solution. The experiments demonstrate that it is possible to find such linearly connected no-barrier regions.", "questions": "1. The posterior predictive likelihood computed on test points in Eq. (5) is also used to define the functional loss barrier in Definition 3. Don't we want the loss barrier to be defined on training points? \n\t- Is this simply a loose usage of the definition, or did the authors specifically imply using test points to construct barrier loss? It looks like from Figure 3 that all kinds of points are used.\n\t- I think it is completely fine to check using the test points as a diagnostic for correlation between behavior at test and train time.\n\t- But then, the matching doesn't seem to impact train landscape at all. Can the authors comment on this, since if I understand this correctly, this only remains a diagnostic method and not a method to actually generate samples from the posterior.\n2. Did the authors try interpolating with a convex mixture of two distributions? There's no need to report results for this, but generally curious if such an interpolation provided some reasonable results since the authors only claim that this choice is trivial and non-informative.\n3. What would using a KL-divergence instead of Eq. (16) do in practice? Would the distances be always too large to meaningfully distinguish for the LAP problem?  \n4. Have the authors considered accounting for such functional symmetries during the training itself, instead of a post-hoc matching procedure?", "limitations": "Yes. Also see, weaknesses and questions.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The authors take the ideas of mode connectivity in the deep learning literature, and try to find an equivalent formulation at the level of distributions over parameters as in Bayesian neural networks. It is a very interesting concept. \n- The concept becomes even more effective by a successful demonstration of an algorithm that finds such a low-loss barrier. However, I do want to say that the existence is not surprising given that the deep learning literature already demonstrates the existence and the ELBO is just optimizing for \"another\" marginalized likelihood instead of an un-marginalized likelihood typically used in DL.\n", "weaknesses": "- I believe that the title of the paper should qualify variational BNNs instead of just BNNs. Approximate VI remains practically distinctive in properties from other approximate Bayesian inference methods like MCMC/HMC. It also helps contextualize the scope for the reader, and something that the authors readily align with in the discussion as well.\n- Definition 3 about barrier loss seems to be connected to $\\mathcal{L}$ instead of $\\mathcal{L}_{\\mathrm{ELBO}}$. We are interested in the loss computed by the marginalization of the model parameters under the approximate posterior. The authors choose this to be computed on the test data, which seems like jumping ahead to information that should not be used by the modeler. See also Question 1. Figure 3 shows such an algorithm does not benefit train landscape at all. Or did I misunderstand?\n- In Line 159, the authors claim to argue that Wasserstein distance is better. Is the argument that covariance information is lost? I think the readers would appreciate a tiny note on using KL-divergence in the appendix and some results showing failure or becoming non-informative. See also Question 3.\n- A key missing feature in the experiments is that the identification of low barrier region is not followed by the construction of posterior predictive distribution to compute the generalization error. Only likelihoods are reported, and I am wondering if the benefits of Bayesian model averaging shine as well as or better than reported in earlier works. Can the authors report those numbers? At least a basic SGD, VI, and the Aligned VI proposed in this work.\n- The design choice of skipping data augmentation severely restricts the applicability of these results, especially with vision problems where data augmentation is still commonly used. I do, however, empathize with the authors on this one and do not consider this to be a big limitation to discount the contributions, since the rest of the community suffers with this limitation too.\n\n### Minor\n- Please use `\\citet` instead of `\\citep` when referring directly to papers. For instance, in Line 134.\n- A very small description of the original LAP problem (introduced in Linear 167) in the appendix would be much appreciated.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the detailed review and valuable feedback provided by the reviewer. We have carefully considered each point raised and addressed them below.\n\n* **Scope of the claim**: The reviewer suggests that the title should qualify \"variational BNNs\" instead of \"BNNs\" to better contextualize the scope of the paper. Given that this point was raised by Reviewer 2rWU as well, we will be changing  Conjecture 1 to better reflect the focus on the variational approximation.\n\n* **Definition of Barrier Loss and Use of Test Points**: The reviewer correctly points out that Definition 3, which defines the barrier, is connected to the posterior predictive likelihood computed on test points, which may seem to incorporate information not available at train time. We acknowledge this concern and apologize for any confusion. In practice, the barrier loss can be computed on both training and testing points. Indeed, all plots report both train and test barriers. We will clarify this in the revised manuscript (line 102 will be changed from \"$\\{x_\\star, y_\\star\\}$ are respectively the test point and its corresponding label\" to  \"$\\{x_\\star, y_\\star\\}$ are respectively the input point under evaluation and its corresponding label\". We also agree with the reviewer's observation that the matching procedure may not significantly impact the train landscape. The primary aim of our method is to establish linear connectivity between approximate Bayesian solutions rather than improving the train landscape.\n\n* **Use of Wasserstein Distance**: The reviewer requests further clarification on why we argue that the Wasserstein distance is better than using KL-divergence. For this we refer the reviewer to the Appendix, where we show that the KL-divergence simply reduces to a distance between means, which disregards any information regarding covariances. In the appendix, we also show a simple motivating example of such failure.\n\n* **Data augmentation**: Skipping data augmentation in our training setup was purely a choice of convenience. We know that with data augmentation we need to be careful once moving to Bayesian inference for a multiple of reasons (e.g., cold posterior effect [92], re-weighting of the likelihood due to the increase of the effective sample size [68]). For this reason, we didn't want to \"pollute\" the results with spurious effects coming from DA, rather than the phenomenon under analysis. Nonetheless, we want also to emphasize that nowhere during the development of the method we make an assumption on DA, which makes our method still applicable in both cases. During this rebuttal week, we were able to run some comparisons with data augmentation for the ResNet20 models. Figure 8 in the rebuttal PDF summarizes this experiment: we see that in both cases (with and without DA), we are still able to recover similar low barrier solutions in terms of likelihood (and accuracy---not reported for space reasons) when following our proposal to align the distributions. As discussed in the general comments, we are planning tp add an additional paragraph in Sec. 5 to better comment on the effect of DA.\n\n* **Benefits of BMA**: In terms of generalization on accuracy, indeed BMA is  beneficial w.r.t. point estimates. See the table below, where we report the accuracy of the two models, as well as the interpolated one $\\tau=0.5$.\n\n**ResNet/CIFAR10 (Accuracy)**\n\n|                | Model 0 | Model interpolated | Model 1 |\n| ---------------|---------|--------------------|---------|\n|             VI | 0.8580  |       0.1025       | 0.8556  |\n| **VI aligned** | 0.8580  |       0.7413       | 0.8556  |\n|            SGD | 0.8558  |       0.1432       | 0.8516  |\n\n**MLP/CIFAR10 (Accuracy)**\n\n|                | Model 0 | Model interpolated | Model 1 |\n| ---------------|---------|--------------------|---------|\n|             VI | 0.5718  |       0.2589       | 0.5734  |\n| **VI aligned** | 0.5718  |       0.5647       | 0.5734  |\n|            SGD | 0.5546  |       0.2500       | 0.5545  |\n\nThese numbers are in line with previous work in the literature.\nWe also report the reference comparison with the likelihood below\n\n**ResNet/CIFAR10 (Likelihood)**\n\n|                |   Model 0 | Model interpolated |  Model  1 |\n| ---------------| ----------|--------------------| ----------|\n|             VI | -0.417142 |      -2.32982      | -0.427943 |\n| **VI aligned** | -0.417142 |      -0.71424      | -0.427599 |\n|            SGD | -0.702219 |      -2.43758      | -0.731933 |\n\n* **Mixture of distributions**: as we discussed for Reviewer 2rWU, we argue that a mixture is not informative for studying the geometry of the posterior. To visualize this argument, please check Figure 1 in the rebuttal PDF. With mixtures we are essentially re-weighting the two (good) solutions, without \"transport\" of distribution mass between the two extremes. As we said in the general comment, this means that if we look at the barrier, we don't see any not because they don't exist, but simply because we are interpolating in such a way that prevents us from exploring the geometry of the posterior.\n\n* **KL divergence**: Thanks for the question. In the appendix we actually show how using the KL divergence in practice falls back to a distance between means (without accounting for the actual shape of the distributions). We also have a simple visualization where this fails in practice, i.e. the LAP with the KL objective fails to recover a clear symmetry.\n\n* **Accounting for Functional Symmetries during Training**: While the idea of accounting for functional symmetries during training is intriguing, it presents a challenging optimization problem. We have not explored this approach in the current paper, but it could be an interesting direction for future research. We will mention this possibility in the discussion section to encourage further investigations."}, {"review_id": "31ykZtfKmW", "reviewer": "Reviewer_hffp", "summary": "This paper considers mode-connectivity in the context of Bayesian neural networks.  It shows roughly what we'd expect, based on the results of Entezari et al. (2022).", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "It shows roughly what we'd expect, based on the results of Entezari et al. (2022).", "weaknesses": "That's my main issue with the paper.  Bayesian neural networks really aren't that different from neural networks.  So we would definitely expect the results of Entenzari et al. (2022) to apply in the BNN context.  I can't see any interesting contributions on top of that, so I can't recommend acceptance.  I would recommend submission to a more specialised venue (e.g. AABI or UAI).\n\nOther points:\n* An important reference for permutations: Aitchison, Laurence, Adam Yang, and Sebastian W. Ober. \"Deep kernel processes.\" International Conference on Machine Learning. PMLR, 2021.\n* Legend for Fig. 8 is _way_ too small.  In general the Figures feel a bit crammed in. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's feedback, but we respectfully disagree with their assessment of our paper.\nBelow is our best attempt at addressing the reviewers' concerns:\n\n* **The paper shows expected results**: The reviewer expresses their main issue with the paper, stating that Bayesian neural networks are not significantly different from neural networks, and therefore, the results are expected based on Entezari et al. (2022). However, while Bayesian neural networks share some similarities with neural networks, their probabilistic nature and approximate inference techniques introduce complexities not present in standard neural networks. As such, zero-barrier connectivity in the context of Bayesian neural networks is a non-trivial extension and a novel contribution to the field. Our work demonstrates that connectivity can indeed be achieved with approximate Bayesian neural networks, offering insights into their behavior and potential applications in uncertainty quantification.\n\n* **Lack of interesting contributions**: This point is very much linked to the previous one. The reviewer states that they cannot see any interesting contributions beyond the prior work. We strongly disagree with this assessment. Our paper introduces the notion of linear connectivity in Bayesian neural networks, providing a formalism as well a novel perspective on understanding their posterior landscapes. Additionally, we propose a tractable algorithm based on the Linear Assignment Problem to connect approximate solutions efficiently. These contributions offer valuable insights and open up new avenues for exploration in Bayesian neural networks.\n\n* **Recommendation for Specialized Venues**: The reviewer suggests that our work would be more suitable for a specialized venue, such as AABI or UAI. While we appreciate the suggestion, we believe that our work addresses an important question in the field of deep learning and Bayesian neural networks, making it relevant to a broader audience. Mode connectivity is a topic of significant interest, and demonstrating its applicability in the context of BNNs, along with the proposed algorithm, contributes to the understanding of optimization landscapes and uncertainty modeling in deep learning. We respectfully maintain that our paper is well-suited for consideration in the current venue.\n\nFor the remaining points:\n\n* **Additional Reference**: We appreciate the reference provided by the reviewer on permutations. We will carefully consider it and include it in our revised manuscript.\n\n* **Figures and Presentation**: See the general comment. In the revised version, we will ensure that the figures are appropriately sized and presented for better clarity and readability.\n"}]

Author rebuttal:
# General comments

First, we would like to thank the Reviewers for their comments and helpful feedback.
With this paper we analyze the geometry of the Bayesian posterior in deep neural networks by considering the permutation symmetries raising from the neural network parameterization. We do it by extending previous analysis on loss-optimized models to the Bayesian setting: this requires, among other things, to extend the concept of barrier and solutions interpolation.
In the variational inference setting, we move then to propose a methodology to align distributions, by framing this problem as a combinatorial optimization one. With this setup, we can find distributions that once interpolated exhibit low/zero barriers even for ResNet-scale models.

We are delighted to see that Reviewers agree on the novelty and quality of presentation (bvHV, Kavz, Kqgj, T6Qt, 2rWU), that our proposal is principled (2rWU) with rigorous experiments (bvHV, 2rWU) and interesting results (2rWU, Kqgj).

We use this space to address some common points:

* **Figures presentation**: several Reviewers (Kavz, 2rWU, hffp) commented on the readability of some of the plots in the submission. Indeed we had to squeeze the figures more than what we would have liked as to remain in the 9 pages, while presenting all the important results. We plan to remake the figures to address this point, possibly by converting the plots to LaTeX using tikzplotlib. Additionally, if allowed one additional page for the camera ready, we will be less constrained in the positioning of the figures, which should help.

* **Scope of the work and extensions beyond variational inference**: Reviewers 2rWU and T6Qt highlights how our initial conjecture is too broad, as it can cover all possible approximations (VI, Laplace, MCMC) while our work primarily focuses on VI. The idea for such a conjecture is to encourage further investigations with other setups and approximations. Nonetheless, we generally agree with this opinion and we will re-frame the Conjecture to better align with our context.

* **Data augmentation (DA)**: Reviewers bvHV and T6Qt expressed a limitation in our choice of not using data augmentation for our experimental setup. This choice was merely a consequence of the difficult interpretation of DA in the Bayesian context (as many previous works in the literature have shown, [e.g. from the paper 92,68,40]). Indeed, without DA (and without BatchNorm for normalization layers), we are provably targeting a correct Bayesian posterior. Having said that, for a sanity check we have run a ResNet20 on CIFAR10 with and without data augmentation, showing similar behaviour in both cases (Figure 8 in the rebuttal PDF).

* **Mixtures vs OT interpolation**: Reviewers 2rWU and T6Qt requested additional clarification on the choice of using optimal transport for the interpolation. While we address this point individually for both Reviewers, we want to give a generic intuition for this choice. We are interested in studying the properties of the approximate Bayesian posterior while we align the solutions w.r.t. permutations symmetries. With mixtures we essentially re-weighting the two solutions, without continuous "transport" of distribution mass between the two extremes. This means that if we look at the barrier, we don't see any not because they don't exist, but simply because we are interpolating in such a way that prevents us from exploring the geometry of the posterior. Said differently, mixtures do not lead us to the construction of low barrier solutions from permutation symmetries.

Below, a summary for the additional experiments run during the rebuttal and plots that you will find in the rebuttal PDF:

* **Figure 1**: Visualization of mixtures vs OT interpolation
* **Figures 2/3**: Two new visualizations of the posterior for aligned and not aligned solutions
* **Figure 4**: Analysis of features variance collapse in interpolated models
* **Figure 5**: Replication of Figure 3 from the paper, but using accuracy rather than likelihood
* **Figure 6**: Zoom of Figure 3 from the paper, with focus on the aligned models only
* **Figure 7**: Analysis of prior variance and posterior temperature on the solution connectivity for ResNet20
* **Figure 8**: Analysis of the effect of data augmentation on ResNet20

Here are the changes planned for the next version/camera-ready of the paper:

1. Change the conjecture in Sec. 1 to contextualize better our contributions
2. Include arguments for using optimal transport to interpolate solutions in Sec. 3
3. Fix the readability of some plots
4. Fix various typos and add a couple of missing references pointed out by the Reviewers
5. Add and comment the new experiments presented in this rebuttal on prior and temperature using ResNet20 models in Sec. 5
6. Add and comment the new experiments presented in this rebuttal on data augmentation in Sec. 5

Furthermore, the remaining discussion point of this rebuttal will be added to the Appendix.

Finally, in the threads below we will address the Reviewers' comments and questions in detail.

Meta-review: The paper proposes an informal conjecture of linear connectivity modulo permutations for Bayesian neural networks trained with variational inference. The authors present some experiments supporting the conjecture, as well as an algorithm for aligning a pair of Bayesian neural networks. 
I expect the authors to implement the suggestions made by the reviewers around Conjecture 1. In addition to weakening the conjecture to be more specific to mean field variational Bayesian inference, I would also suggest changing the language around “functionally transparent permutations”, which is undefined jargon. While conjecture 1 is already informal, this jargon makes the conjecture hard to interpret. 
The paper could be further improved by adding the discussion on the connection between SGD and VI (as in the response to Reviewer 2rWU).
Decision: Accept (poster)