Optimizing Loss Landscape Connectivity via Neuron Alignment

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 . 1 INTRODUCTION . 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 . 2 BACKGROUND ON CONNECTIVITY AND ALIGNMENT . 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 .

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