OPTIMIZING LOSS LANDSCAPE CONNECTIVITY
Year: 2019

Abstract:
The loss landscapes of deep neural networks are poorly understood due to their high nonconvexity. Empirically, the local minima of these loss functions can be connected by a learned curve in model space, along which the loss remains nearly constant. Yet, current path finding algorithms do not consider the influence of symmetry in the loss surface caused by weight permutations of the networks corresponding to the minima. We propose a framework to investigate the effect of symmetry on landscape connectivity by directly optimizing the weight permutations of the networks being connected. To learn a locally optimal permutation, we introduce a proximal alternating minimization scheme and prove some convergence guarantees. Using an neuron alignment technique, we derive an inexpensive heuristic for approximating a locally optimal weight permutation. Empirically, optimizing the permutation is critical for efficiently learning a simple, planar, low-loss curve between networks that successfully generalizes. As an application, we also show that there is a performance gain, particularly for models with less parameters, in the accuracy of ensembles constructed from the learned curve with neuron alignment permutation.

text: The loss landscapes of deep neural networks are poorly understood due to their high nonconvexity. Empirically, the local minima of these loss functions can be connected by a learned curve in model space, along which the loss remains nearly constant. Yet, current path finding algorithms do not consider the influence of symmetry in the loss surface caused by weight permutations of the networks corresponding to the minima. We propose a framework to investigate the effect of symmetry on landscape connectivity by directly optimizing the weight permutations of the networks being connected. To learn a locally optimal permutation, we introduce a proximal alternating minimization scheme and prove some convergence guarantees. Using an neuron alignment technique, we derive an inexpensive heuristic for approximating a locally optimal weight permutation. Empirically, optimizing the permutation is critical for efficiently learning a simple, planar, low-loss curve between networks that successfully generalizes. As an application, we also show that there is a performance gain, particularly for models with less parameters, in the accuracy of ensembles constructed from the learned curve with neuron alignment permutation.

heading: 1 INTRODUCTION
text: Loss surfaces of neural networks have been of recent interest in the deep learning community both from a numerical (Dauphin et al. (2014); Sagun et al. (2014)) and a theoretical (Choromanska et al. (2014); Safran & Shamir (2015)) perspective. Their optimization yields interesting examples of a high-dimensional non-convex problem, where counter-intuitively gradient descent methods successfully converge to non-spurious minima. Practically, recent advancements in several applications have used insights on loss surfaces to justify their approaches. For instance, Moosavi-Dezfooli et al. (2019) investigates regularizing the curvature of the loss surface to increase the adversarial robustness of trained models.
One interesting question about these non-convex loss surfaces is to what extent trained models, which correspond to local minima, are connected. Here, connection denotes the existence of a path between the models, parameterized by their weights, along which loss is nearly constant. There has been conjecture that such models are connected asymptotically, with respect to the width of hidden layers. Recently, Freeman & Bruna (2016) proved this for rectified networks with one hidden layer.
When considering the connection between two neural networks, it is important for us to consider what properties of the neural networks are intrinsic. There is a permutation ambiguity in the indexing of units in a given hidden layer of a neural network, and as a result, this ambiguity extends to the network weights themselves. Thus, there are numerous equivalent points in model space that correspond to a given neural network. This creates weight symmetry in the loss landscape. It is possible that the minimal loss paths between a network and all networks equivalent to a second network could be quite different. If we do not consider the best path among this set, we could fail to see to what extent models are intrinsically connected. Therefore, in this work we propose to develop a technique for more robust model interpolation / optimal connection finding by investigating the effect of weight symmetry in the loss landscape.The analyses and results will give us insight into the geometry of level sets of the loss surfaces of deep networks that are often hard to study theoretically.
Related Work Freeman & Bruna (2016) is one of the first studies to rigorously prove that one hidden layer rectified networks are asymptotically connected and established relevant bounds. Several recent numerical works have shown that parameterized curves along which loss is nearly constant can be successfully learned. Concurrently, Garipov et al. (2018) proposed learning Bezier curves and polygonal chains and Draxler et al. (2018) proposed learning a curve using nudged elastic band energy between two models. Gotmare et al. (2018) showed that these algorithms work even for models trained using different hyperparameters, excluding network architecture. Recently, Kuditipudi et al. (2019) analyzed the connectivity between -dropout stable networks. This body of work can be seen as the extension of the linear averaging of models studied in (Goodfellow et al., 2014). In fact, concurrent with our work, Singh & Jaggi (2019) applied neuron alignment to model averaging.
The symmetry groups in neural network weight space have long been formally studied (Chen et al., 1993). While permutation ambiguity in the weights has been acknowledged, ostensibly ambiguity due to scaling in the weights has received more attention in research. Numerous regularization approaches based on weight scaling such as in (Cho & Lee, 2017) have been proposed to improve the performance of learned models. Recently, Brea et al. (2019) studied the existence of permutation plateaus in which the neurons in the layer of a network can all be permuted at the same cost.
A second line of work studies network similarity. Kornblith et al. (2019) gives a comprehensive review on the topic while introducing centered kernel alignment (CKA) for comparing the behavior of different neural networks. CKA is an improvement over the canonical correlation analysis (CCA) technique introduced in (Raghu et al., 2017) and explored further in (Morcos et al., 2018). Another contribution in this direction is the neuron alignment algorithm from (Li et al., 2016), which showed empirically that two networks of the same architecture learn a subset of similar features.
Contributions We summarize the main contributions of this work as follows: 1. We formalize this problem, and apply a proximal alternating minimization (PAM) scheme to split the problem into iteratively optimizing the permutation of the second model weights and optimizing the curve parameters. We prove convergence of this scheme to a local critical point for feed-forward neural networks which are piece-wise analytic functions and continuously differentiable. 2. Motivated by the neuron alignment technique of (Li et al., 2016) and our PAM framework, we introduce a heuristic for approximating the optimal weight permutation in order to learn aligned curves connecting networks up to a symmetry in their weights. 3. We perform experiments on 3 datasets and 4 architectures affirming that more optimal curves can be learned faster with neuron alignment. We observe that this aligned permutation is close to a locally optimal permutation that PAM converges to under the same initialization. 4. For simple networks, we observe a notable improvement in ensemble accuracy when constructing them by sampling the aligned as opposed to the unaligned curve or a set of independent models.
For the structure of this paper, we first review pertinent background on curve finding and neuron alignment in Section 2. Then, we introduce our proposed optimization models and algorithms for curve finding up to a weight permutation in Section 3. In Section 4, we discuss our experiments in detail. In Section 5, we explore the effect of alignment on the performance of model ensembles constructed through sampling along the curve.

heading: 2 BACKGROUND ON CONNECTIVITY AND ALIGNMENT
text: In this section we review the existing approaches for loss optima connectivity and neuron alignment.
Loss Optima Connectivity To learn the minimal loss path connecting two N -dimensional neural networks, θ1 and θ2, we utilize the curve finding approach introduced in (Garipov et al., 2018). Here we search for the path, r : [0, 1] 7→ RN , that connects the two models while minimizing the average of the loss function, L, along the path. This problem is formalized in equation 1.
r∗ = arg min r
∫ t∈[0,1] L(r(t))‖r
′(t)‖dt∫ t∈[0,1] ‖r′(t)‖dt
subject to r(0) = θ1, r(1) = θ2. (1)
For tractability, r∗ can be approximated by a parameterized curve rφ, where φ denotes the curve parameters. For instance, as described in Section 4, this paper will be using the quadratic Bezier
curve. Computationally, an arclength parameterization, that is ||r′(t)|| = 1 for all t, is assumed to make optimization more computationally feasible. Note that if the endpoint networks are global minima and a flat loss path does exist, then the optimal objective of equation 1 is unchanged.
An equivalent view under the arclength parameterization is that we are minimizing Et∼U [L(rφ(t))], where U is the uniform distribution on the unit interval. This view is taken in Algorithm 2 in Appendix C. For clarity, we emphasize that rφ denotes the curve on the loss surface between two networks while rφ(t) denotes a point on that curve which is a neural network.
Neuron Alignment We give an overview of the neuron alignment framework in (Li et al., 2016). Given input d drawn from the input data distribution D, let X(1)l,i,:(d) ∈ Rk represent the activation values of channel i in layer l of network θ1, where k is the number of units in the channel. As an example, a channel could correspond to one unit in a hidden state or one filter output by a convolutional layer, where k would be 1 or the number of pixels in the filter respectively.
Given networks, θ1 and θ2, we define the channel-wise mean and standard deviation for θ1 in equation 2. We also define the cross-correlation matrix, C(1,2)l , denoting the cross-correlation between each channel in θ1 and θ2 in layer l.
µ (1) l,i = Ed∼D[
1
k k∑ a=1 X (1) l,i,a(d)] σ (1) l,i = √√√√Ed∼D[ 1 k k∑ a=1 ( X (1) l,i,a(d)− µ (1) l,i )2 ]
C (1,2) l,i,j =
Ed∼D[ ∑k a=1 1 k ( X (1) l,i,a(d)− µ (1) l,i )( X (2) l,j,a(d)− µ (2) l,j ) ]
σ (1) l,i σ (2) l,j
(2)
To align the activations in layer l between networks θ1 and θ2, the neuron alignment algorithm maximizes the sum of cross-correlation between aligned activations. Equivalently, this finds the permutation,Pl, that maximizes the trace ofP Tl C (1,2) l,:,: , which is an instance of the linear assignment problem. We formalize this optimization model in equation 3 below, where Kl represents the index set of activations in layer l. We recommend Burkard & Cela (1999) as a reference for those unfamiliar with the assignment problem.
max Pl
trace(P Tl C (1,2) l,:,: ) (3)
subject to Pl1 = 1,P Tl 1 = 1, Pl ∈ Z |Kl|×|Kl| +
The alignment technique is visualized in Figure 1a. This displays the cross-correlation matrix for the TinyTen network and CIFAR100 dataset that we discuss later in Section 4. It is clear that the values along the diagonal are much stronger after alignment. Figure 1b displays the mean cross-correlation at each layer between corresponding neurons. This figure also shows the standard deviation of this signal over a set of 3 network pairs. With this correlation signature being consistent over different pairs and being increased highly with alignment, we can feel confident that some subset of highly correlated features are being matched.

heading: 3 OPTIMA CONNECTIVITY CONSIDERING WEIGHT SYMMETRY
text: We clarify the idea of weight symmetry in a neural network. θ1 is a neural network on the loss surface parameterized by its weights. A permutation Pl is in Π|Kl|, the set of permutations on Kl, the index set of channels in layer l. For simplicity suppose we have an L layer feed-forward network with activation function σ, weights {Wl}Ll=1, and input X0. Then the weight permutation ambiguity becomes clear when we introduce the following set of permutations to the feedforward equation:
Y := WLP T L−1 ◦ σ ◦ PL−1WL−1P TL−2 ◦ σ ◦ PL−2WL−2P TL−3 ◦ ... ◦ σ ◦ P1W1X0 (4)
Then we can define the network weight permutation P as the block diagonal matrix, blockdiag(P1,P2, ...,PL−1). Additionally, Pθ denotes the network parameterized by the weights [P1W1,P2W2P T 1 , ...,WLP T L−1]. Note that we omit permutationsP0 andPL, as the input and output channels of neural networks have a fixed ordering, so they correspond to the identity I . Without much difficulty this framework generalizes for more complicated architectures. We discuss this for residual networks in Appendix E.

heading: 3.1 CURVE FINDING UP TO SYMMETRY
text: From equation 4, it becomes clear that the networks θ1 and Pθ1 share the same structure and intermediate outputs up to indexing. Taking weight symmetry into account, we can find the optimal curve connecting two networks up to symmetry with the model in equation 5.
φ∗,P ∗ = arg min φ,P Et∼U [L(rφ(t))] (5)
subject to rφ(0) = θ1, rφ(1) = Pθ2, P = blockdiag(P1,P2, ...,PL−1) Pl ∈ Π|Kl| for l ∈ {1, 2, ..., L− 1}

heading: 3.1.1 PROXIMAL ALTERNATING MINIMIZATION AS A FRAMEWORK
text: We first introduce a framework to solve the generalized problem in equation 5. Theoretically, this problem is fairly complicated and hard to analyze. Numerically, approaching the problem directly with first order methods could be computationally intensive as we need to store gradients of φ andP simultaneously. The problem can be more easily addressed using the method of proximal alternating minimization (PAM) (Attouch et al., 2010). The PAM scheme involves iteratively solving the two subproblems in equation 6. Here we let Q(φ,P ) denote the objective function in equation 5. We only consider parameterized forms of r that satisfy the endpoint constraints for all φ and P . For generality, we letR denote a regularization term on φ. P k+1 = arg minP Q(φ k,P ) + 12νP ||P − P k||22 such that Pl ∈ Π|Kl| for l ∈ {1, 2, ..., L− 1}
P = blockdiag(P1,P2, ...,PL−1) φk+1 = arg minφ Q(φ,P k+1) +R(φ) + 12νφ ||φ− φ k||22
(6)
Computing the unaligned curve is equivalent to solving the PAM scheme with a very large value of νP . In fact, we are able to prove local convergence results for a certain class of networks. Theorem 3.1 (Convergence). Let {φk+1,P k+1} be the sequence produced by equation 6. Assume that rφ(t) corresponds to a feed-forward neural network with activation function σ for t ∈ [0, 1]. Assume that L, rφ, and σ are all piece-wise analytic functions in C1 and locally Lipschitz differentiable in φ and P . Additionally, assume R is piece-wise analytic in the primal variables and bounded below. Lastly, assume that the input data is bounded and the norm of the network weights are constrained to be bounded above. Then the following statements hold:
1. Q(φk+1,P k+1) +R(φk+1) + 12νφ ||φ k+1 − φk||22 + 12νP ||P k+1 −P k||22 ≤ Q(φk,P k) + R(φk),∀k ≥ 0
2. {φk,P k} converges to a critical point of Q(φ,P ) +R(φ)
Proof See Appendix D
Remark Theorem 3.1 does not extend to neural networks with ReLU activation functions. In Appendix D, we address a technique utilizing this theorem for learning a curve connecting rectified networks while still generating a sequence of iterates with monotonic decreasing objective value.

heading: 3.1.2 NEURON ALIGNMENT AS AN INITIALIZATION
text: In spite of convergence guarantees, PAM still requires a good initialization as the loss landscape is nonconvex. This is critical for avoiding convergence to non-global optima. Conceptually, neuron alignment introduced in (Li et al., 2016) is able to match subsets of similar feature representations. Thus, we believe that the permutation on the network weights induced by neuron alignment could be meaningful enough to provide a good initialization of P .
In practice, we solve the linear sum assignment problem formulated in equation 3 using the Hungarian algorithm. See Kuhn (1955) for further reading on the Hungarian algorithm. Algorithm 1 summarizes the process for efficiently computing a permutation of the network weights from neuron alignment. For an L layer network with a maximum layer width of M , we compute P using a subset of the training data. Then the cost of computing the cross-correlation matrices for all layers is dominated by the forward propogation through the network to accumulate the activations. The running time needed to compute all needed linear assignments is O(LM3). This is on the order of the running time associated with one iteration of forward propagation. Then neuron alignment is relatively cheap as the time complexity of computing curves using neuron alignment is on the same order as traditional curve finding. We refer to these different curves as aligned and unaligned.
Data: Trained Neural Networks θ1 and θ2, Subset of Training Data X0 Result: Aligned Neural Networks θ1 and Pθ2 Initialize Pθ2 := [Ŵ 21 , Ŵ 2 2 , ..., Ŵ 2 L] as [W 2 1 ,W 2 2 , ...,W 2 L]; for each layer l in {1, 2, ..., L− 1} do for each network j in {1, 2} do
compute activations, X(j)l = σ ◦W j l X (j) l−1 ; for each element in the batch, vectorize X(j)l if applicable ; compute, Z(j)l , the Z-score normalization of the activations ;
end compute the cross-correlation matrix, C(1,2)l = Z (1) l Z (2)T l ; compute Pl by solving the assignment problem associated with C (1,2) l using the Hungarian
algorithm ; update Ŵ 2l → PlŴ 2l , Ŵ 2l+1 → Ŵ 2l+1P Tl
end Algorithm 1: Computing Permutation via Neuron Alignment

heading: 4 EXPERIMENTS
text: Datasets In our experiments, we trained neural networks to classify images from CIFAR10 and CIFAR100 (Krizhevsky et al., 2009), as well as STL10 (Coates et al., 2011). The loss function is the cross entropy loss on the softmax of the logits output by the networks. 20% of the images in the training set are used for computing alignments between pairs of models. We augment the data using color normalization, random horizontal flips, random rotation, and random cropping to prevent the models from overfitting on the training set.
Architectures Four different model architectures are used. Table 1 summarizes relevant properties of these architectures. The first architectures considered were the TinySix and TinyTen models. TinyTen, introduced in (Kornblith et al., 2019), is a narrow 10 layer convolutional neural network that uses batch-normalization, ReLU activations, and global average pooling. TinySix is equivalent
to TinyTen with layers 2, 4, 5, and 7 removed. These are useful models for concept testing and allow us to gain insight to networks that are underparameterized. We also include ResNet32 (He et al., 2016) in our experiments to understand the effect of skip connections on curve finding with alignment. VGG16-BN is the third architecture that we considered in our experiments (Simonyan & Zisserman, 2014). VGG16 has significantly more parameters compared to other models. We chose this set of architectures for its varying properties and because of their prevalence in numerical experiments in related literature.
All models used as curve endpoints are trained using stochastic gradient descent. We set a learning rate of 1E−1 that decays by a factor of 0.5 every 20 epochs. Weight decay of 5E−4 was used for regularization. Each model was trained for 250 epochs, and all models were seen to converge.
Quadratic Bezier curves All curves are parameterized as quadratic Bezier curves. Bezier curves are popular in computer graphics as they can be defined by their control points. In the current study, we refer to endpoint models as θ1 and θ2 as well as the control point, θc. Then r is defined in equation 7 rφ(t) = (1− t)2θ1 + 2(1− t)tθc + t2θ2. (7) Then θc is the learnable parameter in φ. Of course one could consider more complicated curve parameterizations. In practice, we find a simple curve to be enough for our experiments, and consider the learning of a planar curve along which loss is nearly constant to be significant in itself.
We explicitly define this curve for use in the PAM algorithm in equation 8. Here the curve has been reparameterized so that the control point is a function of the permutation P . θ̃c captures the deviation of the control point from the linear midpoint between θ1 and Pθ2. For PAM, θ̃c is the learnable curve parameter in φ. It is zero initialized so that the initial curve is a linear interpolation between models as in traditional curve finding. This coupling of the control point with the permutation was seen to be critical for the success of PAM.
r(t; θ̃c,P ) = (1− t)2θ1 + 2(1− t)t ( θ1 + Pθ2
2 + θ̃c
) + t2Pθ2. (8)

heading: 4.1 TRAINING CURVES
text: For each architecture, we train 12, 6, and 6 different models using different random initializations for CIFAR10, CIFAR100, and STL10 respectively. Thus we have 6 or 3 independent model pairs for a dataset. We learn four classes of curves:
• Unaligned: Solution to algorithm 2 given networks θ1 and θ2 • PAM Unaligned: Solution to equation 6 given networks θ1 and θ2 with P (0) = I • PAM Aligned: Solution to equation 6 given networks θ1 and θ2 with P (0) = PAl • Aligned: Solution to algorithm 2 given networks θ1 and PAlθ2
where PAl denotes the permutation learned by neuron alignment (algorithm 1).
We learn PAM curves for all architectures except VGG16, as its size made this computationally prohibitive. We train two sets of each curve class. One set involves the curves learned when the random seed for curve finding is fixed for all model pairs. The other set consists of the curves learned when the random seed is different for each model pair. We find that the learned curves for
different seeds are similar up to reindexing the endpoints. For Figures 2, 3, and 4, we use the first set of curves so that interesting geometric features on the loss surface are not averaged out. For tables and other figures, we use the second set of curves as they are more general.

heading: 4.1.1 NEURON ALIGNMENT
text: First, we investigate the effects of using neuron alignment as a heuristic for curve finding up to symmetry. That is, we are determining some weight permutation PAl and then finding the curve between networks θ1 and PAlθ2. The unaligned and aligned curves were both trained for 200 epochs using stochastic gradient descent with an annealing learning rate. The training of these curves share the same hyperparameters as the training of the individual models.
The test accuracy can be seen for each dataset and curve class in Table 2. Clearly, the aligned curves outperform the unaligned curve. In many cases, the average accuracy along the aligned curves in
1Strictly speaking, the algorithm converges to a local optima in the convex relaxation of the domain. The learned permutation is the projection of this optima to the feasible set.
comparable to the trained models used as endpoints. Table 4 contains the minimum accuracy along the curve, indicating that while aligned curves do not suffer from the same generalization gap that unaligned curves are prone to. Finally, Table 5 contains the training loss for each case at convergence. Overall, it is clear that the strongest gains from using alignment are in the case of underparameterized networks. As seen in Table 2, the largest increase in performance is for TinySix on CIFAR100 while the smallest gain is made for STL10 on VGG16. This is inline with observations by (Freeman & Bruna, 2016).
The test loss and accuracy along the learned curves for CIFAR100 are shown in Figure 2. The corresponding Fourier transform of the loss along the curve for assessing curve smoothness is displayed in Figure 11. We observe that, as expected, the accuracy at each point along the aligned curve exceeds that of the unaligned curve, while the loss along the curve is also smoother with neuron alignment. We are comparing loss and accuracy at the curve parameter, t. Noteworthy is the prominent presence of the accuracy barrier along the unaligned curve around t at 0.8 for all models. This accuracy barrier corresponds to a clear loss barrier for Tiny-10 and ResNet32. In contrast, for VGG16 there is a valley in the loss function at this point on the unaligned curve with worse generalization performance. Overall, we find that loss along the aligned curves varies more smoothly as seen in Figure 11, and this leads to better generalization of the interpolated models.
Figure 3 displays the planes which contains the initializations for curve finding. It is clear that the aligned initialization has better objective value. This can also be seen for the other datasets in Figure 7. The planes containing the learned curves are displayed in Figures 4 and 13. These are the planes containing θ1, Pθ2, and θc, although the control point is out of bounds of the figure. The axis is determined by Gram-Schmidt orthonormalization. The loss displayed on the planes containing the linear initializations and the Bezier curves can be seen in Figures 6 and 12.
Practically, the neuron alignment heurisitc for determining the permutation P may be enough and avoids performing more complicated optimization. We see this by noting the relative flatness of the accuracy along the aligned curves in Figure 2b. Additionally, Figure 5 indicates much faster convergence when learning φ using neuron alignment, which is quite impressive. For example, the aligned curve takes 100 epochs less to achieve the training accuracy that the unaligned curve converges to, when TinyTen is used on CIFAR100. Even for VGG16, the aligned curve reaches the milestone 40 epochs earlier. Figures 8 and 9 display these training curves for the additional datasets.
While these observations are promising, we intend to provide insight into why neuron alignment works, whih is addressed in Appendix F where we investigate how the alignment is preserved along the learned curve. We find that the midpoints of the unaligned curves are highly aligned to each endpoint even though the endpoints themselves are weakly aligned at best. Then, we find that curve finding is essentially trying to smoothly interpolate similar feature representations. Sensibly, neuron alignment of the endpoints makes this task easier.

heading: 4.1.2 PROXIMAL ALTERNATING MINIMIZATION
text: Proximal alternating minimization provides a comprehensive formulation for learning the weight permutation P directly, coupled with some convergence guarantees. We find that curves learned using PAM perform better than the unaligned curves as seen in Table 2. As was the case for the aligned curves, this performance gain is more notable in underparameterized models. Notably, the aligned curves perform comparably to PAM aligned. This indicates that PAl is already close to the locally optimal permutation when PAl is chosen as the initialization for PAM. Additionally, the performance gain of PAM Aligned over PAM Unaligned shows that this permutation is not easy to learn when P (0) is not necessarily close. Then training aligned curves is an inexpensive way to approximate the solution to a rigorous optimization method with good initialization. We stress that this observation shows that the gain from neuron alignment is not trivial.
We now address some technical points in the training of the PAM curves. To learn each curve, we perform 4 iterations of PAM. The permutation subproblem entails 20 epochs of projected stochastic gradient descent to the set of doubly stochastic matrices. This is done as the set of doubly stochastic matrices is the convex relaxation of the set of permutations. This projection is accomplished through 20 iterations of alternating projection of the updated permutation to the set of nonnegative matrices and the set of matrices with row and column sum of 1. After the 20 epochs of projected gradient descent, each layer permutation is projected to the set of permutations, Π|Kl|. The curve

heading: 5 APPLICATION: CONSTRUCTING ENSEMBLES ALONG ALIGNED CURVES
text: Curve Ensembles (%) Independent Ensembles (%) curve parameter t {0.0, 1.0} {0.0, 0.5, 1.0} Unaligned Aligned
TinySix 41.80± 0.15 42.66± 0.48 43.82± 0.78 42.05± 0.26 TinyTen 61.23± 0.36 61.39± 0.25 62.40± 0.30 61.76± 0.01 ResNet32 71.35± 0.30 72.13± 0.12 72.33± 0.12 72.03± 0.08 VGG16 74.00± 0.17 74.69± 0.22 74.91± 0.10 74.88± 0.06

heading: 6 DISCUSSION AND FUTURE WORK
text: We generalize the curve finding problem by removing the weight symmetry ambiguity associated with the endpoint models. The optimal permutation of these weights can be approximated using neuron alignment. We find empirically that this approximation performs comparably to a proximal alternating scheme with the same initialization which learns a locally optimal permutation. Additionally, we prove that this PAM scheme has some convergence guarantees. Empirically, we show that neuron alignment can be used to successfully and efficiently learn optimal connections between neural nets. Addressing the ambiguity of weight symmetry is critical for learning planar curves on the loss surface along which accuracy is mostly constant. Our results hold true over a range of datasets and network architectures. With neuron alignment, these curves can be trained in less epochs and to higher accuracy. We also see a modest to notable increase in the performance of ensembling.
Future work will include gaining a deeper theoretical understanding of how neuron alignment affects curve finding dynamics. We plan to further explore how similar feature representations between models are smoothly interpolated by these learned curves, and how this relates to network training.

heading: A ADDITIONAL FIGURES
text: A.1 PLANES CONTAINING LINEAR INITIALIZATIONS
A.2 TRAINING
A.3 PLANES CONTAINING LEARNED BEZIER CURVES

heading: B ADDITIONAL RESULTS
text: Table 4: The minimum accuracy along the curve with standard deviation is reported for each combination of dataset, network architecture, and curve class.
Model Endpoints CIFAR10 CIFAR100 STL10 TinySix Unaligned 71.6± 0.6 35.9± 2.3 62.6± 0.6 PAM Unaligned 74.8± 0.5 39.8± 0.2 63.4± 0.4 PAM Aligned 75.7± 0.4 39.9± 0.3 64.8± 0.1 Aligned 75.6± 0.5 39.9± 0.2 64.8± 0.1 TinyTen Unaligned 85.0± 1.2 53.2± 1.1 72.3± 0.6 PAM Unaligned 87.2± 0.3 56.1± 0.1 72.6± 0.5 PAM Aligned 88.0± 0.1 57.7± 0.3 73.5± 0.3 Aligned 88.6± 0.1 57.7± 0.4 73.4± 0.3 ResNet32 Unaligned 91.2± 0.7 64.7± 0.4 75.7± 0.3PAM Unaligned 91.7± 0.3 65.1± 1.0 76.0± 0.4 PAM Aligned 91.8± 0.2 65.3± 0.6 76.1± 0.3 Aligned 92.6± 0.2 66.6± 0.1 76.0± 0.2 VGG16 Unaligned 92.6± 0.1 69.6± 0.4 71.3± 1.8Aligned 93.0± 0.1 70.8± 0.1 71.3± 1.8
Table 5: The training loss with standard deviation is reported for each combination of dataset, network architecture, and curve class.
Model Endpoints CIFAR10 CIFAR100 STL10 TinySix Unaligned 0.896± 0.008 2.636± 0.008 0.931± 0.014 PAM Unaligned 0.826± 0.006 2.501± 0.015 0.914± 0.010 PAM Aligned 0.813± 0.008 2.489± 0.018 0.863± 0.005 Aligned 0.817± 0.006 2.490± 0.018 0.864± 0.004 TinyTen Unaligned 0.460± 0.004 1.839± 0.010 0.479± 0.018 PAM Unaligned 0.419± 0.005 1.733± 0.013 0.469± 0.000 PAM Aligned 0.387± 0.003 1.662± 0.007 0.393± 0.007 Aligned 0.392± 0.003 1.693± 0.008 0.389± 0.007 ResNet32 Unaligned 0.212± 0.005 1.124± 0.005 0.292± 0.007PAM Unaligned 0.191± 0.002 1.093± 0.011 0.283± 0.001 PAM Aligned 0.192± 0.002 1.086± 0.011 0.283± 0.001 Aligned 0.180± 0.002 1.011± 0.002 0.256± 0.002 VGG16 Unaligned 0.212± 0.002 0.710± 0.004 0.372± 0.029Aligned 0.198± 0.001 0.676± 0.005 0.366± 0.015

heading: C ALGORITHMS
text: This section contains algorithms described in Section 2.
Data: Two trained models, θ1 and θ2 Result: A parameterized curve, rφ, connecting θ1 and θ2 along which loss is flat Initialize rφ(t) as θ1 + t(θ2 − θ1); while not converged do
for batch in dataset do sample point t0 in [0, 1]; compute loss L(rφ(t0)) ; optimization step on network rφ(t0) to update φ ;
end end
Algorithm 2: Curve Finding (Garipov et al., 2018)
In this algorithm, the optimization step can correspond to a variety of techniques. In this paper, we use traditional stochastic gradient descent to update the curve parameters φ. Notice that stochasticity is introduced by the sampling of t as well as the training data.
For the purpose of computing validation loss and test loss for rφ, important care must be given for networks that contain batch normalization layers. This is because batch normalization aggregates running statistics of the network output that are used when evaluating the model. Though, rφ(t0) gives the weights for the model at point t0, the running statistics need to be aggregated for each normalization layer. In practice, this can be done by training the model for one epoch, while freezing all learnable parameters of the model. Since batch statistics would need to be computed for each point sampled along the curve, it happens that computing the validation or test loss of the curve rφ is more expensive than an epoch of training.

heading: D PROOFS
text: For the following proofs, we first establish and more rigorously define some terminology. We first discuss an important abuse of notation. For clarity the parameterized curve connecting networks under some permutation P that has been written as rφ(t) will now sometimes be referred to as r(t;φ,P ).
Feed-forward neural networks In this section, we will be analyzing feed-forward neural networks. We let X0 ∈ Rm0×d be the input to the neural network, d samples of dimension m0. Then we let Wi ∈ Rmi×mi−1 denote the network weights mapping from layer l − 1 to layer l. Additionally, σ denotes the pointwise activation function. Then we can express the output of a feed-forward neural network, Y , as:
Y := WLσ ◦WL−1σ ◦WL−2...σ ◦W1X0 (9)
To include biases, {bi}Li=1, we simply convert to homogeneous coordinates,
X̂0 = [ X0 1 ] , Ŵi = [ Wi bi 0 1 ] , Ŷ = [ Y 1 ] (10)
In all proofs, these terms are interchangeable.
Huberized ReLU The commonly used ReLU function is defined as σ(t) := max(0, t). However, this function is not in C1 and hence not locally Lipschitz differentiable. This makes conducting analysis with this function difficult. Thus, we will approach studying it through the lens of the huberized ReLU function, defined as:
σδ(t) :=  0 for t ≤ 0 1 2δ t
2 for 0 ≤ t ≤ δ t− δ2 for δ ≤ t
(11)
It is clear that σδ is a C1 approximation of σ such that ||σ − σδ||∞ = δ2 . Using huberized forms of loss functions for analysis is a fairly common technique such as in (Xu et al., 2016) which studies huberized support vector machines.
Kurdyka-Lojasiewicz property The function f is said to have the Kurdyka-Lojasiewics (KL) property at x̄ if there exist ν ∈ (0,+∞], a neighborhood U of x̄ and a continuous concave function ψ : [0, ν)→ R+ such that:
• ψ(0) = 0 • ψ is C1 on (0, ν) • ∀s ∈ (0, ν), ψ′(s) > 0 • ∀x ∈ U ∩ [f(x̄) < f < f(x̄) + ν], the Kurdyka-Lojasiewics inequality holds
ψ′(f(x)− f(x̄))dist(0, ∂f(x)) ≥ 1. (12)
Here ∂f denotes the subdifferential of f . Informally, a function that satisfies this inequality is one whose range can be reparameterized such that a kink occurs at its minimum. More intuitively, if ψ has the form, s1−θ, and f is differentiable on (0, ν), then the inequality reduces to
1
(1− θ) |f(x)|θ ≤ ||∇f(x)|| (13)
Semialgebraic function A subset of Rn is semialgebraic if it can be written as a finite union of sets of the form
{x ∈ Rn : pi(x) = 0, qi(x) < 0, i = {1, 2, ..., p}}
where pi and qi are real polynomial functions. A function f : Rn → R ∪ {+∞} is said to be semialgebraic if its graph is a semialgebraic subset of Rn+1.
Subanalytic function Globally subanalytic sets are sets of the form {(x, t) ∈ [−1, 1]n × R : f(x) = t} where f : [−1, 1]n → R is an analytic function that can be extended analytically on a neighborhood of the interval [−1, 1]n. A function is subanalytic if its graph is a globally subanalytic set.
D.1 PROOF OF THEOREM 3.1
To prove this, we need that our problem meets the conditions required for local convergence of proximal alternating minimization (PAM) described in (Attouch et al., 2010). This requires the following:
1. Each term in the objective function containing only one primal variable is bounded below and lower semicontinuous.
2. Each term in the objective function which contains both variables is in C1 and is locally Lipschitz differentiable.
3. The objective function satisfies the Kurdyka-Lojasiewicz (KL) property.
First we reformulate the problem so that it becomes unconstrained. Let χ denote the indicator function, where:
χC(t) := { 0, for t ∈ C +∞, otherwise (14)
This problem contains two hard constraints. First, each permutation matrix, Pl, must clearly be restricted to the set of permutation matrices of size |Kl|, Π|Kl|. Additionally, it is assumed that the norm of the weights are bounded above. Without loss of generality, let KW denote an upper bound valid for all the weights. We denote the set of weights that satisfy the norm constraint as {A : ||A||22 ≤ KW }. Then equation 5 with added regularization is equivalent to:
φ∗,P ∗ = arg min φ,P Q(φ,P ) +R(φ) + L−1∑ l=1 χΠ|Kl|(Pl) + L∑ l=1 χ{A:||A||22<KW }(Wl) (15)
We now address each requirement for local convergence.
1. By the theorem, R is assumed to be bounded below and lower semicontinuous. It is easy to see from equation 14 that the sum of indicator functions are bounded below and lower semicontinuous.
2. Now we consider the form of the function, Q(φ,P ). It has been defined as∫ 1 t=0 L(r(t;φ,P ))dt
We know that r(t;φ,P ) corresponds to a feed-forward neural network. Then Q can be expressed as:∫ 1
t=0
L (WL(t;φ,P )σ ◦WL−1(t;φ,P )...σ ◦W1(t;φ,P )X0) dt (16)
with weight matrices Wi and activation function σ. It becomes clear that for Q(φ,P ) to be in C1 and locally Lipschitz differentiable, the same must be true for L, σ, and {Wi}Li=1. The first two are true as they are assumptions of the theorem. Since, rφ is in C1 and locally Lipschitz differentiable in the primal variables, then this is also true for all Wi. Thus, Q(φ,P ) is in C1 and locally Lipschitz differentiable.
3. To satisfy the KL property, the objective function must be a tame function (Attouch et al., 2010). Rigorously, this means that the graph of the function belongs to an o-minimal structure, a concept from algebraic geometry. We refer curious readers to (van den Dries & Speissegger, 2002) for further reference.
First, we note that Q(φ,P ) is piece-wise analytic. This is because Q is a composition of piece-wise analytic functions, L, σ, and rφ. Additionally, because the input data is bounded and the norm of the weight matrices are bounded, it follows that the domain of Q is bounded. Since, Q is a piece-wise analytic function with bounded domain, it follows that Q is a subanalytic function. The boundedness of the domain is an important detail here. This is because analytic functions are not necessarilly subanalytic unless their domain is bounded; a popular example of such a function is the exponential function.
The regularization function, R, is assumed to be a piece-wise analytic function. It follows from the previous reasoning thatR has bounded domain. Thus,R is a subanalytic function.
We now consider the constraints associated with this problem, which have been reexpressed as indicator functions in the objective. The set of permutation matrices, Π|Kl|, is finite and thus it is clearly a semi-algebraic set. Notice that the set of weight matrices satisfying the norm bound is equivalent to {A : ||A||22 − KW < 0}. The function that defines this set is a polynomial, so it is a semi-algebraic set. Indicator functions on semi-algebraic sets are semi-algebraic functions. Thus, the indicator functions in the objective are semi-algebraic.
The graphs of semi-algebraic functions and subanalytic functions both belong to the logarithmic-exponential structure, an o-minimal structure. A basic algebraic property of o-minimal structures is that the graphs of addition and multiplication are also elements of the structure (van den Dries & Speissegger, 2002). Since our objective function is a linear combination of semi-algebraic functions and subanalytic functions, it follows that the graph of our objective function is an element of the logarithmic-exponential structure. Therefore, our objective function is a tame function and it satisfies the KL property.
D.2 CONSIDERING RECTIFIED NETWORKS
Theorem 3.1 does not extend to the class of rectified networks. However, we are still interested in contructing a sequence of iterates {φk,P k} such that the objective value, Et∼U [L(r(t;φk,P k))], is monotonic decreasing. The following theorem will introduce a technique for constructing such a sequence.
Lemma D.1 (L restricted to possible network outputs is Lipschitz continuous). For a feed-forward neural network, assume that L is continuous and that the neural network input, X0, is bounded. Additionally, assume that the spectral norm of all weights, {Wi}Li=1, is bounded above by KW , and the activation function, σ, is continous with ||σ|| ≤ 1. Let SY denote the set of Y where
Y = WLσ ◦WL−1σ ◦WL−2...σ ◦W1X0 (17) such that ||Wi||2 ≤ KW ∀i ∈ {1, 2, ..., L}
Then L restricted to the set SY is Lipschitz continuous with some Lipschitz constant K.
Proof. Since X0 is bounded, it follows that there exists some constant KX such that ||X0|| ≤ KX . Since, the spectral norm of W1 is bounded above by KW , it is easy to see that ||W1X0|| ≤ KWKX . Now since the pointwise activation function is a non-expansive map, it immediately follows that ||σ ◦W1X0|| ≤ KWKX . Following this process inductively, we see that the network output, Y , is bounded and that: ||Y || ≤ KLWKX (18) Since Y is arbitrary, it follows that this is a bound for SY . Then we can restrict L to the ball in RmL×d of radius KLWKX . This ball is compact and L is continuous, so it follows that L restricted to this ball is Lipschitz continuous. Thus, there exists some Lipschitz constant K. Clearly, SY is contained in this ball. Therefore, L is Lipschitz continuous on the set of all possible network outputs with Lipschitz constant K.
Let θ1 and θ2 be feed-forward neural networks with ReLU activation function. Assume that L and rφ are piece-wise analytic functions in C1 and locally Lipschitz differentiable. Assume that R is piece-wise analytic, lower semi-continuous, and bounded below. Assume that the maximum network width at any layer isM units. Additionally, assume that the network weights have a spectral norm bounded above by KW , and that this is a hard constraint when training the networks. Finally, any point on rφ must be equivalent to an affine combination of neural networks (Bezier curves, polygonal chains, etc.) satisfying the previously stated spectral norm bound.
Create the parameterized curve rδ(t;φ,P ) by substituting the huberized ReLU function, σδ , into all ReLU functions in r(t;φ,P ). We refer to the objective values associated with these curves as Qδ(φ,P ) and Q(φ,P ) respectively.
Theorem D.2 (Monotonic Decreasing Sequence for Rectified Networks). For a feed-forward network, assume the above assumptions have been met. Additionally, assume that X0 is bounded, so thatL restricted to the set of possible network outputs is Lipschitz continuous with Lipschitz constant KL by Lemma D.1. Now generate the sequence {φk,P k} by solving equation 6 for rδ(t;φ,P ). On this sequence impose the additional stopping criteria that
1
2νφ ||φk+1 − φk||22 +
1
2νP ||P k+1 − P k||22 ≥ KL
√ M δ
2 L−1∑ i=1 KiW ∀k ≥ 0. (19)
Then, the sequence of curves r(t;φk,P k) connecting rectified networks has monotonic decreasing objective value.
Proof. First we consider the approximation error from replacing σ with σδ . It is straightforward to see that
max t |σ(t)− σδ(t)| ≤
δ 2 . (20)
Then it follows that for any input x,
||σ ◦W1x− σδ ◦W1x||2 ≤ √ M δ
2 .
Since the spectral norm of Wi are bounded above by KW , then we see that
||W2σ ◦W1x−W2σδ ◦W1x||2 ≤ KW √ M δ
2 .
Now notice that
||σ ◦W2σ ◦W1x− σδ ◦W2σδ ◦W1x|| ≤ ||σ ◦W2σ ◦W1x− σ ◦W2σδ ◦W1x|| (21) + ||σ ◦W2σδ ◦W1x− σδ ◦W2σδ ◦W1x||.
Since the ReLU function is a non-expansive map, it must be that the first term is bounded above by the previous error, KW √ M δ2 . The second term corresponds once again to the error associated
with the huberized form of the ReLU function, √ M δ2 . Thus the total error can be bounded by
(KW + 1) √ M δ2 .
Following this inductively, it can be seen that the this error grows geometrically with the number of layers. Additionally, the loss function is Lipschitz continuous when restricted to the set of possible network outputs. So we find the following bounds:
||Y − Yδ|| ≤ √ M δ
2 L−1∑ i=1 KiW
||L(Y )− L(Yδ)|| ≤ KL √ M δ
2 L−1∑ i=1 KiW (22)
Since any point on the curve is an affine combination of networks with theKW bound on the spectral norm of their weights, it immediately follows this spectral norm bound also holds for the weights for any point on the curve. Then ||Q(φ,P ) − Qδ(φ,P )|| is also bounded above by the bound in equation 22.
Then let {φk,P k} be the sequence generated by solving equation 6 using the curve rδ . σδ is a piecewise analytic function in C1 and is locally Lipschitz differentiable. Additionally, the spectral norm constraint on the weights is semi-algebraic and bounded below, so Theorem 3.1 can be applied. It then follows that
Q(φk+1,P k+1) +R(φk+1) + 1 2νφ ||φk+1 − φk||22 + 1 2νP ||P k+1 − P k||22 (23)
≤ Q(φk,P k) +R(φk) +KL √ M δ
2 L−1∑ i=1 KiW , ∀k ≥ 0
Thus, r(t;φk,P k) is a sequence of curves, connecting rectified networks, with monotonic decreasing objective value as long as
1
2νφ ||φk+1 − φk||22 +
1
2νP ||P k+1 − P k||22 ≥ KL
√ M δ
2 L−1∑ i=1 KiW ∀k ≥ 0
Since the above equation is a stopping criterion introduced in the theorem statement, it follows that we have constructed a sequence of curves, connecting rectified networks, with monotonic decreasing objective value.

heading: E RESIDUAL NETWORK ALIGNMENT
text: Algorithm 1 applies to networks with a typical feed-forward structure. In this section, we discuss how we compute alignments for the ResNet32 architecture as it is more complicated. It is important to align networks such that the network structure is preserved and network activations are not altered. In the context of residual networks, special consideration must be given to skip connections.
Consider the formulation of a basic skip connection,
Xk+1 = σ ◦ (Wk+1Xk) +Xk−1 (24)
In this equation, we can see that Xk+1 and Xk−1 share the same indexing of their units. This becomes clear when you consider permuting the hidden units inXk−1 without permuting the hidden units ofXk+1. It is impossible to do so without breaking the structure of the equation above, where there is essentially the use of an identity mapping fromXk−1 toXk+1.
We consider a traditional residual network that is decomposed into residual blocks. In each block the even layers have skip connections while the odd layers do not. So, we compute the alignment as usual for odd layers. For all even layers within a given residual block, we determine a shared alignment. We do this by solving the assignment problem for the average of the cross-correlation matrix over the even layers in that residual block.

heading: F ALIGNMENT ALONG CURVES
text: Clearly, alignment is a useful method for learning better flat loss curves between models. An interesting question is how curve finding itself relates to alignment. Until now, we have only considered the alignment between the endpoint models, r(0) and r(1). Now, we consider how points along the curve, r(t), align to the endpoints. To study this numerically, we will use the curve midpoint r(0.5). From Figure 4, we see that this is the point on the quadratic Bezier curve that is roughly linearly connected to both endpoints.
F.1 CORRELATION SIGNATURE
First, we consider how the correlation signature changes along the curve. Figure 14 displays the correlation signature between the curve midpoint and each endpoint in blue. To gain a better understanding of this signature, we require some context. Thus, the correlation signature between the linear midpoint and each endpoint is displayed in green. This allows us to understand how the correlation signature changes over the curve finding optimization. Additionally, we display the correlation signature between the curve midpoint and each endpoint, where the midpoint has been aligned to the given endpoint, in yellow. This essentially gives us context on how highly the midpoint is aligned to each endpoint. This is because the yellow curve acts as an upper bound for the blue curve.
There are several observations to be made about Figure 14. The correlation signature between the endpoint and the curve midpoint is fairly high. For unaligned endpoints, the correlation is only slightly lower than the signature computed when the curve midpoint is aligned to the endpoint. In the case where the endpoints are aligned, the signatures are seen to coincide. This suggests that the curve finding algorithm is finding the quadratic curve along which similar feature representations are being interpolated.
Concerning the linear midpoint, the correlation at the linear midpoint decays to 0 when endpoints are unaligned as the network goes deeper. When endpoints are aligned, the correlation signature at the linear midpoint is similar to the correlation signature at the curve midpoint. Since these linear connections between the endpoints are the initializations for the curve finding algorithm, this gives some intuition on how alignment works to give a better initialization.

Review summaries:
Given two parameters theta_1 and theta_2 of the same architecture, the authors propose to learn a minimal loss curve between theta_1 and P theta_2, where P is a permutation matrix yielding another equivalent parameterization of the same network. The authors show that either by initializing P with neuron alignment or by optimizing P jointly with the curve, one can find a parameter-space with better accuracy and loss than by just learning a curve between theta_1 and theta_2. The authors also show that initializing P is sufficient to obtain these gains, avoiding the complexity of also optimizing for P. Furthermore, they show that ensembles across models from these curves have a very mild gain in accuracy to those of non-aligned curves.

The paper investigates the connection between symmetries in the neural network architectures and the loss landscape of the corresponding neural networks. In the previous works, there was shown that the two local minima of a neural network can be connected by a curve with the low validation/train loss along the curve. Despite the loss on the curve being close to the loss at the end points, there are segments of the curve on which the loss in higher than loss at the local minima. To overcome this problem, the authors proposed two-step procedure: 1. Align weights between two neural networks 2. Launch the path finding algorithm between aligned weights. In other words, the authors proposed to connect not original local minima but local minima that parametrize the same functions as the original ones, but have a different order of neurons. The authors also proposed PAM algorithm where they iteratively apply path finding algorithm and weights alignment algorithm.