Title: On permutation symmetries in Bayesian neural network posteriors: a variational perspective
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

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

1 Introduction

Throughout the last decade, deep neural networks (DNNs) have achieved significant success in a wide range of practical applications, becoming the fundamental ingredient for e.g., computer vision [e.g., 57, 25, 41, 64], language models [e.g., 24, 82, 17] and generative models [e.g., 54, 96, 97, 98, 102, 34]. Despite recent important advancements, understanding the loss landscape of 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 [e.g., 3, 5, 7, 8, 18, 109], leading the loss landscape to have many optima due to symmetries and over-parameterization [76, 26, 111, 94]. At the same time, the role of the depth of a model in relation with its width is far less understood [80], despite wide neural networks exhibiting important theoretical properties in their infinite limit behavior [e.g., 73, 23, 29, 53, 48, 37, 20].

Two notions that have been useful to shed light on the geometry of loss landscapes are that of loss barriers and mode connectivity [36, 26]. The mode connectivity hypothesis states that given two points in the landscape, there exists a path connecting them such that the loss is constant or near constant (or, said differently, the loss barrier is null). We refer to linear mode connectivity when the path connecting the two solutions is linear [32]. Recently, evidence has surfaced that stochastic gradient descent (SGD) solutions to the loss minimization problem can be linearly connected. Indeed, Entezari et al. [30] discuss the role of permutation symmetries from a loss connectivity viewpoint, conjecturing the possibility that mode connectivity is actually linear once accounting for all permutation invariances.

-2

0

2

x

-2

0

2

y

Solution 1

-2

0

2

x

Solution 2

-2

0

2

x

Naive interpolation

-2

0

2

x

Aligned interpolation

Additionally, [2] gathers compelling empirical evidence across several network architectures and tasks, that under such a permutation symmetry the loss landscape often contains a single, nearly convex basin.

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

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

While being similar to the one in [30, 2], if this 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 to allow for a more general discussion. More specifically, in this paper we will analyze and focus our discussion on the variational inference framework, making a more specific conjecture:

Conjecture 1. Variational inference solutions for approximate Bayesian inference in neural networks are linearly connected after accounting for functionally equivalent permutations.

Contributions. With this work, we aim at studying the linear connectivity properties of approximate solutions to the Bayesian inference problem and we make several contributions. (i) We extend the formalism of loss barrier and solution interpolation to BNNs. (ii) For the variational inference setting, propose a matching algorithm to search for linearly connected solutions by aligning the distributions of two independent solutions with respect to permutation matrices. Inspired by [2], we frame the problem as a combinatorial optimization problem using approximation to the linear sum assignment problem. (iii) We then experiment on a variety of architectures and datasets, finding nearly zero-loss barriers for linearly connected solutions. In Fig. 2 we present a sneak-peek and a visualization of our findings, where we show that after weight distribution alignment we can find a permutation map P of the solution q 1 such that it can be linearly connected through high density regions to q 0 .

2 Preliminaries on Bayesian deep learning

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

where a ( · ) is a non-linearity, W l is a D l × D l -1 weight matrix and b l the corresponding bias vector. We shall refer to the parameters of the layer l as θ l = { b l , W l } , and the union of all trainable parameters as θ = { θ l } L l =1 .

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

where the normalization constant Z is the marginal likelihood p ( Y | X ) . As usually done, we assume that the likelihood factorizes over observations, i.e. p ( Y | θ , X ) = ∏ N i =1 p ( y i | θ , x i ) .

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

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

3 Loss barriers

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

Log-posterior for CIFAR10

P

#

q

1

q

0

q

1

likelihood, a proper scoring method for probabilistic models [83], defined as

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

Let's assume two models trained with VI with two different initializations, random seeds, and batch ordering. Variational inference in the classic inverse sense KL [ q ‖ p ] is mode seeking, thus we expect the two runs to converge to different solutions, say q 0 and q 1 . To test the loss barrier as we interpolate between the two solutions we need to decide on the interpolation rule. We decide to interpolate the solutions following the Wasserstein geodesics between q 0 and q 1 . First, let's start with a few definitions. Let q ∈ P ( R d ) be a probability measure on R d and T : R d → R d a measurable map; we denote T # q the push-forward measure of q through T . Now we can introduce the Wasserstein geodesics , as follows.

Definition 1. The Wasserstein geodesics between q 0 and q 1 is defined as the path

where Id is the identity map and T q 1 q 0 is the optimal transport map between q 0 and q 1 , which for Brenier's theorem [16], is unique.

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

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 ( R d ) → R , L is λ geodesics convex with λ > 0 if for any q 0 , q 1 ∈ P ( R d ) it holds that

where W 2 2 ( q 0 , q 1 ) is the Wasserstein distance defined as [104, 51, 52]

with Π( · , · ) being the space of measure with q 0 and q 1 as marginals.

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

Definition 3. The functional loss barrier along the Wasserstein geodesics from q 0 and q 1 is defined as the highest difference between the marginal loss computed when interpolating two solutions q 0 and q 1 and the linear interpolation of the loss at q 0 and q 1 :

where q τ follows the definition in Eq. (6) .

This definition is a more general than the ones in [2, 30, 32] but we can recover [30] by assuming delta posteriors q i = δ ( θ -θ i ) and we can further recover [2, 32] by also assuming L ( q 0 ) = L ( q 1 ) .

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

0.0

0.2

0.4

0.6

0.8

1.0

Interpolation coefficient

τ

-2.00

-1.75

-1.50

-1.25

Predictive likelihood

Method

Mixture

OT

0

2

1

2

q

0

0

2

1

2

q

1

0

2

1

2

(1

-

t

)

q

0 +

tq

1

0

2

1

2

((1

-

t

)

Id

+

tT q

1

q

0

)

#

q

0

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 × d permutation matrices. Given a generic distribution q ( θ ) , we can apply a permutation matrix P ∈ S ( D l ) to a hidden layer output at layer l , and if we define θ ′ to be equivalent to θ with the exception of

then P # q is the equivalent push-forward distribution for θ ′ , where P is the associated permutation map. Let us define the distribution over the functional output of the model as

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

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

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

where D is a generic measure of discrepancy. 2

4.1 Problem setup for permutation of vectors

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

Gaussian distributions, we know that if q = N ( m , S ) , then P # q = N ( Pm , PSP /latticetop ) and if q = N ( m , diag ( s 2 )) then P # q = N ( Pm , diag ( Ps 2 )) [11]. With the KL divergence KL [ P # q 1 ‖ q 0 ] it's easy to verify that it leads to just a distance between means, disregarding any covariance information. While certainly this represents a valid choice, we argue that we can find a better solution by using the Wasserstein distance. For Gaussian measures, the Wasserstein distance has analytic solution:

where ‖·‖ F denotes the Frobenius norm, ‖ A ‖ = ∑ ij a ij 2 , and where the second line is valid only if the covariances commute ( S 1 S 0 = S 0 S 1 ). In our case, then, we can simplify as follows:

To summarize, the problem now can be written as:

where the expression 〈 A | B 〉 F is the Frobenius inner product, 〈 A | B 〉 F = ∑ ij A ij B ij . Note that the r.h.s. of Eq. (18) is a valid instantiation of the linear assignment problem (LAP) [10], which can be solved in polynomial time.

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 to explicitly change our main objective by enforcing the functional equivalence constraint as follows:

where the notation ( P l ◦ P /latticetop l -1 ) represents the composition of the two permutation maps applied to rows and columns of the random weight matrices. More conveniently, this can be rewritten in terms of means and standard deviations. To leave the notation uncluttered, let's collect the means and the standard deviations for the layer l in M ( l ) and S ( l ) , which are now both D l × D l -1 matrices, so that q ( l ) = ∏ ij N ( M ( l ) ij , S ( l ) ij ) . Now we can write,

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

As discussed in [2], going through each layer, and greedily selecting its best P l , leads to a coordinate descent algorithm which guarantees to end in finite time. We present a pseudo-code in Algorithm 1.

Algorithm 1: Algorithm to align variational inference solutions

0.00

0.25

0.50

0.75

1.00

Interpolation coefficient

τ

-1.0

-0.5

0.0

Predictive likelihood

MLP on MNIST

0.00

0.25

0.50

0.75

1.00

Interpolation coefficient

τ

-2.0

-1.5

-1.0

Predictive likelihood

MLP on CIFAR10

0.00

0.25

0.50

0.75

1.00

Interpolation coefficient

τ

-2

-1

0

Predictive likelihood

VI (Train)

VI with distr. alignment (Train)

VI (Test)

VI with distr. alignment (Test)

ResNet20 on CIFAR10

5 Experiments

Now, we present some supporting evidence to Conjecture 1. We start by training two replicas of BNNs with variational inference (we refer to the Appendix for additional details on the experimental setup). We then compute the marginalized barrier as B ( q 0 , q 1 ) = max τ L ( q τ ) -((1 -τ ) L ( q 0 ) + τ L ( q 1 )) where L ( · ) is the predictive likelihood and τ ∈ [0 , 1] , from which we take 25 evenly distributed points. In particular, we seek to understand what happens to the VI solutions first for the naive interpolation from q 0 and q 1 , and then for the interpolation after aligning q 0 and P # q 1 . We experiment with MLPs with three layers and ResNet20 [41] with various widths on MNIST [60], Fashion-MNIST [108] and CIFAR10 [56]. All models are trained without data augmentation [106] and with filter response normalization (FRN) layers instead of BatchNorm. Finally, we set the prior to be Gaussian N ( 0 , α 2 I ) , with the flexibility of choosing the variance.

5.1 Low-barrier interpolations

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

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

2 6

2 9

Number of features

0

1

2

Barrier in the

predictive likelihood

MLP - MNIST

2 5

2 7

2 9

Number of features

1

2

Barrier in the

predictive likelihood

ResNet20 - CIFAR10

VI (Train)

VI (Test)

VI with align. (Train)

VI with align. (Test)

Posterior without alignment

3.966

3.962

3.958

10 7

(

a

)

Posterior with alignment

3.966

3.962

3.958

10 7

(

b

)

(

a

)

(

b

)

q

0

q

0

q

τ

P

#

q

τ

q

1

P

#

q

1

Note:

τ

= 0

.

5

×

×

10

2

10

1

10 0

Prior variance

α

2

0.0

0.5

1.0

Barrier in the

predictive likelihood

MLP - MNIST

10

2

10

1

10 0

Prior variance

α

2

0.0

0.5

1.0

Barrier in the

predictive likelihood

MLP - FMNIST

10

2

10

1

10 0

Prior variance

α

2

1

2

Barrier in the

predictive likelihood

ResNet20 - CIFAR10

VI (Train)

VI with distr. alignment (Train)

VI (Test)

VI with distr. alignment (Test)

Wespeculate that this might be due to the limiting behavior of Bayesian neural networks, which makes the posterior landscape Gaussian-like [48, 40]. While this does not fully explain the phenomenon observed, the existing connections between 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 of symmetries in weight space [80].

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

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

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

Additionally, several previous works have analyzed the effect of tempering the posterior in BNNs [106, 112, 110, 47]. Specifically, we are interested in the distribution p T ( θ | Y ) ∝ ( p ( Y | θ , X ) p ( θ )) 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

ResNet20 - CIFAR10

10

1

10 0

10 1

Posterior temperature

T

0.0

0.5

1.0

1.5

Barrier in the

predictive likelihood

MLP - MNIST

10

1

10 0

10 1

Posterior temperature

T

0

1

2

Barrier in the

predictive likelihood

MLP - FMNIST

10

1

10 0

Posterior temperature

T

1

2

Barrier in the

predictive likelihood

VI (Train)

VI with distr. alignment (Train)

VI (Test)

VI with distr. alignment (Test)

sharper than the true Bayesian posteriors, while for T > 1 we have warm posterior, which are more diffused. In Fig. 8 we see that barriers for cold posteriors with alignment are marginally closer to zero than for warm posteriors. Note that cold temperatures concentrate the distribution around the MAP, which motivates a further comparison with a non-Bayesian approach.

5.3 Comparison with SGD solutions

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

2 5

2 7

2 9

Features multiplier

0.5

1.0

1.5

2.0

Barrier in the

predictive likelihood

ResNet20 - CIFAR10

VI

SGD

VI with alignment

SGD with weight matching

5.4 Effect of normalization layers and data augmentation

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

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

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

3

6

9

12

15

18

21

Layer index

0.0

0.5

Variance ratio

ResNet20 - CIFAR10

Interpolation after alignment

Naive interpolation

0.00

0.25

0.50

0.75

1.00

Interpolation coefficient

τ

-2

-1

Predictive likelihood

ResNet20 - CIFAR10

Norm. layer

Alignment

FRN

True

LN

False

0.00

0.25

0.50

0.75

1.00

Interpolation coefficient

τ

-2

-1

Predictive likelihood

ResNet20 - CIFAR10

DA

Alignment

True

True

False

False

assumption on data augmentation and in Fig. 12 we experiment both with and without augmentation, showing that we are still able to recover similar low barrier solutions in both cases. Having said that, we advocate caution when using data augmentation in Bayesian neural networks, as it changes the shape of the posterior.

6 Related work

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

7 Conclusions

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

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

M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Jozefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Man´ e, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, I. Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Vasudevan, F. Vi´ egas, 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.

S. Ainsworth, J. Hayase, and S. Srinivasa. Git Re-Basin: Merging Models modulo Permutation Symmetries. In The Eleventh International Conference on Learning Representations , 2023.

Z. Allen-Zhu and Y. Li. What Can ResNet Learn Efficiently, Going Beyond Kernels? In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alch´ e-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

L. Ambrosio, N. Gigli, and G. Savare. Gradient Flows: In Metric Spaces and in the Space of Probability Measures . Lectures in Mathematics. ETH Z¨ urich. Birkh¨ auser Basel, 2008. ISBN 9783764387228.

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´ eBuc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer Normalization, 2016.

Y. Bai and J. D. Lee. Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks. In International Conference on Learning Representations , 2020.

Y. Bengio, O. Delalleau, and N. Le Roux. The Curse of Dimensionality for Local Kernel Machines. Technical Report TR-1258, March 2005.

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 Machine Learning Research , pages 769-779. PMLR, 18-24 Jul 2021.

D. Bertsekas. Network Optimization: Continuous and Discrete Methods . Athena scientific optimization and computation series. Athena Scientific, 1998. ISBN 9788865290279.

C. M. Bishop. Pattern recognition and machine learning . Springer, 1st ed. 2006. Corr. 2nd printing 2011 edition, 2006. ISBN 0387310738.

D. Blei, S. Mohamed, and R. Ranganath. Variational inference: Foundations and modern methods. 2016.

D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational Inference: A Review for Statisticians. 112(518):859-877, 2017.

C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight Uncertainty in Neural Network. 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 1613-1622. PMLR, 2015.

J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: Composable Transformations of Python+NumPy Programs, 2018.

Y. Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on Pure and Applied Mathematics , 44(4):375-417, 1991.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, 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.

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.

T. Chen, E. Fox, and C. Guestrin. Stochastic Gradient Hamiltonian Monte Carlo. In E. P. Xing and T. Jebara, editors, Proceedings of the 31st International Conference on Machine Learning , Proceedings of Machine Learning Research, pages 1683-1691. PMLR, 2014.

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.

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.

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.

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.

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.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations , 2021.

F. Draxler, K. Veschgini, M. Salmhofer, and F. 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.

S. Duane, A. Kennedy, B. J. Pendleton, and D. Roweth. Hybrid Monte Carlo. 195(2):216 222, 1987.

M. M. Dunlop, M. A. Girolami, A. M. Stuart, and A. L. Teckentrup. How Deep Are Deep Gaussian Processes? 19(1):2100-2145, 2018-01.

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.

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.

S. Fort, C. H. Hu, and B. Lakshminarayanan. Deep Ensembles: A Loss Landscape Perspective, 2020.

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 119 of Proceedings of Machine Learning Research , pages 3259-3269. PMLR, 13-18 Jul 2020.

G. 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.

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.

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.

T. 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-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31 , pages 8789-8798. Curran Associates, Inc., 2018.

A. Garriga-Alonso, C. E. Rasmussen, and L. Aitchison. Deep Convolutional Networks as shallow Gaussian Processes. In International Conference on Learning Representations , 2019.

C. Godfrey, D. Brown, T. Emerson, and H. Kvinge. On the Symmetries of 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.

A. Graves. Practical Variational Inference for Neural Networks. In J. Shawe-Taylor, R. S. Zemel, P. L. Bartlett, F. Pereira, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 24 , pages 2348-2356. Curran Associates, Inc., 2011.

B. He, B. Lakshminarayanan, and Y. W. Teh. Bayesian Deep Ensembles via the Neural Tangent Kernel. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33 , 2020.

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.

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.

M. D. Hoffman, D. M. Blei, C. Wang, and J. W. Paisley. Stochastic Variational Inference. 14 (1):1303-1347, 2013.

T. Huix, S. Majewski, A. Durmus, E. Moulines, and A. Korba. Variational Inference of overparameterized Bayesian Neural Networks: a theoretical and empirical study, 2022.

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.

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.

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.

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-8580. Curran Associates, Inc., 2018.

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.

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.

L. V. Kantorovich. On the transfer of masses. 37:227-229, 1942.

L. V. Kantorovich. On a problem of Monge. 3:225-226, 1948.

M. E. E. Khan, A. Immer, E. Abedi, and M. Korzepa. Approximate Inference Turns Deep Networks into Gaussian Processes. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alch´ e Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

D. P. Kingma and M. Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations , 2014.

D. P. Kingma, T. Salimans, and M. Welling. Variational Dropout and the Local Reparameterization Trick. In Advances in Neural Information Processing Systems 28 , pages 2575-2583. Curran Associates, Inc., 2015.

A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.

A. Krizhevsky, I. Sutskever, and G. E. Hinton. ImageNet Classification with Deep Convolutional Neural Networks. In F. Pereira, C. J. C. Burges, L. Bottou, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 25. Curran Associates, Inc., 2012.

R. Kurle, R. Herbrich, T. Januschowski, B. Wang, and J. Gasthaus. On the detrimental effect of invariances in the likelihood for variational inference. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022.

B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30 , pages 6402-6413. Curran Associates, Inc., 2017.

Y. LeCun, C. Cortes, and C. Burges. MNIST handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist , 2, 2010.

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.

Y. Li and Y. Gal. Dropout Inference in Bayesian Neural Networks with Alpha-divergences. 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 2052-2061. PMLR, 2017.

Q. Liu and D. Wang. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems 29 , pages 2378-2386. Curran Associates, Inc., 2016.

Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. A ConvNet for the 2020s. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , 2022.

D. J. C. MacKay. Bayesian Model Comparison and Backprop Nets. In J. E. Moody, S. J. Hanson, and R. Lippmann, editors, Advances in Neural Information Processing Systems 4 , pages 839-846. Morgan Kaufmann, 1991.

D. J. C. MacKay. A Practical Bayesian Framework for Backpropagation Networks. 4(3): 448-472, 1992-05.

D. 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 211-254. Springer, 1994.

D. J. C. MacKay. Choice of Basis for Laplace Approximation. 33(1):77-86, 1998.

C. J. Maddison, A. Mnih, and Y. W. Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. In Proceedings of the 5th International Conference on Learning Representations . OpenReview.net, 2017.

S. Mandt, , M. D. Hoffman, and D. M. Blei. Stochastic Gradient Descent as Approximate Bayesian Inference. Journal of Machine Learning Research , 18(134):1-35, 2017.

A. G. d. G. Matthews, J. Hensman, R. Turner, and Z. Ghahramani. On Sparse Variational Methods and the Kullback-Leibler 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.

N. Metropolis and S. Ulam. The monte carlo method. 44(247):335-341, 1949. PMID: 18139350.

R. M. Neal. Bayesian Learning for Neural Networks . PhD thesis, 1994.

R. M. Neal. Priors for infinite networks. Technical report crg-tr-94-1, University of Toronto, 1994.

R. M. Neal. MCMC Using Hamiltonian Dynamics , chapter 5. CRC Press, 2011.

B. Neyshabur, Z. Li, S. Bhojanapalli, Y. LeCun, and N. Srebro. The role of over-parametrization in generalization of neural networks. In International Conference on Learning Representations , 2019.

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 Processing Systems , volume 32, pages 4287-4299. Curran Associates, Inc., 2019.

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. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alch´ e Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

F. Pittorino, A. Ferraro, G. Perugini, C. Feinauer, C. Baldassi, and R. Zecchina. Deep Networks on 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 on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 17759-17781. PMLR, 17-23 Jul 2022.

G. Pleiss and J. P. Cunningham. The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems , 2021.

Y. Polyanskiy and Y. Wu. Strong Data-Processing Inequalities for Channels and Bayesian Networks. In E. Xarlen, M. Madiman, and E. M. Werner, editors, Convexity and Concentration , pages 211-249, New York, NY, 2017. Springer New York.

A. Radford and K. Narasimhan. Improving Language Understanding by Generative PreTraining. 2018.

C. E. Rasmussen and C. K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning) . The MIT Press, 2005. ISBN 026218253X.

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.

H. Robbins and S. Monro. A Stochastic Approximation Method. 22(3):400-407, 1951.

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.

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.

S. Rossi, S. Marmin, and M. Filippone. Walsh-Hadamard Variational Inference for Bayesian Deep Learning. In Advances in Neural Information Processing Systems 33 , 2020.

T. G. J. Rudner, Z. Chen, Y . W. Teh, and Y . Gal. Tractable Function-Space Variational Inference in Bayesian Neural Networks. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022.

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.

D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning Representations by Backpropagating Errors. 323(6088):533-536, 1986-10.

T. Salimans and D. A. Knowles. Fixed-Form Variational Posterior Approximation through Stochastic Linear Regression. 8(4):837-882, 2013.

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.

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.

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.

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

J. Song, C. Meng, and S. Ermon. Denoising Diffusion Implicit Models. In International Conference on Learning Representations , 2021.

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.

S. Sun, G. Zhang, J. Shi, and R. B. Grosse. Functional Variational Bayesian Neural Networks. In Proceedings of the 7th International Conference on Learning Representations . OpenReview.net, 2019.

A. Takatsu. Wasserstein geometry of Gaussian measures. Osaka Journal of Mathematics , 48 (4):1005 - 1026, 2011.

Tishby, Levin, and Solla. Consistent inference of probabilities in layered networks: predictions and generalizations. In International 1989 Joint Conference on Neural Networks , pages 403-409 vol.2, 1989.

B.-H. Tran, S. Rossi, D. Milios, P. 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.

B.-H. Tran, S. Rossi, D. Milios, and M. Filippone. All You Need is a Good Functional Prior for Bayesian Deep Learning. In The Journal of Machine Learning Research , 2022.

C. Villani. Optimal Transport: Old and New , volume 338. Springer Science & Business Media, 2008.

M. Welling and Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning , pages 681-688, Madison, WI, USA, 2011. Omnipress. ISBN 9781450306195.

F. Wenzel, K. Roth, B. S. Veeling, J. ´ Swia ¸tkowski, L. Tran, S. Mandt, J. Snoek, T. Salimans, R. Jenatton, and S. Nowozin. How Good is the Bayes Posterior in Deep Neural Networks Really? In Proceeding of the 37th International Conference on Machine Learning , 2020.

J. G. Wiese, L. Wimmer, T. Papamarkou, B. Bischl, S. G¨ unnemann, and D. R¨ ugamer. Towards Efficient MCMC Sampling in Bayesian Neural Networks by Exploiting Symmetry, 2023.

H. Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms. CoRR , abs/1708.07747, 2017.

G. Yehudai and O. Shamir. On the Power and Limitations of Random Features for Understanding Neural Networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alch´ e-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

C. Zeno, I. 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.

C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations , 2017.

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.

A Future work and open questions

Besides the specific open questions discussed above, we anticipate some possible future work. The framework we proposed to analyze the loss barriers in BNNs is general and can be applied to approximations other than VI. Future work could undertake to extend this analysis to the Laplace approximation [68]. However, this raises methodological challenges since Algorithm 1 is not directly applicable, due to the use of a dense covariance matrix (the inverse of the Hessian). Additionally, it would be worth extending these results to sample-based inference, like stochastic gradient Hamiltonian Monte Carlo (SGHMC) [19], or particle-based inference, like Stein variational inference [63]. This would call for a 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 [ P # q 1 ‖ q 0 ] between Gaussians we have

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.

-2

-1

0

1

2

-2

-1

0

1

2

Original distributions

q

1

q

2

-2

-1

0

1

2

LAP with KL cost

-2

-1

0

1

2

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 ( q 0 , q 1 ) = max τ L ( q τ ) -((1 -τ ) L ( q 0 ) + τ L ( q 1 )) where L ( · ) is the predictive likelihood and τ ∈ [0 , 1] , from which we take 25 evenly distributed points. In particular, we seek to understand what happens to the VI solutions with and without alignment applied to one of the two distribution ( q 1 in our case). We experiment with MLPs with three layers and ResNet20 [41] with various widths on MNIST [60], Fashion-MNIST [108] and CIFAR10 [56]. All models are trained without data augmentation, 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 N ( 0 , α 2 I ) , with the flexibility of choosing the variance. All VI 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 MLPs and convolutional neural networks (CNNs) base architectures used in our experimental campaign, while Table 1 reports the hyperparameters used in the experiments. Note that differently from Entezari et al. [30] and Ainsworth

et 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.

C.1 Computing platform

The experiments have been performed using JAX [15] and run on two AWS p4d.24xlarge instances with 8 NVIDIA A100 GPUs. Experiments were conducted using in the eu-west-1 region, which has a carbon efficiency of 0.62 kgCO 2 eq/kWh. A 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 a pessimistic estimation of the true utilization. Total emissions are estimated to be 1007.5 kgCO 2 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 the time for the deterministic case [2]. In Fig. A14 we show the results for MLP and ResNet20 architectures, varying the model width. It is evident that, in the majority of cases, the algorithm completes within a minute. Moreover, as anticipated, in case of VI 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 of performance as in 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 et 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 a number of questions: the Bayesian posterior is the product of two ingredients, the prior and the likelihood, conditioned to observing a dataset.

16

32

64

128

256

512 1024 2048

Model width

10

1

10 0

10 1

10 2

Time [sec.]

MLP

Mode

VI

SGD

16

32

64

128

256

512

Model width

10 0

10 1

Time [sec.]

ResNet20

Mode

VI

SGD

0.0

0.2

0.4

0.6

0.8

1.0

Interpolation coefficient

τ

-4

-2

0

Predictive likelihood

ResNet20 on CIFAR100

VI (Train)

VI with distr. alignment (Train)

VI (Test)

VI with distr. alignment (Test)

Finally, as an additional check, we analyze the logposterior 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 of the form

10 6

where a, b ∈ R and θ a , θ b and θ c are the means of q 0 , q 1 and P # q 1 . With this configuration, all three solutions lie on this hyper-plane. In Fig. 6, we present the visualization of ResNet20 trained on CIFAR10. We see that the distributions q 0 and P # q 1 are connected by higher density regions than the ones between q 0 and q 1 . Also, as expected the symmetries arise from the form of the likelihood, and the prior has a 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-posterior 3.54 3.60 3.66 ×

×

Log-posterior

Log-likelihood

3.54

3.60

3.66

3.72

×

10 6

175000

100000

25000

Model

q

0

Model

q

1

Model

P

#

q

1

Log-likelihood

Log-prior

175000

100000

25000

3.680

3.696

3.712

3.728

10 6

3.680

E A primer on variational inference for Bayesian neural networks

VI is a classic tool to tackle intractable Bayesian inference [50, 13]. VI casts the inference problem into an optimization-based procedure to compute a tractable approximation of the true posterior. Assume a generic parametric model f parameterized by some unknown parameters θ (i.e. f ( · , θ ) ) and a collection of data y ∈ R N corresponding to some input points X = { x i | x i ∈ R D in } i =1 ,...,N . In our setting, we have a probabilistic model p ( y | f ( X ; θ )) with parameters θ , a prior distributions on them p ( θ ) and a 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 q ( θ ) ∈ Q and the true posterior p ( θ | y ) ; and finally (iii) providing a programmatic way to find the distribution ˜ q ( θ ) that minimizes such distance. In practice, q ( θ ) has some free parameters ν (also known as variational parameters ), which are optimized such that

3.72

Log-prior

3.696

3.712

×

3.728

10 6

the approximating distribution q ( θ ; ν ) is as closer as possible to the true posterior p ( θ | y ) . We can derive the variational objective starting from the definition of the KL,

Rearranging we have that

The r.h.s. of the equation defines our variational objective, also known as ELBO, that can be arranged as follows,

This formulation highlights the property of this objective, which is made of two components: the first one is the expected log-likelihood under the approximate posterior q and measures how the model fits the data. The second term, on 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 θ i is taken to be independent with respect to the remaining θ -i . Effectively, this imposes a factorization of the posterior,

where ν i is the set of variational parameters for the parameter θ i . On top of this approximation, q ( θ i ) is often chosen to be Gaussian,

Now, the collection of all means and variances { µ i , σ 2 i } K i =1 defines the set of variational parameters to optimize.

For BNNs the analytic evaluation of the ELBO (and its gradients) is always untractable due the nonlinear 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 N MC times from q ν ,

In practice, this is as simple as re-sampling the weights and the biases for all the layers N MC times and computing the output for each new sample.

We now have a tractable objective that needs to be optimized with respect to the variational parameters ν . Very often the KL 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 reparameterization trick [92, 54]. The reparameterization trick aims at constructing θ as an invertible function T of the variational parameters ν and of another random variable ε , so that θ = T ( ε ; ν ) . Generally, a 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]. ε is chosen such that its marginal p ( ε ) 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 σ 2 , T corresponds to as simple scale-location transformation of an isotropic Gaussian noise,

This simple transformation ensures that p ( ε ) = N (0 , 1) does not depends on the variational parameters ν = { µ, σ 2 } . The gradients of the ELBO can be therefore computed as

The gradient ∇ θ log p ( y | θ ) depends on the model and it can be derived with automatic differentation tools [1, 78], while ∇ ν T ( ε ; ν ) doesn't have any stochastic components and therefore can be known deterministically. Note that the reparameterization trick can be also used when the KL is not analitically available. In that case, we would end up with,

Roeder et al. [86] argue that when we believe that q ( θ ; ν ) ≈ p ( y | θ ) , Eq. (30) should be prefered over Eq. (29) even if computing analitically the KL is possible. Note that this case is 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 to the evaluation of the likelihood and its gradients N MC times. Assuming factorization of the likelihood,

this quantity can be approximated using mini-batching [39, 43]. Recalling y as the set of labels of our dataset with N examples, by taking B ⊂ y as a random subset of y , the likelihood term can be estimated in an unbiased way as

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 a local optimum provided with a decreasing learning rate and sufficient gradient updates [85].

Peer reviews:
comments: 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).

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).

Other points:
* An important reference for permutations: Aitchison, Laurence, Adam Yang, and Sebastian W. Ober. "Deep kernel processes." International Conference on Machine Learning. PMLR, 2021.
* Legend for Fig. 8 is _way_ too small.  In general the Figures feel a bit crammed in.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes. Also see, weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe authors have addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nBringing 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).\n\nWEAKNESSES:\nIn 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.\n\nQUESTIONS:\nGiven 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nUnderstanding 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.\n\nWEAKNESSES:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.*\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.*\n\nQUESTIONS:\nEmpirically, 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.\n\nLIMITATIONS:\nThis 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}