Neural networks with late-phase weights

1 INTRODUCTION . Neural networks trained with SGD generalize remarkably well on a wide range of problems . A classic technique to further improve generalization is to ensemble many such models ( Lakshminarayanan et al. , 2017 ) . At test time , the predictions made by each model are combined , usually through a simple average . Although largely successful , this technique is costly both during learning and inference . This has prompted the development of ensembling methods with reduced complexity , for example by collecting models along an optimization path generated by SGD ( Huang et al. , 2017 ) , by performing interpolations in weight space ( Garipov et al. , 2018 ) , or by tying a subset of the weights over the ensemble ( Lee et al. , 2015 ; Wen et al. , 2020 ) . An alternative line of work explores the use of ensembles to guide the optimization of a single model ( Zhang et al. , 2015 ; Pittorino et al. , 2020 ) . We join these efforts and develop a method that fine-tunes the behavior of SGD using late-phase weights : late in training , we replicate a subset of the weights of a neural network and randomly initialize them in a small neighborhood . Together with the stochasticity inherent to SGD , this initialization encourages the late-phase weights to explore the loss landscape . As the late-phase weights explore , the shared weights accumulate gradients . After training we collapse this implicit ensemble into a single model by averaging in weight space . Building upon recent work on ensembles with shared parameters ( Wen et al. , 2020 ) we explore a family of late-phase weight models involving multiplicative interactions ( Jayakumar et al. , 2020 ) . We focus on low-dimensional late-phase models that can be ensembled with negligible overhead . Our experiments reveal that replicating the ubiquitous batch normalization layers ( Ioffe & Szegedy , 2015 ) is a surprisingly simple and effective strategy for improving generalization1 . Furthermore , we find that late-phase weights can be combined with stochastic weight averaging ( Izmailov et al. , 2018 ) , a complementary method that has been shown to greatly improve generalization . 1We provide code to reproduce our experiments at https : //github.com/seijin-kobayashi/ late-phase-weights 2 METHODS AND MODELS . 2.1 LEARNING WITH LATE-PHASE WEIGHTS . Late-phase weights . To apply our learning algorithm to a given neural network model fw we first specify its weights w in terms of two components , base and late-phase ( θ and φ , resp. ) . The two components interact according to a weight interaction function w = h ( θ , φ ) . Base weights are learned throughout the entire training session , and until time step T0 both θ and φ are learned and treated on equal grounds . At time step T0 , a hyperparameter of our algorithm , we introduce K late-phase components Φ = { φk } Kk=1 , that are learned together with θ until the end . This procedure yields a late-phase ensemble of K neural networks with parameter sharing : reusing the base weights θ , each late-phase weight φk defines a model with parameters wk = h ( θ , φk ) . Late-phase weight averaging at test time . Our ensemble defined by the K late-phase weight configurations in Φ is kept only during learning . At test time , we discard the ensemble and obtain a single model by averaging over the K late-phase weight components . That is , given some input pattern x , we generate a prediction y ( x ) using the averaged model , computed once after learning : y ( x ) = fw ( x ) , w ≡ h ( θ , 1 K K∑ k=1 φk ) . ( 1 ) Hence , the complexity of inference is independent ofK , and equivalent to that of the original model . Late-phase weight initialization . We initialize our late-phase weights from a reference base weight . We first learn a base parameter φ0 from time step t = 0 until T0 , treating φ0 as any other base parameter in θ . Then , at time t = T0 , each configuration φk is initialized in the vicinity of φ0 . We explore perturbing φ0 using a symmetric Gaussian noise model , φk = φ0 + σ0 Z ( φ0 ) k , ( 2 ) where k is a standard normal variate of appropriate dimension and σ0 is a hyperparameter controlling the noise amplitude . We allow for a φ0-dependent normalization factor , which we set so as to ensure layerwise scale-invariance , which helps finding a single σ0 that governs the initialization of the entire network . More concretely , for a given neural network layer l with weights φ ( l ) 0 of dimension D ( l ) , we choose Z ( φ ( l ) 0 ) = √ D ( l ) /‖φ ( l ) 0 ‖ . Our perturbative initialization ( Eq . 2 ) is motivated by ongoing studies of the nonconvex , highdimensional loss functions that arise in deep learning . Empirical results and theoretical analyses of simplified models point to the existence of dense clusters of connected solutions with a locallyflat geometry ( Hochreiter & Schmidhuber , 1997a ) that are accessible by SGD ( Huang et al. , 2017 ; Garipov et al. , 2018 ; Baldassi et al. , 2020 ) . Indeed , the eigenspectrum of the loss Hessian evaluated at weight configurations found by SGD reveals a large number of directions of low curvature ( Keskar et al. , 2017 ; Chaudhari et al. , 2019 ; Sagun et al. , 2018 ) . For not yet completely understood reasons , this appears to be a recurring phenomenon in overparameterized nonlinear problems ( Brown & Sethna , 2003 ; Waterfall et al. , 2006 ) . Based on these observations , we assume that the initial parameter configuration φ0 can be perturbed in a late phase of learning without leading to mode hopping across the different models wk . While mode coverage is usually a sought after property when learning neural network ensembles ( Fort et al. , 2020 ) , here it would preclude us from taking the averaged model at the end of learning ( Eq . 1 ) . Stochastic learning algorithm . Having decomposed our weights into base and late-phase components , we now present a stochastic algorithm which learns both θ and Φ . Our algorithm works on the standard stochastic ( minibatch ) neural network optimization setting ( Bottou , 2010 ) . Given a loss function L ( D , w ) = 1|D| ∑ x∈D L ( x , w ) to be minimized with respect to the weights w on a set of data D , at every round we randomly sample a subsetM from D and optimize instead the stochastic loss L ( M , w ) . However , in contrast to the standard setting , in late stages of learning ( t > T0 ) we simultaneously optimize K parameterizationsW : = { wk | wk = h ( θ , φk ) } Kk=1 , instead of one . We proceed by iteration over W . At each step k , we sample a minibatch Mk and immediately update the late-phase weights φk , while accumulating gradients over the shared base weights θ . Such gradient accumulation has been previously used when learning ensembles ( Lee et al. , 2015 ; Wen et al. , 2020 ) and multi-task models ( Rebuffi et al. , 2017 ) with shared base parameters . A single iteration is finally concluded by changing the base weights in the direction opposite of the accumulated gradient . We scale the accumulated gradient by γθ ; setting γθ = 1/K recovers the original step size in θ , but other choices are possible . In particular , we find that a large γθ of unit size is in practice often tolerated , resulting in accelerated learning . Algorithm 1 : Late-phase learning Require : Base weights θ , late-phase weight set Φ , dataset D , gradient scale factor γθ , loss L Require : Training iteration t > T0 for 1 ≤ k ≤ K do Mk ← Sample minibatch from D ∆θk ← ∇θ L ( Mk , θ , φk ) φk ← Uφ ( φk , ∇φk L ( Mk , θ , φk ) ) θ ← Uθ ( θ , γθ ∑K k=1 ∆θk ) We summarize an iteration of our method in Algorithm 1 , where the loss L ( M , θ , φ ) is now seen as a function of θ and φ . We opt for a general presentation using unspecified gradient-based update operators Uφ and Uθ . These operators can be set to optimizers of choice . For instance , our method might benefit from additional noise injection onto parameter updates ( Welling & Teh , 2011 ) . Furthermore , late-phase optimizers need not coincide with the optimizer used in the early phase . In our work we typically set Uφ and Uθ to a single step of SGD with Nesterov momentum ( Nesterov , 2004 ) , and explore Adam ( Kingma & Ba , 2015 ) and plain SGD in a smaller set of experiments . 2.2 LATE-PHASE WEIGHT MODELS . As detailed next , we consider a number of distinct late-phase weight models in our experiments . In particular , we explore weight interaction functions h in which late-phase weights have low dimensionality , to avoid a large increase in complexity with the ensemble size K. To counteract this reduced dimensionality , we make extensive use of multiplicative base-late weight interactions . This design choice is motivated by the large expressive power of multiplicative interactions despite low dimensionality , which has been demonstrated in a wide range of settings ( Jayakumar et al. , 2020 ) . Late-phase batch normalization layers . Batch normalization layers ( BatchNorm ; Ioffe & Szegedy , 2015 ) are a staple of current deep neural network models . Besides standardizing the activity of the layer they are applied to , BatchNorm units introduce a learnable multiplicative ( scale ) parameter γ and an additive ( shift ) parameter β . While being low-dimensional , these additional parameters have large expressive power : it has been shown that learning only γ and β keeping the remaining weights frozen can lead to significantly lower loss than when learning random subsets of other weights of matching dimensionality ( Frankle et al. , 2020 ; Mudrakarta et al. , 2019 ) . We take the scale and shift parameters of BatchNorm layers as our first choice of late-phase weights ; the base weights are the remaining parameters of the model . Batch statistics are also individually estimated for each model in W . This late-phase weight parameterization is motivated by ( i ) the expressive power of γ and β discussed above , and by ( ii ) practical considerations , as BatchNorm layers are generally already present in feedforward neural network models , and are otherwise easy to implement efficiently . More concretely , let us consider an affine transformation layer l which maps an input vector r ( l−1 ) to θ ( l ) w r ( l−1 ) + θ ( l ) b , where the early-phase weight matrix θ ( l ) w and bias vector θ ( l ) b are already standardized using the respective batch statistics . For this standard layer , our model introduces a multiplicative interaction between base and late-phase weights , diag ( γ ( l ) ) θ ( l ) w , and an additive interaction between base and late-phase bias parameters , θ ( l ) b + β ( l ) . Late-phase rank-1 matrix weights . We also study a closely related late-phase weight model , where existing weight matrices – the base components , as before – are multiplied elementwise by rank-1 matrices ( Wen et al. , 2020 ) . For a given affine layer l , we define a late-phase weight matrix with resort to a pair of learnable vectors , φ ( l ) = u ( l ) v ( l ) T . Taking the Hadamard product with the base weight matrix yields the effective weights W ( l ) = φ ( l ) ◦ θ ( l ) . With this parameterization , we recover the ensemble proposed by Wen et al . ( 2020 ) , except that here it is generated late in training using our perturbative initialization ( Eq . 2 ) . Unlike BatchNorm layers , which include the shift parameter , rank-1 late-phase weights interact in a purely multiplicative manner with base weights . We study this model since it is easy to implement on neural networks which do not feature BatchNorm layers , such as standard long short-term memories ( LSTMs ; Hochreiter & Schmidhuber , 1997b ) . Hypernetworks with late-phase weight embeddings . Additionally , we generalize the late-phase weight models described above using hypernetworks ( Ha et al. , 2017 ) . A hypernetwork generates the parameters w of a given target neural network fw based on a weight embedding . In our framework , we can use a hypernetwork to implement the interaction function w = h ( θ , φ ) directly , with parameters θ corresponding to base weights and embeddings φ to late-phase weights . We experiment with linear hypernetworks and use the same hypernetwork to produce the weights of multiple layers , following Savarese & Maire ( 2019 ) ; Ha et al . ( 2017 ) ; von Oswald et al . ( 2020 ) . In this scheme , the weight embedding input specifies the target layer whose parameters are being generated . More specifically , the weight matrix for some layer l belonging to a group of layers g which share a hypernetwork is given byW ( g , l ) = θ ( g ) φ ( g , l ) , where θ ( g ) and φ ( g , l ) are appropriatelysized tensors . Sharing θ ( g ) over a layer group g allows countering an increase in the overall number of parameters . We parameterize our hypernetworks such that the weight embedding vectors φ ( g , l ) are small , and therefore cheap to ensemble . Late-phase classification layers . Finally , inspired by Lee et al . ( 2015 ) , in classification experiments we take the weights of the last linear layer as late-phase weights by default . In modern neural network architectures these layers do not usually comprise large numbers of parameters , and our architecture explorations indicated that it is typically beneficial to ensemble them . We therefore include W ( L ) in our late-phase weights φ , where W ( L ) denotes the weights of the final layer L .

Summary: This work suggests a variant of ensembling that is more compute-efficient. Specifically, it involves forking an ensemble only in the late stage of training, and forming this ensemble via a "low-dimentional" family. That is, instead of maintaining independent networks, maintain only "low-rank"-style perturbations of the base network (for various instanciations of "low-rank").