H OW MANY DEGREES OF FREEDOM DO WE NEED TO TRAIN DEEP NETWORKS : A LOSS LANDSCAPE PERSPECTIVE
A BSTRACT
A variety of recent works, spanning pruning, lottery tickets, and training within random subspaces, have shown that deep neural networks can be trained using far fewer degrees of freedom than the total number of parameters.
We analyze this phenomenon for random subspaces by ﬁrst examining the success probability of hitting a training loss sublevel set when training within a random subspace of a given training dimensionality.
We ﬁnd a sharp phase transition in the success probability from 0 to 1 as the training dimension surpasses a threshold.
This threshold training dimension increases as the desired ﬁnal loss decreases, but decreases as the initial loss decreases.
We then theoretically explain the origin of this phase transition, and its dependence on initialization and ﬁnal desired loss, in terms of properties of the high dimensional geometry of the loss landscape.
In particular, we show via Gordon’s escape theorem, that the training dimension plus the Gaussian width of the desired loss sub-level set, projected onto a unit sphere surrounding the initialization, must exceed the total number of parameters for the success probability to be large.
In several architectures and datasets, we measure the threshold training dimension as a function of initialization and demonstrate that it is a small fraction of the total parameters, implying by our theory that successful training with so few dimensions is possible precisely because the Gaussian width of low loss sub-level sets is very large.
Moreover, we compare this threshold training dimension to more sophisticated ways of reducing training degrees of freedom, including lottery tickets as well as a new, analogous method: lottery subspaces.
I NTRODUCTION
How many parameters are needed to train a neural network to a speciﬁed 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 ﬁrst 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 afﬁne 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 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 afﬁne subspace can sufﬁce 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 understanding the geometry of high dimensional error landscapes has been a subject of intense interest in deep learning, see e.g. Dauphin et al.
; Goodfellow et al.
(2014); Fort & Jastrzebski (2019);
Ghorbani et al.
; Sagun et al.
(2016; 2017); Yao et al.
; Fort & Scherlis (2019); Papyan (2020);
Gur-Ari et al.
; 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 efﬁcacy 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 ﬁnd 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.
A N 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 ﬁrst deﬁne 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 =
f w ( x ) be a neural network with weights w ∈ R D and inputs x ∈ R k .
For a given training set { x n , y n } Nn =1 and loss function (cid:96) , the empirical loss landscape is given by
L ( w ) = 1 N (cid:80) Nn =1 (cid:96) (cid:16) f w ( x n ) , y n (cid:17) .
Though our theory is general, we focus on classiﬁcation 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 (cid:96) (ˆ y , y ) is the cross-entropy loss.
In general, the loss sublevel set S ( (cid:15) ) at a desired value of loss (cid:15) is the set of all points for which the loss is less than or equal to (cid:15) : S ( (cid:15)
[Equation]
:= { w
[Equation]
[Equation]
[Equation]
Random afﬁne subspace.
Consider a d dimensional random afﬁne hyperplane contained in D dimensional weight space, parameterized by θ ∈ R d : w ( θ ) =
A θ + w 0 .
Here A ∈ R D ×
d is a random Gaussian matrix with columns normalized to 1 and w 0 ∈ R D a random weight initialization by standard methods.
To train within this subspace, we initialize θ = 0 , which corresponds to randomly initializing the network at w 0 , and we minimize L (cid:0) w ( θ ) (cid:1) with respect to θ .
Burn-in afﬁne subspace.
Alternatively, we can initialize the network with parameters w 0 and train the network in the full space for some number of iterations t , arriving at the parameters w t .
We can then construct the random burn-in subspace
[Equation]
with A chosen randomly as before, and then subsequently train within this subspace by minimizing
L (cid:0) w ( θ ) (cid:1) with respect to θ .
The random afﬁne subspace is identical to the burn-in afﬁne subspace but with t = 0 .
Exploring the properties of training within burn-in as opposed to random afﬁne 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 ( θ ))
= (cid:15) implies that the intersection between our random or burn-in afﬁne subspace and the loss sub-level set
S ( (cid:15) (cid:48) ) is non-empty for all (cid:15) (cid:48) ≥ (cid:15) .
As both the subspace A and the initialization w 0 leading to w t are random, we are interested in the success probability P s ( d, (cid:15), t ) that a burn-in (or random when t = 0 ) subspace of training dimension d actually intersects a loss sub-level set S ( (cid:15) ) :
[Equation]
Here, span( A ) denotes the column space of A .
Note in practice we cannot 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 ﬁnd that for any ﬁxed t , the success probability P s ( d, (cid:15), t ) in the (cid:15) by d plane undergoes a sharp phase transition.
In particular for a desired (not too low) loss (cid:15) it transitions sharply from 0 to 1 as the training dimension d increases.
To capture this transition we deﬁne:
Deﬁnition 2.1.
[Threshold training dimension]
The threshold training dimension d ∗ ( (cid:15), t, δ ) is the minimal value of d such that P s ( d, (cid:15), t ) ≥
1 − δ for some small δ > 0 .
For any chosen criterion δ (and ﬁxed t ) we will see that the curve d ∗ ( (cid:15), t, δ ) forms a phase boundary in the (cid:15) by d plane separating two phases of high and low success probability.
This deﬁnition also gives an operational procedure to approximately measure the threshold training dimension: run either the random or burn-in afﬁne subspace method repeatedly over a range of training dimensions d and record the lowest loss value (cid:15) found in the plane when optimizing via Adam .
We can then construct the empirical probability across runs of hitting a given sublevel set S ( (cid:15) ) and the threshold training dimension is lowest value of d for which this probability crosses 1 − δ (where we employ δ = 0 . 1 ).
A N EMPIRICAL DEMONSTRATION OF A TRAINING PHASE TRANSITION
In this section, we carry out this operational procedure, comparing random and burn-in afﬁne subspaces across a range of datasets and architectures.
We examined 3 architectures: 1) Conv-  which is a simple 2-layer CNN with 16 and 32 channels, ReLU activations and maxpool after each convolution followed by a fully connected layer; 2) Conv-3 which is a 3-layer CNN with 32, 64, and 64 channels but otherwise identical setup to Conv-2 ; and 3) ResNet20v1 as described in He et al.
(2016) with on-the-ﬂy batch normalization (Ioffe & Szegedy, 2015).
We perform experiments on 5 datasets: MNIST (LeCun et al., 2010), Fashion MNIST (Xiao et al., 2017), CIFAR-10 and
CIFAR-100 (Krizhevsky et al., 2014), and SVHN (Netzer et al., 2011).
Baselines and experiments were run for the same number of epochs for each model and dataset combination; further details on architectures, hyperparameters, and training procedures are provided in the appendix.
The code for the experiments was implemented in JAX (Bradbury et al., 2018).
Figure 2 shows results on the training loss for 4 datasets for both random and burn-in afﬁne subspaces with a Conv-2 .
We obtain similar results for the two other architectures (see Appendix).
Figureexhibits several broad and important trends.
First, for each training method within a random subspace, there is indeed a sharp phase transition in the success probability P s ( d, (cid:15), t ) in the (cid:15) (or equivalently accuracy) by d plane from 0 (white regions) to 1 (black regions).
Second, the threshold training dimension d ∗ ( (cid:15), t, δ ) (with δ = 0 . 1 ) does indeed track the tight phase boundary separating these two regimes.
Third, broadly for each method, to achieve a lower loss, or equivalently higher accuracy, the threshold training dimension is higher; thus one needs more training dimensions to achieve better performance.
Fourth, when comparing the threshold training dimension across all 4 methods on the same dataset (ﬁnal column of Figure 2) we see that at high accuracy (low loss (cid:15) ), increasing the amount of burn in lowers the threshold training dimension.
To see this, pick a high accuracy for each dataset, and follow the horizontal line of constant accuracy from left to right to ﬁnd the threshold training dimension for that accuracy.
The ﬁrst method encountered with the lowest threshold training dimension is burn-in with t = 16 .
Then burn-in with t = 8 has a higher threshold training dimension and so on, with random afﬁne having the highest.
Thus the main trend is, for some range of desired accuracies, burning more information into the initialization by training on the training data reduces the number of subsequent training dimensions required to achieve the desired accuracy.
Figure 3 shows the threshold training dimension for each accuracy level for all three models on
MNIST, Fashion MNIST and CIFAR-10, not only for training accuracy, but also for test accuracy.
The broad trends discussed above hold robustly for both train and test accuracy for all 3 models.
A THEORY OF THE PHASE TRANSITION IN TRAINING SUCCESS
Here we aim to give a theoretical explanation for the major trends observed empirically above, namely: (1) there exists a phase transition in the success probability P s ( d, (cid:15), t ) yielding a phase boundary given by a threshold training dimension d ∗ ( (cid:15), t, δ ) ; (2) at ﬁxed t and δ this threshold increases as the desired loss (cid:15) decreases (or desired accuracy increases), indicating more dimensions are required to perform better; (3) at ﬁxed (cid:15) and δ , this threshold decreases as the burn-in time t increases, indicating fewer training dimensions are required to achieve a given performance starting from a better burned-in initialization.
Our theory will build upon several aspects of high dimensional geometry which we ﬁrst review.
In particular we discuss, in turn, the notion of the Gaussian width of a set, then Gordon’s escape theorem, and then introduce a notion of local angular dimension of a set about a point.
Our ﬁnal result, stated informally, will be that the threshold training dimension plus the local angular dimension of a desired loss sub-level set about the initialization must equal the total number of parameters D .
As we will see, this succinct statement will conceptually explain the major trends observed empirically.
First we start with the deﬁnition of Gaussian width:
Deﬁnition 3.1 (Gaussian Width) .
The Gaussian width of a subset S ⊂ R D is given by (see Figure 4):
[Equation]
As a simple example, let S be a solid l 2 ball of radius r and dimension d (cid:28) D embedded in R D .
Then its Gaussian width for large D is well approximated by w ( S ) =
r √ d .
Gordon’s escape theorem.
The Gaussian width w ( S ) of a set S , at least when that set is contained in a unit sphere around the origin, in turn characterizes the probability that a random subspace intersects that set, through Gordon’s escape theorem (Gordon, 1988):
Theorem 3.1.
[Escape Theorem] Let S be a closed subset of the unit sphere in R D .
If k > w ( S ) 2 , then a d = D − k dimensional subspace Y drawn uniformly from the Grassmannian satisﬁes (Gordon, 1988):
[Equation]
A clear explanation of the proof can be found in Mixon (2014).
Thus, the bound says when k > w 2 ( S ) , the probability of no intersection quickly goes to 1 − (cid:15) for any (cid:15) > 0 .
Matching lower bounds which state that the intersection occurs with high probability when k ≤ w ( S ) 2 have been proven for spherically convex sets (Amelunxen et al., 2014).
Thus, this threshold is sharp except for the subtlety that you are only guaranteed to hit the spherical convex hull of the set (deﬁned on the sphere) with high probability.
When expressed in terms of the subspace dimension d = D − k , rather than its co-dimension k , these results indicate that a d dimensional subspace will intersect a closed subset S of the unit sphere around the origin with high probability if and only if d + w ( S ) 2 ≥ D , with a sharp transition at the threshold d ∗ =
D − w ( S ) 2 .
This is a generalization of the result that two random subspaces in R D of dimension d and d 2 intersect with high probability if and only if d + d 2 > D .
Thus we can think of w ( S ) 2 as playing a role analogous to dimension for sets on the centered unit sphere.
3. I NTERSECTIONS OF RANDOM SUBSPACES WITH GENERAL SUBSETS
To explain the training phase transition, we must now adapt Gordon’s escape theorem to a general loss sublevel set S in R D , and we must take into account that the initialization w t is not at the origin in weight space.
To do so, we ﬁrst deﬁne the projection of a set S onto a unit sphere centered at w t :
[Equation]
Then we note that any afﬁne subspace Y of the form in eq.
(2.2) centered at w t intersects S if and only if it intersects proj w t ( S ) .
Thus we can apply Gordon’s escape theorem to proj w t ( S ) to compute the probability of the training subspace in eq.
(2.2) intersecting a sublevel set S .
Since the squared Gaussian width of a set in a unit sphere plays a role analogous to dimension, we deﬁne:
Deﬁnition 3.2 (Local angular dimension) .
The local angular dimension of a general set S ⊂ R D about a point w t is deﬁned as
[Equation]
An escape theorem for general sets S and afﬁne subspaces now depends on the initialization w t also, and follows from the above considerations and Gordon’s original escape theorem:
Theorem 3.2.
[Main Theorem] Let S be a closed subset of R D .
If k > w (proj w t ( S )) 2 , then a d = D − k dimensional afﬁne subspace drawn uniformly from the Grassmannian and centered at w t satisﬁes:
[Equation]
To summarise this result in the context of our application, given an arbitrary loss sub-level set S ( (cid:15) ) , a training subspace of training dimension d starting from an initialization w t will hit the (convex hull) of the loss sublevel set with high probability when d + d local ( S ( (cid:15) ) , w t ) > D , and will miss it (i.e have empty intersection) with high probability when d + d local ( S ( (cid:15) ) , w t ) < D .
This analysis thus establishes the existence of a phase transition in the success probability P s ( d, (cid:15), t ) in eq.
(2.3), and moreover establishes the threshold training dimension d ∗ ( (cid:15), t, δ ) for small values of δ in deﬁnition 2.1:
[Equation]
Our theory provides several important insights on the nature of threshold training dimension.
Firstly, small threshold training dimensions can only arise if the local angular dimension of the loss sublevel set S ( (cid:15) ) about the initialization w t is close to the ambient dimension.
Second, as (cid:15) increases,
S ( (cid:15) ) becomes larger, with a larger d local ( S ( (cid:15) ) , w t ) , and consequently a smaller threshold training dimension.
Similarly, if w t is closer to S ( (cid:15) ) , then d local ( S ( (cid:15) ) , w t ) will be larger, and the threshold training dimension will also be lower (see ﬁg. 4).
This observation accounts for the observed decrease in threshold training dimension with increased burn-in time t .
Presumably, burning in information into the initialization w t for a longer time t brings the initialization closer to the sublevel set S ( (cid:15) ) , making it easier to hit with a random subspace of lower dimension.
This effect is akin to staring out into the night sky in a single random direction and asking with what probability we will see the moon; this probability increases the closer we are to the moon [Figure]
3. A PARADIGMATIC LOSS LANDSCAPE EXAMPLE : THE QUADRATIC WELL
[Figure]
To illustrate our theory we work out the paradigmatic example of a quadratic loss function L ( w ) =2 w T Hw where w ∈ R d and H ∈ R D
× D is a symmetric, positive deﬁnite Hessian matrix.
A sublevel set S ( (cid:15) ) of the quadratic well is an ellipsoidal body with principal axes along the eigenvectors ˆ e i of H .
The radius r i along principal axis ˆ
e
i obeys 12 λ i r 2 i = (cid:15) where λ i is the eigenvalue.
Thus r i = (cid:112) 2 (cid:15)/λ i , and so a large (small) Hessian eigenvalue leads to a narrow (wide) radius along each principal axis of the ellipsoid.
The overall squared Gaussian width of the sublevel set obeys w 2 ( S ( (cid:15) )) ∼ 2 (cid:15)
Tr( H − 1 ) =
(cid:80) i r 2 i , where ∼ denotes bounded above and below by this expression times positive constants (Vershynin, 2018).
We next consider training within a random subspace of dimension d starting from some initialization w 0 .
To compute the probability the subspace hits the sublevel set S ( (cid:15) ) , as illustrated in Fig.
4, we must project this ellipsoidal sublevel set onto the surface of the unit sphere centered at w 0 .
The Gaussian width of this projection proj w
( S ( (cid:15) )) depends on the distance R ≡ || w
|| from the initialization to the global minimum at w =
(i.e. it should increase with decreasing R ).
We can develop a crude approximation to this width as follows.
Assuming D (cid:29) 1 , the direction ˆ e
i will be approximately orthogonal to w 0 , so that | ˆ e i · x 0 | (cid:28) R .
The distance between the tip of the ellipsoid at radius r i along principal axis e i and the initialization w 0 is therefore ρ i = (cid:112) R 2 + r 2 i .
The ellipse’s radius r i then gets scaled down to approximately r i / (cid:112) R 2 + r 2 i when projected onto the surface of the unit sphere.
Note the subtlety in this derivation is that the point actually projected onto the sphere is where a line through the center of the sphere lies tangent to the ellipse rather than the point of fullest extent.
As a result, r i / (cid:112) R 2 + r
i provides a lower bound to the projected extent on the circle.
This is formalized in the appendix along with an explanation as to why this bound becomes looser with decreasing R .
Taken together, a lower bound on the local angular dimension of
S ( (cid:15) ) about w 0 is:
[Equation]
where again r i = (cid:112) 2 (cid:15)/λ i .
In Fig.
5, we plot the corresponding upper bound on the threshold training dimension, i.e. D − d local ( (cid:15), R ) alongside simulated results for two different Hessian spectra.
C HARACTERIZING AND COMPARING THE SPACE OF PRUNING METHODS
Training within random subspaces is primarily a scientiﬁc tool to explore loss landscapes.
It further has the advantage that we can explain theoretically why the number of degrees of freedom required to train can be far fewer than the number of parameters, as described above.
However, there are many other pruning methods of interest.
For example, the top row of Table 1 focuses on pruning to axis aligned subspaces, starting from random weight pruning, to lottery tickets which use information from training to prune weights, and/or choose the initialization if not rewound to init.
As one moves from left to right, one achieves better pruning (fewer degrees of freedom for a given accuracy).
Our analysis can be viewed as relaxing the axis-aligned constraint to pruning to general subspaces (second row of Table 1), either not using training at all (random afﬁne subspaces), or using information from training to only to choose the init (burn in afﬁne subspaces).
This analogy naturally leads to the notion of lottery subspaces described below (an analog of lottery tickets with axis-alignment relaxed to general subspaces) either rewound to init or not (last two entries of Table 1).
We compare the methods we have theoretically analyzed (random and burn-in afﬁne subspaces) to popular methods like lottery tickets rewound to init, and our new method of lottery subspaces, in an effort understand the differential efﬁcacy of various choices like axis-alignment, initialization, and the use of full training information to prune.
A full investigation of table 1 however is the subject of future work.
[Table]
tickets, rewound to step t
Generalsubspaces
Random afﬁne subspaces
Burn-in afﬁne at step t
Lottery subspaces
Lottery subspaces at step t
Lottery subspaces.
We ﬁrst train the network in the full space starting from an initialization w 0 .
We then form the matrix U d ∈ R D × d whose d columns are the top d principal components of entire the training trajectory w 0: T (see Appendix for details).
We then train within the subspace w ( θ ) =
U d θ + w t starting from a rewound initialization w t ( t = 0 is rewinding to the original init).
Since the subspace is optimized to match the top d dimensions of the training trajectory, we expect lottery subspaces to achieve much higher accuracies for a given training dimension than random or potentially even burn-in afﬁne subspaces.
This expectation is indeed borne out in Fig.
3 (purple lines above all other lines).
Intriguingly, very few lottery subspace training dimensions (in the range ofto 60 depending on the dataset and architecture) are required to attain full accuracy, and thus lottery subspaces can set a (potentially optimistic) target for what accuracies might be attainable by practical network pruning methods as a function of training dimension.
Empirical comparison of pruning methods.
Figure 6 presents empirical results comparing a subset of the methods in table 1: random afﬁne subspaces, burn-in afﬁne subspaces, lottery subspaces, and lottery tickets plotted against model compression ratio (deﬁned as parameters in full model over parameters, or training dimension, in restricted model).
The lottery tickets were constructed by training for 2 epochs, performing magnitude pruning of weights and biases, rewinding to initialization, and then training for the same number of epochs as the other methods.
Note that lottery tickets are created by pruning the full model (increasing compression ratio) in contrast to all other methods which are built up from a single dimension (decreasing compression ratio).
We observe lottery subspaces signiﬁcantly outperform random subspaces and lottery tickets at low training dimensions (high compression ratios), and we explore the spectrum of these spaces in more detail in the Appendix.
The comparison to lottery tickets at low compression ratios is limited by the fact that it is computationally expensive to project to higher dimensional subspaces and thus the highest training dimension we used was 4096 .
In the regions where the experiments overlap, the lottery tickets do not outperform random afﬁne subspaces, indicating that they are not gaining an advantage from the training information they utilize.
A notable exception is Conv-2 on CIFAR-10 in which the lottery tickets do outperform random afﬁne subspaces.
Finally, we note lottery tickets do not perform well at high compression ratios due to the phenomenon of layer collapse, where an entire layer gets pruned.
C ONCLUSION
The surprising ability of pruning methods like lottery tickets to achieve high accuracy with very few well chosen parameters, and even higher accuracy if not rewound to init, but to a later point in training, has garnered great interest in deep learning, but has been hard to analyze.
In this paper we focused on gaining theoretical insight into when and why training within a random subspace starting at different inits (or burn-ins) along a full training trajectory can achieve a given low loss (cid:15) .
We ﬁnd that this can occur only when the local angular dimension of the loss sublevel set S ( (cid:15) ) about the init is high, or close to the ambient dimension D .
Our theory also explains geometrically why longer burn-in lowers the number of degrees of freedom required to train to a given accuracy.
This is analogous to how rewinding to a later point in training reduces the size of lottery tickets, and indeed may share a similar mechanism.
Overall, these theoretical insights and comparisons begin to provide a high dimensional geometric framework to understand and assess the efﬁcacy of a wide range of network pruning methods at or beyond initialization. A CKNOWLEDGEMENTS
B.W.L. was supported by the Department of Energy Computational Science Graduate Fellowship program (DE-FG02-97ER25308).
S.G. thanks the Simons Foundation, NTT Research and an NSF
Career award for funding while at Stanford.
A E XPERIMENT SUPPLEMENT
The core experiment code is available on Github: https://github.com/ganguli-lab/ degrees-of-freedom .
The three top-level scripts are burn_in_subspace.py , lottery_subspace.py , and lottery_ticket.py .
Random afﬁne experiments were run by setting the parameter init_iters to 0 in the burn-in subspace code.
The primary automatic differentiation framework used for the experiments was JAX Bradbury et al.
The code was developed and tested using JAX v0.1.74, JAXlib v0.1.
and Flax v0.2.0 and run on an internal cluster using NVIDIA TITAN
Xp GPU’s.
Figures 7 and 8 show the corresponding empirical probability plots for the two other models considered in this paper: Conv-3 and ResNet20 .
These plots are constructed in the same manner as ﬁg. 2 except a larger value of δ was used since fewer runs were conducted ( δ was always chosen such that all but one of the runs had to successfully hit a training accuracy super-level set).
The data in these plots is from the same runs as ﬁgs.
3 and 6.
A. C OMPARISON TO L INEARIZED N ETWORKS (N EURAL T ANGENT K ERNEL )
For general neural networks, we do not expect to be able bound the local angular dimension; instead, we use the relationship between the threshold training dimension and local angular dimension to empirically probe this important property of the loss landscape as in the experiments of ﬁg. 3.
For a single basin, we can consider the second-order approximation to the landscape at the optimum which yields a quadratic well based on the spectrum of the Hessian at this point, corresponding to the experiments presented in ﬁg. 5 using a well with the appropriate spectrum.
In this section, we consider how linearizing via the neural tangent kernel (NTK) can be used as a tool to better approximate the landscape of the network around a single basin while being potentially more amenable to theoretical characterization than the full network.
[Figure]
For this experiment we ﬁrst train in the full network starting from initialization w 0 ∈ R D until we ﬁnd a local optimum w opt .
Instead of using the second-order approximation around this optimum given by the Hessian, we linearize the network around w opt via the NTK (Jacot et al., 2018).
In essence, if f ( w , x ) is the function that outputs the i th logit for a given input x we instead consider the following approximation which is a linear function in w :
[Equation]
At w opt , the full and linearized network are identical; however, in the linearized network there is only one basin which is around w opt .
We then compare these networks by returning to the initialization point w 0 and perform the experiment training within random afﬁne subspaces across a range of dimensions in both the full and linearized network.
Figure 9 shows the results of this experiment for both MNIST and Fashion MNIST using the model
Conv-2 .
In these two settings, the threshold training dimension of the linearized model approximates this property of the full model fairly well, indicating promise as a useful approximation to the true loss landscape around a basin.
Thus, we consider theoretically characterizing the local angular dimension of these linearized models interesting future work.
A. S PECTRA OF LOTTERY SUBSPACES
In our experiments, we formed lottery subspaces by storing the directions traveled during a full training trajectory and then ﬁnding the singular value decomposition of this matrix.
As we increased the subspace dimension, directions were added in order of descending singular values.
Figure 10 and the left panel of ﬁg. 11 show the associated spectra for the results presented in ﬁgs.
3 and 6.
Note that in the main text ﬁgures we plot the accuracies as a running max over the current and smaller dimensions.
This is because the subspaces are nested such that if we increase the dimension and ﬁnd a point of lower accuracy, it indicates a failure of optimization to ﬁnd the intersection as the higher accuracy point is still in the subspace.
In these supplement ﬁgures, we plot the recorded accuracies without this processing step for completeness.
We see that in several cases this optimization failure did occur as we moved to higher dimensions; we suspect this is related to how quickly the singular values fall off meaning the higher dimensions we add are much less informative.
The spectra are aligned with the train and test accuracy plots such that the value directly below a point on the curve corresponds to the singular value of the last dimension added to the subspace.
There were 10 runs for Conv-2 , 5 for Conv-3 , and 3 for ResNet20 .
Only the ﬁrst 5 out ofruns are displayed for the experiments with Conv-2 .
No signiﬁcant deviations were observed in the remaining runs.
From these plots, we observe that the spectra for a given dataset are generally consistent across architectures.
In addition, the decrease in accuracy after a certain dimension (particularly for CIFAR10) corresponds to the singular values of the added dimensions falling off towards 0.
The right panel of ﬁg. 11 shows a tangential observation that lottery subspaces for CIFAR-10 display a sharp transition in accuracy at d = 10 .
This provides additions evidence for the conjecture explored by Gur-Ari et al.
(2018), Fort & Ganguli (2019), and Papyan (2020) that the sharpest directions of the
Hessian and the most prominent logit gradients are each associated with a class.
Very little learning happens in these directions, but during optimization you bounce up and down along them so that the are prominent in the SVD of the gradients.
This predicts exactly the behavior observed.
A. A CCURACY OF B URN - IN I NITIALIZATION
Figure 12 shows a subset of the random afﬁne and burn-in afﬁne subspace experiments with a value plotted at dimension 0 to indicate the accuracy of the random or burn-in initialization.
This is to give context for what sublevel set the burn-in methods are starting out, enabling us to evaluate whether they are indeed reducing the threshold training dimension of sublevel sets with higher accuracy.
In most cases, as we increase dimension the burn-in experiments increase in accuracy above their initialization and at a faster pace than the random afﬁne subspaces.
A notable exception is Convon MNIST in which the burn-in methods appear to provide no advantage.
A. H YPERPARAMETERS
Random hyperplanes were chosen by sampling a D × d matrix of independent, standard Gaussians and then normalizing the columns to 1.
This is equivalent to sampling uniformly from the Grassmanian as required by theorem 3.2.
Optimization restricted to an afﬁne subspace was done using Adam Kingma & Ba (2014) with β 1 = 0 .
9 , β 2 = 0 .
999 , and (cid:15) =
10 − 7 .
We explored using 5 · 10 − 2 and 10 − 2 for the learning rate but 5 · 10 − 2 worked substantially better for this restricted optimization and was used in all experiments; a batch size of 128 was used.
The full model runs used the better result of 5 · 10 − 2 and 10 − 2 for the learning rate.
ResNet20v1 was run with on-the-ﬂy batch normalization
Ioffe & Szegedy (2015), meaning we simply use the mean and variance of the current batch rather than maintaining a running average.
Table 2 shows the number of epochs used for each dataset and architecture combination across all experiments.
3 epochs was chosen by default and then increased if the full model was not close to convergence.
[Table]
B T HEORY SUPPLEMENT
In this section, we provide additional details for our study of the threshold training dimension of the sublevel sets of quadratic wells.
We also derive the threshold training dimension of afﬁne subspaces to provide further intuition.
B. P ROOF : G AUSSIAN WIDTH OF SUBLEVEL SETS OF THE QUADRATIC WELL
In our derivation of eq.
(3.4), we employ the result that the Gaussian width squared of quadratic well sublevel sets is bounded as w 2 ( S ( (cid:15) ))
∼ 2 (cid:15)
Tr( H − 1 ) =
(cid:80) i r
i , i.e. bounded above and below by this expression times positive constants.
This follows from well-established bounds on the Gaussian width of an ellipsoid which we now prove.
In our proof, we will use an equivalent expression for the Gaussian width of set S :
[Equation]
Lemma B.1 (Gaussian width of ellipsoid) .
Let E be an ellipsoid in R D deﬁned by the vector r ∈ R D with strictly positive entries as:
[Equation]
Then w ( E ) 2 or the Gaussian width squared of the ellipsoid satisﬁes the following bounds:
[Equation]
Proof.
Let g ∼ N ( 0 , I D × D ) .
Then we upper-bound w ( E ) by the following steps:
[Equation]
Cauchy-Schwarz inequality
[Equation]
Deﬁnition of E
[Equation]
Jensen’s inequality
[Equation]
E [ w
[Equation]
giving the upper bound in the lemma.
For the lower bound, we will begin with a general lower bound for Gaussian widths using two facts.
The ﬁrst is that if (cid:15) i are i.i.d.
Rademacher random varaibles and, then (cid:15) i | g
i | ∼ N (0 , 1) .
Second, we have:
[Equation]
Then for the Gaussian width of a general set:
[Equation]
[Equation]
(cid:35)(cid:35)
Using (cid:15) i | g i | ∼ N (0 , 1)
[Equation]
(cid:34) sup
[Equation]
D (cid:88) i = (cid:15)
[Equation]
[Equation]
(cid:35)
Jensen’s Inequality = (cid:114)
[Equation]
E (cid:34) sup
[Equation]
∈
[Equation]
D (cid:88) i = (cid:15) i x i (cid:35)
All that remains for our lower bound is to show that for the ellipsoid E (cid:104) sup x ∈E (cid:80) Di =1 (cid:15)
i x i (cid:105) =
(cid:16)(cid:80)
Di =1 r 2 i (cid:17) 1 / 2 .
We begin by showing it is an upper-bound:
[Equation]
Using E is symmetric
[Equation]
Cauchy-Schwarz inequality
[Equation]
Deﬁnition of E
In the ﬁrst line, we mean that E is symmetric about the origin such that we can use (cid:15) i
= 1 for all i without loss of generality.
Finally, consider x such that x i =
r
i
(cid:16)(cid:80)
Di =1 r 2 i (cid:17) 1 / 2 .
For this choice we have x ∈ E and:
[Equation]
showing that equality is obtained in the bound.
Putting these steps together yields the overall desired lower bound:
[Equation]
With this bound in hand, we can immediately obtain the following corollary for a quadratic well deﬁned by Hessian H .
The Gaussian width is invariant under afﬁne transformation so we can shift the well to the origin.
Then note that S ( (cid:15) ) is an ellipsoid with r i = (cid:112) 2 (cid:15)/λ i and thus (cid:80) i r 2 i = (cid:15)
Tr( H − 1 ) .
Corollary B.1 (Gaussian width of quadratic sublevel sets) .
Consider a quadratic well deﬁned by
Hessian H ∈ R D × D .
Then the Gaussian width squared of the associated sublevel sets S ( (cid:15) ) obey the following bound:
[Equation]
B. D ETAILS ON THRESHOLD TRAINING DIMENSION UPPER BOUND
In section 3.2, we consider the projection of ellipsoidal sublevel sets onto the surface of a unit sphere centered at w 0 .
The Gaussian width of this projection proj w
( S ( (cid:15) )) will depend on the distance
R ≡ || w
|| from the initialization to the global minimum at w =
(i.e. it should increase with decreasing R ).
We used a crude approximation to this width as follows.
Assuming D (cid:29) 1 , the direction ˆ e
i will be approximately orthogonal to w 0 , so that | ˆ e i · x 0 | (cid:28) R .
The distance between the tip of the ellipsoid at radius r i along principal axis e i and the initialization w 0 is therefore ρ i = (cid:112) R 2 + r 2 i .
The ellipse’s radius r i then gets scaled down to approximately r i / (cid:112) R 2 + r 2 i when projected onto the surface of the unit sphere.
We now explain why this projected size is always a lower bound by illustrating the setup in two dimensions in ﬁg. 13.
As shown, the linear extent of the projection will always result from a line that is tangent to the ellipse.
For an ellipse ( x/a ) 2 + (( y − R ) /b ) 2 = 1 and a line y = cx in a two-dimensional space (we set the origin at the center of the unit circle), a line tangent to the ellipse must satisfy c = a/ √ R 2 − b 2 .
That means that the linear extent of the projection on unit circle will be a/ √ a 2 + R 2 − b 2 .
For a = 2 (cid:15)/λ i and R =
R , this is exactly Eq.
3.4 provided b = 0 .
The b (cid:54) =will always make the linear projections larger , and therefore Eq.
3.4 will be a lower bound on the projected Gaussian width.
Furthermore, this bound will be looser with decreasing R .
We then obtain a corresponding upper bound on the threshold training dimension, i.e. D − d local ( (cid:15), R ) .
B. T HRESHOLD TRAINING DIMENSION OF AFFINE SUBSPACES
In Section 3.2, we considered the threshold training dimension of the sublevel sets of a quadratic well and showed that it depends on the distance from the initialization to the set, formalized in eq.
As a point of contrast, we include a derivation of the threshold training dimension of a random afﬁne subspace in ambient dimension D and demonstrate that this dimension does not depend on distance to the subspace.
Intuitively this is because any dimension in the subspace is of inﬁnite or zero extent, unlike the quadratic sublevel sets which have dimensions of ﬁnite extent.
Let us consider a D -dimensional space for which we have a randomly chosen d -dimensional afﬁne subspace A deﬁned by a vector offset x 0 ∈ R D and a set of d orthonormal basis vectors { ˆ v i } di = that we encapsulate into a matrix M ∈ R d × D .
Let us consider another random n -dimensional afﬁne subspace B .
Our task is to ﬁnd a point x ∗ ∈ A that has the minimum (cid:96) 2 distance to the subspace B , i.e.:
[Equation]
[Table]
In words, we are looking for a point in the d -dimensional subspace A that is as close as possible to its closest point in the n -dimensional subspace B .
Furthermore, points within the subspace A can be parametrized by a d -dimensional vector θ ∈ R d
as x ( θ ) =
θ M + x 0 ∈ A ; for all choices of θ , the associated vector x is in the subspace A .
Without loss of generality, let us consider the case where the n basis vectors of the subspace B are aligned with the dimensions D − n, D − n + 1 , . . .
, D of the coordinate system (we can rotate our coordinate system such that this is true).
Call the remaining axes s = D − n the short directions of the subspace B .
A distance from a point x to the subspace B now depends only on its coordinates 1 , 2 , . .
, s .
Under our assumption of the alignment of subspace B we then have:
[Equation]
The only coordinates inﬂuencing the distance are the ﬁrst s values, and thus let us consider a R s subspace of the original R D only including those without loss of generality.
Now θ ∈ R d , M ∈ R
d × s and x 0 ∈ R d , and the distance between a point within the subspace A parameterized by the vector θ is given by:
[Equation]
The distance l attains its minimum for
[Equation]
yielding the optimality condition θ ∗ M = − x 0 .
There are 3 cases based on the relationship between d and s .
The overdetermined case, d > s .
In case d > s = D − n , the optimal θ ∗ = − x 0 M − 1 belongs to a ( d − s =
d + n − D )-dimensional family of solutions that attain 0 distance to the plane B .
In this case the afﬁne subspaces A and B intersect and share a ( d + n − D )-dimensional intersection.
A unique solution case, d = s .
In case of d = s = D − n , the solution is a unique θ ∗ = − x 0 M − 1 .
After plugging this back to the distance equation, we obtain θ is
[Equation]
The matrix M is square in this case and cancels out with its inverse M − 1 .
An underdetermined case, d < s .
In case of d < s , there is generically no intersection between the subspaces.
The inverse of M is now the Moore-Penrose inverse M + .
Therefore the closest distance θ :
[Equation]
Before our restriction from D → s dimensions, the matrix M consisted of d D -dimensional, mutually orthogonal vectors of unit length each.
We will consider these vectors to be component-wise random, each component with variance 1 / √ D to satisfy this condition on average.
After restricting our space to s dimensions, M ’s vectors are reduced to s components each, keeping their variance 1 / √ D .
They are still mutually orthogonal in expectation, but their length are reduced to √ s/ √ D .
The transpose of the inverse M + consists of vectors of the same directions, with their lengths scaled up to √ D/ √ s .
That means that in expectation, MM + is a diagonal matrix with d diagonal components set to 1 , and the remainder being 0 .
The matrix ( I − M + M ) contains ( s − d ) ones on its diagonal.
The projection (cid:107) x 0 ( I − M + M ) (cid:107) 2 is therefore of the expected value of (cid:107) x 0 (cid:107) 2 ( s − d ) 2 /D .
The expected distance between the d -dimensional subspace A and the d -dimensional subspace B is:
[Equation]
To summarize, for a space of dimension D , two afﬁne subspaces generically intersect provided that their dimensions d A and d B add up to at least the ambient (full) dimension of the space.
The exact condition for intersection is d A + d B ≥ D , and the threshold training dimension of subspace B is
D − d .
This result provides two main points of contrast to the quadratic well:
• Even extended directions are not inﬁnite for the quadratic well.
While in the case of the afﬁne subspaces even a slight non-coplanarity of the target afﬁne subspace and the random training subspace will eventually lead to an intersection, this is not the case for the sublevel sets of the quadratic well.
Even its small eigenvalues, i e. shallow directions, will still have a ﬁnite extent for all ﬁnite (cid:15) .
• Distance independence of the threshold training dimension.
As a result
of the dimensions having ﬁnite extent, the distance independence of threshold training dimension for afﬁne subspaces does not carry over to the case of quadratic wells.
In the main text, this dependence on distance is calculated by projecting the set onto the unit sphere around the initialization enabling us to apply Gordon’s Escape Theorem.