O PTIMIZING L OSS L ANDSCAPE C ONNECTIVITY VIA N EURON A LIGNMENT
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁnding algorithms do not consider the inﬂuence 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 efﬁciently 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.
I NTRODUCTION
Loss surfaces of neural networks have been of recent interest in the deep learning community both from a numerical (Dauphin et al.
; Sagun et al.
and a theoretical (Choromanska et al.
; 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 rectiﬁed 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 ﬁnding 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 ﬁrst studies to rigorously prove that one hidden layer rectiﬁed 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 (cid:15) -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:
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.
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.
We perform experiments on 3 datasets and 4 architectures afﬁrming 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.
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 ﬁrst review pertinent background on curve ﬁnding and neuron alignment in Section 2.
Then, we introduce our proposed optimization models and algorithms for curve ﬁnding 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.
B ACKGROUND ON C ONNECTIVITY AND A LIGNMENT
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 ﬁnding approach introduced in (Garipov et al., 2018).
Here we search for the path, r : [0 , 1] (cid:55)→ R N , that connects the two models while minimizing the average of the loss function, L , along the path.
This problem is formalized in equation 1.
[Equation]
dt
[Equation]
subject to r
[Equation]
[Equation]
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 (cid:48)
( t ) ||
= 1 for all t , is assumed to make optimization more computationally feasible.
Note that if the endpoint networks are global minima and a ﬂat 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 E t ∼ 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 )
∈ R
k 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 ﬁlter output by a convolutional layer, where k would be 1 or the number of pixels in the ﬁlter respectively.
Given networks, θ 1 and θ 2 , we deﬁne the channel-wise mean and standard deviation for θ 1 in equation 2.
We also deﬁne the cross-correlation matrix, C (1 , 2) l , denoting the cross-correlation between each channel in θ 1 and θ 2 in layer l .
[Equation]
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 ﬁnds the permutation, P l , that maximizes the trace of P Tl C (1 , 2) l, : , : , which is an instance of the linear assignment problem.
We formalize this optimization model in equation 3 below, where K l represents the index set of activations in layer l .
We recommend Burkard & Cela (1999) as a reference for those unfamiliar with the assignment problem.
[Equation]
subject to
[Equation]
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 ﬁgure 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 conﬁdent that some subset of highly correlated features are being matched.
O PTIMA C ONNECTIVITY C ONSIDERING W EIGHT S YMMETRY
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 P l is in Π | K l | , the set of permutations on K l , the index set of channels in layer l .
For simplicity suppose we have an L layer feed-forward network with activation function σ , weights { W l } Ll =1 , and input X 0 .
Then the weight permutation ambiguity becomes clear when we introduce the following set of permutations to the feedforward equation:
[Equation]
Then we can deﬁne the network weight permutation P as the block diagonal matrix, blockdiag ( P 1 , P 2 , ..., P L − 1 ) .
Additionally, P θ denotes the network parameterized by the weights [ P 1 W 1 , P 2 W 2 P T 1 , ..., W L P TL − 1 ] .
Note that we omit permutations P 0 and P L , as the input and output channels of neural networks have a ﬁxed ordering, so they correspond to the identity I .
Without much difﬁculty this framework generalizes for more complicated architectures.
We discuss this for residual networks in Appendix E.
[Figure]
3. C URVE F INDING UP TO S YMMETRY
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 ﬁnd the optimal curve connecting two networks up to symmetry with the model in equation 5.
[Equation]
5) subject to
[Equation]
3.1. P ROXIMAL A LTERNATING M INIMIZATION AS A F RAMEWORK
We ﬁrst 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 ﬁrst order methods could be computationally intensive as we need to store gradients of φ and P 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 let R denote a regularization term on φ .
[Equation]
such that
[Equation]
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 C 1 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:
Q ( φ k +1 , P k +1 )
+ R ( φ k +1 )
+ 12 ν φ ||
φ k +1 −
φ k ||
+ 12 ν P
|| P
k +1 −
P k
||
22 ≤ Q ( φ k , P k )
+ R ( φ k ) ,
∀ k ≥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 rectiﬁed networks while still generating a sequence of iterates with monotonic decreasing objective value.
3.1. N EURON A LIGNMENT AS AN I NITIALIZATION
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.
Algorithmsummarizes the process for efﬁciently 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 ( LM 3 ) .
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 ﬁnding.
We refer to these different curves as aligned and unaligned.
Data: Trained Neural Networks θ 1 and θ 2 , Subset of Training Data X
Result: Aligned Neural Networks θ 1 and P θ
Initialize P θ 2 :=
[ ˆ W 21 , ˆ W 22 , ..., ˆ W 2 L ] as [ W 21 , W 22 , ..., W 2 L ] ; for each layer l in { 1 , 2 , ..., L − 1 } do for each network j in { 1 , 2 } do compute activations, X
E XPERIMENTS
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 ﬂips, random rotation, and random cropping to prevent the models from overﬁtting on the training set.
Architectures Four different model architectures are used.
Table 1 summarizes relevant properties of these architectures.
The ﬁrst 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 ﬁnding with alignment.
VGG16-BN is the third architecture that we considered in our experiments (Simonyan & Zisserman, 2014).
VGG16 has signiﬁcantly 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.
[Table]
Quadratic Bezier curves All curves are parameterized as quadratic Bezier curves.
Bezier curves are popular in computer graphics as they can be deﬁned 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 deﬁned in equation
[Equation]
Then θ c is the learnable parameter in φ .
Of course one could consider more complicated curve parameterizations.
In practice, we ﬁnd 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 signiﬁcant in itself.
We explicitly deﬁne 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 ﬁnding.
This coupling of the control point with the permutation was seen to be critical for the success of PAM.
[Equation]
4. T RAINING C URVES
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 θ• 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) =
P Al •
Aligned: Solution to algorithm 2 given networks θ 1 and P Al θ
where P Al 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 ﬁnding is ﬁxed for all model pairs.
The other set consists of the curves learned when the random seed is different for each model pair.
We ﬁnd that the learned curves for different seeds are similar up to reindexing the endpoints.
For Figures 2,
[Table]
and 4, we use the ﬁrst set of curves so that interesting geometric features on the loss surface are not averaged out.
For tables and other ﬁgures, we use the second set of curves as they are more general.
[Table]
[Figure]
[Table]
[Figure]
[Table]
N EURON A LIGNMENT
First, we investigate the effects of using neuron alignment as a heuristic for curve ﬁnding up to symmetry.
That is, we are determining some weight permutation P Al and then ﬁnding the curve between networks θ 1 and P Al θ 2 .
The unaligned and aligned curves were both trained forepochs 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.
[Figure]
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 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 CIFAR 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 ﬁnd 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 ﬁnding.
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 ﬁgure.
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.
[Figure]
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 ﬂatness 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 ﬁnd 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 ﬁnd that curve ﬁnding is essentially trying to smoothly interpolate similar feature representations.
Sensibly, neuron alignment of the endpoints makes this task easier.
4.1. P ROXIMAL A LTERNATING M INIMIZATION
Proximal alternating minimization provides a comprehensive formulation for learning the weight permutation P directly, coupled with some convergence guarantees.
We ﬁnd 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 P Al is already close to the locally optimal permutation when P Al 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, Π | K l | .
The curve parameter subproblem, which optimizes ˜ θ c from equation 8, entails 40 epochs of stochastic gradient descent.
The same hyperparameters are used as in training the endpoint models.
The learning rates are annealed with each iteration of PAM.
This training can be seen for CIFAR100 in Figure 10.
The outlier case in Table 2 is for ResNet32 where Aligned notably outperforms PAM Aligned.
This result is due to the integral nature of the problem.
Essentially, during PAM Aligned, the scheme learns a more optimal doubly stochastic matrix that actually projects to a worse permutation at the end of the permutation subproblem.
This is a known possible negative, when solving a problem in the convex hull and then projecting back to the feasible set.
A PPLICATION : C ONSTRUCTING E NSEMBLES A LONG A LIGNED C URVES
We investigate if the diversity of models along the curve suffers due to alignment through the lens of constructing network ensembles.
We consider the simple ensemble that performs classiﬁcation by averaging the probability distributions output by the individual networks.
In our experiment, we look at four cases.
The ensemble formed by considering the curve endpoints, the ensembles formed by the curve endpoints and the midpoint, and an ensemble of the curve endpoints and a third independent model.
We consider ensembling only on the CIFAR100 dataset with results for the aligned and unaligned case being summarized in Table 3.
We see that at best, ensembles constructed by sampling along the unaligned curve perform as well as the independent ensemble.
When constructing ensembles by sampling along the aligned curve, those outperform the independent ensemble for the TinySix, TinyTen, and Resnet32 case and show comparable performance for VGG16.
The enhanced performance is most obvious on TinyTen.
This result makes sense since ensembling has been reported to lead to better performance increase for simpler networks (Ju et al., 2018).
This result is encouraging as ensembles of simple models are common in practice.
We also note that the performance gain from the aligned ensemble is comparable in magnitude to the gains of Fast
Geometric Ensembling in (Garipov et al., 2018) over Snapshot Ensembling (Huang et al., 2017)
[Table]
D ISCUSSION AND F UTURE W ORK
We generalize the curve ﬁnding 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 ﬁnd 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 efﬁciently 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 [Table]
Future work will include gaining a deeper theoretical understanding of how neuron alignment affects curve ﬁnding 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.
A DDITIONAL F IGURES
[Figure] [Figure]
[Figure] [Figure]
A. P LANES CONTAINING LEARNED B EZIER CURVES
[Figure] [Figure]
B A DDITIONAL R ESULTS
[Table] [Table]
366 ± 0 .
C A LGORITHMS
This section contains algorithms described in Section 2.
Data: Two trained models, θ 1 and θ
Result: A parameterized curve, r φ , connecting θ 1 and θ 2 along which loss is ﬂat
Initialize r φ ( t ) as θ 1 + t ( θ 2 − θ 1 ) ; while not converged do for batch in dataset do sample point t 0 in [0 , 1] ; compute loss L ( r φ ( t 0 )) ; optimization step on network r φ ( t 0 ) to update φ ; end end et al.,
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 φ ( t 0 ) gives the weights for the model at point t 0 , 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.
D P ROOFS
For the following proofs, we ﬁrst establish and more rigorously deﬁne some terminology.
We ﬁrst 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 X 0 ∈ R m
×
d be the input to the neural network, d samples of dimension m 0 .
Then we let W i
∈
R m i ×
m i − 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:
[Equation]
To include biases, { b i } Li =1 , we simply convert to homogeneous coordinates,
[Equation]
In all proofs, these terms are interchangeable.
Huberized ReLU
The commonly used ReLU function is deﬁned as σ ( t ) :
= max(0 , t ) .
However, this function is not in C 1 and hence not locally Lipschitz differentiable.
This makes conducting analysis with this function difﬁcult.
Thus, we will approach studying it through the lens of the huberized ReLU function, deﬁned as:
[Equation]
It is clear that σ δ is a C 1 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) =
• ψ is C 1 on (0 , ν ) • ∀ s ∈ (0 , ν ) , ψ (cid:48)
( s ) >• ∀ x ∈ U ∩ [ f (¯ x ) <
f < f (¯ x ) + ν ]
, the Kurdyka-Lojasiewics inequality holds
[Equation]
Here ∂f denotes the subdifferential of f .
Informally, a function that satisﬁes this inequality is one whose range can be reparameterized such that a kink occurs at its minimum.
More intuitively, if ψ has the form, s 1 − θ , and f is differentiable on (0 , ν ) , then the inequality reduces to
[Equation]
Semialgebraic function A subset of R n is semialgebraic if it can be written as a ﬁnite union of sets of the form
[Equation]
where p i and q i are real polynomial functions.
A function f : R n → R ∪ { + ∞} is said to be semialgebraic if its graph is a semialgebraic subset of R n +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. P ROOF OF T HEOREM 3.
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:
Each term in the objective function containing only one primal variable is bounded below and lower semicontinuous.
Each term in the objective function which contains both variables is in C 1 and is locally Lipschitz differentiable.
The objective function satisﬁes the Kurdyka-Lojasiewicz (KL) property.
First we reformulate the problem so that it becomes unconstrained.
Let χ denote the indicator function, where:
[Equation]
This problem contains two hard constraints.
First, each permutation matrix, P l , must clearly be restricted to the set of permutation matrices of size | K l | , Π | K l | .
Additionally, it is assumed that the norm of the weights are bounded above.
Without loss of generality, let K W denote an upper bound valid for all the weights.
We denote the set of weights that satisfy the norm constraint as { A
[Equation]
|
[Equation]
||
[Equation]
≤ K W } .
Then equation 5 with added regularization is equivalent to:
[Equation]
We now address each requirement for local convergence.
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.
Now we consider the form of the function, Q ( φ, P ) .
It has been deﬁned as
[Equation]
We know that r ( t ; φ, P ) corresponds to a feed-forward neural network.
Then Q can be expressed as:
[Equation]
with weight matrices W i and activation function σ .
It becomes clear that for Q ( φ, P ) to be in C 1 and locally Lipschitz differentiable, the same must be true for L , σ , and { W i } Li =1 .
The ﬁrst two are true as they are assumptions of the theorem.
Since, r φ is in C 1 and locally
Lipschitz differentiable in the primal variables, then this is also true for all W i .
Thus,
Q ( φ, P ) is in C 1 and locally Lipschitz differentiable.
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 that R 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, Π | K l | , is ﬁnite 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 ||
− K W < 0 } .
The function that deﬁnes 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 satisﬁes the KL property.
D. C ONSIDERING R ECTIFIED N ETWORKS
Theorem 3.1 does not extend to the class of rectiﬁed networks.
However, we are still interested in contructing a sequence of iterates { φ k , P k } such that the objective value, E t ∼ 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, X 0 , is bounded.
Additionally, assume that the spectral norm of all weights, { W i } Li =1 , is bounded above by K W , and the activation function, σ , is continous with || σ || ≤ 1 .
Let S Y denote the set of Y where
[Equation]
such that || W i || 2 ≤ K W ∀ i ∈ { 1 , 2 , ..., L }
Then L restricted to the set S Y is Lipschitz continuous with some Lipschitz constant K .
Proof.
Since X 0 is bounded, it follows that there exists some constant K X such that || X 0 ||
≤ K X .
Since, the spectral norm of W 1 is bounded above by K W , it is easy to see that
|| W 1 X
||
≤ K W K X .
Now since the pointwise activation function is a non-expansive map, it immediately follows that ||
σ ◦ W 1 X
|| ≤ K W K X .
Following this process inductively, we see that the network output, Y , is bounded and that:
[Equation]
Since Y is arbitrary, it follows that this is a bound for S Y .
Then we can restrict L to the ball in R m L × d of radius K LW K
X .
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, S Y 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 C 1 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 is M units.
Additionally, assume that the network weights have a spectral norm bounded above by K W , and that this is a hard constraint when training the networks.
Finally, any point on r φ must be equivalent to an afﬁne 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 Rectiﬁed Networks) .
For a feed-forward network, assume the above assumptions have been met.
Additionally, assume that X 0 is bounded, so that L restricted to the set of possible network outputs is Lipschitz continuous with Lipschitz constant
K L 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
[Equation]
Then, the sequence of curves r ( t ; φ k , P k ) connecting rectiﬁed networks has monotonic decreasing objective value.
Proof.
First we consider the approximation error from replacing σ with σ δ .
It is straightforward to see that
[Equation]
Then it follows that for any input x ,
[Equation]
Since the spectral norm of W i are bounded above by K W , then we see that
[Equation]
Now notice that
[Equation]
Since the ReLU function is a non-expansive map, it must be that the ﬁrst term is bounded above by the previous error, K W √ 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
[Equation]
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 ﬁnd the following bounds:
[Equation]
Since any point on the curve is an afﬁne combination of networks with the K W 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 C 1 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
[Equation]
Thus, r ( t ; φ k , P k ) is a sequence of curves, connecting rectiﬁed networks, with monotonic decreasing objective value as long as
[Equation]
Since the above equation is a stopping criterion introduced in the theorem statement, it follows that we have constructed a sequence of curves, connecting rectiﬁed networks, with monotonic decreasing objective value.
E R ESIDUAL N ETWORK A LIGNMENT
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,
[Equation]
In this equation, we can see that X k +1 and X k − 1 share the same indexing of their units.
This becomes clear when you consider permuting the hidden units in X k − 1 without permuting the hidden units of X k +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 from X k − 1 to X k +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.
F A LIGNMENT A LONG C URVES
Clearly, alignment is a useful method for learning better ﬂat loss curves between models.
An interesting question is how curve ﬁnding 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. C ORRELATION S IGNATURE
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 ﬁnding 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 ﬁnding algorithm is ﬁnding 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 ﬁnding algorithm, this gives some intuition on how alignment works to give a better initialization.