Gradient Descent Happens in a Tiny Subspace

1 INTRODUCTION . Stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) and its variants are used to train nearly every large-scale machine learning model . Its ubiquity in deep learning is connected to the efficiency at which gradients can be computed ( Rumelhart et al. , 1985 ; 1986 ) , though its success remains somewhat of a mystery due to the highly nonlinear and nonconvex nature of typical deep learning loss landscapes ( Bottou et al. , 2016 ) . In an attempt to shed light on this question , this paper investigates the dynamics of the gradient and the Hessian matrix during SGD . In a common deep learning scenario , models contain many more tunable parameters than training samples . In such “ overparameterized ” models , one expects generically that the loss landscape should have many flat directions : directions in parameter space in which the loss changes by very little or not at all ( we will use “ flat ” colloquially to also mean approximately flat ) .1 Intuitively , this may occur because the overparameterization leads to a large redundancy in configurations that realize the same decrease in the loss after a gradient descent update . One local way of measuring the flatness of the loss function involves the Hessian . Small or zero eigenvalues in the spectrum of the Hessian are an indication of flat directions ( Hochreiter & Schmidhuber , 1997 ) . In Sagun et al . ( 2016 ; 2017 ) , the spectrum of the Hessian for deep learning crossentropy losses was analyzed in depth.2 These works showed empirically that along the optimization trajectory the spectrum separates into two components : a bulk component with many small eigenvalues , and a top component of much larger positive eigenvalues.3 Correspondingly , at each point in parameter space the tangent space has two orthogonal components , which we will call the bulk subspace and the top subspace . The dimension of the top subspace is k , the number of classes in the classification objective . This result indicates the presence of many flat directions , which is consistent with the general expectation above . In this work we present two novel observations : • First , the gradient of the loss during training quickly moves to lie within the top subspace of the Hessian.4 Within this subspace the gradient seems to have no special properties ; its direction appears random with respect to the eigenvector basis . 1 Over parameterization suggests many directions in weight space where the loss does not change . This implies that the curvature of the loss , captured through the hessian spectrum , vanishes in these directions . In the remainder of the paper , we use the term flat , as is common in the literature , in a slightly broader sense to describe this curvature of the loss surface , not necessarily implying vanishing of the gradient . 2For other recent work on the spectrum of the Hessian as it relates to learning dynamics , see Pascanu et al . ( 2014 ) ; Dauphin et al . ( 2014 ) ; Chaudhari et al . ( 2016 ) . 3We provide our own evidence of this in Appendix B and provide some additional commentary . 4 This is similar to Advani & Saxe ( 2017 ) , who found that a large fraction of the weights in overparameter- ized linear models remain untrained from their initial values ( thus the gradient in those directions vanishes ) . • Second , the top Hessian eigenvectors evolve nontrivially but tend not to mix with the bulk eigenvectors , even over hundreds of training steps or more . In other words , the top subspace is approximately preserved over long periods of training . These observations are borne out across model architectures , including fully connected networks , convolutional networks , and ResNet-18 , and data sets ( Figures 1 , 2 , Table 1 , Appendices C-D ) . Taken all together , despite the large number of training examples and even larger number of parameters in deep-learning models , these results seem to imply that learning may happen in a tiny , slowly-evolving subspace . Indeed , consider a gradient descent step−ηg where η is the learning rate and g the gradient . The change in the loss to leading order in η is δL = −η ‖g‖2 . Now , let gtop be the projection of g onto the top subspace of the Hessian . If the gradient is mostly contained within this subspace , then doing gradient descent with gtop instead of g will yield a similar decrease in the loss , assuming the linear approximation is valid . Therefore , we think this may have bearing on the question of how gradient descent can traverse such a nonlinear and nonconvex landscape . To shed light on this mechanism more directly , we also present a toy model of softmax regression trained on a mixture of Gaussians that displays all of the effects observed in the full deep-learning scenarios . This isn ’ t meant as a definitive explanation , but rather an illustrative example in which we can understand these phenomenon directly . In this model , we can solve the gradient descent equations exactly in a limit where the Gaussians have zero variance.5 We find that the gradient is concentrated in the top Hessian subspace , while the bulk subspace has all zero eigenvalues . We then argue and use empirical simulations to show that including a small amount of variance will not change these conclusions , even though the bulk subspace will now contain non-zero eigenvalues . Finally , we conclude by discussing some consequences of these observations for learning and optimization , leaving the study of improving current methods based on these ideas for future work . 2 THE GRADIENT AND THE TOP HESSIAN SUBSPACE . In this section , we present the main empirical observations of the paper . First , the gradient lies predominantly in the smaller , top subspace . Second , in many deep learning scenarios , the top and bulk Hessian subspaces are approximately preserved over long periods of training . These properties come about quickly during training . In general , we will consider models with p parameters denoted by θ and a cross-entropy loss function L ( θ ) . We will generally use g ( θ ) ≡ ∇L ( θ ) for the gradient and H ( θ ) ≡ ∇∇TL ( θ ) for the Hessian matrix of the loss function at a point θ in parameter space . A gradient descent update with learning rate η at step t is θ ( t+1 ) = θ ( t ) − η g ( θ ( t ) ) , ( 1 ) and for stochastic gradient descent we estimate the gradient using a mini-batch of examples . 2.1 THE GRADIENT CONCENTRATES IN THE TOP SUBSPACE . For a classification problem with k classes , consider a point θ in parameter space where the Hessian spectrum decomposes into a top and a bulk subspace as discussed above.6 Now , let Vtop be the subspace of tangent space spanned by the top k eigenvectors of the Hessian ; we will call this the top subspace . Let Vbulk be the orthogonal subspace . The gradient at this point can be written as a sum g ( θ ) = gtop + gbulk where gtop ( gbulk ) is the orthogonal projection of g onto Vtop ( Vbulk ) . The fraction of the gradient in the top subspace is then given by ftop ≡ ‖gtop‖2 ‖g‖2 . ( 2 ) 5Other works where the dynamics of gradient descent were analyzed directly include Fukumizu ; Saxe et al . ( 2013 ) ; Arora et al . ( 2018 ) . 6As we have mentioned , this decomposition was originally found in Sagun et al . ( 2016 ; 2017 ) , and we provide additional discussion of the Hessian spectrum in Appendix B . Figure 1 shows this fraction for common datasets and network architectures during the early stages of training . The fraction starts out small , but then quickly grows to a value close to 1 , implying that there is an underlying dynamical mechanism that is driving the gradient into the top subspace . For these experiments , training was carried out using vanilla stochastic gradient descent on a variety of realistic models and dataset combinations . However , measurements of the gradient and Hessian were evaluated using the entire training set . Additionally , all of our empirical results have been replicated in two independent implementations . ( See Appendix A for further details on the numerical calculation . ) In the next subsection we provide evidence that this effect occurs in a broader range of models . 2.2 HESSIAN-GRADIENT OVERLAP . In this section , we consider the overlap between the gradient g and the Hessian-gradient productHg during training , defined by overlap ( g , Hg ) ≡ g THg ‖g‖ · ‖Hg‖ . ( 3 ) The overlap takes values in the range [ −1 , 1 ] . Computing the overlap is computationally much more efficient than computing the leading Hessian eigenvectors . We argue below that the overlap becomes big ( of order 1 ) if the gradient is contained in the top subspace of the Hessian . We can use the overlap as a proxy measurement : if the overlap is large , we take that to be evidence that the gradient lives mostly in the top subspace . We measured the overlap in a range of deep learning scenarios , and the results are shown in Table 1 . In these experiments we consider fully-connected networks , convolutional networks , a ResNet-18 ( He et al. , 2016 ) , as well as networks with no hidden layers , models with dropout ( Srivastava et al. , 2014 ) and batch-norm ( 201 ) , models with a smooth activation function ( e.g . softplus instead of ReLU ) , models trained using different optimization algorithms ( SGD and Adam ) , models trained using different batch sizes and learning rates , models trained on data with random labels ( as was considered by Zhang et al . ( 2016 ) ) , and a regression task . The overlap is large for the gradient and Hessian computed on a test set as well ( except for the case where the labels are randomized ) . In addition , we will see below that the effect is not unique to models with cross-entropy loss ; a simpler version of the same effect occurs for linear and deep regression models . In all the examples that we checked , the overlap was consistently close to one after some training . Let us now show that the overlap tends to be large for a random vector in the top Hessian subspace . Let λi be the Hessian eigenvalues in the top subspace of dimension k , with corresponding eigenvectors vi . Let w be a vector in this subspace , with coefficients wi in the vi basis . To get an estimate for the overlap equation 3 , we choose w to be at a random vertex on the unit cube , namely choosing wi = ±1 at random for each i . The overlap is then given by overlap ( w , Hw ) = ∑k i λiw 2 i√ ( ∑k j w 2 j ) ( ∑k l λ 2 lw 2 l ) = ∑ki λi√ k ∑k j λ 2 j . ( 4 ) As discussed above , in typical scenarios the spectrum will consist of k positive eigenvalues where k is the number of classes and all the rest close to zero . To get a concrete estimate , we approximate this spectrum by taking λi ∝ i ( a rough approximation , empirically , when k = 10 ) , and take k large so that we can compute the sums approximately . This estimate for the overlap is √ 3/4 ≈ 0.87 , which is in line with our empirical observations . This should compared with a generic random vector not restricted to the top subspace , which would have an overlap much less than 1 . We have verified empirically that a random unit vector w in the top Hessian subspace will have a large overlap with Hw , comparable to that of the gradient , while a random unit vector in the full parameter space has negligible overlap . Based on these observations , we will take the overlap equation 3 to be a proxy measurement for the part of the gradient that lives in the top Hessian subspace .

Summary: The authors build on recent works that study the spectrum of the Hessian of deep networks (e.g. Sagun et al). Previous work argues that Hessian is approximately low-rank i.e. there are few large eigenvalues and many small eigenvalues. This work argues that after some training, large eigenvectors of the Hessian converges to a subspace and stays there.