SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\underset {\boldsymbol {P} \in \mathbb {S} (d)} {\arg \min } \| \boldsymbol {m} _ {0} - \boldsymbol {P} \boldsymbol {m} _ {1} \| _ {2} ^ {2} + \| \boldsymbol {s} _ {0} - \boldsymbol {P} \boldsymbol {s} _ {1} \| _ {2} ^ {2} = \underset {\boldsymbol {P} \in \mathbb {S} (d)} {\arg \max } \left\langle \boldsymbol {P} | \boldsymbol {m} _ {0} \boldsymbol {m} _ {1} ^ {\top} + \boldsymbol {s} _ {0} \boldsymbol {s} _ {1} ^ {\top} \right\rangle_ {F}, \tag {18}
$$
where the expression  $\langle A|B\rangle_F$  is the Frobenius inner product,  $\langle A|B\rangle_F = \sum_{ij}A_{ij}B_{ij}$ . Note that the r.h.s. of Eq. (18) is a valid instantiation of the linear assignment problem (LAP) [10], which can be solved in polynomial time.
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:
$$
\underset {\{P _ {i} \}} {\arg \min} \mathcal {W} _ {2} ^ {2} \left(P _ {1 \#} q _ {1} ^ {(1)}, q _ {0} ^ {(1)}\right) + \mathcal {W} _ {2} ^ {2} \left(\left(P _ {2} \circ P _ {1} ^ {\top}\right) _ {\#} q _ {1} ^ {(2)}, q _ {0} ^ {(2)}\right) + \dots + \mathcal {W} _ {2} ^ {2} \left(\left(P _ {L - 1} ^ {\top}\right) _ {\#} q _ {1} ^ {(L)}, q _ {0} ^ {(L)}\right),
$$
where the notation  $\left(P_{l} \circ P_{l-1}^{\top}\right)$  represents the composition of the two permutation maps applied to rows and columns of the random weight matrices. More conveniently, this can be rewritten in terms of means and standard deviations. To leave the notation uncluttered, let's collect the means and the standard deviations for the layer  $l$  in  $\boldsymbol{M}^{(l)}$  and  $\boldsymbol{S}^{(l)}$ , which are now both  $D_{l} \times D_{l-1}$  matrices, so that  $q^{(l)} = \prod_{ij} \mathcal{N}(M_{ij}^{(l)}, S_{ij}^{(l)})$ . Now we can write,
$$
\begin{array}{l} \operatorname * {a r g   m a x} _ {\left\{\boldsymbol {P} _ {i} \right\} _ {i = 1} ^ {L}} \left\langle \boldsymbol {M} _ {0} ^ {(1)} \middle | \boldsymbol {P} _ {1} \boldsymbol {M} _ {1} ^ {(1)} \right\rangle_ {F} + \left\langle \boldsymbol {S} _ {0} ^ {(1)} \middle | \boldsymbol {P} _ {1} \boldsymbol {S} _ {1} ^ {(1)} \right\rangle_ {F} + \left\langle \boldsymbol {M} _ {0} ^ {(2)} \middle | \boldsymbol {P} _ {2} \boldsymbol {M} _ {1} ^ {(2)} \boldsymbol {P} _ {1} ^ {\top} \right\rangle_ {F} + \left\langle \boldsymbol {S} _ {0} ^ {(2)} \middle | \boldsymbol {P} _ {2} \boldsymbol {S} _ {1} ^ {(2)} \boldsymbol {P} _ {1} ^ {\top} \right\rangle_ {F} + \\ + \dots + \left\langle \boldsymbol {M} _ {0} ^ {(L)} \middle | \boldsymbol {M} _ {1} ^ {(L)} \boldsymbol {P} _ {L - 1} ^ {\top} \right\rangle_ {F} + \left\langle \boldsymbol {S} _ {0} ^ {(L)} \middle | \boldsymbol {S} _ {1} ^ {(L)} \boldsymbol {P} _ {L - 1} ^ {\top} \right\rangle_ {F}. \\ \end{array}
$$
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.
$$
\begin{array}{l} \underset {\boldsymbol {P} _ {l}} {\arg \max} \left\langle \boldsymbol {M} _ {0} ^ {(l)} \Big | \boldsymbol {P} _ {l} \boldsymbol {M} _ {1} ^ {(l)} \boldsymbol {P} _ {l - 1} ^ {\top} \right\rangle_ {F} + \left\langle \boldsymbol {M} _ {0} ^ {(l + 1)} \Big | \boldsymbol {P} _ {(l + 1)} \boldsymbol {M} _ {1} ^ {(l + 1)} \boldsymbol {P} _ {l} ^ {\top} \right\rangle_ {F} + \\ \left\langle \boldsymbol {S} _ {0} ^ {(l)} \right| \boldsymbol {P} _ {l} \boldsymbol {S} _ {1} ^ {(l)} \boldsymbol {P} _ {l - 1} ^ {\top} \bigg \rangle_ {F} + \left\langle \boldsymbol {S} _ {0} ^ {(l + 1)} \right| \boldsymbol {P} _ {(l + 1)} \boldsymbol {S} _ {1} ^ {(l + 1)} \boldsymbol {P} _ {l} ^ {\top} \bigg \rangle_ {F} = \\ = \underset {\boldsymbol {P} _ {l}} {\arg \max } \left\langle \boldsymbol {P} _ {l} \middle | \boldsymbol {M} _ {0} ^ {(l)} \boldsymbol {P} _ {l - 1} \left(\boldsymbol {M} _ {1} ^ {(l)}\right) ^ {\top} + \left(\boldsymbol {M} _ {0} ^ {(l + 1)}\right) ^ {\top} \boldsymbol {P} _ {l + 1} \boldsymbol {M} _ {1} ^ {(l + 1)} + \right. \\ \left. \boldsymbol {S} _ {0} ^ {(l)} \boldsymbol {P} _ {l - 1} \left(\boldsymbol {S} _ {1} ^ {(l)}\right) ^ {\top} + \left(\boldsymbol {S} _ {0} ^ {(l + 1)}\right) ^ {\top} \boldsymbol {P} _ {l + 1} \boldsymbol {S} _ {1} ^ {(l + 1)} \right\rangle_ {F}. \tag {19} \\ \end{array}
$$
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
Data: Variational inference solutions  $q_{0}$  and  $q_{1}$  Result: Permutation matrices  $\pmb{P}_{i}$    
1  $\pmb {P}_i\gets \pmb {I},\forall i\in \{1,\dots ,L\} ;$    
2 while not converged do   
3 for  $i\in$  RandomPerm(1,..,L-1) do   
4  $\begin{array}{r}\big{\lfloor}\pmb {P}_i = \arg \max \langle \pmb {P}_i|\star \rangle ,\mathrm{where}\star \mathrm{is the r.h.s. of Eq. (19)} \end{array}$
Figure 4: Zero barrier solutions. Comparison of loss barriers for standard VI (gray) and VI with alignment (orange). While loss barriers always appear between two solutions in the standard VI approach, in the case of VI with alignment there is no noticeable loss barrier for MLPs and a nearly-zero loss barrier for ResNet20.
Now, we present some supporting evidence to Conjecture 1. We start by training two replicas of BNNs with variational inference (we refer to the Appendix for additional details on the experimental setup). We then compute the marginalized barrier as  $\mathcal{B}(q_0,q_1) = \max_{\tau}\mathcal{L}(q_\tau) - ((1 - \tau)\mathcal{L}(q_0) + \tau \mathcal{L}(q_1))$  where  $\mathcal{L}(\cdot)$  is the predictive likelihood and  $\tau \in [0,1]$ , from which we take 25 evenly distributed points. In particular, we seek to understand what happens to the VI solutions first for the naive interpolation from  $q_0$  and  $q_1$ , and then for the interpolation after aligning  $q_0$  and  $P_{\#}q_1$ . We experiment with MLPs with three layers and ResNet20 [41] with various widths on MNIST [60], Fashion-MNIST [108] and CIFAR10 [56]. All models are trained without data augmentation [106] and with filter response normalization (FRN) layers instead of BatchNorm. Finally, we set the prior to be Gaussian  $\mathcal{N}(\mathbf{0},\alpha^2\mathbf{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.
Figure 5: Effect of width. After distribution alignment, wider models exhibit lower likelihood barrier.
Figure 6: Posterior density visualization. Analysis of the log-posterior computed for ResNet20. Samples from  $q_{0}$  and  $q_{1}$  are connected by lower density regions, while  $q_{0}$  and  $P_{\#}q_{1}$  are not.
Figure 7: Effect of prior variance. After distribution alignment, prior variance has low effect in finding zero-loss barriers, while with naive interpolation we see a decreasing trend the higher the variance is.
We speculate that this might be due to the limiting behavior of Bayesian neural networks, which makes the posterior landscape Gaussian-like [48, 40]. While this does not fully explain the phenomenon observed, the existing connections between 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 = \{\pmb{\theta} \in \mathbb{R}^d \mid \pmb{\theta} = a\pmb{\theta}_a + b\pmb{\theta}_b + (1 - a - b)\pmb{\theta}_c\}$ , where  $a, b \in \mathbb{R}$  and  $\pmb{\theta}_a, \pmb{\theta}_b$  and  $\pmb{\theta}_c$  are the samples either from  $q_0, q_1$  and  $q_{\tau}$  without alignment or from  $q_0, P_{\#}q_1$  and  $P_{\#}q_{\tau}$  with alignment. With this configuration, all three samples always lie on this hyper-plane. In Fig. 6, we present the visualization of ResNet20 trained on CIFAR10. We see that the samples from  $q_0$  and  $P_{\#}q_1$  are connected by higher density regions than the ones between  $q_0$  and  $q_1$ . This is in line with the results in Fig. 4, where we see that the loss barrier is lower after alignment.
5.2 Analyzing the effect of the prior and testing the cold posterior effect
In all previous experiments we used a Gaussian prior  $\mathcal{N}(\mathbf{0},\alpha^2\mathbf{I})$  with fixed  $\alpha^2$ ; now we study the effect of a varying prior variance  $\alpha^2$  on the behavior of the loss barriers. We experiment this on a MLP trained on MNIST and Fashion-MNIST and on a ResNet20 (width x8) on CIFAR10. We report the results in Fig. 7. We can appreciate two behaviors: with alignment, there is no measurable effect of using different variances in finding zero-barrier solutions; on the contrary, without alignment we see that naive VI solutions are easier to interpolate with lower barrier when the prior is more diffused. At the same time, we see that higher variances produce bigger gaps between train barriers and test barriers. We speculate that this is due to overfitting happening with more relaxed priors, which makes low-barrier (but low-likelihood) solutions easier to find.
Additionally, several previous works have analyzed the effect of tempering the posterior in BNNs [106, 112, 110, 47]. Specifically, we are interested in the distribution  $p_T(\pmb{\theta} \mid \pmb{Y}) \propto (p(\pmb{Y} \mid \pmb{\theta}, \pmb{X}) p(\pmb{\theta}))^{1 / T}$ , where  $T$  is known as the temperature. Note that starting from the above definition, we can write an equivalent ELBO for VI which takes into account  $T$ . For  $T < 1$ , we have cold posteriors, which are
Figure 8: Effect of temperature. After alignment, cold posteriors makes barriers marginally closer to zero
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.
Figure 9: VI versus SGD. VI and SGD behave equivalently after permutation alignment in narrow models. For wider models, SGD solutions reach lower barriers than VI.
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
Figure 10: Variance collapse. Variance collapse is present for the naive interpolation, but it is reduced after alignment.
Figure 11: Effect of normalization layers. Both FRN and LN can be aligned, and LN exhibits lower barriers.
Figure 12: Effect of data augmentation. With and without DA, we are able to recover similar low barrier solutions.
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.
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.
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.
The authors want to thank David Bertrand and the AIAO team atstellantis for their work on the software and hardware infrastructure used for this work.
[105] 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.
[106] F. Wenzel, K. Roth, B. S. Veeling, J. Światkowski, 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 Proceedings of the 37th International Conference on Machine Learning, 2020.
[107] J. G. Wiese, L. Wimmer, T. Papamarkou, B. Bischl, S. Gunnemann, and D. Rügamer. Towards Efficient MCMC Sampling in Bayesian Neural Networks by Exploiting Symmetry, 2023.
[108] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms. CoRR, abs/1708.07747, 2017.
[109] 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é-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[110] 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.
[111] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
[112] R. Zhang, C. Li, J. Zhang, C. Chen, and A. G. Wilson. Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning. In Proceedings of the 8th International Conference on Learning Representations. OpenReview.net, 2020.
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  $\mathrm{KL}\left[P_{\#}q_1\parallel q_0\right]$  between Gaussians we have
$$
\operatorname {K L} \left[ P _ {\#} q _ {1} \| q _ {0} \right] = \log \det \operatorname {d i a g} \left(\boldsymbol {s} _ {0}\right) - \log \det \operatorname {d i a g} \left(\boldsymbol {P} \boldsymbol {s} _ {1}\right) + \operatorname {T r} \left(\operatorname {d i a g} \left(\boldsymbol {P} \boldsymbol {s} _ {1} \boldsymbol {s} _ {0} ^ {- 1}\right)\right) + \tag {20}
$$
$$
\left(\boldsymbol {m} _ {0} - \boldsymbol {P} \boldsymbol {m} _ {1}\right) ^ {\top} \operatorname {d i a g} \left(\boldsymbol {s} _ {0} ^ {- 1}\right) \left(\boldsymbol {m} _ {0} - \boldsymbol {P} \boldsymbol {m} _ {1}\right) \tag {21}
$$
It's easy to verify that the first three terms do not depend on  $P$ , leading to just a distance between means and disregarding any covariance information. In the figure below, we visualize the difference between doing LAP with the KL cost and LAP with the Wasserstein cost.
Figure A13: Alignment using different objectives. Given two distributions symmetrical w.r.t. the  $y = x$  plane, using the KL cost LAP results in the identity permutation (which fails to recover the symmetry), while the Wasserstein cost better aligns the two distributions
If not otherwise stated, we start by training two replicas of BNNs with variational inference and we compute the marginalized barrier as  $\mathcal{B}(q_0,q_1) = \max_{\tau}\mathcal{L}(q_\tau) - ((1 - \tau)\mathcal{L}(q_0) + \tau \mathcal{L}(q_1))$  where  $\mathcal{L}(\cdot)$  is the predictive likelihood and  $\tau \in [0,1]$ , from which we take 25 evenly distributed points. In particular, we seek to understand what happens to the VI solutions 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  $\mathcal{N}(0,\alpha^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