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.
Paper
Similar papers
Peer review
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
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
2 fair
Limitations
N/A
Summary
The authors conjecture that after accounting for permutation symmetries in overparametrized neural networks that lead to same functional behavior, the low-loss solutions are linearly connected. In the context of Bayesian neural networks with variational inference, the authors use this conjecture to propose a simple matching algorithm based on Linear Assignment Problem, to find an equivalent variational distribution but attains linear connectivity with another solution. The experiments demonstrate that it is possible to find such linearly connected no-barrier regions.
Strengths
- The authors take the ideas of mode connectivity in the deep learning literature, and try to find an equivalent formulation at the level of distributions over parameters as in Bayesian neural networks. It is a very interesting concept. - 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.
Weaknesses
- I believe that the title of the paper should qualify variational BNNs instead of just BNNs. Approximate VI remains practically distinctive in properties from other approximate Bayesian inference methods like MCMC/HMC. It also helps contextualize the scope for the reader, and something that the authors readily align with in the discussion as well. - 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? - 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. - 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. - 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. ### Minor - Please use `\citet` instead of `\citep` when referring directly to papers. For instance, in Line 134. - A very small description of the original LAP problem (introduced in Linear 167) in the appendix would be much appreciated.
Questions
1. The posterior predictive likelihood computed on test points in Eq. (5) is also used to define the functional loss barrier in Definition 3. Don't we want the loss barrier to be defined on training points? - 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. - I think it is completely fine to check using the test points as a diagnostic for correlation between behavior at test and train time. - 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. 2. 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. 3. 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? 4. Have the authors considered accounting for such functional symmetries during the training itself, instead of a post-hoc matching procedure?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
4 excellent
Presentation
3 good
Contribution
3 good
Limitations
Yes. Also see, weaknesses and questions.
Summary
The authors extend the recent idea of linear mode connectivity up to permutation symmetry to the setting of Bayesian neural networks. They demonstrate that two different variational approximations to the Bayes posterior enjoy mode connectivity along the Wasserstein geodesic of one distribution, and a suitably permuted version of the other. Such a permutation is discovered by replacing the L2 distance in previous work with the more suitable Wasserstein distance, and similarly, the objective can then be relaxed to a layer-wise linear assignment problem, leading to a tractable coordinate descent algorithm. The authors verify their results numerically, showing that two approximate solutions can indeed be connected for modern networks such as ResNet20 on CIFAR10.
Strengths
1. Studying mode connectivity for approximate Bayesian inference is a natural follow-up question to previous work, while at the same time requiring non-trivial extensions such as the Wasserstein geodesic and Wasserstein distance. It is very nice and somewhat surprising that the resulting objective (which is seemingly involved) can be relaxed in a very similar spirit, leading to a tractable problem. 2. The experimental setup is quite carefully created and a lot of ablations for different parameters such as the prior variance, the temperature and the width of the network are performed, giving a very complete picture.
Weaknesses
1. The proposed algorithm and the setup seem to heavily rely on the specific approximation method, namely variational inference with a Gaussian distribution of diagonal covariance as the variational family. Even if the diagonal covariance assumption is relaxed, it is not obvious to me how one can guarantee tractability, let alone moving to multi-modal approximations such as MCMC where not even the Wasserstein geodesic or distance are known in closed form. This is somewhat unsatisfying, since the power of BNNs and the Bayesian posterior in general only really starts to unfold once multiple modes are leveraged. It is thus somewhat unclear how much a unimodal approximation to the Bayes posterior really captures and how much it is really different from a simple point estimate. 2. 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. 3. 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.
Questions
1. Is data augmentation employed for the experiments involving the cold posterior effect? Data augmentation has been observed to be the main driver of the CPE [1, 2, 3] and hence using it would probably lead to a more visible effect. It has also been recently justified that tempering is a principled way to use data augmentation in Bayesian frameworks [4], so it would not affect the validity of the approximation. 2. Are there gains in terms of test accuracy of the (tempered) approximate Bayesian posteriors versus a standard SGD baseline? [1] What are Bayesian Neural Network Posteriors Really Like?, Izmailov et al. [2] Data augmentation in Bayesian neural networks and the cold posterior effect, Nabarro et al. [3] Disentangling the Roles of Curation, Data-Augmentation and the Prior in the Cold Posterior Effect, Noci et al. [4] How Tempering Fixes Data Augmentation in Bayesian Neural Networks, Bachmann et al.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
2 fair
Limitations
The authors have addressed the limitations of their work.
Summary
This work permutes together the distributions of bayesian neural network (BNN) parameters in the context of variational inference (VI) so that they are linearly connected. This is done by adapting recent work on permuting together SGD solutions to be linearly connected, by optimizing for similarity of means and variances from VI in the permutation assignment objective, instead of similarity of parameters (as in Ainsworth et al. 2022). The problem setup and permutation objective are described and derived in detail. Experiments show that this method is comparable in loss barrier to directly permuted SGD solutions.
Strengths
Bringing permutation alignment into a probabilistic setting is novel, and the formalization of the methodology is satisfying. The paper is well-organized, and a number of variations (variance of prior, temperature) are considered in the experiments. Despite its focus on BNNs, this work also has relevant implications for the SGD solution setting, since it would be nice to be able to extend linear mode connectivity from individual solutions to families of SGD solutions whose parameters are stochastic (e.g. due to randomness in initialization/training).
Weaknesses
In both the temperature and prior variance experiments, a comparison of barriers at the limit of 0 variance/temperature (corresponding to direct alignment of the MAP solution) for CIFAR-10 would be interesting, as it would put the observed difference between VI alignment and SGD alignment in figure 8 into context. Specifically, it would be nice to see figures 6 and 7 replicated on CIFAR-scale networks, and a reference line included in each figure to indicate the barrier achieved by SGD alignment. From 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.
Questions
Given that barrier is higher when aligning VI distributions versus aligning the MAP parameters (akin to directly aligning SGD solutions), what are the potential benefits of optimizing alignment from a distributional standpoint? If the goal is to minimize barrier, why not simply align the MAP? Some motivation would be helpful here.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
4 excellent
Presentation
4 excellent
Contribution
3 good
Limitations
The main limitation (which is briefly discussed in section 7, but not strongly emphasized throughout) is that the assumptions on the posterior distribution are too rigid: namely that parameters are independent Gaussians. The authors conjecture (line 239) that this may be the cause of reduced performance relative to Ainsworth et al. (2023). Given the derivation in (16) already admits arbitrary covariances, it seems very interesting to consider the case where the covariance of the posterior is non-diagonal (lines 275-276). This could also lead to significant innovations over the existing alignment algorithm of Ainsworth et al. (2023), which is largely unchanged in this work.
Summary
The authors study the geometry of SGD-trained Gaussian mean-field variational approximations to the posteriors of Bayesian neural networks (BNN). In large part, the authors propose extensions of the method and analysis of Ainsworth et al. [1] from MAP-estimated neural networks to BNNs. Notably, the authors informally conjecture that a permutation symmetry exists in the solutions that approximate Bayesian methods can find, similar to the permutation symmetry of MAP solutions SGD finds. To test this hypothesis, the authors first use optimal transport theory to extend the linear interpolation and the loss barrier framework of [1] to the variational posterior approximation setting. Then, they show that they can align individual mean-field Gaussian weight posteriors by approximately solving a bilinear assignment problem using an algorithm analogous to the one proposed by [1] for MAP solutions. The authors demonstrate empirically that zero-barrier interpolations exist for permutation-aligned variational posteriors of non-trivial architectures (ResNets) trained on non-trivial datasets (MNIST, FashionMNIST and CIFAR-10). ## References [1] Ainsworth, S. K., Hayase, J., & Srinivasa, S. (2022). Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836.
Strengths
Understanding the properties of approximate solutions to the posteriors of BNNs is one of the central challenges of Bayesian deep learning; hence extending the method and analysis of Ainsworth et al. to the variational BNN setting is an important step towards this goal. The proposed extension of the interpolation framework and the alignment procedure using optimal transport theory seems natural, and the experimental methodology is sound and reasonably thorough. Besides the experiments analogous to the ones of Ainsworth et al., the authors also investigate the effect of prior variance and the cold posterior effect, the limit of which is the MAP solution. The paper is mostly well-written and easy to follow.
Weaknesses
While I did not find any significant weaknesses in the work, there are a couple of points that, if improved/clarified, could significantly strengthen the paper. ### Conjecture 1 First, 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: "Solutions of approximate Bayesian inference for neural networks are linearly connected after accounting for functionally transparent permutations." First, 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. Related 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 "SGD-based solutions of mean-field variational Bayesian inference for neural networks are linearly connected after accounting for functionally transparent permutations." Thus, this is also a more direct generalization of the conjecture given in [1]. What do the authors think? ### Connection between the interpolation and the alignment process The 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)? Relatedly, the authors state the following on L116: "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." What 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. In 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. ### Figure 5 Given 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? ### Miscellaneous Eq (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. Line 4 in Algorithm 1 should be broken up into two or three lines, and the font size should be increased. The font size in Figures 3-8 is too small and should be increased to match at least the font size of the captions. The training procedure needs clarification. Did the authors use Bayes by Backprop [3] or the local reparameterization trick [4]? ## References [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. [2] Ainsworth, S. K., Hayase, J., & Srinivasa, S. (2022). Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836. [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. [4] Kingma, D. P., Salimans, T., & Welling, M. (2015). Variational dropout and the local reparameterization trick. Advances in neural information processing systems, 28.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Summary
This paper paper extends linear mode connectivity modulo permutation to Bayesian neural networks posteriors. Authors do this by imposing a Wasserstein metric on the space of distributions and look at how log likelihood changes along the geodesic between two distributions obtained via approximate Bayesian inference. For experiments, authors study BNNs parametrized by multivariate Gaussians with diagonal covariance. Analogous to weight matching algorithm proposed in [1], authors propose to find permutation that minimizes Wasserstein distance between the two distributions. For Gaussian distribution with diagonal covariance, authors propose an algorithm similar to weight matching that also takes into account the variances of the Gaussian. In experiments authors show that this heuristic for finding permutation that leads to small change in log likelihood along the geodesic. *[1] Git Re-Basin: Merging Models modulo Permutation Symmetries. Ainsworth et al.*
Strengths
- Overall a well-written paper that extends previously studied linear mode connectivity modulo permutation of neural networks to Bayesian neural networks. - Experiments demonstrate that the heuristic proposed to compute permutations works for network architectures and datasets studied in the paper.
Weaknesses
- Limitations of the proposed heuristic are not extensively discussed in the paper. For instance model architectures with batch norm typically fails without additional fixes [2]. - 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]. *[2] REPAIR: REnormalizing Permuted Activations for Interpolation Repair. Jordan et al.*
Questions
Empirically, how important is the role of adding covariance matrix in the objective when computing the permutations? For instance, using standard weight-matching in this setting just assumes that the covariance matrix is identity. it would be useful to see a table with comparison of different approaches.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
This paper limits itself on studying connectivity modulo permutation for approximate Bayesian NNs posteriors found via gradient based algorithms. It’s not clear if this phenomenon holds for other approaches which should be interesting future work.
Response to the Authors
I thank the authors for their rebuttal. After reading the other reviews and the authors' rebuttals, I maintain that the paper should be accepted and I raise my confidence score to reflect this. I thank the authors for providing a very nice explanation of why the mixture interpolation is unhelpful. I think they should also include Figure 1 in the appendix of the camera-ready version of the paper, as it would help readers who do not have a strong background in OT (like myself) to understand the authors' choices much easier. Though they did not address my point regarding geometric mixtures, I am now much more convinced that the OT formulation is more appropriate, as it captures the intuition of linear interpolation analogous to the point-estimate setting. *Regarding the disconnect between the interpolation and alignment procedure:* I think, in a certain way, the missing puzzle piece for me regarding the OT-based formulation is how SGD is connected to it. In particular, I wonder if using SGD to train the VI solution could be construed as some sort of OT procedure to transport the initial posterior guess to some optimal one. For example, is it connected to the Wasserstein gradient flow (Salim et al., 2020)? I think some form of a positive answer would conclusively put my worries to rest since this would shed light on how the VI setting generalizes the point-estimate setting. If SGD on VI was performing OT, it would be much more intuitively obvious why we should expect the OT interpolation and alignment procedures to be the right things to look at. Perhaps this could also help establish a connection between the two, though this is more of a shot in the dark on my part. ## References Salim, A., Korba, A., & Luise, G. (2020). The Wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 33, 12356-12366.
Thanks for reply. Regarding the geometric mixture, we apologize for missing this point earlier. We are not sure how the interpolation paths will look like with a geometric mixture, especially at the extremes ($\tau=\{0,1\}$). Indeed if $q$ is Gaussian, we have $$ q^\tau \propto \exp\left(-\frac \tau 2 \left(\frac{x-\mu}{\sigma}\right)^2 \right) = \exp\left(-\frac 1 2 \left(\frac{x-\mu}{\sigma/\sqrt{\tau}}\right)^2 \right) $$ which is the un-normalized PDF of a Gaussian RV with variance $\sigma^2/\tau$ (with $\tau \rightarrow 0$ we have a degenerate distribution). Regarding the connection with the Wasserstein gradient flow, an interesting reference connecting VI and OT could be [1]. In this paper, the authors show how the dynamics of the mean and (co)variance of Gaussian VI follows the gradient flow of the Kullback–Leibler (KL) divergence $KL(\cdot \| \pi)$ on a submanifold of the Wasserstein space of Gaussian distributions on $\mathcal P_2(\mathbb R^d)$ (known as Bures–Wasserstein manifold), which is equipped with the 2-Wasserstein distance. Additionally, the authors also show an SGD version for the discretization of the Bures–Wasserstein gradient flow, which might represent an interesting results to further characterize the properties of the approximate neural network posterior. Thanks again for this discussion. [1] Marc Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, Philippe Rigollet. Variational inference via Wasserstein gradient flows. NeurIPS 2022
Response to Rebuttal
I thank the authors for running additional experiments in such a short amount of time! **Tractability:** I agree that the Sinkhorn algorithm could be used to approximate the Wasserstein distance, but it would be tricky to reduce the permutation problem to anything tractable in this case, right? I might be missing something. **Precision of Approximation and Interpretability:** I agree that the exact posterior remains elusive and understanding approximate methods is thus very valuable. I simply remain unsure whether permutations really say anything fundamental about approximate posteriors. While of course SGD is also not an exact minimizer of its objective, even if it were, the permutation problem still remains! Moreover, here the permutation symmetry really reveals something fundamental about the problem. **Evaluation Metric:** Thank you for adding this! The precise definition of linear connectivity is always an issue in this line of work. What do the authors consider as connected here? Strictly speaking, none of the results are linearly-connected as the values do worsen, as evident now in the zoomed plot. I'm aware that similar results might also have passed as being "connected" in the literature, so I don't want to impose a new standard here, but it would be great if the authors could at least compare their connectivity values (for accuracy) with [1]. Especially the ResNet seems to worsen by almost 10% on the training data (removing the baseline here too would be helpful), which seems somewhat drastic. I hope the authors can clarify this. **Data augmentation:** Thank you for clarifying this! It's nice that data augmentation does not affect results to drastically. [1] Git Re-Basin: Merging Models modulo Permutation Symmetries, Ainsworth et al
We thank the reviewer for these additional comments. Here's our reply: * **Tractability**: We agree with the reviewer on this point. Extending this to sample-based inference is not going to be trivial nor easy, hence our comment on leaving this as future work. * **Precision of Approximation and Interpretability**: We can give the reviewer one possible application where such permutation analysis can be beneficial for approximate inference: convergence analysis of SG-MCMC. We know by experience that classic convergence statistics (like R-Hat) are not robust to assess convergence behavior of MCMC chains in large models. In Figure 2 in [1] we see that despite unrealistic compute availability, the R-Hat statistics severely underestimates the convergences of the chains, due to the exploration by HMC of permutation equivalent modes. While it is possible to compute such statistic in function-space, technically this is not a standard practice, since the R-Hat becomes also a function of the inputs, rather than just being a property of the chains/samples. On the other hand, by accounting for permutation symmetries we could derive more appropriate convergence statistics for MCMC methods to be used in this context. * **Evaluation Metric**: If we look at Figure 9 in [2] we see that for the MLP the results are actually very similar with ours (if not marginally worse, see CIFAR10/MLP). The biggest difference is indeed with ResNet20, where the only difference is the normalization layer (LayerNorm for [2] and FRN for us). We speculate this being the cause. Nonetheless, we want to highlight how our alignment method decreases the (accuracy) barrier by 85.89% on the train set and by 87.55% on the test set. [1] Pavel et al. What Are Bayesian Neural Network Posteriors Really Like? [2] Ainsworth et al. Git Re-Basin: Merging Models modulo Permutation Symmetries.
**Interpretability:** I'm not sure if I completely understand. Wouldn't the algorithm from *Ainsworth et al.* suffice to perform a permutation check between two MCMC samples (as far as I know, no Gaussian is built around MCMC samples)? **Evaluation Metric:** Thanks for the comparisons! FRN vs LayerNorm could be the reason that the ResNet20 is less connected but of course it could also be the novel alignment procedure proposed in this work. The only way to know for sure is to either check the novel procedure with the same LayerNorm-based ResNet20 or compare against other connectivity results for the FRN ResNet20 in the literature.
**Interpretability**: with MCMC, we would not be interested in single MCMC samples but rather on populations of samples (e.g. comparing different chains). In this sense, Ainsworth et al. is not directly applicable here. **Evaluation Metric**: this was exactly what we were running this past few days. It's not clear whether we can add links to external services but if you are interested here ([https://imgur.com/a/0iHJxiY](https://imgur.com/a/0iHJxiY)) you'll find a comparison between FRN and LN (hopefully we are not braking any rules, here). If you prefer not to follow the link, the gist of the figure is confirming that LN is performing better than FRN, further reducing the barrier by ~35%. | Normalization | Barrier Accuracy | Barrier Likelihood | | ----------------- | ------------------- | -------------------- | | FRN (with align.) | 0.163350 | 0.519373 | | LN (with align.) | 0.105150 (-35.629%) | 0.336379 (-35.2336%) |
**Interpretability:** I see, so that's kind of the multi-modal case though that's not doable now. But this work might serve as a first step towards this. At least the mathematical formulation might be helpful here. **Evaluation Metric:** Thanks a lot for running this! Are those test scores? And for accuracy, is the barrier for LN 10%? What was the value for *Ainsworth et al.*?
**Evaluation Metric**: Yes, those are the test score. Note that the width of the model is different: for sake of compute time, we are showing the x8, while they only have the x32 (sorry for the missing information above). Also Ainsworth et al. doesn't report this number
Thank you for all the clarifications. They mostly address my concerns and I have raised my score. I think for the final version of this paper, it would be great to have a 1-1 correspondence with the setup in *Ainsworth et al.* so that connectivity values can actually be compared. I still do not feel completely certain whether a 10% drop qualifies for enough connectivity.
Thank you for the additional results and the clarification regarding tractable posteriors. Given the new results show the advantage of including covariance in the alignment objective, I will raise my score.
Response to the rebuttal
Thank you for running additional experiments. It does appear that adding covariance to the objective indeed leads to relatively smaller loss barrier. I will keep my original score.
Decision
Accept (poster)