Original: REPAIR: RE NORMALIZING P ERMUTED A CTIVATIONS FOR I NTERPOLATION R EPAIR
Revised: REPAIR: RE NORMALIZING P ERMUTED A CTIVATIONS FOR I NTERPOLATION R EPAIR

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

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

Original: (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.
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: 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.
Revised: First, we observe that neuron alignment methods alone are insufﬁcient to establish lowbarrier 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: Next, we propose REPAIR (REnormalizing Permuted Activations for Interpolation Repair) which mitigates variance collapse by rescaling the preactivations of such interpolated networks.
Revised: Next, we propose REPAIR (REnormalizing Permuted Activations for Interpolation Repair) which mitigates variance collapse by rescaling the preactivations of such interpolated networks.

Original: 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.
Revised: 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.

Original: In particular, we report a 74% barrier reduction for ResNet50 on ImageNet and 90% barrier reduction for ResNet on CIFAR10.
Revised: In particular, we report a 74% barrier reduction for ResNet50 on ImageNet and 90% barrier reduction for ResNet18 on CIFAR10.

Original: Our code is available at https://anonymous.4open.science/ r/Network-Permutations-2BEC/README.md .
Revised: Our code is available at https://github.com/KellerJordan/REPAIR .

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 that behave similarly on the training data (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 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).
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 the loss landscape with generalization and optimization, a large body of recent works (Li et al., 2017;
Revised: Due to the relationship of the loss landscape with generalization and optimization, a large body of recent works (Li et al., 2017;

Original: Mei et al., 2018; Geiger et al., 2019; Nguyen et al., 2018; Fort et al., 2019; S¸ims¸ek et al., 2021;
Revised: Mei et al., 2018; Geiger et al., 2019; Nguyen et al., 2018; Fort et al., 2019; S¸ims¸ek et al., 2021;

Original: 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).
Revised: 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 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).
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 such a linear path between SGD solutions if the permutation invariance of neural networks’ weight 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.

Original: 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.
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: This 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 for 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 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;
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;

Original: Guo et al., 2022), and constructively merge models trained on separate data splits (Wang et al., 2020; Ainsworth et al., 2022).
Revised: Guo et al., 2022), and constructively merge models trained on separate data splits (Wang et al., 2020; Ainsworth et al., 2022).

Original: Therefore, any improvements toward reducing the obstacles to interpolation between trained models has the potential to lead to empirical progress in the above areas.
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.

Original: Prior and concurrent works on linear interpolation (Singh and Jaggi, 2020; Entezari et al., 2021;
Revised: Prior and concurrent works on linear interpolation (Singh and Jaggi, 2020; Entezari et al., 2021;

Original: 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.
Revised: 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: Singh and
Revised: Singh and

Original: 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%.
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: This constitutes a “barrier” of 16% relative to the original endpoint networks which achieve over 93% accuracy on the CIFAR-10 test set.
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: Entezari et al.
Revised: Entezari et al.

Original: (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.
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: Ainsworth et al.
Revised: Ainsworth et al.

Original: (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: (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.

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

Original: 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.
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: 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.
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: Our contributions are as follows:
Revised: Our contributions are as follows:

Original: • 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.
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 further identify and explain the reason behind this variance collapse.
Revised: We further identify and explain the reason behind this variance collapse.

Original: (Figure 1 (left) and Section 3).
Revised: (Figure 1 (left) and Section 3).

Original: • 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.
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: (Section 4).
Revised: (Section 4).

Original: • 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)).
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:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this section we give preliminary deﬁnitions and algorithms to be used throughout the paper.
Revised: In this section we give preliminary deﬁnitions and algorithms to be used throughout the paper.

Original: 2. L INEAR INTERPOLATION OF NEURAL NETWORKS
Revised: 2. L INEAR INTERPOLATION OF NEURAL NETWORKS

Original: We consider the problem of interpolating between independently trained neural networks.
Revised: We consider the problem of interpolating between independently trained neural networks.

Original: 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 θ α
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: =
Revised: =

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

Original: θ 1 + α θ 2 for 0 < α < 1 .
Revised: θ 1 + α θ 2 for 0 < α < 1 .

Original: We refer to such networks θ α as interpolated networks , and to θ 1 , θ 2 as the endpoint networks .
Revised: We refer to such networks θ α as interpolated networks , and to θ 1 , θ 2 as the endpoint networks .

Original: 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.
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.

Original: In notation:
Revised: In notation:

Original: [Equation]
Revised: [Equation]

Original: The barrier between pairs of unmodiﬁed, independent SGD solutions θ 1 , θ 2 is typically large.
Revised: The barrier between pairs of unmodiﬁed, independent SGD solutions θ 1 , θ 2 is typically large.

Original: We next discuss permutation invariance in neural networks.
Revised: We next discuss permutation invariance in neural networks.

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

Original: 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.
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: 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.
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: 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.
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, then all SGD solutions for a sufﬁciently wide network trained on the same task become linearly mode connected.
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.

Original: For completeness we provide the formal statement of the conjecture below.
Revised: For completeness we provide the formal statement of the conjecture below.

Original: Conjecture 1.
Revised: Conjecture 1.

Original: (Entezari et al., 2021)
Revised: (Entezari et al., 2021)

Original: For a given neural architecture, let P be the set of all valid permutations of hidden units, and P : R k ×P
Revised: For a given neural architecture, let P be the set of all valid permutations of hidden units, and P : R k ×P

Original: →
Revised: →

Original: R k be the function that applies a given permutation to a weight vector and returns the permuted version.
Revised: R k be the function that applies a given permutation to a weight vector and returns the permuted version.

Original: Then for sufﬁciently wide networks, there exists a high-probability set of SGD solutions S ⊆ R k and a function Q : S
Revised: Then for sufﬁciently wide networks, there exists a high-probability set of SGD solutions S ⊆ R k and a function Q : S

Original: → P such that for any θ 1 , θ 2 ∈ S , we have small interpolation barrier B ( P ( θ 1 , Q ( θ 1 )) , θ 2 ) ≈ 0 .
Revised: → P such that for any θ 1 , θ 2 ∈ S , we have small interpolation barrier B ( P ( θ 1 , Q ( θ 1 )) , θ 2 ) ≈ 0 .

Original: 2. N EURON ALIGNMENT ALGORITHMS
Revised: 2. N EURON ALIGNMENT ALGORITHMS

Original: A number of works have proposed methods of ﬁnding such alignments between the hidden units of a pair of neural networks.
Revised: A number of works have proposed methods of ﬁnding such alignments between the hidden units of a pair of neural networks.

Original: Li et al.
Revised: Li et al.

Original: (2015) propose to maximize the sum of correlations between the activations of paired neurons across a batch of training data.
Revised: (2015) propose to maximize the sum of correlations between the activations of paired neurons across a batch of training data.

Original: 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 layer (across a batch of training data), then we optimize the permutation P l to minimize following cost:
Revised: 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 layer (across a batch of training data), then we optimize the permutation P l to minimize following cost:

Original: [Equation]
Revised: [Equation]

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

Original: Recent works have proposed alternative approaches:
Revised: Recent works have proposed alternative approaches:

Original: He et al.
Revised: He et al.

Original: (2018) compute a
Revised: (2018) compute a

Original: Hessian approximation to align functionally similar neurons, and Singh and Jaggi (2020) develop an optimal transport-based method of soft-alignment.
Revised: Hessian approximation to align functionally similar neurons, and Singh and Jaggi (2020) develop an optimal transport-based method of soft-alignment.

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

Original: (2022) compare three methods, including one based on Li et al.
Revised: (2022) compare three methods, including one based on Li et al.

Original: (2015) and two novel approaches.
Revised: (2015) and two novel approaches.

Original: Tatro et al.
Revised: Tatro et al.

Original: (2020) also perform alignment based on minimizing Equation 2, in order to reduce the barrier to non-linear interpolation.
Revised: (2020) also perform alignment based on minimizing Equation 2, in order to reduce the barrier to non-linear interpolation.

Original: They show using a proximal alternating minimization scheme that alignments found using this method are nearly optimal for their purposes.
Revised: They show using a proximal alternating minimization scheme that alignments found using this method are nearly optimal for their purposes.

Original: In this paper, we continue to use this alignment method which was originally introduced by Li et al.
Revised: In this paper, we continue to use this alignment method which was originally introduced by Li et al.

Original: 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 does not change.
Revised: 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 does not change.

Original: In particular, the same permutation of hidden units must be applied to all layers which feed into a single residual stream.
Revised: In particular, the same permutation of hidden units must be applied to all layers which feed into a single residual stream.

Original:  V ARIANCE COLLAPSE IN INTERPOLATED NETWORKS
Revised:  V ARIANCE COLLAPSE IN INTERPOLATED NETWORKS

Original: We begin our study of interpolated networks by picking up where previous works left off:
Revised: We begin our study of interpolated networks by picking up where previous works left off:

Original: • The barrier between aligned networks decreases with width, including to nearly zero for very wide MLPs trained on MNIST.
Revised: • The barrier between aligned networks decreases with width, including to nearly zero for very wide MLPs trained on MNIST.

Original: The barrier also increases sharply with depth, becoming large for MLPs or simple CNNs of more than a few layers (Entezari et al., 2021).
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: • 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).
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: 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.
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.

Original: (2021), which we hypothesize to be the same phenomenon that causes a high barrier to interpolation for standard ResNets.
Revised: (2021), which we hypothesize to be the same phenomenon that causes a high barrier to interpolation for standard ResNets.

Original: 3. I DENTIFYING THE PROBLEM : V ARIANCE C OLLAPSE
Revised: 3. I DENTIFYING THE PROBLEM : V ARIANCE C OLLAPSE

Original: What causes this rapid drop-off in the performance of interpolated networks which are deeper than a few layers?
Revised: What causes this rapid drop-off in the performance of interpolated networks which are deeper than a few layers?

Original: To answer this question, we investigate the internal behavior of such networks, focusing on the statistics of hidden units (Figure 2).
Revised: To answer this question, we investigate the internal behavior of such networks, focusing on the statistics of hidden units (Figure 2).

Original: 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 .
Revised: 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: That is, the variance of their activations progressively decays as we move deeper into the network, with activations in later layers becoming nearly constant.
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.

Original: For each layer, we quantify this decay as follows.
Revised: For each layer, we quantify this decay as follows.

Original: First, we measure the variance of each neuron across a batch of training data.
Revised: First, we measure the variance of each neuron across a batch of training data.

Original: Then we sum these variances across the neurons in the layer.
Revised: Then we sum these variances across the neurons in the layer.

Original: 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 .
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 .

Original: 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.
Revised: 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.

Original: For an example of the set of per-neuron variances in a single layer of an interpolated ResNet18, see
Revised: For an example of the set of per-neuron variances in a single layer of an interpolated ResNet18, see

Original: Figure 3 (left).
Revised: Figure 3 (left).

Original: 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.
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.

Original: This effect seems to be further exacerbated when directly interpolating between unaligned networks.
Revised: This effect seems to be further exacerbated when directly interpolating between unaligned networks.

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

Original: 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.
Revised: 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: 3. W HY DOES THIS PHENOMENON OCCUR ?
Revised: 3. W HY DOES THIS PHENOMENON OCCUR ?

Original: We argue that this phenomenon can be understood through the following statistical calculation.
Revised: We argue that this phenomenon can be understood through the following statistical calculation.

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

Original: Such a unit will be functionally equivalent to the linear interpolation between the respective units in the endpoint networks.
Revised: Such a unit will be functionally equivalent to the linear interpolation between the respective units in the endpoint networks.

Original: That is, if we represent the unit’s preactivation by X α in the interpolated network, and
Revised: That is, if we represent the unit’s preactivation by X α in the interpolated network, and

Original: X 1 , X 2 in the two endpoint networks (as random variables over the input data distribution), then the equality X α
Revised: 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: X 1 + αX 2 holds.
Revised: X 1 + αX 2 holds.

Original: We will argue that the variance of X α is typically reduced as compared to that of X 1 or X 2 .
Revised: We will argue that the variance of X α is typically reduced as compared to that of X 1 or X 2 .

Original: If the two endpoint networks are perfectly aligned and have learned the same features, then we should have corr( X 1 , X 2 ) = 1 .
Revised: If the two endpoint networks are perfectly aligned and have learned the same features, then we should have corr( X 1 , X 2 ) = 1 .

Original: 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 )
Revised: 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: ≈ 0 .
Revised: ≈ 0 .

Original: When considering the midpoint interpolated network ( α = 0 . 5 ), the variance of X α is given by
Revised: When considering the midpoint interpolated network ( α = 0 . 5 ), the variance of X α is given by

Original: [Equation]
Revised: [Equation]

Original: We typically have std( X 1 )
Revised: We typically have std( X 1 )

Original: ≈ std
Revised: ≈ std

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

Original: = (0 . 

Revised: = (0 .

Original: + 0 . 

Revised: + 0 .

Original: · corr( X 1 , X 2 ))
Revised: · corr( X 1 , X 2 ))

Original: · Var( X 1 ) .
Revised: · Var( X 1 ) .

Original: With our typical value of corr( X 1 , X 2 ) ≈ 0 .
Revised: With our typical value of corr( X 1 , X 2 ) ≈ 0 .

Original: 4 for pairs of neurons in aligned networks, this yields Var( X α )
Revised: 4 for pairs of neurons in aligned networks, this yields Var( X α )

Original: = 0 .
Revised: = 0 .

Original: 7 · Var( X 1 ) , a 30% reduction compared to the endpoint networks.
Revised: 7 · Var( X 1 ) , a 30% reduction compared to the endpoint networks.

Original: This analysis cannot be rigorously extended to deeper layers of the interpolated network, but intuitively we expect this decay to compound with depth.
Revised: This analysis cannot be rigorously extended to deeper layers of the interpolated network, but intuitively we expect this decay to compound with depth.

Original: This intuition matches our experiments, where we see that variance collapse becomes worse as we progress through the layers of
Revised: This intuition matches our experiments, where we see that variance collapse becomes worse as we progress through the layers of

Original: MLP, VGG, and ResNet50 networks (Figure 2).
Revised: MLP, VGG, and ResNet50 networks (Figure 2).

Original:  REPAIR
Revised:  REPAIR

Original: We now propose a method which mitigates variance collapse by rescaling the preactivations of hidden units in interpolated networks.
Revised: We now propose a method which mitigates variance collapse by rescaling the preactivations of hidden units in interpolated networks.

Original: We call this method REPAIR (REnormalizing Permuted
Revised: We call this method REPAIR (REnormalizing Permuted

Original: Activations for Interpolation Repair).
Revised: Activations for Interpolation Repair).

Original: 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.
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.

Original: For example, for VGG networks we correct the outputs (preactivations) of every convolutional layer.
Revised: For example, for VGG networks we correct the outputs (preactivations) of every convolutional layer.

Original: For ResNets we correct both these convolutional outputs and the outputs of each residual block.
Revised: For ResNets we correct both these convolutional outputs and the outputs of each residual block.

Original: 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.
Revised: 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: Let us consider a particular channel, e.g. , the 45th convolutional channel of the 8th layer in an interpolated ResNet18.
Revised: Let us consider a particular channel, e.g. , the 45th convolutional channel of the 8th layer in an interpolated ResNet18.

Original: Similar to the analysis in the last section, let X 

Revised: Similar to the analysis in the last section, let X

Original: 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.
Revised: 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.

Original: Then after rescaling, we want the following two conditions to hold:
Revised: Then after rescaling, we want the following two conditions to hold:

Original: [Equation]
Revised: [Equation]

Original: Whereas before any correction, we typically have std( X α ) (cid:28) min(std( X 1 ) , std( X 2 )) due to variance collapse.
Revised: Whereas before any correction, we typically have std( X α ) (cid:28) min(std( X 1 ) , std( X 2 )) due to variance collapse.

Original: 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.
Revised: 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.

Original: We also propose an approximate algorithm which avoids the use of forward passes in the interpolated network in Appendix C.
Revised: We also propose an approximate algorithm which avoids the use of forward passes in the interpolated network in Appendix C.

Original: This algorithm depends ﬁrst upon the computation of the statistics
Revised: This algorithm depends ﬁrst upon the computation of the statistics

Original: E
Revised: E

Original: [ X 1 ] , E [ X 2 ] , std( X 1 ) , std( X 2 ) for each selected channel in the endpoint networks.
Revised: [ X 1 ] , E [ X 2 ] , std( X 1 ) , std( X 2 ) for each selected channel in the endpoint networks.

Original: We present a PyTorch-based approach to perform this computation in Appendix B.2.
Revised: We present a PyTorch-based approach to perform this computation in Appendix B.2.

Original: Given these values, the algorithm proceeds as follows.
Revised: Given these values, the algorithm proceeds as follows.

Original: 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.
Revised: 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: PyTorch pseudocode for this wrapper can also be found in Appendix B.2.
Revised: PyTorch pseudocode for this wrapper can also be found in Appendix B.2.

Original: Each such added BatchNorm layer contains afﬁne ( i.e. , per-channel rescale-and-shift) parameters.
Revised: Each such added BatchNorm layer contains afﬁne ( i.e. , per-channel rescale-and-shift) parameters.

Original: 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).
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: 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).
Revised: With the added BatchNorm layers in training mode, during execution they ﬁrst renormalize their inputs to have zero mean and unit variance per channel, and then apply the given afﬁne transformation, such that our statistical conditions (3), (4) are exactly induced with respect to each batch of input data.

Original: 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.
Revised: Next, we pass a set of training data through the network ( ∼ 5,000 examples sufﬁces) so that the running mean and variance parameters of these BatchNorm layers will be accurately estimated.

Original: During this pass, any BatchNorm layers which already existed in the original network are kept frozen.
Revised: During this pass, any BatchNorm layers which already existed in the original network are kept frozen.

Original: Finally, we set the added BatchNorm layers to evaluation mode, so that they behave as afﬁne layers which do not recompute statistics.
Revised: Finally, we set the added BatchNorm layers to evaluation mode, so that they behave as afﬁne layers which do not recompute statistics.

Original: 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.
Revised: 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: If we wish to generate a new parameter vector θ (cid:48)
Revised: If we wish to generate a new parameter vector θ (cid:48)

Original: α 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.
Revised: α 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 bias-shifting values are computed from each added BatchNorm layer, and then applied to the preceding convolutional ﬁlters.

Original: The networks resulting from this process have their internal statistics corrected, so that all selected channels satisfy conditions (3), (4).
Revised: The networks resulting from this process have their internal statistics corrected, so that all selected channels satisfy conditions (3), (4).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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%.
Revised: 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: After correction, all checkpoints along the linear path have signiﬁcantly boosted accuracy, with the midpoint performing at 94.1%.
Revised: After correction, all checkpoints along the linear path have signiﬁcantly boosted accuracy, with the midpoint performing at 94.1%.

Original: In comparison, Singh and Jaggi (2020) report a linear midpoint accuracy of 77.0% using strong optimal-transport based alignment methods.
Revised: In comparison, Singh and Jaggi (2020) report a linear midpoint accuracy of 77.0% using strong optimal-transport based alignment methods.

Original: Throughout the rest of the paper, we refer to the above method as REPAIR.
Revised: Throughout the rest of the paper, we refer to the above method as REPAIR.

Original: 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.
Revised: 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.

Revised: We explore this in Appendix Figure 15, where we report barrier curves for REPAIR applied to interpolations between unaligned networks as well.

Original: In the remainder of the paper, we demonstrate the effectiveness of REPAIR across a wide range of scenarios.
Revised: In the remainder of the paper, we demonstrate the effectiveness of REPAIR across a wide range of scenarios.

Original:  E XPERIMENTAL RESULTS
Revised:  E XPERIMENTAL RESULTS

Original: In this section we report the results of the following experiments.
Revised: In this section we report the results of the following experiments.

Original: • 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.
Revised: • 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.

Original: (Section 5.1) •
Revised: (Section 5.1) •

Original: We apply REPAIR to the challenging scenario of interpolating between aligned ResNets trained on ImageNet.
Revised: We apply REPAIR to the challenging scenario of interpolating between aligned ResNets trained on ImageNet.

Original: (Section 5.2) •
Revised: (Section 5.2) •

Original: Finally, we perform an improved replication using REPAIR of an experiment from Ainsworth et al.
Revised: Finally, we perform an improved replication using REPAIR of an experiment from Ainsworth et al.

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

Original: (Section 5.3) • We further explore the effect of width and depth across a variety of datasets.
Revised: (Section 5.3) • We further explore the effect of width and depth across a variety of datasets.

Original: (Figure 5).
Revised: (Figure 5).

Original: Extra ﬁgures replicating these experiments in terms of alternate metrics can be found in Appendix D.
Revised: Extra ﬁgures replicating these experiments in terms of alternate metrics can be found in Appendix D.

Original: 5. N ORMALIZATION LAYER AND NETWORK WIDTH
Revised: 5. N ORMALIZATION LAYER AND NETWORK WIDTH

Original: 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).
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 on the effect of width.
Revised: We combine this ablation with a study on the effect of width.

Original: For each class of networks, we vary the width multiplier from 1 × , where the ﬁnal block haschannels, up to 16 × or 1024 channels.
Revised: 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 4 we report barriers to interpolation, both with and without REPAIR.
Revised: In Figure 4 we report barriers to interpolation, both with and without REPAIR.

Original: 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 × .
Revised: 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: The LayerNorm-based ResNets are a special case that deserves comment.
Revised: The LayerNorm-based ResNets are a special case that deserves comment.

Original: We ﬁrst note that LayerNorm is nonstandard in ResNets, and has worse performance when compared to BatchNorm.
Revised: We ﬁrst note that LayerNorm is nonstandard in ResNets, and has worse performance when compared to BatchNorm.

Original: With our training conﬁguration, LayerNorm-based ResNet20s achieve 92.4% CIFAR-10 test set accuracy, vs. 93.6% for the BatchNorm baseline.
Revised: With our training conﬁguration, LayerNorm-based ResNet20s achieve 92.4% CIFAR-10 test set accuracy, vs. 93.6% for the BatchNorm baseline.

Original: But these networks also have a unique property: when interpolating between aligned LayerNorm-ResNets, the midpoints already perform relatively well before
Revised: But these networks also have a unique property: when interpolating between aligned LayerNorm-ResNets, the midpoints already perform relatively well before

Original: REPAIR is applied, see Figure 4 (middle and right).
Revised: REPAIR is applied, see Figure 4 (middle and right).

Original: The midpoint between our pair of 8 × -width
Revised: The midpoint between our pair of 8 × -width

Original: LayerNorm-ResNet20s attains 90.3% test-set accuracy, which is boosted to 93.2% by an application of REPAIR.
Revised: LayerNorm-ResNet20s attains 90.3% test-set accuracy, which is boosted to 93.2% by an application of REPAIR.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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 is functionally just an afﬁne layer during testtime, so that variance reductions do compound.
Revised: In contrast, BatchNorm is functionally just an afﬁne layer during testtime, so that variance reductions do compound.

Original: We claim that this observation regarding LayerNorm-based ResNets explains the contradiction between the results of Singh and Jaggi (2020) and Ainsworth et al.
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.

Original: In the former, the authors develop a strong optimal-transport based method of aligning networks, which contains the approach of Li et al.
Revised: In the former, the authors develop a strong optimal-transport based method of aligning networks, which contains the approach of Li et al.

Original: as a special case.
Revised: as a special case.

Original: 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.
Revised: 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.

Original: In contrast, Ainsworth et al.
Revised: In contrast, Ainsworth et al.

Original: (2022) report that a variety of alignment methods, including the method of Li et al.
Revised: (2022) report that a variety of alignment methods, including the method of Li et al.

Original: (2015) which they call “activation matching”, sufﬁce to establish nearly zero-barrier connectivity between wide ResNets trained on CIFAR-10.
Revised: (2015) which they call “activation matching”, sufﬁce to establish nearly zero-barrier connectivity between wide ResNets trained on CIFAR-10.

Original: How can both results be true?
Revised: How can both results be true?

Original: We claim that this contradiction is resolved by the observation that Ainsworth et al.
Revised: We claim that this contradiction is resolved by the observation that Ainsworth et al.

Original: (2022) replace standard BatchNorm layers with LayerNorm in their ResNets, as is evident in the code release 2 .
Revised: (2022) replace standard BatchNorm layers with LayerNorm in their ResNets, as is evident in the code release 2 .

Original: In our experiments, the use of LayerNorm in
Revised: In our experiments, the use of LayerNorm in

Original: 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.
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 impact of REPAIR on the barrier to interpolation between standard ResNet models trained from scratch 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 test ResNet18, ResNet50, and a double-width variant of ResNet50 in Figure 6 (right).
Revised: We test ResNet18, ResNet50, and a double-width variant of ResNet50 in Figure 6 (right).

Original: Without REPAIR, the interpolated midpoints between each aligned pair of networks perform at below 1% (top-1) accuracy on the ImageNet validation set.
Revised: Without REPAIR, the interpolated midpoints between each aligned pair of networks perform at below 1% (top-1) accuracy on the ImageNet validation set.

Original: After REPAIR, the midpoint ResNet18 improves to 41.1%, ResNet50 to 56.5%, and double-width ResNet50 to 64.2%.
Revised: After REPAIR, the midpoint ResNet18 improves to 41.1%, ResNet50 to 56.5% and double-width ResNet50 to 64.2%.

Original: 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.
Revised: 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: 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%).
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%).

Original: 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.
Revised: 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.

Original: Performing such a reset on averaged networks from a single training trajectory goes back to Izmailov et al.
Revised: Performing such a reset on averaged networks from a single training trajectory goes back to Izmailov et al.

Original: (2018), but
Revised: (2018), but

Original: as far as we are aware, has not been applied to interpolations between independently trained networks until now.
Revised: as far as we are aware, has not been applied to interpolations between independently trained networks until now.

Original: We provide a comparison between the ImageNet results of our work and those of Ainsworth et al.
Revised: We provide a comparison between the ImageNet results of our work and those of Ainsworth et al.

Original: (2022) in
Revised: (2022) in

Original: Appendix Figure 9.
Revised: Appendix Figure 9.

Original: In general, the barriers for these architectures on ImageNet are still relatively high.
Revised: In general, the barriers for these architectures on ImageNet are still relatively high.

Original: The standard
Revised: The standard

Original: ResNet50 architecture has a ﬁnal-block bottleneck width of 1024, and we measure the barrier after
Revised: ResNet50 architecture has a ﬁnal-block bottleneck width of 1024, and we measure the barrier after

Original: REPAIR to be 20.1% in terms of test error.
Revised: REPAIR to be 20.1% in terms of test error.

Original: The double-width ResNet50 variant has a ﬁnal-block bottleneck width of 2048, reducing the barrier to 12.9%.
Revised: The double-width ResNet50 variant has a ﬁnal-block bottleneck width of 2048, reducing the barrier to 12.9%.

Original: In comparison, the widest ResNet20 we studied on CIFAR-10 had ﬁnal-layer width of 1024, and a barrier of nearly zero.
Revised: In comparison, the widest ResNet20 we studied on CIFAR-10 had ﬁnal-layer width of 1024, and a barrier of nearly zero.

Original: Therefore, it may be the case that for more difﬁcult datasets, larger widths are required in order to reach low barriers.
Revised: Therefore, it may be the case that for more difﬁcult datasets, larger widths are required in order to reach low barriers.

Original: [Figure]
Revised: 5. S PLIT DATA TRAINING

Original: In this section we study a setting where the two endpoint networks are trained on disjoint splits of the training dataset.
Revised: In this section we study a setting where the two endpoint networks are trained on disjoint splits of the training dataset.

Original: We aim to improve the corresponding experiment of Ainsworth et al.
Revised: We aim to improve the corresponding experiment of Ainsworth et al.

Original: (2022), by using
Revised: (2022), by using

Original: REPAIR applied to standard BatchNorm networks.
Revised: REPAIR applied to standard BatchNorm networks.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We then train one network on each split.
Revised: We then train one network on each split.

Original: The result is that the ﬁrst network is more accurate on the ﬁrst 50 classes, and the other more accurate on the second.
Revised: The result is that the ﬁrst network is more accurate on the ﬁrst 50 classes, and the other more accurate on the second.

Original: Both networks underperform either their ensemble or a network trained on the full training set.
Revised: Both networks underperform either their ensemble or a network trained on the full training set.

Original: We next align the hidden units of these two networks, and generate a series of weight-space interpolations, applying REPAIR to each.
Revised: We next align the hidden units of these two networks, and generate a series of weight-space interpolations, applying REPAIR to each.

Original: 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).
Revised: 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 sense, the two networks can be said to have been constructively merged.
Revised: In this sense, the two networks can be said to have been constructively merged.

Original: The best interpolated network of the corresponding experiment in Ainsworth et al.
Revised: The best interpolated network of the corresponding experiment in Ainsworth et al.

Original: (2022) was reported to obtain a loss of 1.73 with mixing coefﬁcient α ≈ 0 . 3 .
Revised: (2022) was reported to obtain a loss of 1.73 with mixing coefﬁcient α ≈ 0 . 3 .

Original: Using REPAIR our best interpolated network achieves a loss of 1.30, also with α = 0 .
Revised: Using REPAIR our best interpolated network achieves a loss of 1.30, also with α = 0 .

Original: We attribute this improvement partially to REPAIR, and partially to the increased performance of standard ResNets compared to the LayerNorm-based variants used in
Revised: We attribute this improvement partially to REPAIR, and partially to the increased performance of standard ResNets compared to the LayerNorm-based variants used in

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

Original: In Appendix D, we compare these results against a strong baseline.
Revised: In Appendix D, we compare these results against a strong baseline.

Original:  D ISCUSSION AND FUTURE WORK
Revised:  D ISCUSSION AND FUTURE WORK

Original: In this paper we proposed REPAIR, a method of mitigating the variance collapse phenomenon that we identiﬁed in interpolated networks.
Revised: In this paper we proposed REPAIR, a method of mitigating the variance collapse phenomenon that we identiﬁed in interpolated networks.

Original: We demonstrated that REPAIR improves the performance of interpolated networks across a wide variety of architectures and datasets.
Revised: We demonstrated that REPAIR improves the performance of interpolated networks across a wide variety of architectures and datasets.

Original: For example, we used
Revised: For example, we used

Original: 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
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

Original: ResNet50s from below 1% to 56.5% on ImageNet.
Revised: ResNet50s from below 1% to 56.5% on ImageNet.

Original: REPAIR is effective for networks trained with many choices of normalization, including BatchNorm, LayerNorm, and normalization-free.
Revised: REPAIR is effective for networks trained with many choices of normalization, including BatchNorm, LayerNorm, and normalization-free.

Original: To explain these results, we provided an analysis of the variance collapse phenomenon and how
Revised: To explain these results, we provided an analysis of the variance collapse phenomenon and how

Original: REPAIR mitigates it.
Revised: REPAIR mitigates it.

Original: 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.
Revised: 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: In so far as we establish low-barrier permuted interpolation for further scenarios, these results provide 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: 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.
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.

Original: Our correction method REPAIR was developed as a generalization of the method of resetting BatchNorm statistics of averaged networks (Izmailov et al., 2018).
Revised: Our correction method REPAIR was developed as a generalization of the method of resetting BatchNorm statistics of averaged networks (Izmailov et al., 2018).

Original: 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.
Revised: 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 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.
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.

Original: We hope that further such applied results will be possible, including improvements to weight-space ensembling, checkpoint averaging, and robust ﬁnetuning.
Revised: We hope that further such applied results will be possible, including improvements to weight-space ensembling, checkpoint averaging, and robust ﬁnetuning.

Revised: A CKNOWLEDGMENTS

Revised: Thanks to Luke Johnston for his insightful comments on the manuscript.

Revised: This collaboration was facilitated by ML Collective

Original: A PPENDIX A F URTHER DISCUSSION OF RELATED WORK
Revised: A PPENDIX A F URTHER DISCUSSION OF RELATED WORK

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

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

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

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

Original: Prior work was unable to achieve low-barrier for many standard architectures and/or challenging tasks ( e.g. , ImageNet classiﬁcation).
Revised: Prior work was unable to achieve low-barrier for many standard architectures and/or challenging tasks ( e.g. , ImageNet classiﬁcation).

Original: He et al.
Revised: He et al.

Original: (2018) propose a neuron alignment method in the context of multi-task model compression.
Revised: (2018) propose a neuron alignment method in the context of multi-task model compression.

Original: The algorithm leverages layer-wise Hessian approximation to match neurons by computing a similarity measure based on their functional difference.
Revised: The algorithm leverages layer-wise Hessian approximation to match neurons by computing a similarity measure based on their functional difference.

Original: 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: 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.

Original: Tatro et al.
Revised: Tatro et al.

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

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

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

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

Original: A recent work by Pittorino et al.
Revised: A recent work by Pittorino et al.

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

Original: B I MPLEMENTATION DETAILS B. T RAINING HYPERPARAMETERS
Revised: B I MPLEMENTATION DETAILS B. T RAINING HYPERPARAMETERS

Original: Table 1 summarizesthe
Revised: Table 1 summarizes

Original:  hyperparameters we used to train the neural 
Revised: the hyperparameters we used to train the neural networks which appear in this work.

Original: networks which appear

Original: in this

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

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

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

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

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

Original: [Table]
Revised: [Table]

Original:  P SEUDOCODE FOR P Y T ORCH MODULES
Revised: B. P SEUDOCODE FOR P Y T ORCH MODULES

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

Original: 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: 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.

Original: 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: 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:

Original: class TrackLayer(nn.

Original: Module): def __init__(self

Original: [Equation]

Original: layer):

Original: super

Original: [Equation]

Original: [Equation]

Original: __init

Original: [Equation]

Original: self

Original: [Equation]
Revised: [Equation]

Original: self
Revised: (self

Original: [Equation]
Revised: [Equation]

Original: bn
Revised: bn

Original: [Equation]
Revised: [Equation]

Original: self

Original: [Equation]
Revised: [Equation]

Original: self
Revised: self

Original: [Equation]
Revised: [Equation]

Original: eval
Revised: bn

Original: [Equation]
Revised: [Equation]

Original:  defget
Revised: _var.sqrt()) def

Original: [Equation]
Revised: [Equation]

Original: return

Original: [Equation]
Revised: [Equation]

Original:  running def [Equation] [Equation]
Revised: : outputs = self.layer(inputs) # Apply BatchNorm so the running variance are updated; the output self.bn(outputs) returnoutputs

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

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

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

Original: class ResetLayer(nn.
Revised: class ResetLayer

Original:  Module): def__init__(self, layer):
Revised: [Equation] nn [Equation] Module): def

Original: [Equation]
Revised: [Equation]

Original: def
Revised: init

Original: [Equation]
Revised: [Equation]

Original:  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: self [Equation] [Equation] 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.

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

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

Original: C C LOSED - FORM APPROXIMATE VARIANT OF REPAIR
Revised: C C LOSED - FORM APPROXIMATE VARIANT OF REPAIR

Original: [Table]
Revised: [Table]

Original: were trained
Revised: were trained

Original: [Table]
Revised: [Table]

Original: 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: 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.

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

Original: 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: 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.

Original: The variant proceeds as follows.
Revised: The variant proceeds as follows.

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

Original: 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: 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.

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

Original: X α =
Revised: X α =

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

Original: X 1 + α · X 2 (this holds only because the unit is in the ﬁrst layer).
Revised: X 1 + α · X 2 (this holds only because the unit is in the ﬁrst layer).

Original: 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
Revised: 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

Original: α exactly according to the formula
Revised: α exactly according to the formula

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: 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.
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 ﬁnd that this rescaling signiﬁcantly improves the performance of such interpolated networks.
Revised: We ﬁnd that this rescaling signiﬁcantly improves the performance of such interpolated networks.

Original: 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.
Revised: 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.

Original: For this case, the approximate rescaling performs similarly to the full exact method presented in Section 4.
Revised: For this case, the approximate rescaling performs similarly to the full exact method presented in Section 4.

Original: D A DDITIONAL PLOTS
Revised: D A DDITIONAL PLOTS

Revised: It appears that REPAIR is effective towards the endpoints, and neuron-alignment becomes essential near the midpoint.

Revised: Both methods are necessary in order to have a high- performing midpoint.