Original: L INEAR M ODE C ONNECTIVITY OF D EEP N EURAL N ETWORKS VIA P ERMUTATION I NVARIANCE AND R ENORMALIZATION
Revised: REPAIR: RE NORMALIZING P ERMUTED A CTIVATIONS FOR I NTERPOLATION R EPAIR

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In this paper we empirically investigate the conjecture from Entezari et al.
Revised: In this paper we look into the conjecture of Entezari et al.

Original: (2021) which states that if  permutation invariance  is taken into account, then there should be no loss barrier to the linear interpolation between SGD solutions.
Revised: (2021) which states that if the permutation invariance of neural networks is taken into account, then there is likely no loss barrier to the linear interpolation between SGD solutions.

Original: We conduct our investigation using standard computer vision architectures trained on CIFAR10 and ImageNet. First, we observe  ageneral phenomenon in which interpolated deep networks suffer a collapse in the variance of their activations.
Revised: First, we observe that neuron alignment methods alone are insufﬁcient to establish low-barrier linear connectivity between SGD solutions due to a phenomenon we call variance collapse : interpolated deep networks suffer a collapse in the variance of their activations, causing poor performance.

Original: We demonstrate that an appropriate rescalingof the pre-activations of the interpolated networks ameliorates this problem and signiﬁcantly reduces the barrier. Second, by combining this with an algorithm for ﬁnding permutations based on maximizing correlations between  the activations of matched neurons, we are able to reduce the interpolation barrier for a standard ResNet18 trained on CIFAR-10 to 1.5% absolute test error.
Revised: Next, we propose REPAIR (REnormalizing Permuted Activations for Interpolation Repair) which mitigates variance collapse by rescaling the preactivations of such interpolated networks. We explore the interaction between our method and the choice of normalization layer, network width, and depth, and demonstrate that using REPAIR on top of neuron alignment methods leads to 60%-100% relative barrier reduction across a wide variety of architecture families and tasks. In particular, we report a 74% barrier reduction for ResNet50 on ImageNet and 90% barrier reduction for ResNet on CIFAR10.

Original: We explore the interaction between our method and the choice of normalization layer, and demonstrate its robustness across a variety of architectures and training sets.
Revised: Our code is available at https://anonymous.4open.science/ r/Network-Permutations-2BEC/README.md .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Training a neural network corresponds to optimizing a highly non-linear function by navigating a complex loss landscape with numerous minima, symmetries and saddles (Zhang et al., 2017; Keskar et al., 2017; Draxler et al., 2018; S¸ims¸ek et al., 2021).
Revised: Training a neural network corresponds to optimizing a highly non-linear function by navigating a complex loss landscape with numerous minima, symmetries and saddles (Zhang et al., 2017; Keskar et al., 2017; Draxler et al., 2018; S¸ims¸ek et al., 2021).

Original: Overparameterization is one of the reasons behind the abundance of minima leading to different functions  behave similarly on the training databut different on other inputs (Neyshabur et al., 2017;
Revised: Overparameterization is one of the reasons behind the abundance of minima leading to different functions that behave similarly on the training data (Neyshabur et al., 2017;

Original: Nguyen et al., 2018; Li et al., 2018; Liu et al., 2020).
Revised: Nguyen et al., 2018; Li et al., 2018; Liu et al., 2020).

Original: Another reason are permutation and scaling invariances which lead to functionally identical minima that differ in the weight space (Brea et al., 2019; Entezari et al., 2021).
Revised: Another reason is the existence of permutation and scaling invariances which lead to functionally identical minima that differ in the weight space (Brea et al., 2019; Entezari et al., 2021).

Original: Due to the relationship of  loss landscape with generalization and optimization, a large body of recent work (Li et al., 2017; Mei et al., 2018; Geiger et al., 2019; Nguyen et al., 2018; Fort et al., 2019; S¸ims¸ek et al., 2021; Juneja et al., 2022) study theproperties of the loss landscape of DNNs with the goal to navigate the optimizer to agood region with desired properties, e.g. , with respect to ﬂatness around the SGD solution (Baldassi et al., 2020; Pittorino et al., 2020).
Revised: Due to the relationship of the loss landscape with generalization and optimization, a large body of recent works (Li et al., 2017; Mei et al., 2018; Geiger et al., 2019; Nguyen et al., 2018; Fort et al., 2019; S¸ims¸ek et al., 2021; Juneja et al., 2022) study the loss landscape of deep neural networks with the goal of navigating the optimizer to a region with desired properties, e.g. , with respect to ﬂatness around the SGD solution (Baldassi et al., 2020; Pittorino et al., 2020).

Original: Early works conjectured the existence of a non-linear path of non-increasing loss between solutions found by SGD (Freeman & Bruna, 2016; Draxler et al., 2018) and empirically showed how to ﬁnd it (Garipov et al., 2018;  Pittorino et al., 2022).
Revised: Early work conjectured the existence of a non-linear path of non-increasing loss between solutions found by SGD (Freeman and Bruna, 2016; Draxler et al., 2018) and empirically showed how to ﬁnd it (Garipov et al., 2018; Tatro et al., 2020; Pittorino et al., 2022).

Original: Recently, Entezari et al.
Revised: Recently, Entezari et al.

Original: (2021) conjectured the existence of  a linear path between SGD solutions if  permutation invariance of neural networks function space is taken into account.
Revised: (2021) conjectured the existence of such a linear path between SGD solutions if the permutation invariance of neural networks’ weight space is taken into account.

Revised: That is, with high probability over SGD solutions, for each pair of trained networks A and B there exists a permutation of the hidden units in each layer of B such that the linear path between A and the permuted network B (B’) is of non-increasing loss relative to the endpoints.

Original: The conjecture is important from both theoretical and empirical perspectives.
Revised: This conjecture is important from both theoretical and empirical perspectives.

Original: Theoretically, it leads to a drastic simpliﬁcation of the loss landscape, reducing the complexity obstacle fortheoretically analyzing deep neural networks.
Revised: Theoretically, it leads to a drastic simpliﬁcation of the loss landscape, reducing the complexity obstacle for analyzing deep neural networks.

Original: Empirically, linear interpolation between  network weights has become an important tool, having recently been used to set state of the art accuracy on ImageNet (Wortsman et al., 2022a), improve the robust accuracy of ﬁnetuned models (Wortsman et al., 2022b; Ilharco et al., 2022),and build effective weight-space model ensembles (Izmailov et al., 2019; Frankle et al., 2020; 
Revised: Empirically, linear interpolation between neural network weights has become an important tool, having recently been used to set state of the art accuracy on ImageNet (Wortsman et al., 2022a), improve robustness of ﬁnetuned models (Wortsman et al., 2022b; Ilharco et al., 2022), build effective weight-space model ensembles (Izmailov et al., 2019; Frankle et al., 2020; Guo et al., 2022), and constructively merge models trained on separate data splits (Wang et al., 2020; Ainsworth et al., 2022).

Original: Guo et al., 2022). New insights about the behavior of interpolated networks may lead tofurther empirical progress in utilizing linear interpolation  to achieve better results.
Revised: Therefore, any improvements toward reducing the obstacles to interpolation between trained models has the potential to lead to empirical progress in the above areas. Prior and concurrent works on linear interpolation (Singh and Jaggi, 2020; Entezari et al., 2021; Ainsworth et al., 2022) have focused on improving the algorithms used to bring the hidden units of two networks into alignment, in order to reduce the barrier to interpolation between them.

Original: Entezari et al.
Revised: Singh and

Original: (2021)’s conjecture has been shown to hold true in special cases both theoretically and empirically.
Revised: Jaggi (2020) develop a strong optimal transport-based method which allows linear interpolation between a pair of ResNet18 (He et al., 2016) networks such that the minimum accuracy attained along the path is 77%.

Original: In particular, the loss barrier approaches zero for sufﬁciently wide overparameterized shallow networks.
Revised: This constitutes a “barrier” of 16% relative to the original endpoint networks which achieve over 93% accuracy on the CIFAR-10 test set.

Original: S¸ims¸ek et al.
Revised: Entezari et al.

Original: (2021) show that adding one extra neuron to each MLP layer is sufﬁcient to connect previously discrete minima into a single manifold. Singh & Jaggi (2020) explore linear mode connectivity withsoft-alignment between ResNet18s, achieving a barrier of 14% .
Revised: (2021) use an approach based on simulated annealing (Zhan et al., 2016) in order to ﬁnd permutations such that wide multi-layer perceptrons (MLPs) (Rosenblatt, 1958) trained on MNIST (LeCun, 1998) can be linearly interpolated with a barrier of nearly zero.

Original: Pittorino et al.
Revised: Ainsworth et al.

Original: (2022) normalize networks by projecting them onto a hyper-toroidal manifold, which they ﬁnd reduces the barrier for VGG16 on CIFAR10 to 20% .
Revised: (2022) make the ﬁrst demonstration of zero-barrier connectivity between wide ResNets trained on CIFAR-10 by replacing 1 the standard Batch Normalization (Ioffe and Szegedy, 2015) layers with Layer Normalization (Ba et al., 2016), and develop several novel alignment methods.

Revised: Further discussion of related work can be found in Appendix A.

Original: Earlier works are unable to establish low-barrier linear mode connectivity between permuted solutions for standard deep networks.
Revised: In this paper we are interested in understanding why alignment of the endpoint networks alone has so far been insufﬁcient to reach low-barrier linear connectivity between standard deep networks.

Original:  In this work we investigate and resolve the cause of thetypically poor performance of  deep networks which have been  interpolated from independently-trained weights.
Revised: Contributions In this work we focus on understanding the source of the poor performance of standard deep networks (ResNet18, VGG11) whose weights have been linearly interpolated from between pairs of networks with aligned neurons.

Original: In summary, the contributionsof this paper are as follows:
Revised: Our contributions are as follows:

Original: • We empirically investigate the behavior of interpolated deep networks in order to determine why permutation invariance alone is insufﬁcient to allow low-barrier linear interpolation between SGD solutions to standard deep architectures.

Original:  We ﬁnd that  interpolated networks suffer from a phenomenon  in which their hidden units have signiﬁcantly smaller activation variance relative to the corresponding units of the trained networks.
Revised: • We ﬁnd that such interpolated networks suffer from a phenomenon of variance collapse in which their hidden units have signiﬁcantly smaller activation variance compared to the corresponding units of the original networks from which they were interpolated.

Original: We refer to thisphenomenon as variance collapse (Section 3).
Revised: We further identify and explain the reason behind this variance collapse. (Figure 1 (left) and Section 3).

Original: • We propose  a method which correctsfor variance collapse by rescaling hidden units in the interpolated network such that their statistics match those of the original trained networks (Section 4).
Revised: • We propose REPAIR (REnormalizing Permuted Activations for Interpolation Repair), a method that corrects variance collapse by rescaling hidden units in the interpolated network such that their statistics match those of the original networks.

Original: • We show that after applying this correction, it is possible to linearly interpolate between permuted SGD solutions to standard deep architectures with a low barrier. For example, we demonstrate (Figure 1) that when statistically corrected, an interpolated ResNet18 reaches 94.1% test accuracy on CIFAR-10, only 1.5% lower than the original trained networks.
Revised: (Section 4).

Original: To our knowledge, this is a 10 × improvement relative to previous approaches in terms of the size of the barrier to linear interpolation for a standard ResNet18. Wefurther demonstrate the utility of our method to  a  variety of other architectures and datasets (Section 5).
Revised: • We demonstrate that applying REPAIR to such interpolated networks leads to signiﬁcant barrier reductions across a wide variety of architectures, datasets, normalization techniques, and network width/depth (Section 5 and Figure 1 (middle and right)).

Original:  1. R ELATED W ORK
Revised:  P RELIMINARIES

Original: Neuron alignment is a technique of ﬁnding a match between neurons in the same layer of two neural networks, so that these networks are linear mode connected. Early works (Li et al., 2015) investigated the extent to which representations learned by independently trained networks converge to a similar set of features. They use neuron alignment based on linear correlation and estimated mutual information metrics. Since brute force is ineffective when searching for winning permutations due to a factorial time complexity, several recent works proposed heuristics for neuron alignment.
Revised: In this section we give preliminary deﬁnitions and algorithms to be used throughout the paper.

Original: Entezari et al. (2021) implements neuron alignment based on Simulated Annealing (Zhan et al., 2016) in order to remove the barrier between independently trained models. The method is successful only in a very limited number of cases. He et al.
Revised: 2. L INEAR INTERPOLATION OF NEURAL NETWORKS

Original: (2018) proposed a neuron alignment method in the context of multi-task model compression.
Revised: We consider the problem of interpolating between independently trained neural networks.

Original: The algorithm leverages layer-wise Hessian approximation to match neuron by computing a similarity measure based on their functional difference. Singh & Jaggi (2020) propose a model fusion algorithm which leverages optimal transport to perform neuron alignment in the Wasserstein space extensively studied by Agueh & Carlier (2011). The method allows building high-quality ensembles of models trained on heterogeneous datasets.
Revised: That is, if we let θ 1 , θ 2 be the weight vectors of two such networks, then we are interested in networks whose weights are of the form θ α

Original: Tatro et al. (2020) point out that special care must be taken for networks that contain batch normalization layers when it comes to neuron alignment. They match neurons by optimizing a curve in the weight space between two models by aligning correlations of post-activations. The running statistics are normalized by training the model for one epoch, while freezing all learnable parameters of the model.
Revised: =

Original: However, the crucial role of statistics reset in the context of linear interpolation of the models is not investigated. A recent work by Pittorino et al. (2022) uses LayerNorm as part of symmetry removal, so that the network behavior remains unchanged. Finally, concurrently to this work,
Revised: (1 − α )

Original: Ainsworth et al. (2022) develop a method which achieves zero-barrier interpolation between sufﬁciently wide ResNet models which have been modiﬁed to replace batch normalization with layer normalization. We note that this detail of the normalization layer was out of the scope of their study.  P RELIMINARIES 2. L INEAR I NTERPOLATION OF N EURAL N ETWORKS
Revised: θ 1 + α θ 2 for 0 < α < 1 .

Original: Let θ be a weight vector which parameterizes a neural network, and L ( θ ) be a given loss metric (e.g., train or test error). We consider the problem of interpolating between a pair of neural networkswith parameters θ 0 ,  θ 1  .
Revised: We refer to such networks θ α as interpolated networks , and to θ 1 , θ 2 as the endpoint networks .

Original: The loss barrier B ( θ 0 , θ 1 )  is deﬁned (Entezari et al., 2021) as the maximum increase in loss along the linear path from θ 0 to θ 1  , relative to the linear interpolation of the loss values at each of them.
Revised: The loss barrier B ( θ 1 , θ 2 ) between a pair of networks is deﬁned (Entezari et al., 2021) as the maximum increase in loss along the linear path between θ 1 and θ 2 , relative to the corresponding convex combination of the two endpoint losses. In notation:

Original: [Equation]
Revised: [Equation]

Original: Thetwo networks θ 0 and θ 1 are called linear mode connected if the barrier between them is roughly zero (Frankle et al., 2020). When referring to weights θ α produced by interpolating between the two weights θ 0 and θ 1 ,we refer to the network parameterized by θ α as the interpolated network, and the networks parameterized by θ 0 and θ 1 as the trained or parent networks.
Revised: The barrier between pairs of unmodiﬁed, independent SGD solutions θ 1 , θ 2 is typically large. We next discuss permutation invariance in neural networks.

Original: 2. P ERMUTATION I NVARIANCE
Revised: 2. P ERMUTATION INVARIANCE

Original: The class of neural networks with ReLU activation satisﬁes two invariance properties. Namely, scale invariance or unit-rescaling due to positive homogeneity of ReLU activations (Neyshabur et al., 2015) and the other is permutation of hidden units. Unit-rescaling has been well-studied and empirical evidence suggests that implicit bias of SGD would make the solution converge to a stage where the weights are more balanced (Neyshabur et al., 2015; Wu et al., 2019). Since permutations are equally likely for all SGD solutions, we follow theline of work that focuses on permutation invariance.
Revised: For typical neural architectures, the neurons in each layer can be permuted without functionally changing the network; this is known as the permutation invariance property.

Original: Permutation invariance is considered in form of permutations of hidden units in each layer of the network, i.e. , each layer i with parameters W i is replaced with P i W i Pi − 1 where P i is a permutation matrixand P l = P 0 is the identity matrix. For the case of networks with skip connections such as ResNets, care must be taken to apply a permutation in such a manner that the function represented by the network does not change.
Revised: In a simple feedforward network, this amounts to the observation that we can replace the L th weight matrix by PW L and the ( L + 1) th weight by W L +1 P − 1 for a permutation matrix P without changing the function represented by the network.

Original: In particular, the same permutation of hidden units must be applied to all convolutional layers which feed into a single residual stream.
Revised: Therefore, even if two networks θ 1 , θ 2 learn functionally identical sets of neurons at each layer, it is possible for their neurons to be arbitrarily permuted or misaligned.

Original: Entezari et al.
Revised: Entezari et al.

Original: (2021) conjecture that if permutation invariance is taken into account,  SGD solutions for a  network trained on the same training data become linear mode connected, i.e. , there will be no barrier between them.
Revised: (2021) conjecture that if permutation invariance is taken into account, then all SGD solutions for a sufﬁciently wide network trained on the same task become linearly mode connected. For completeness we provide the formal statement of the conjecture below. Conjecture 1.

Original: This conjecture is the subject of  the study performed in this paper.
Revised: (Entezari et al., 2021) For a given neural architecture, let P be the set of all valid permutations of hidden units, and P : R k ×P → R k be the function that applies a given permutation to a weight vector and returns the permuted version.

Original: [Equation]
Revised: Then for sufﬁciently wide networks, there exists a high-probability set of SGD solutions S ⊆ R k and a function Q : S → P such that for any θ 1 , θ 2 ∈ S , we have small interpolation barrier B ( P ( θ 1 , Q ( θ 1 )) , θ 2 ) ≈ 0 .

Original:  N EURON A LIGNMENT A LGORITHMS
Revised: 2. N EURON ALIGNMENT ALGORITHMS

Original: (Li et al., 2015) explore the notion that neurons of one trained network might be aligned with those of another. To perform this alignment for a given layer, they compute the correlation matrix between all pairs of hidden units ofthe two networks, and then solve the linear assignment problem corresponding to ﬁnding a permutation which maximizes the diagonal of this matrix.
Revised: A number of works have proposed methods of ﬁnding such alignments between the hidden units of a pair of neural networks.

Original: In notation, let
Revised: Li et al.

Original:  X (0) l,i and X (1) l,i be random variables corresponding to the  i -th hidden units(or channels depending on terminology) of the l -th [-layers in the networks parameterized by θ 
 and θ 1 , respectively, and where randomness is over the data distribution. Then-] we want to ﬁnd the permutation P l which maximizes the following the sum of correlations between matched units, i.e. , the following cost:
Revised: (2015) propose to maximize the sum of correlations between the activations of paired neurons across a batch of training data. That is, if we let X (0) l,i and X (1) l,i be random variables corresponding to the activations of the i -th hidden units of the l -th [-layers in the networks parameterized by θ 
 and θ 1 , respectively, and where randomness is over the data distribution. Then-]layer (across a batch of training data), then we optimize the permutation P l to minimize following cost:

Original: [Equation]
Revised: [Equation]

Original: In Figure 2, we display the outcome of applying this algorithm to the ﬁrst layer of convolutional ﬁlters in  two ResNet50s trained on ImageNet.
Revised: This amounts to a linear sum assignment problem corresponding to the matrix of correlations between pairs of hidden units in the two networks which can be solved via the Hungarian algorithm (Kuhn, 1955). Recent works have proposed alternative approaches:

Original: The alighment results in two sets of 64 7 × 7 convolutional ﬁlters, in which many of the matched pairs appear visually similar.
Revised: He et al.

Original: In the view of the conjecture, this alignment method is a tool for increasing linear mode connectivity of two networks.
Revised: (2018) compute a Hessian approximation to align functionally similar neurons, and Singh and Jaggi (2020) develop an optimal transport-based method of soft-alignment.

Original: Speciﬁcally, consider two parameterizations θ 0 and θ 1 of a network.
Revised: Ainsworth et al. (2022) compare three methods, including one based on Li et al. (2015) and two novel approaches. Tatro et al.

Original: Then applying the above alignment method will produce a modiﬁed version of θ 1 which has been permuted such that hidden units of the two networks are brought into alignment.
Revised: (2020) also perform alignment based on minimizing Equation 2, in order to reduce the barrier to non-linear interpolation. They show using a proximal alternating minimization scheme that alignments found using this method are nearly optimal for their purposes.

Original: The idea of the conjecture is that  network [-weights which are interpolated between θ 0 and this modiﬁed θ 1 should have better performance than those which are directly interpolated between θ 
-]
Revised: In this paper, we continue to use this alignment method which was originally introduced by Li et al. We note that for networks with residual connections, care must be taken to restrict the set of permutations such that the function represented by the network [-weights which are interpolated between θ 0 and this modiﬁed θ 1 should have better performance than those which are directly interpolated between θ 
-]does not change.

Original: and θ 1 . We next move on to understanding the behavior of such interpolated networks.
Revised: In particular, the same permutation of hidden units must be applied to all layers which feed into a single residual stream.

Original: C OLLAPSE OF  V ARIANCE  IN I NTERPOLATED N ETWORKS
Revised:  V ARIANCE COLLAPSE IN INTERPOLATED NETWORKS

Original: In this section we empirically investigate the source of the poor performance of interpolated deep networks. We  study two scenarios: multi-layer perceptrons (MLPs) (Rosenblatt, 1961) trained on 
Revised: We begin our study of interpolated networks by picking up where previous works left off: • The barrier between aligned networks decreases with width, including to nearly zero for very wide MLPs trained on MNIST.

Original: MNIST, and VGG11 (Simonyan & Zisserman, 2014) trained on CIFAR-10 (Krizhevsky et al., 2009). In Figure 3, we evaluate the effect of depth on barrier.
Revised: The barrier also increases sharply with depth, becoming large for MLPs or simple CNNs of more than a few layers (Entezari et al., 2021).

Original: We observe that although applying a suitable permutation to the weights is sufﬁcient toachieve nearly zero-barrier mode connectivity for shallower networks, past a certain depth even the permuted networks are no longer connected.
Revised: • Even strong optimal transport-based methods, which go far beyond alignment by allowing each neuron in network A to be matched to a weighted sum of neurons in network B, are insufﬁcient to achieve low-barrier (below 5% test-error) connectivity between standard ResNets (Singh and Jaggi, 2020).

Original: This reproduces the ﬁndings of  Entezari et al. 
Revised: We focus on understanding the source of this sharp increase in barrier between aligned networks as a function of depth observed in Entezari et al. (2021), which we hypothesize to be the same phenomenon that causes a high barrier to interpolation for standard ResNets. 3. I DENTIFYING THE PROBLEM : V ARIANCE C OLLAPSE What causes this rapid drop-off in the performance of interpolated networks which are deeper than a few layers?

Original: To  investigate the  behavior of interpolated deep MLPs, we focus on the statistics ofthe hidden units. For each hidden unit, we measure the variance of its activation across the MNIST training set, see Figure 4.
Revised: To answer this question, we investigate the internal behavior of such networks, focusing on the statistics of hidden units (Figure 2). We ﬁnd that for deep MLPs, interpolated from pairs of aligned endpoint networks which have high accuracy on the MNIST test-set, hidden units undergo a variance collapse .

Original: For each layer, we take the average of  the variances of all hidden units of that layer, and call this the layerwise variance.
Revised: That is, the variance of their activations progressively decays as we move deeper into the network, with activations in later layers becoming nearly constant. For each layer, we quantify this decay as follows.

Original: We plot in Figure 4 the ratio between the layerwise variances of the interpolated and parent networks for MLP and VGG architectures on MNIST and CIFARrespectively. For the parent network values, we take the layerwise variances to be the average of those of the two parents.
Revised: First, we measure the variance of each neuron across a batch of training data. Then we sum these variances across the neurons in the layer.

Original: To be precise, if we suppose that the l -th layers of the two parent networks have layerwise variances σ 20 and σ 21 , and that thel -th layer of the interpolated network has layerwise variance σ 2 α , then we report the ratio 
Revised: Finally, if we let this sum be denoted v α , v 1 , v 2 for the layer in the interpolated and two endpoint networks, respectively, then we return the ratio v α ( v 1 + v 2 ) / 2 . We report this ratio for each layer, yielding the sequence of values in Figure 2 (left), where we see variance collapse occur in MLPs, VGG networks, and ResNets. For an example of the set of per-neuron variances in a single layer of an interpolated ResNet18, see

Original: [Equation]
Revised: Figure 3 (left).

Original: We observe a striking trend in the statistics of the interpolated networks: by the ﬁnal layer, the activation variance within the directly interpolated network has become 10 4 times smaller than the activation variance of  the parent networks. We refer to this phenomenon as variance collapse . This progressive decrease in variance seems to occur at a roughly exponential rate, in that the relative variances within each layer of the interpolated networks are roughly a constant factor smaller than those of the previous layer.
Revised: We observe that variance decays to nearly zero by the ﬁnal layer of an interpolated 35-layer MLP, indicating that the activations in these last layers have become nearly constant. This effect seems to be further exacerbated when directly interpolating between unaligned networks.

Original: In the case of the VGG network, the trend is not quite so precise, but many layers still have their variance reduced by more than 10 × . In all cases, the application of the permutation to align hidden units before interpolating decreases but does not prevent the rate of collapse.
Revised: We repeat this experiment for VGG (Simonyan and Zisserman, 2014) and ResNet50 architectures, trained on CIFARand ImageNet respectively, and ﬁnd that variance by the ﬁnal layers decays by more than 10 × (Figure 2 (middle) and Figure 2 (right)). This is a problem: if these networks have nearly constant activations in their ﬁnal layers, then they will no longer even be able to differentiate between inputs.

Original: C ORRECTING THE A CTIVATION S TATISTICS
Revised: 3. W HY DOES THIS PHENOMENON OCCUR ?

Original: To correct for thisstatistical phenomenon (variance collapse), we develop the following method.
Revised: We argue that this phenomenon can be understood through the following statistical calculation.

Original: Let θ 0 , θ 1 be the weights we wish to interpolate between, with a suitable permutation already applied to θ 1 in order toalign the hidden units of the two networks.
Revised: Consider a hidden unit or channel in the ﬁrst layer of the interpolated network. Such a unit will be functionally equivalent to the linear interpolation between the respective units in the endpoint networks.

Original: Let α ∈ (0 , 1) be the desired interpolation coefﬁcient, such that the directly interpolated weights are θ α
Revised: That is, if we represent the unit’s preactivation by X α in the interpolated network, and X 1 , X 2 in the two endpoint networks (as random variables over the input data distribution), then the equality X α

Original: =
Revised: =

Original: (1 − α )
Revised: (1 − α )

Original: [-θ 
-]
Revised: [-θ 
-]X 1 + αX 2 holds. We will argue that the variance of X α is typically reduced as compared to that of X 1 or X 2 . If the two endpoint networks are perfectly aligned and have learned the same features, then we should have corr( X 1 , X 2 ) = 1 . But in practice, it is more typical for pairs of aligned units (whose alignment minimizes the cost function given by Equation 2) to have a correlation of corr( X 1 , X 2 )

Original: + αθ 1 .
Revised: ≈ 0 .

Original: As discussed in the previous section, we typically observe that the directly interpolated networks have poor performance, with hidden units in later layers having “collapsed” variance.
Revised: When considering the midpoint interpolated network ( α = 0 . 5 ), the variance of X α is given by [Equation] We typically have std( X 1 ) ≈ std

Revised: ( X 2 ) , so that this simpliﬁes to Var( X α )

Original: [Figure]{+= (0 . 
 + 0 . 
+}
Revised: {+= (0 . 
 + 0 . 
+}

Original: We want to compute a rescaled set of weights θ (cid:48) α in which this phenomenon is corrected.
Revised: · corr( X 1 , X 2 )) · Var( X 1 ) . With our typical value of corr( X 1 , X 2 ) ≈ 0 . 4 for pairs of neurons in aligned networks, this yields Var( X α )

Revised: = 0 .

Original: In particular, let σ 20 be  the pre-activation variance (w.r t. the training data) of some particular hidden unit in the network which is parameterized by θ 0 .
Revised: 7 · Var( X 1 ) , a 30% reduction compared to the endpoint networks. This analysis cannot be rigorously extended to deeper layers of the interpolated network, but intuitively we expect this decay to compound with depth. This intuition matches our experiments, where we see that variance collapse becomes worse as we progress through the layers of MLP, VGG, and ResNet50 networks (Figure 2).

Original: Let σ 21 be the variance of the same hidden unit in the network given by θ 1 .
Revised: We now propose a method which mitigates variance collapse by rescaling the preactivations of hidden units in interpolated networks. We call this method REPAIR (REnormalizing Permuted Activations for Interpolation Repair).

Original: Then we want the same unit in the network parameterized by θ (cid:48)
Revised: Given an interpolated network θ α = (1 − α ) · θ 1 + α · θ 2 for some 0 < α < 1 (between aligned endpoint networks θ 1 , θ 2 ), we select the set of hidden units or channels whose statistics we aim to correct. For example, for VGG networks we correct the outputs (preactivations) of every convolutional layer. For ResNets we correct both these convolutional outputs and the outputs of each residual block. Our goal will be to compute afﬁne (rescale-and-shift) coefﬁcients for every selected channel, such that the statistics of all selected channels are corrected.

Original: α to have pre-activation variance which is equal to the square{+analysis in the last section, let X 
 and X 2 be the values+} of the  interpolated standard deviation:
Revised: Let us consider a particular channel, e.g. , the 45th convolutional channel of the 8th layer in an interpolated ResNet18. Similar to the {+analysis in the last section, let X 
 and X 2 be the values+} of the channel in the two endpoint networks, viewed as random variables over the input training data, and let X α be the same channel in the interpolated network. Then after rescaling, we want the following two conditions to hold:

Original: [Equation]
Revised: [Equation]

Original:  We  propose  the following practical method for achieving this.
Revised: Whereas before any correction, we typically have std( X α ) (cid:28) min(std( X 1 ) , std( X 2 )) due to variance collapse. In the following, we present an algorithm which computes exact afﬁne coefﬁcients for each selected channel such that conditions (3), (4) become true. We also propose an approximate algorithm which avoids the use of forward passes in the interpolated network in Appendix C.

Original: First, we compute the mean and variance ofevery hidden unit in the parent networks θ 0 , θ 1 .
Revised: This algorithm depends ﬁrst upon the computation of the statistics E [ X 1 ] , E [ X 2 ] , std( X 1 ) , std( X 2 ) for each selected channel in the endpoint networks.

Original: Next, we modify the interpolated network by adding a new BatchNorm layer after every hidden convolutional or linear layer. In these added
Revised: We present a PyTorch-based approach to perform this computation in Appendix B.2. Given these values, the algorithm proceeds as follows. For each module in the interpolated network whose outputs we have identiﬁed as targets for statistical correction, we apply a wrapper, which adds a Batch Normalization layer after the wrapped module that is initially set to “train” mode.

Original: BatchNorm layers we also include a channelwise afﬁne transformation (this is the default setting in common learning libraries such as PyTorch and TensorFlow).
Revised: PyTorch pseudocode for this wrapper can also be found in Appendix B.2. Each such added BatchNorm layer contains afﬁne ( i.e. , per-channel rescale-and-shift) parameters.

Original: We set the biases and weights of this transformation to be equal tothe interpolated means (1 − α ) µ 0 + αµ 1  and standard deviations (1 − α ) σ 0 + ασ 1 .
Revised: For a given channel incoming to the BatchNorm layer, we set the respective afﬁne weight to (1 − α ) std( X 1 ) + α std( X 2 ) and the bias to (1 − α ) E [ X 1 ] + α E [ X 2 ] , where X 1 and Xare the respective channels in the endpoint networks as in conditions (3), (4).

Original: We then recompute thetracked normalization statistics of our added
Revised: With the added BatchNorm layers in training mode, these afﬁne parameters will exactly induce our statistical conditions with respect to batches of input data. The reason for this is that during execution, the added BatchNorm layers ﬁrst renormalize their inputs to have zero mean and unit variance per channel, and then apply the given afﬁne transformation which sets the statistics of the output to be that of conditions (3), (4). Next, we pass a set of training data through the network ( ∼ 5, examples sufﬁces) so that the running mean and variance parameters of these BatchNorm layers will be accurately estimated.

Original: BatchNorm layers; this can be done by a forward pass over the training data in “train mode” in the standard libraries. Finally, we perform “fusion” of the  BatchNorm layer back into the previous hidden convolutional or linear layer.
Revised: During this pass, any BatchNorm layers which already existed in the original network are kept frozen. Finally, we set the added BatchNorm layers to evaluation mode, so that they behave as afﬁne layers which do not recompute statistics. At this point, the resulting network is functionally equivalent to one in which the weights of our selected set of channels have been rescaled and biases shifted.

Original: This amounts to extracting rescaling and shift coefﬁcients from the  BatchNorm statistics and afﬁne parameters, which are respectively multiplied into the weights and  addedto the bias of the previous layer. We note that the technique of recomputing BatchNorm statistics in order toboost the accuracy of interpolated networks goes back to (Izmailov et al., 2018).
Revised: If we wish to generate a new parameter vector θ (cid:48) α which is compatible with the original network architecture ( i.e. , lacks these added BatchNorm layers), then we can perform BatchNorm layer fusion (Marku ˇ s, 2018) in which appropriate rescaling and biasshifting values are computed from each added BatchNorm layer, and then applied to the preceding convolutional ﬁlters.

Original: After this process, we are left with a new set of parameters θ (cid:48) α in which weights have been rescaled and biases shifted such that hidden units have activation statistics which match the desired interpolated values.
Revised: The networks resulting from this process have their internal statistics corrected, so that all selected channels satisfy conditions (3), (4). In Figure 2 and Figure 3 (left) we observe that REPAIR resolves variance collapse, bringing the statistics of the interpolated units back up to the level of the endpoint networks. In Figure 3 (left) we apply REPAIR to MLPs and in Figure 3 (right) we apply REPAIR to networks where the weights are linearly interpolated between a pair of ResNet18s whose hidden units have been brought into alignment. Before the correction, networks near the midpoint have a reduced accuracy of 66.0% on the CIFAR-10 test set, while the endpoints accuracy is 95.5%.

Original: R ESULTS
Revised: After correction, all checkpoints along the linear path have signiﬁcantly boosted accuracy, with the midpoint performing at 94.1%. In comparison, Singh and Jaggi (2020) report a linear midpoint accuracy of 77.0% using strong optimal-transport based alignment methods. Throughout the rest of the paper, we refer to the above method as REPAIR. We note that the utility of this method is orthogonal to any improvements in the algorithm used to align the hidden units of the endpoint networks.

Original: In this section we apply our correction method to various scenarios.
Revised: In the remainder of the paper, we demonstrate the effectiveness of REPAIR across a wide range of scenarios.

Original:  Speciﬁcally we look into MLP, ResNet and VGG architectures and cover MNIST, CIFAR-10, and ImageNet datasets.
Revised:  E XPERIMENTAL RESULTS

Original:  We investigate the effect of depth and width, as well as studying the choice of ResNet architecture normalization layer between BatchNorm (Ioffe & Szegedy, 2015) and LayerNorm (Ba et al., 2016)), as well as no normalization layers via Fixup (Zhang et al., 2019) initialization.
Revised: In this section we report the results of the following experiments. • We investigate the effectiveness of REPAIR for ResNets of varying width, trained on CIFAR-10 using three different normalization layers: standard BatchNorm, normalizationfree using FixUp (Zhang et al., 2019), and LayerNorm. (Section 5.1) • We apply REPAIR to the challenging scenario of interpolating between aligned ResNets trained on ImageNet.

Original: 5. MLP S
Revised: (Section 5.2) •

Original: First we establish the effect of our correction on the MLPs which were discussed in Section 3.
Revised: Finally, we perform an improved replication using REPAIR of an experiment from Ainsworth et al.

Original: In
Revised: (2022), in which a pair of models trained on disjoint subsets of data is constructively merged.

Original: Figure we plot the barrier as a function of number of layers (depth) for four different interpolation scenarios.
Revised: (Section 5.3) • We further explore the effect of width and depth across a variety of datasets. (Figure 5).

Original: vanilla refers to direct interpolation between two SGD solutions. permute refers to the barrier between two SGD solutions when we permute one of them such that the barrier is reduced; these ﬁrst two scenarios reproduce Figure 3. renorm refers to the renormalization correction of variance collapse discussed in Section 4 as applied to the directly interpolated networks and permute-renorm refers to using neuron alignment and renormalization together. We observe that both permutation and renormalization lead to barrier reduction separately and using them together reduces the barrier leading to almost zero barrier interpolation for a greater range of network depths.
Revised: Extra ﬁgures replicating these experiments in terms of alternate metrics can be found in Appendix D.

Original: 5.E FFECT OF N ORMALIZATION L AYER IN R ES N ETS
Revised: 5. N ORMALIZATION LAYER AND NETWORK WIDTH

Original: Next, we consider ResNet networks trained on CIFAR-10 across three different types of normalization layer used during training: BatchNorm, LayerNorm, and  Fixup initialization which does not use a normalization layer.
Revised: We begin by considering ResNet20s trained on CIFAR-10 using three different types of normalization layer: BatchNorm (Ioffe and Szegedy, 2015), LayerNorm (Ba et al., 2016), and normalizationfree via Fixup (Zhang et al., 2019).

Original: We combine this ablation with a study of the effect of width – for each class of networks, we vary the width multiplier from 1 × , where the ﬁnal block has 64 channels, to 16 × which gives the ﬁnal block 1024 channels.
Revised: We combine this ablation with a study on the effect of width. For each class of networks, we vary the width multiplier from 1 × , where the ﬁnal block haschannels, up to 16 × or 1024 channels.

Original: In Figure 6 we reportthe permuted interpolation barriers  for  three metrics: the training set (cross-entropy) loss, test set loss, and test set accuracy.
Revised: In Figure 4 we report barriers to interpolation, both with and without REPAIR. We ﬁnd that REPAIR is effective for all three cases, shrinking the loss barrier to below 0.05 for BatchNorm and LayerNorm-based networks when width is scaled to 16 × .

Original:  We ﬁnd that our correction method is able to signiﬁcantly reduce the barrier in all cases.
Revised: The LayerNorm-based ResNets are a special case that deserves comment. We ﬁrst note that LayerNorm is nonstandard in ResNets, and has worse performance when compared to BatchNorm. With our training conﬁguration, LayerNorm-based ResNet20s achieve 92.4% CIFAR-10 test set accuracy, vs. 93.6% for the BatchNorm baseline.

Original: For the case of LayerNorm, we ﬁnd that the interpolated networks  have reasonably high performance even without applying our statistical correction.
Revised: But these networks also have a unique property: when interpolating between aligned LayerNorm-ResNets, the midpoints already perform relatively well before REPAIR is applied, see Figure 4 (middle and right).

Original: We hypothesize that his has to do with the fact that LayerNorm  applies  renormalization at evaluation-time, which can partially correct the variance collapse.
Revised: The midpoint between our pair of 8 × -width LayerNorm-ResNet20s attains 90.3% test-set accuracy, which is boosted to 93.2% by an application of REPAIR. In contrast, for standard BatchNorm ResNets of the same width, the interpolated network obtains only 32.3% accuracy, which is boosted to 94.4% by REPAIR. We hypothesize that LayerNorm partially mitigates variance collapse because it applies test-time renormalization after each layer, which may prevent per-layer reductions in variance from compounding between layers.

Original: In contrast, BatchNorm layers do not compute fresh normalization statistics at evaluation time, instead using running mean and variance statistics that are computed during train time.
Revised: In contrast, BatchNorm is functionally just an afﬁne layer during testtime, so that variance reductions do compound.

Original: (Ainsworth et al., 2022) studied interpolation for the case of ResNet and VGG networks which use LayerNorm as their normalization layer. We note that our reported barriers for (uncorrected)
Revised: We claim that this observation regarding LayerNorm-based ResNets explains the contradiction between the results of Singh and Jaggi (2020) and Ainsworth et al. In the former, the authors develop a strong optimal-transport based method of aligning networks, which contains the approach of Li et al.

Original: LayerNorm interpolations are higher than theirs.
Revised: as a special case. They show that even this method is insufﬁcient to attain lowbarrier connectivity between standard ResNets before ﬁne-tuning; their best result is a barrier of 16% error between aligned ResNet18s trained on CIFAR-10. In contrast, Ainsworth et al.

Original: This may be due to a difference in LayerNorm implementation – we normalize with a single mean and variance computed across each entire layer of activations, but it is also common to normalize each spatial location separately.
Revised: (2022) report that a variety of alignment methods, including the method of Li et al. (2015) which they call “activation matching”, sufﬁce to establish nearly zero-barrier connectivity between wide ResNets trained on CIFAR-10. How can both results be true?

Original: Regardless, we note that with our proposed correction, the barriers for standard BatchNorm-based ResNets are lower across all widths than those reported by (Ainsworth et al., 2022) for uncorrected LayerNorm-based networks.
Revised: We claim that this contradiction is resolved by the observation that Ainsworth et al. (2022) replace standard BatchNorm layers with LayerNorm in their ResNets, as is evident in the code release 2 .

Original: InFigure 7, we report the results of applying our proposed correction to VGG networks of varying depth. In the left plot, the barrier is measured in terms of train and test loss. In the right plot, the barrier is measured interms of the respective accuracies.
Revised: In our experiments, the use of LayerNorm in

Original: We observe that the uncorrected interpolations have a large barrier which increases with the depth. In contrast, the barrier for interpolated networks with our correction applied stays roughly constant across all depths we consider.
Revised: ResNets uniquely allows for low-barrier linear connectivity without the requirement of a statistical correction such as REPAIR, at the cost of reduced performance on most datasets.

Original: 5. I MAGE N ET
Revised: 5. I MAGE N ET

Original: Next we explore the interpolation barriers of  ResNet models trained  on ImageNet (Deng et al., 2009).
Revised: Next, we explore the impact of REPAIR on the barrier to interpolation between standard ResNet models trained from scratch on ImageNet (Deng et al., 2009).

Original: We present interpolated accuracy curves of ResNet18 and  ResNet50architectures in Figure 
Revised: We test ResNet18, ResNet50, and a double-width variant of ResNet50 in Figure 6 (right).

Original: We observe that both the direct and permuted interpolations appear to have maximum barrier in  that the midpoints do not perform better than random chance.
Revised: Without REPAIR, the interpolated midpoints between each aligned pair of networks perform at below 1% (top-1) accuracy on the ImageNet validation set. After REPAIR, the midpoint ResNet18 improves to 41.1%, ResNet50 to 56.5%, and double-width ResNet50 to 64.2%. We ﬁnd in Figure 6 (left) that it is important to apply REPAIR not just to all convolutional layer outputs, but also to the outputs of every residual block.

Original: With thecorrection applied, performance at the midpoint increases to 26.4% and 27.3% respectively.
Revised: For the case of ResNet50, REPAIRing these these extra channels boosted the performance of the midpoint from 53.2% to 56.5%, which is a 14% reduction in the barrier (from 23.4% to 20.1%). We note that for architectures which contain BatchNorm after convolutional layers, applying REPAIR only to convolutional layer outputs is equivalent to resetting the BatchNorm statistics. Performing such a reset on averaged networks from a single training trajectory goes back to Izmailov et al. (2018), but

Original: To the best of our knowledge this is the ﬁrst method that leads to nonzero barrier reduction for the ImageNet dataset.
Revised: as far as we are aware, has not been applied to interpolations between independently trained networks until now. We provide a comparison between the ImageNet results of our work and those of Ainsworth et al.

Original: We note however that the barrier in this scenario is still high at 41.7% and 50.4% for ResNet18 and ResNet50 respectively.
Revised: (2022) in Appendix Figure 9. In general, the barriers for these architectures on ImageNet are still relatively high.

Original: As a future direction, it is of interest to look for extra interventions which may lead to this barrier being further reduced.
Revised: The standard ResNet50 architecture has a ﬁnal-block bottleneck width of 1024, and we measure the barrier after

Original:  In Figure 9, we visualize the relationship between test error and (corrected) barrierfor all of our models which were trained on CIFAR-10.
Revised: REPAIR to be 20.1% in terms of test error. The double-width ResNet50 variant has a ﬁnal-block bottleneck width of 2048, reducing the barrier to 12.9%. In comparison, the widest ResNet20 we studied on CIFAR-10 had ﬁnal-layer width of 1024, and a barrier of nearly zero.

Original: There is generally a relationship between the test error and size of the barrier in that our models with smaller test error also tend to have smaller barrier.
Revised: Therefore, it may be the case that for more difﬁcult datasets, larger widths are required in order to reach low barriers.

Original: Comparing this to Figure 4(c) of (Entezari et al., 2021), we note the same relationship between barrier values to test error.
Revised: [Figure] In this section we study a setting where the two endpoint networks are trained on disjoint splits of the training dataset. We aim to improve the corresponding experiment of Ainsworth et al. (2022), by using

Original: It is therefore reasonable to conclude that the high barrier of ImageNet models stems inpart simply from the relatively higher test error of said models, as compared to our models which were trained on CIFAR-10.
Revised: REPAIR applied to standard BatchNorm networks. We ﬁrst split the CIFAR-100 training set, consisting of 50,000 images distributed across 100 classes, into two disjoint sets of 25,000 images. The ﬁrst split contains a random 80% of the images in the ﬁrstclasses, and 20% of the second 50 classes; the second split has the proportions reversed. We then train one network on each split.

Original: C ONCLUSION AND F UTURE W ORK
Revised: The result is that the ﬁrst network is more accurate on the ﬁrst 50 classes, and the other more accurate on the second. Both networks underperform either their ensemble or a network trained on the full training set. We next align the hidden units of these two networks, and generate a series of weight-space interpolations, applying REPAIR to each. We ﬁnd that many of these interpolated networks signiﬁcantly outperform either of the two endpoints in terms of loss on the full CIFAR-100 test set (Figure 7).

Original: In this paper we investigated the source of the typically high barrier when interpolating between neuron-aligned deep networks.
Revised: In this sense, the two networks can be said to have been constructively merged. The best interpolated network of the corresponding experiment in Ainsworth et al. (2022) was reported to obtain a loss of 1.73 with mixing coefﬁcient α ≈ 0 . 3 .

Original: We observed a phenomenon of variance collapse in which the hidden units of interpolated networks have signiﬁcantly reduced variance relative to those of their trained parent networks.
Revised: Using REPAIR our best interpolated network achieves a loss of 1.30, also with α = 0 . We attribute this improvement partially to REPAIR, and partially to the increased performance of standard ResNets compared to the LayerNorm-based variants used in Ainsworth et al. In Appendix D, we compare these results against a strong baseline.

Original: We showed that this phenomenon is general and scales with depth, and proposed  a method of correcting for it by rescaling the activations of the  interpolated network.
Revised: In this paper we proposed REPAIR, a method of mitigating the variance collapse phenomenon that we identiﬁed in interpolated networks. We demonstrated that REPAIR improves the performance of interpolated networks across a wide variety of architectures and datasets. For example, we used

Original: We demonstrated experimentally that this correction signiﬁcantly improves the performance of intepolated networks, including reducing theinterpolation barrier  for a pair of ResNet18s trained on 
Revised: REPAIR to reduce the barrier to permuted interpolation for standard ResNet18s trained on CIFAR10 from 16% (Singh and Jaggi, 2020) to 1.5%, and to improve the performance of interpolated ResNet50s from below 1% to 56.5% on ImageNet. REPAIR is effective for networks trained with many choices of normalization, including BatchNorm, LayerNorm, and normalization-free.

Original: CIFAR-10 to 1.5% absolute test error.
Revised: To explain these results, we provided an analysis of the variance collapse phenomenon and how REPAIR mitigates it. We also demonstrated that LayerNorm-based ResNets are unique in attaining a relatively low barrier to aligned interpolation before the use of REPAIR, resolving the contradiction between the results of Singh and Jaggi (2020) and Ainsworth et al.

Original: Our work provides support for the conjecture of Entezari et al.
Revised: In so far as we establish low-barrier permuted interpolation for further scenarios, these results provide support for the conjecture of Entezari et al.

Original: by showing that once permutation invariance and  the statistical correction of variance collapse are taken into account, it is possible to move all SGD solutions into a single basin with a low barrier between all pairs.
Revised: On the other hand, to practically do so we needed to rescale the preactivations of the interpolated network, moving us out of the realm of strictly permuted linear interpolation. Our correction method REPAIR was developed as a generalization of the method of resetting BatchNorm statistics of averaged networks (Izmailov et al., 2018). Such a correction appears to be necessary in order to establish low-barrier permuted linear connectivity, at practical widths and without the use of test-time normalization layers.

Original: On the empirical side, we hope that a better understanding of the neural network loss landscape and in particular linear interpolation will lead to improvements in applications including  weight-space ensembling, checkpoint averaging, and robust ﬁnetuning. 
Revised: On the applications side, we demonstrated that REPAIRed interpolations between two networks which were trained on disjoint, biased dataset splits are able to outperform either network from which they were interpolated. We hope that further such applied results will be possible, including improvements to weight-space ensembling, checkpoint averaging, and robust ﬁnetuning. A PPENDIX A F URTHER DISCUSSION OF RELATED WORK Entezari et al.

Revised: (2021) try to ﬁnd a low-barrier permutation between two SGD solutions by searching in the set of all permutations using simulated annealing.

Revised: Neuron alignment methods use efﬁcient heuristics to ﬁnd a matching that maximizes a deﬁned similarity measure given two neural networks (Li et al., 2015; Ashmore and Gashler, 2015; Singh and Jaggi, 2020; Tatro et al., 2020; Pittorino et al., 2022).

Revised: This similarity measure is often based on correlation between weights or activations of the neurons in the same layer.

Original: R EPRODUCIBILITY AND E THICS S TATEMENT
Revised: Prior work was unable to achieve low-barrier for many standard architectures and/or challenging tasks ( e.g. , ImageNet classiﬁcation). He et al. (2018) propose a neuron alignment method in the context of multi-task model compression. The algorithm leverages layer-wise Hessian approximation to match neurons by computing a similarity measure based on their functional difference.

Revised: Singh and Jaggi (2020) propose a model fusion algorithm which leverages optimal transport to perform neuron alignment in the Wasserstein space (Agueh and Carlier, 2011), achieving a barrier of 16% for ResNet18 trained on CIFAR-10; furthermore they show that by ﬁnetuning such interpolated networks, the original performance is recovered.

Revised: Tatro et al.

Revised: (2020) point out that special care must be taken for networks that contain batch normalization layers when it comes to neuron alignment.

Revised: They match neurons by optimizing a curve in the weight space between two models by aligning correlations of post-activations.

Revised: The running statistics are normalized by training the model for one epoch, while freezing all learnable parameters of the model.

Revised: However, the crucial role of statistics reset in the context of linear interpolation of the models is not investigated.

Revised: A recent work by Pittorino et al.

Revised: (2022) uses LayerNorm as part of symmetry removal, so that the network behavior remains unchanged.

Original: All the datasets in this paper are public data and cited in the paper.
Revised: Table 1 summarizes the hyperparameters we

Revised: used to

Original: The code can be found at this anonymous link https://anonymous.4open.science/r/
Revised: train the neural networks which appear in this work

Revised: • We train all networks using SGD with momentum 0.9.

Revised: The weight decay and learning rates differ for each task, and are speciﬁed below.

Revised: • For our MLP trainings, we keep the hyperparameters below constant across varying widths, depths, and datasets.

Revised: • When training on MNIST and SVHN, we remove cutout and horizontal ﬂip from the list of data augmenations used by our ResNet20 training.

Revised: Otherwise, we keep the below ResNet hyperparameters constant across different choices of width, dataset, and normalization layer.

Revised: [Table]

Revised: P SEUDOCODE FOR P Y T ORCH MODULES

Revised: The ﬁrst step of REPAIR is to measure the statistics of identiﬁed channels in the endpoint networks.

Revised: There are many ways to do this, but we found that the following was efﬁcient and had low code-complexity in a Pytorch environment.

Revised: For each module in the interpolated network whose outputs we wish to REPAIR, we wrap the corresponding modules in the endpoint networks with the following:

Revised: class TrackLayer(nn.

Revised: Module): def __init__(self

Revised: [Equation]

Revised: layer):

Revised: super

Revised: [Equation]

Revised: [Equation]

Revised: __init

Revised: [Equation]

Revised: self

Revised: [Equation]

Revised: self

Revised: [Equation]

Revised: bn

Revised: [Equation]

Revised: self

Revised: [Equation]

Revised: self

Revised: [Equation]

Revised: eval

Revised: [Equation]

Revised: def get

Revised: [Equation]

Revised: return

Revised: [Equation]

Revised: running def

Revised: [Equation]

Revised: [Equation]

Revised: It then sufﬁces to pass a small set of training data ( ∼ 5,000 examples) through the tracked endpoint networks.

Revised: At this point, the statistics of each tracked module’s output can be retrieved from the wrapping TrackLayer.

Revised: Next, we wrap each module in the interpolated network that we wish to REPAIR with the following:

Original: Network-Permutations-2BEC/README.md
Revised: class ResetLayer(nn. Module): def __init__(self, layer): [Equation] def

Revised: [Equation]

Revised: The next step is to iterate over each triple of corresponding (TrackLayer, TrackLayer, ResetLayer) modules coming from the two endpoint networks and the interpolated network.

Revised: For each triple, we set the statistics of the ResetLayer to be the interpolation of the statistics from the two TrackLayers.

Revised: Working and minimal code to accomplish this can be found in our code release.

Revised: C C LOSED - FORM APPROXIMATE VARIANT OF REPAIR

Original:  We believe that our submission raises no questions regarding the Code of Ethics as we use open public datasets and model.
Revised: [Table] were trained [Table] We present an alternative correction algorithm which generates channel-wise afﬁne coefﬁcients using a closed-form approximation, eliminating the need for forward passes in the interpolated network.

Revised: This alternative is effective for MLP networks, but we found that it did not work well for deep convolutional architectures.

Revised: In Figure 8 we compare the afﬁne coefﬁcients produced by this variant with those of the primary method in Section 4, ﬁnding substantial correlation.

Revised: The variant proceeds as follows.

Revised: Consider a hidden unit in the ﬁrst layer of the interpolated network.

Revised: As in Section 4, let X α represent the unit in the interpolated network, and X 1 , X 2 the same unit in the two endpoint networks, respectively.

Revised: Condition (3) will already be satisﬁed for this unit by virtue of the equation

Revised: X α =

Revised: (1 − α )

Original: Our study also does not involve human subjects, potentially harmful insights, methodologies and applications, potential conﬂicts of interest and sponsorship, discrimination/bias/fairness concerns, privacy and security issues, legal compliance, and research integrity issues.
Revised: X 1 + α · X 2 (this holds only because the unit is in the ﬁrst layer). Given the values Var( X 1 ) , Var( X 2 ) , and Cov( X 1 , X 2 ) , which can be computed only using forward-passes in the endpoint networks, it is possible to compute the variance of X α exactly according to the formula [Equation]

Revised: Therefore, to satisfy condition (4) for this unit, the rescaling coefﬁcient β must be

Revised: [Equation]

Revised: which is simply the desired standard deviation divided by the standard deviation of X α .

Revised: For each unit in the ﬁrst layer, this factor is exactly correct in order to obtain the desired statistics.

Revised: For deeper layers, this factor is an approximation which we empirically test.

Revised: In Figure 3 (right), we apply this rescaling to every hidden unit of MLPs of depth between 2 andhidden layers, which are linearly interpolated ( α = 0 . 5 ) between aligned endpoints networks trained on MNIST.

Original: We also referred to close related works and compare our results wherever applicable
Revised: We ﬁnd that this rescaling signiﬁcantly improves the performance of such interpolated networks. In particular, we obtain interpolated checkpoints of up to 27 layers that achieve over 90% accuracy, whereas without a correction, we hit this limit after only 6 layers. For this case, the approximate rescaling performs similarly to the full exact method presented in Section 4. D A DDITIONAL PLOTS