How many degrees of freedom do we need to train deep networks: a loss landscape perspective

1 INTRODUCTION . How many parameters are needed to train a neural network to a specified accuracy ? Recent work on two fronts indicates that the answer for a given architecture and dataset pair is often much smaller than the total number of parameters used in modern large-scale neural networks . The first is successfully identifying lottery tickets or sparse trainable subnetworks through iterative training and pruning cycles ( Frankle & Carbin , 2019 ) . Such methods utilize information from training to identify lowerdimensional parameter spaces which can optimize to a similar accuracy as the full model . The second is the observation that constrained training within a random , low-dimension affine subspace , is often successful at reaching a high desired train and test accuracy on a variety of tasks , provided that the training dimension of the subspace is above an empirically-observed threshold training dimension ( Li et al. , 2018 ) . These results , however , leave open the question of why low-dimensional training is so successful and whether we can theoretically explain the existence of a threshold training dimension . In this work , we provide such an explanation in terms of the high-dimensional geometry of the loss landscape , the initialization , and the desired loss . In particular , we leverage a powerful tool from high-dimensional probability theory , namely Gordon ’ s escape theorem , to show that this threshold training dimension is equal to the dimension of the full parameter space minus the squared Gaussian width of the desired loss sublevel set projected onto the unit sphere around initialization . This theory can then be applied in several ways to enhance our understanding of neural network loss landscapes . For a quadratic well or second-order approximation around a local minimum , we derive an analytic bound on this threshold training dimension in terms of the Hessian spectrum and the distance of the initialization from the minimum . For general models , this relationship can be used in reverse to measure important high dimensional properties of loss landscape geometry . For example , by performing a tomographic exploration of the loss landscape , i.e . training within random subspaces of varying training dimension , we uncover a phase transition in the success probability of hitting a given loss sub-level set . The threshold-training dimension is then the phase boundary in this transition , and our theory explains the dependence of the phase boundary on the desired loss sub-level set and the initialization , in terms of the Gaussian width of the loss sub-level set projected onto a sphere surrounding the initialization . Motivated by lottery tickets , we furthermore consider training not only within random dimensions , but also within optimized subspaces using information from training in the full space . Lottery tickets can be viewed as constructing an optimized , axis-aligned subspace , i.e . where each subspace dimension corresponds to a single parameter . What would constitute an optimized choice for general subspaces ? We propose two new methods : burn-in subspaces which optimize the offset of the subspace by taking a few steps along a training trajectory and lottery subspaces determined by the span of gradients along a full training trajectory ( Fig . 1 ) . Burn-in subspaces in particular can be viewed as lowering the threshold training dimension by moving closer to the desired loss sublevel set . For all three methods , we empirically explore the threshold training dimension across a range of datasets and architectures . Related Work : An important motivation of our work is the observation that training within a random , low-dimensional affine subspace can suffice to reach high training and test accuracies on a variety of tasks , provided the training dimension exceeds a threshold that was called the intrinsic dimension ( Li et al. , 2018 ) and which we call the threshold training dimension . However Li et al . ( 2018 ) provided no theoretical explanation for this threshold and did not explore the dependence of this threshold on the quality of the initialization . Our primary goal is to provide a theoretical explanation for the existence of this threshold in terms of the geometry of the loss landscape and the quality of initialization . Indeed understand- ing the geometry of high dimensional error landscapes has been a subject of intense interest in deep learning , see e.g . Dauphin et al . ( 2014 ) ; Goodfellow et al . ( 2014 ) ; Fort & Jastrzebski ( 2019 ) ; Ghorbani et al . ( 2019 ) ; Sagun et al . ( 2016 ; 2017 ) ; Yao et al . ( 2018 ) ; Fort & Scherlis ( 2019 ) ; Papyan ( 2020 ) ; Gur-Ari et al . ( 2018 ) ; Fort & Ganguli ( 2019 ) ; Papyan ( 2019 ) ; Fort et al . ( 2020 ) , or Bahri et al . ( 2020 ) for a review . But to our knowledge , the Gaussian width of sub-level sets projected onto a sphere surrounding initialization , a key quantity that determines the threshold training dimension , has not been extensively explored in deep learning . Another motivation for our work is contextualizing the efficacy of diverse more sophisticated network pruning methods like lottery tickets ( Frankle & Carbin , 2019 ; Frankle et al. , 2019 ) . Further work in this area revealed the advantages obtained by pruning networks not at initialization ( Frankle & Carbin , 2019 ; Lee et al. , 2018 ; Wang et al. , 2020 ; Tanaka et al. , 2020 ) but slightly later in training ( Frankle et al. , 2020 ) , highlighting the importance of early stages of training ( Jastrzebski et al. , 2020 ; Lewkowycz et al. , 2020 ) . We find empirically , as well as explain theoretically , that even when training within random subspaces , one can obtain higher accuracies for a given training dimension if one starts from a slightly pre-trained , or burned-in initialization as opposed to a random initialization . 2 AN EMPIRICALLY OBSERVED PHASE TRANSITION IN TRAINING SUCCESS . We begin with the empirical observation of a phase transition in the probability of hitting a loss sub-level set when training within a random subspace of a given training dimension , starting from some initialization . Before presenting this phase transition , we first define loss sublevel sets and two different methods for training within a random subspace that differ only in the quality of the initialization . In the next section we develop theory for the nature of this phase transition . Loss sublevel sets . Let ŷ = fw ( x ) be a neural network with weights w ∈ RD and inputs x ∈ Rk . For a given training set { xn , yn } Nn=1 and loss function ` , the empirical loss landscape is given by L ( w ) = 1N ∑N n=1 ` ( fw ( xn ) , yn ) . Though our theory is general , we focus on classification for our experiments , where y ∈ { 0 , 1 } C is a one-hot encoding of C class labels , ŷ is a vector of class probabilities , and ` ( ŷ , y ) is the cross-entropy loss . In general , the loss sublevel set S ( ) at a desired value of loss is the set of all points for which the loss is less than or equal to : S ( ) : = { w ∈ RD : L ( w ) ≤ } . ( 2.1 ) Random affine subspace . Consider a d dimensional random affine hyperplane contained in D dimensional weight space , parameterized by θ ∈ Rd : w ( θ ) = Aθ + w0 . Here A ∈ RD×d is a random Gaussian matrix with columns normalized to 1 and w0 ∈ RD a random weight initialization by standard methods . To train within this subspace , we initialize θ = 0 , which corresponds to randomly initializing the network at w0 , and we minimize L ( w ( θ ) ) with respect to θ. Burn-in affine subspace . Alternatively , we can initialize the network with parameters w0 and train the network in the full space for some number of iterations t , arriving at the parameters wt . We can then construct the random burn-in subspace w ( θ ) = Aθ +wt , ( 2.2 ) with A chosen randomly as before , and then subsequently train within this subspace by minimizing L ( w ( θ ) ) with respect to θ . The random affine subspace is identical to the burn-in affine subspace but with t = 0 . Exploring the properties of training within burn-in as opposed to random affine subspaces enables us to explore the impact of the quality of the initialization , after burning in some information from the training data , on the success of subsequent restricted training . Success probability in hitting a sub-level set . In either training method , achieving L ( w ( θ ) ) = implies that the intersection between our random or burn-in affine subspace and the loss sub-level set S ( ′ ) is non-empty for all ′ ≥ . As both the subspace A and the initialization w0 leading to wt are random , we are interested in the success probability Ps ( d , , t ) that a burn-in ( or random when t = 0 ) subspace of training dimension d actually intersects a loss sub-level set S ( ) : Ps ( d , , t ) ≡ P [ S ( ) ∩ { wt + span ( A ) } 6= ∅ ] . ( 2.3 ) Here , span ( A ) denotes the column space of A . Note in practice we can not guarantee that we obtain the minimal loss in the subspace , so we use the best value achieved by Adam ( Kingma & Ba , 2014 ) as an approximation . Thus the probability of achieving a given loss sublevel set via training constitutes an approximate lower bound on the probability in ( 2.3 ) that the subspace actually intersects the loss sublevel set . Threshold training dimension as a phase transition boundary . We will find that for any fixed t , the success probability Ps ( d , , t ) in the by d plane undergoes a sharp phase transition . In particular for a desired ( not too low ) loss it transitions sharply from 0 to 1 as the training dimension d increases . To capture this transition we define : Definition 2.1 . [ Threshold training dimension ] The threshold training dimension d∗ ( , t , δ ) is the minimal value of d such that Ps ( d , , t ) ≥ 1− δ for some small δ > 0 . For any chosen criterion δ ( and fixed t ) we will see that the curve d∗ ( , t , δ ) forms a phase boundary in the by d plane separating two phases of high and low success probability . This definition also gives an operational procedure to approximately measure the threshold training dimension : run either the random or burn-in affine subspace method repeatedly over a range of training dimensions d and record the lowest loss value found in the plane when optimizing via Adam . We can then construct the empirical probability across runs of hitting a given sublevel set S ( ) and the threshold training dimension is lowest value of d for which this probability crosses 1− δ ( where we employ δ = 0.1 ) .

Summary: The paper studies DNNs training when restricted to a random affine subspace (centered either at the initialization point $x_0$ or at a point $x_t$ obtained by training an unrestricted network for $t$ steps, which they call burn-in), and the probability that a loss lower than $\epsilon$ can be reached with a subspace of dimension $d$. They observe a sharp transition between pairs $(\epsilon,d)$ where this probability is either almost $0$ or $1$. They observe that the lower the dimension, the less likely one is to reach a loss $\epsilon$.  The authors show that this transition can be understood in terms of the Gaussian width of the projection of the sublevel set $S(\epsilon)$ to the unit ball around the center point $x_0$ (or $x_t$), which they call the local angular dimension. Gordon's escape theorem makes it possible to bound the transition between the two phases as a function of the local angular dimension.  Calculating the local angular dimension of the sublevel sets of the loss of DNNs is difficult, instead the authors study a quadratic loss, giving an approximation for the local angular dimension in this case. They then compare the empirical transition and the bound on the transition in terms of the local angular dimension between the two phases and find a good agreement with the theory.  Finally they propose a notion of lottery subspaces inspired by the lottery tickets paper and compare their results.