A N ANALYTIC THEORY OF GENERALIZATION DYNAMICS AND TRANSFER LEARNING IN DEEP LINEAR NETWORKS
A BSTRACT
Much attention has been devoted recently to the generalization puzzle in deep learning: large, deep networks can generalize well, but existing theories bounding generalization error are exceedingly loose, and thus cannot explain this striking performance.
Furthermore, a major hope is that knowledge may transfer across tasks, so that multi-task learning can improve generalization on individual tasks.
However we lack analytic theories that can quantitatively predict how the degree of knowledge transfer depends on the relationship between the tasks.
We develop an analytic theory of the nonlinear dynamics of generalization in deep linear networks, both within and across tasks.
In particular, our theory provides analytic solutions to the training and testing error of deep networks as a function of training time, number of examples, network size and initialization, and the task structure and
SNR.
Our theory reveals that deep networks progressively learn the most important task structure ﬁrst, so that generalization error at the early stopping time primarily depends on task structure and is independent of network size.
This suggests any tight bound on generalization error must take into account task structure, and explains observations about real data being learned faster than random data.
Intriguingly our theory also reveals the existence of a learning algorithm that proveably out-performs neural network training through gradient descent.
Finally, for transfer learning, our theory reveals that knowledge transfer depends sensitively, but computably, on the SNRs and input feature alignments of pairs of tasks.
I NTRODUCTION
Many deep learning practitioners closely monitor both training and test errors, hoping to achieve both a small training error and a small generalization error, or gap between testing and training errors.
Training is usually stopped early, before overﬁtting sets in and increases the test error.
This procedure often results in large networks that generalize well on structured tasks, raising an important generalization puzzle (Zhang et al., 2016): many existing theories that upper bound generalization error (Bartlett & Mendelson, 2002; Neyshabur et al., 2015; Dziugaite & Roy, 2017;
Golowich et al., 2017; Neyshabur et al., 2017; Bartlett et al., 2017; Arora et al., 2018, e.g) in terms of various measures of network complexity yield very loose bounds.
Therefore they cannot explain the impressive generalization capabilities of deep nets.
In the absence of any such tight and computable theory of deep network generalization error, we develop an analytic theory of generalization error for deep linear networks.
Such networks exhibit highly nonlinear learning dynamics (Saxe et al., 2013a;b) including many prominent phenomena like learning plateaus, saddle points, and sudden drops in training error.
Moreover, theory developed for the learning dynamics of deep linear networks directly inspired better initialization schemes for nonlinear networks (Schoenholz et al., 2016; Pennington et al., 2017; 2018).
Here we show that deep linear networks also provide a good theoretical model for generalization dynamics.
In particular we develop an analytic theory for both the training and test error of a deep linear network as a function of training time, number of training examples, network architecture, initialization, and task structure and SNR.
Our theory matches simulations and reveals that deep networks with small weight initialization learn the most important aspects of a task ﬁrst.
Thus the optimal test error at the early stopping time depends largely on task structure and SNR, and not on network architecture, as long as the architecture is expressive enough to attain small training error.
Thus our exact analysis of generalization dynamics reveals the important lesson that any theory that seeks to upper bound generalization error based only on network architecture, and not on task structure, is likely to yield exceedingly loose upper bounds.
Intriguingly our theory also reveals a non-gradient-descent learning algorithm that proveably out-performs neural network training through gradient descent.
We also apply our theory to multi-task learning, which enables knowledge transfer from one task to another, thereby further lowering generalization error (Dong et al., 2015; Rusu et al., 2015; Luong et al., 2016, e.g.).
Moreover, knowledge transfer across tasks may be key to human generalization capabilities (Hansen et al., 2017; Lampinen et al., 2017).
We provide an analytic theory for how much knowledge is transferred between pairs of tasks, and we ﬁnd that it displays a sensitive but computable dependence on the relationship between pairs of tasks, in particular, their SNRs and feature space alignments.
We note that a related prior work (Advani & Saxe, 2017) studied generalization in shallow and deep linear networks, but that work was limited to networks with a single output, thereby precluding the possibility of addressing the issue of transfer learning.
Moreover, analyzing networks with a single output also precludes the possibility of addressing interesting tasks that require higher dimensional outputs, for example in language (Dong et al., 2015, e.g.), generative models (Goodfellow et al., 2014, e.g), and reinforcement learning (Mnih et al., 2015; Silver et al., 2016, e.g).
T HEORETICAL FRAMEWORK
We work in a student-teacher scenario in which we consider an ensemble of low rank, noisy teacher networks that generate training data for a potentially more complex student network, and deﬁne the training and test errors whose dynamics we wish to understand.
2. A N ENSEMBLE OF LOW - RANK NOISY TEACHERS
We ﬁrst consider an ensemble of 3-layer linear teacher networks with N i units in layer i , and weight matrices W 21 ∈ R N 2 × N 1 and W 32 ∈ R N 3 × N 2 between the input to hidden, and hidden to output layers, respectively.
The teacher network thus computes the composite map y
=
Wx , where
W ≡ W 32 W 21 .
Of critical importance is the singular value decomposition (SVD) of W :
[Equation]
Where U ∈ R N 3 × N 2 and V ∈ R N 1 × N 2 are both matrices with orthonormal columns and S is an
N 2 × N 2 diagonal matrix.
We construct a random teacher by picking U and V to be random matrices with orthonormal columns and choosing O (1) values for the diagonal elements of S .
We work in the limit N 1 , N 3 → ∞ with an O (1) aspect ratio A = N 3 /N 1 ∈
0 , 1] so that the teacher has fewer outputs than inputs.
Also, we hold N 2 ∼ O (1) , so the teacher has a low, ﬁnite rank, and we study generalization performance as a function of the N 2 teacher singular values.
We further assume the teacher generates noisy outputs from a set of N 1 orthonormal inputs:
[Equation]
This training set yields important second-order training statistics that will guide student learning:
[Equation]
Here the input covariance Σ 11 is assumed to be white (a common pre-processing step), the inputoutput covariance Σ 31 is simpliﬁed using (2), and Z ∈ R N 3 × N 1 is the noise matrix, whose µ ’th column is z µ .
Its matrix elements z µ i are drawn iid.
from a Gaussian with zero mean and variance σ 2 z /N 1 .
The noise scaling is chosen so the singular values of the teacher W and the noise Z are both
O (1) , leading to non-trivial generalization effects.
As generalization performance will depend on the ratio of teacher singular values to the noise variance parameter σ 2 z , we simply set σ z = 1 in the following.
Thus we can think of teacher singular values as signal to noise ratios (SNRs).
Finally, we note that while we focus for ease of exposition in the main paper on the case of one hidden layer networks and a full orthonormal basis of P = N 1 training inputs in the main paper, neither of these assumptions are essential to our theory.
Indeed in Section 3.4 and App.
A we extend our theory to networks of arbitrary depth, and in App.
G we extend our theory to the case of white inputs with P (cid:54) =
N 1 , obtaining a good match between theory and experiment in both cases.
S TUDENT TRAINING AND TEST ERROR
Now consider a student network with N i units in each layer.
We assume the ﬁrst and last layers match the teacher (i.e. N 1 = N 1 and N 3 = N 3 ) but N 2 ≥ N 2 , allowing the student to have more hidden units than the teacher.
We also consider deeper students (see below and App.
A).
Now consider any student whose input-output map is given by y
= W 32 W 21 ≡ Wx .
Its training error on the teacher dataset in (2) and its test error over a distribution of new inputs are given by
[Equation]
respectively.
Here ˆx µ and ˆy µ are the noisy training set inputs and outputs in (2), whereas x denotes a random test input drawn from zero mean Gaussian with identity covariance, y µ =
Wx µ is noise free teacher output, and (cid:104)·(cid:105) denotes an average w.r.t the distribution of the test input x .
Due to the orthonormality of the training and isotropy of the test inputs, both ε train and ε test can be expressed as
[Equation]
Both ε train and ε test can be further expressed in terms of the student, training data and teacher SVDs, which we denote by W = USV T , Σ 31 = ˆUˆS ˆV T , and W = U S V T respectively.
Speciﬁcally,
[Equation]
Thus as the student learns, its training and test error dynamics depends on the alignment of the time-evolving student singular modes { s α , u α , v α } with the ﬁxed training data { ˆ s α , ˆu α , ˆv α } and teacher { s α , u α , v α } singular modes respectively.
S INGLE TASK GENERALIZATION DYNAMICS : THEORY AND EXPERIMENT
Here we derive and numerically test analytic formulas for both the training and test errors of a student network as it learns from training data generated from a teacher network.
We explore the dependence of these quantitites on the student network size, student initialization, teacher SNR, and training time.
3. S TUDENT TRAINING DYNAMICS AND TRAINING - ALIGNED (TA) NETWORKS
We assume the student weights undergo batch gradient descent with learning rate λ on the training error (cid:80) µ || ˆy µ
[Equation]
[Equation] 21 ˆx µ || 22 , which for small λ is well approximated by the differential equations:
[Equation]
strength.
(a) shows how modes of different singular value are learned, (b) shows that there is a wave of learning that picks up singular dimensions with smaller and smaller singular values as t → ∞ .
(where τ ≡ 1 /λ ), which must be solved from an initial set of student weights at time t =
(Saxe et al., 2013a).
We consider two classes of student initializations.
The ﬁrst initialization corresponds to a random student where the weights W 21 and W 32 are chosen such that the composite map
W = W 32 W 21 has an SVD W = (cid:15) UV T , where U and V are random singular vector matrices and all student singular values are (cid:15) .
As such a random student learns, the composite map undergoes a time dependent evolution W ( t ) =
U ( t ) S ( t ) V ( t ) T
= (cid:80) N 2 α =
s α ( t ) u α ( t ) v α ( t ) T .
For white inputs, as t → ∞ , W → Σ 31 , and so the time-dependent student singular modes { s α ( t ) , u α ( t ) , v ( t ) } converge to the training data singular modes { ˆ s α , ˆu α , ˆv α } .
However, the explicit dynamics of the student singular modes can be difﬁcult to obtain analytically from random initial conditions.
Thus we also consider a special class of training aligned (TA) initial conditions in which W 21 and
W 32 are chosen such that the composite map W = W 32 W 21 has an SVD W = (cid:15) ˆU ˆV T .
That is, the TA network (henceforth referred to simply as the TA) has the same singular vectors as the training data covariance Σ 31 , but has all singular values equal to (cid:15) .
As shown in (Saxe et al., 2013a), as the
TA learns according to (8), the singular vectors of its composite map W remain unchanged, while the singular values evolve as s α ( t )
=
s ( t, ˆ s α ) , where the learning curve function s ( t, ˆ s ) as well as its functional inverse t ( s, ˆ s ) is given by
[Equation]
Here the function s ( t, ˆ s ) describes analytically how each training set singular value
ˆ s drives the dynamics of the corresponding TA singular value s , and for notational simplicity, we have suppressed the dependence of s ( t, ˆ s ) on τ and the initial condition (cid:15) .
As shown in Fig.
1A, for each ˆ s , s ( t, ˆ s ) is a sigmoidal learning curve that undergoes a sharp transition around time t/τ = 12ˆ s ln (ˆ s/(cid:15) − 1) , at which it rises from its small initial value of (cid:15) at t = 0 to its asymptotic value of ˆ s as t/τ → ∞ .
Alternatively, we can plot s ( t, ˆ s ) /
ˆ s as a function of ˆ s for different training times t/τ , as in Fig.
1B.
This shows that TA learning corresponds to a singular mode detection wave which progressively sweeps from large to small singular values.
At any given training time t , training data modes with singular values ˆ s > t/τ have been learned, while those with singular values ˆ s < t/τ have not.
While the TA is more sophisticated than the random student, since it already knows the singular vectors of the training data before learning, we will see that the analytic solution for the TA learning dynamics provides a good approximation to the student learning dynamics, not only for the training error, as shown in (Saxe et al., 2013a), but also for the generalization error as shown below.
The results in this section assume a single hidden layer, but Saxe et al.
(2013a) derived t ( s, ˆ s ) for networks of arbitrary depth and we apply our theory to some deeper networks.
The general differential equation and derivations for deeper networks can be found in Appendix A.
3. H OW THE TEACHER IS BURIED IN THE TRAINING DATA : A RANDOM MATRIX ANALYSIS
In the previous section, we reviewed an exact analytic solution for the composite map of a TA network, namely that its singular modes are related to those of the training data through the relation
[Equation]
However, computation of the generalization error through (5) then requires understanding how the teacher singular modes of W are buried within the noisy training data singular modes of Σ 31 through the relation (3).
Since the input matrix ˆX is orthonormal, Σ 31 is simply a perturbation of the low rank teacher W by a high dimensional noise matrix Z .
The relation between the singular modes of a low rank matrix and its noise perturbed version has been studied extensively in Benaych-Georges &
Nadakuditi (2012), in the high dimensional limit we are working in, namely N 1 , N 3 → ∞ with the aspect ratio A = N 3 /N 1 ∈ (0 , 1] , and N 2 ∼ O (1) .
In this limit, the top N 2 singular values and vectors of Σ 31 converge to ˆ s ( s α ) , where the transfer function from a teacher singular value s to a training data singular value
ˆ s is given by the function
[Equation]
[Equation]
The associated top N 2 singular vectors of Σ 31 can also acquire a nontrivial overlap with the Nmodes of the teacher through the relation | ˆu α · u α | | ˆv α · v α |
= O
( s α ) , where the singular vector overlap function is given by
[Equation]
The rest of the N 3 − N 2 singular vectors of Σ 31 are orthogonal to the top N 2 ones, and their singular values are distributed according to the the Marchenko-Pastur (MP) distribution:
[Equation]
Overall, these equations describe a singular vector phase transition in the training data, as illustrated in Fig.
2BC.
For example in the case of no teacher, the training data is simply noise and the singular values of Σ 31 are distributed as an MP sea spread between 1 ± √A .
When one adds a teacher, how each teacher singular mode is imprinted on the training data depends crucially on the teacher singular value s , and the nature of this imprinting undergoes a phase transition at s =
A 1 / 4 .
For s ≤ A 1 / 4 , the teacher mode SNR is too low and this mode is not imprinted in the noisy training data; the associated training data singular value ˆ s remains at the edge of the MP sea at 1 + √A , and the overlap O ( s ) between training and teacher singular vectors remains zero.
However, when s > A 1 / 4 , this teacher mode is imprinted in the training data; there is an associated training data singular value
ˆ s that pops out of the MP sea (Fig. 2AB).
However, the training data singular value emerges at a position ˆ s > s that is inﬂated by the noise, though the inﬂation effect decreases at larger s , with the ratio ˆ s/s approaching the unity line as s becomes large (Fig. 2B).
Similarly, the corresponding training data singular vectors acquire a non-trivial overlap with the teacher singular vectors when s > A 1 / 4 , and the alignment approaches unity as s increases (Fig.
2C).
3. P UTTING IT ALL TOGETHER : AN ANALYTIC THEORY OF GENERALIZATION DYNAMICS
Based on an analytic understanding of how the singular mode structure { s α , u α , v α } of the teacher
W is imprinted in the modes { ˆ s α , ˆu α , ˆv α } of the training data covariance Σ 31 through (11), (12) and (13), and in turn how this training data singular structure drives the time evolving singular modes of a
TA network { s α ( t ) , ˆu α , ˆv α } of through (9), we can now derive analytic expressions for ε train and ε test in (6) and (7), for a TA network.
We will also show that these learning curves closely approximate those of a random student with time-evolving singular vectors { u α ( t ) , v α ( t ) } , and match on several key aspects.
First, inserting the TA dynamics in (10) into ε train in (6), we obtain ε
[Equation]
[Equation]
Here s α ( t ) =
s (ˆ s α , t ) as deﬁned in (9) are the TA singular values, and ˆ s α
= ˆ s ( s α ) as deﬁned in (11) are the training data singular values associated with the teacher singular values s α .
Also (cid:104)·(cid:105)
R denotes an average with respect to the MP distribution in (13) over a region R .
Two distinct regions contribute to training error.
First R in contains those top N 2 − N 2 training data singular values that do not correspond to the N 2 singular values of the teacher but will be learned by a rank Nstudent.
Second, R out corresponds to the remaining N 3 − N 2 lowest training data singular values that cannot be learned by a rank N 2 student.
In terms of the MP distribution, R out =
[1 − √A , f ] and R in = [ f, 1 + √A ] , where f is the point at which the MP density has 1 − N 2 /N 3 of its mass to the left and N 2 /N 3 of its mass to the right.
In the simple case of a full rank student, f = 1 −√A , and one need only integrate over R in which is the entire range.
Equation (14) for ε train makes it manifest that it will go to zero for a full rank student as its singular values approach those of the training data.
Of course the test error can behave very differently.
Inserting the TA training dynamics in (10) into ε test in (7), and using (11), (12) and (13) to relate training data to the teacher, we ﬁnd ε test ( [Equation]
[Equation]
Together (14) and (15) constitute a complete theory of generalization dynamics in terms of the structure of the data distribution (i.e. the teacher rank N 2 , teacher SNRs { s α } , and the teacher aspect ratio A = N 3 /N 1 ), the architectural complexity of the student (i.e. its rank N 2 , its number of layers
N l , and the norm (cid:15) of its initialization), and the training time t .
They yield considerable insight into the dynamics of good generalization early in learning and overﬁtting later, as we show below.
3. N UMERICAL TESTS OF THE THEORY OF NEURAL NETWORK GENERALIZATION DYNAMICS
Fig.
3 demonstrates an excellent match between the theory and simulations for the TA, and a close match for random students, for single-hidden-layer students and various teacher ranks N 2 .
Intuitively, as time t proceeds, learning corresponds to singular mode detection wave sweeping from large to small training data singular values (i.e. the wave in Fig.
1B sweeps across the training data spectrum in Fig 2A).
Initially, strong singular values associated with large SNR teacher modes are learned and both ε train and ε test drop.
Fig.
3A-D are for a rank 1 teacher, and so in Fig 3AB we see a single sharp drop early on, if the teacher SNR is sufﬁciently high.
By contrast, with a rank 3 teacher in Fig.
3E-H, there are several early drops as the three modes are picked up.
However, as time progresses, the singular mode detection wave penetrates the MP sea, and the student picks up noise structure in the data, so ε train drops but ε test rises, indicating the onset of overﬁtting.
The main difference between the random student and TA learning curves is that the random student learning is slightly delayed relative to the TA, especially late in training.
This is understandable because the TA already knows the singular vectors of the training data, while the random student must learn them.
Nevertheless, two of the most important aspects of learning, namely the optimal early stopping time t optgradient ≡ argmin t ε test ( t ) and the minimal test error achieved at this time ε optgradient ≡ min t ε test ( t ) , match well between TA and random student, as shown in Fig.
3CD.
At low teacher
SNRs, the student takes a little longer to learn than the TA, but their optimal test errors match.
Our theory can also be easily extended to describe the learning dynamics deeper networks.
Saxe et al.
(2013a) derived t ( s, ˆ s ) for networks of arbitrary depth, so we only need to adjust this factor in our formulas, see App.
A for details.
In Fig.
4 we show that again there is an excellent match between TA networks and theory for student networks with N l = 5 layers (i.e. 3 hidden layers).
Randomly-initialized networks show a much longer alignment lag for deeper networks (see App.
B for details), but the curves are qualitatively similar and optimal stopping errors match.
We also demonstrate extensions of our theory to different numbers of training examples (App. G).
Importantly, many of the phenomena we observe in linear networks are qualitatively replicated in nonlinear networks (Fig. 5), suggesting that our theory may help guide understanding of the nonlinear case.
In particular, features such as stage-like initial learning, followed by a plateau if SNR is high, and ﬁnally followed by overﬁtting, are replicated.
However, there are some discrepancies, in particular nonlinear networks (especially deeper ones) begin overﬁtting earlier than linear networks.
This is likely because a mode in a non-linear network can be co-opted by an orthogonal mode, while in a linear network it cannot.
Thus noise modes are able to “stow away” on the strong signal modes once they are learned.
However, overall learning patterns are similar, and we show below that many interesting phenomena in nonlinear networks are understandable in the linear case, such as the (non-)effects of overparameterization, the dynamics of memorization, and the beneﬁts of transfer. [Figure] N [Figure]
N
[Figure] Randomizing the data dilutes the signal singular values, spreading their variance out over many modes, hence randomly labelled data is learned more slowly. ( N 1 = 100 , N 2 = 50 , N 3 = 50 .) 3. R ANDOMIZED DATA VS . REAL DATA : A LEARNING TIME PUZZLE
An intriguing observation that resurrected the generalization puzzle in deep learning was the observation by Zhang et al.
that deep networks can memorize data with the labels randomly permuted.
However, as Arpit et al.
(2017) pointed out, the learning dynamics of training error for randomized labels can be slower than than for structured data.
This phenomenon also arises in deep linear networks, and our theory yields an analytic explanation for why.
We randomize data by choosing orthonormal inputs ˆx µ as in the structured case, but we choose the outputs ˆy µ to be i.i.d.
Gaussian with zero mean and the same diagonal variance as the structured training data generated by the teacher.
For structured data generated by a low rank teacher with singular values s α , the diagonal output variance is given by σ 2 r = 1 N 3 (cid:104)(cid:80) N 2 i
= α ¯ s 2 α (cid:105) + 1 N 1 σ 2 z , where σ z is the noise variance, as before.
Since there is no relation between input and output, Σ 31 is now distributed as a MP distribution whose support is [( σ r (1 − √A ) , σ r (1 + √A )] .
Thus randomization essentially destroys the outlier signal singular values in Σ 31 reﬂecting the teacher, and distributes them across all randomized data modes, yielding this stretched MP distribution (compare 6A top and bottom).
However, even on this stretched MP distribution, the right edge will be much smaller than the signal singular values, since the signal variance will be diluted by spreading it out over many more modes in the randomized data.
Thus the randomized data will lead to slower initial training error drops relative to the structured
N 1 inputs ···
Task1outputs
Task1teacher ···
N 1 inputs ···
Task2outputs
Task2teacher + =
Aligned?
Relativesignalstrength? ···
N 1 inputs ··· ···
Multi-taskoutputs
Multi-taskteacher
Student ··· ··· ··· ···
Training inputs (cid:76) o
u t p u t s t a r g e t s
IIDGaussiannoise(ﬁxedacrosstraining)
Although this task is impossible to learn on its own, with support from another aligned task, especially one with high SNR, learning can occur.
(b) s A = 3 .
Tasks with modest signals will face interference from poorly aligned tasks, but beneﬁts from well aligned tasks.
These effects are ampliﬁed by
SNR.
(c) s A
= 100 .
Tasks with very strong signals will show little effect from other tasks (note y-axis scales), but any impact will be negative unless the tasks are very well aligned.
( N 1 =
N A 2 = N B 2 = 1 , N 2 = 50 , N 3 = 50 .) data (Fig. 6B) since the singular mode detection wave encounters the ﬁrst signal singular values in structured data earlier than it encounters the edge of the stretched MP sea in randomized data.
3. O UT - PERFORMING OPTIMAL EARLY STOPPING THROUGH A NON - GRADIENT ALGORITHM
For the case of a rank 1 teacher, it is straightforward to derive a good analytic approximation to the important quantities ε optgradient and t optgradient .
{+We assume the teacher SNR is beyond the phase transition point so its unique singular value s 1 > A 1 / 
 , yielding a separation between the training data singular value+} ˆ s 1 in (11) and the edge of the MP sea.
In this scenario, optimal early stopping will occur at a time before the detection wave in Fig.
1B penetrates the MP sea, so to minimize test error, we can neglect the ﬁrst term in (15).
Then optimizing the second term yields the optimal student singular value s
= s 1 O
( s 1 ) .
Inserting this value into (15) yields ε optgradient = 1 − O ( s 1 ) 2 , and inserting it into (9) yields t optgradient .
Thus the optimal generalization error with a rank 1 teacher is very simply related to the alignment of the top training data singular vectors with the teacher singular vectors, and it decreases as this alignment increases.
In App.
E, we show this match in the rank 1 case.
With higher rank teachers, ε optgradient and t optgradient must negotiate a more complex trade-off between teacher modes with different SNRs.
For example, as the singular mode detection wave passes the top training data singular value, s 1 ( t ) → ˆ s 1 which is greater than the optimal s 1 = s 1 O ( s 1 ) for mode 1 .
Thus as learning progresses, the student overﬁts on the ﬁrst mode but learns lower modes.
However, this neural generalization dynamics suggests a superior non-gradient training algorithm that simply optimally sets each s α to s α O
( s α ) in (15), yielding an optimal generalization error:
[Equation]
Standard gradient descent learning cannot achieve this low generalization error because it cannot independently adjust all student singular values.
A simple algorithm that achieves ε optnon-gradient is as follows.
From the training data covariance Σ 31 , extract the top singular values
ˆ s α that pop-out of the
MP sea, use the functional inverse of (11) to compute s a (ˆ s α ) , use (12) to compute the optimal s α , and then construct a matrix W with the same top singular vectors as Σ 31 , but with the outlier singular values shrunk from ˆ s α to s α and the rest set to zero.
This non-gradient singular value shrinkage algorithm provably outperforms neural network training with ε optnon-gradient ≤ ε optgradient .
A THEORY FOR THE TRANSFER OF KNOWLEDGE ACROSS MULTIPLE TASKS
Consider two tasks A and B , described by N 3 by N 1 teacher maps W A and W B , of ranks N A 2 and
N B 2 , respectively.
Now two student networks can learn from the two teacher networks separately, each achieving optimal early stopping test errors ε opt A and ε opt B .
Alternatively, one could construct a composite teacher (and student) that concatenates the hidden and output units, but shares the same input units (Fig. 7).
The composite student and teacher each have two heads, one for each task, with
N 3 neurons per head.
Optimal early stopping on each head of the student yields test errors ε opt A ← B and ε opt B ← A .
We deﬁne the transfer beneﬁt that task B confers on task A to be T A ← B ≡
ε opt A − ε opt A ← B .
A postive (negative) transfer beneﬁt implies learning tasks A and B simultaneously yields a lower (higher) optimal test error on task A compared to just learning task A alone.
A foundational question is how the transfer beneﬁt T A ← B depends on the two tasks deﬁned by the teachers W A and W B .
To answer this, consider the SVDs of each teacher alone: W A =
U A S
A V AT and W B = U B S B V BT .
From the above, we know that ε
opt A depends on W A only through S A .
In App.
D we show that the transfer beneﬁt depends on both W A and W B only through
S A , S B , and the N A 2 by N B 2 similarity matrix Q =
V AT V B .
If we think of the columns of each V as spanning a low dimensional feature space in N 1 dimensional input space that is important for each task, then Q reﬂects the input feature subspace similarity matrix.
Interestingly, the transfer beneﬁt is independent of output singular vectors U A and U B .
What matters for knowledge transfer in this setting are the relevant input features, not how you must respond to them.
We describe the transfer beneﬁt for the simple case of two rank one teachers.
Then S A , S B , and Q are simply scalars s A , s B and q , and we explore the function T
A ← B ( s A , s B , q ) in Fig.
5ABC, which reveals several interesting features.
First, knowledge can be transferred from a high SNR task to a low SNR task (Fig.
5A)
and the degree of transfer increases with task alignment q .
This can make it possible to capture signals from task A which would otherwise sink into the MP sea by learning jointly with a related task, even if the tasks are only weakly aligned (Fig. 5A).
However, if task A already has a high SNR, task B must be very well aligned to it for transfer to be beneﬁcial – otherwise there will be interference.
The degree of alignment required increases as the task A SNR increases, but the quantity of beneﬁt or interference decreases correspondingly (Fig.
5BC).
In Appendix D we explain why our theory predicts these results.
Furthermore, in Appendix F we demonstrate these phenomena are qualitatively recapitulated in nonlinear networks, which suggests that our theory may give insight into how to choose auxiliary tasks.
D ISCUSSION
In summary, our analytic theory of generalization dynamics in deep linear networks reveals that many puzzling aspects of generalization in deep learning already arise in the simple linear setting, where the puzzles can be understood analytically.
In particular, deep linear networks learn more important structure in data ﬁrst, leading to generalization errors that depend on task structure much more than network size.
Our theory explains why deep linear networks learn randomized data more slowly than structured data, and provides a non-gradient based learning method that out-performs gradient descent learning in the linear case.
Finally, we provide an analytic theory of how knowledge is transferred from one task to another, demonstrating that the degree of alignment of input features important for each task, but not how one must respond to these features, is critical for facilitating knowledge transfer.
We think these analytic results provide useful insight into the similar generalization and transfer phenomena observed in the nonlinear case.
Among other things, we hope our work will motivate and enable: (1) the search for tighter upper bounds on generalization error that take into account task structure; (2) the design of non gradient based training algorithms that outperform gradient-based learning; and (3) the theory-driven selection of auxiliary tasks that maximize knowledge transfer.
A L EARNING DYNAMICS FOR DEEPER NETWORKS
In the main text, we described the dynamics of how a single-hidden-layer network converges toward the training data singular modes { ˆ s α , ˆu α , ˆv α } , which were originally derived in Saxe et al.
(2013a).
There it was also proven that for a network with N l layers (i.e. N l − 2 hidden layers), the strength of the mode obeys the differential equation:
[Equation]
This equation is separable and can be integrated for any integer number of layers.
In particular, we consider the case of 5 layers (3 hidden), in which case:
[Equation]
(cid:15)
This expression cannot be analytically inverted to ﬁnd s ( t, ˆ s ) , so we numerically invert it where necessary.
B A LIGNMENT LAG IN RANDOMLY INITIALIZED NETWORKS
As noted in the main text, the randomly-initialized networks behave quite similarly to the TA networks, except that the randomly-initialized networks show a lag due to the time it takes for the network’s modes to align with the data modes.
In ﬁg. 9 we explore this lag by plotting the alignment of the modes and the increase in the singular value for several randomly initialized networks.
Notice that stronger modes align more quickly.
Furthermore, the mode alignment is relatively independent – whether the teacher is rank 1 or rank 3, the alignment of the modes is similar for the mode of singular value 2.
Most importantly, note how the deeper networks show substantially slower mode alignment, with alignment not completed until around when the singular value increases.
This explains why deeper networks show a larger lag between randomly-initialized and TA networks – the alignment process is much slower for deeper networks.
C T RAIN AND TEST ERRORS AFTER A PROJECTION
In the case of transfer learning, or more generally when we want to evaluate a network’s loss on a subset of its outputs, we need to use a slight generalization of the train and test error formulas given in the main text.
Suppose we are interested in the train and test errors after applying a projection operator P :
[Equation]
respectively.
As in the main text, we can rexpress these as:
[Equation]
Using the cyclic property of the trace, we can modify these to get:
[Equation]
As before
[Equation]
we express these in terms of the student
[Equation]
training data and teacher SVDs [Equation] W = USV T , Σ 31 = ˆUˆS ˆV T
[Equation]
and W = U S V T respectively.
Speciﬁcally,
[Equation] [Equation]
D T RANSFER LEARNING DERIVATIONS & DETAILS
Thm 1 (Transfer theorem)
The transfer beneﬁt T A ← B :
• Is unaffected by the U A and U B .
• Is completely determined by only σ 2 z , S A , S B , and
the N A 2 by N B 2 similarity matrix Q = V AT V B .
Proof: We deﬁne
[Equation]
Because of the 0 blocks in U AB , the vectors in blocks corresponding to task A and task B are completely orthogonal, so U AB remains orthonormal.
Thus the relationship between the U A and
U A is irrelevant to the transfer.
(In our simulations we use arbitrary orthonormal matrices for
U A and U B .)
Therefore the transfer effects will be entirely driven by the relationship between the matrices V A and V B and the singular values.
We deﬁne N A 2 by N B 2 similarity matrix Q = V AT V B .
If we think of the columns of each V as spanning a low dimensional feature space in N 1 dimensional input space that is important for each task, then Q reﬂects the input feature subspace similarity matrix.
We can now calculate the singular values of W A + B .
First, note that the input singular modes of W A + B are eigenvectors of
W A + BT W A + B , and the associated singular values are square roots of the eigenvalues of W A + B .
Now
[Equation]
Now if (cid:126)c is an eigenvector of this matrix:
[Equation]
This implies that
[Equation]
Hence eigenvalues of V AB S AB 2 V ABT are also eigenvalues of V ABT V AB S AB 2 , with the mapping between the eigenvectors given by V AB .
Furthermore, this mapping must be a bijection for eigenvectors with non-zero eigenvalues, since the matrices have the same rank (the rank of V AB ).
To see this, note that S AB 2 is full rank.
From this, it is clear that
[Equation]
Furthermore, S AB 2 is positive deﬁnite, so
[Equation]
Now that we know the eigenvectors of these matrices are in bijection, note that:
[Equation]
Because the output modes don’t matter (as noted above), the alignment between the eigenvectors of
V AB S AB 2 V ABT and V A , weighted by their respective eigenvalues, gives the transfer beneﬁt.
In other words:
For any given tasks, the transfer beneﬁt can be calculated using our theory.
However, in certain special cases, we can give exact answers.
For example, in the rank one case with equal singular values between the tasks ( s A = s B = s ), the matrix (cid:20) I Q Q T
I (cid:21)   S A
S B  reduces to (cid:20)
[Equation]
q q
[Equation]
with eigenvalues s √ 1 ± q and eigenvectors
[Equation]
Corresponding to the shared structure between the tasks and the differences between them.
We note that the sign of the alignment q is irrelevant as a special case of the fact (noted above) that any orthogonal transformation on the output modes does not affect transfer.
D. M ISALIGNMENT AND INTERFERENCE
Why is there interference between tasks which are not well aligned?
In the rank one case, we are effectively changing the (input) singular dimensions of Y A from V A to V AB .
The two singular modes of V AB correspond to the shared structure between the tasks (weighted by the relative signal strengths), and the differences between them, respectively.
Although we may be improving our estimates of the shared mode if q > 0 (by increasing its singular value relative to s A ), we are actually decreasing its alignment with V A unless q = 1 .
This misalignment is captured by the second mode of V AB , but the increase in the singular value of the ﬁrst mode must come at the cost of a decrease in the singular value of the second mode.
See Fig.
10 for a conceptual illustration of this.
This means that the multi-task setting allows the distinctions between the tasks to sink towards the sea of noise, while pulling out the common structure.
In other words, transferring knowledge from one task always comes at the cost of ignoring differences between the tasks.
Furthermore, incorporating a task B allows its noise to seep into the task A signal.
Together, these two effects help to explain why transfer can be sometimes beneﬁcial but sometimes detrimental.
A ← B , the transfer beneﬁt (or cost) arises from alignment between the task’s input modes.
E N ON - GRADIENT TRAINING ALGORITHM
In Fig.
11 we show the match between the error achieved by training the student by gradient descent and the optimal stopping error predicted by the non-gradient shrinkage algorithm in the case of a rank-1 teacher.
F T RANSFER RESULTS GENERALIZE TO NON - LINEAR NETWORKS
Since most deep learning practitioners do not train linear networks, it is important that our theoretical insights generalize beyond this simple case.
In this section we show that the transfer patterns qualitatively generalize to non-linear networks.
Here we show results from teacher networks with N 1 = 100 N 3 = 50 , N 2 = 4 (thus the task is higher rank) and leaky relu non-linearities at the hidden and output layers.
We train a student with leaky relu units and N 2 = N 3 to solve this task.
Results qualitatively look quite similar to those in
Fig 5.
of the main text for rank one linear teachers, see below.
Thus our insights into transfer may help to understand multi-task beneﬁts in more complicated architectures
G V ARYING THE NUMBER OF TRAINING EXAMPLES
In the main text, we focused on the test error dynamics in the case in which the number of examples equalled the number of inputs.
Here we show how the formula for test error curves is modiﬁed as the number of training examples P is varied.
For simplicity, when P (cid:54) =
N 1 , we focus on the case of a full rank student with aspect ratio A = 1 (so that N 1 = N
=
N 3 ).
The more general case of lower rank students with non-unity aspect ratios can be easily found from this case, but with some additional bookkeeping.
As before, we assume the teacher generates noisy outputs from a set of P inputs:
[Equation]
When P ≥ N 1 , the minimum generalization error is simply determined by SNR (cid:112) P/N 1 , so all curves converge to a single asymptotic line in (d) as P increases.
When P < N 1 , however, the curves for different SNRs separate because the projection and noise effects depend on initial SNR
Optimal stopping error for gaussian vs. orthogonal inputs, showing a strong correlation.
Thus our use of orthogonal inputs in the theory also yields insight into the more general case of approximately unit norm Gaussian inputs
N 2 = N 3 =
This training set yields important second-order training statistics that will guide student learning:
[Equation]
Here ˆX , ˆY , and Z are each N 1 by P , N 3 by P , and N 3 by P matrices respectively, whose µ ’th columns are ˆx µ , ˆy µ , and ˆz µ , respectively.
Σ 11 is an N 1 by N 1 input correlation matrix, and Σ 31 is an N 3 by N 1 the input-output correlation matrix.
We choose the matrix elements z µi of the noise matrix Z to be drawn iid from a Gaussian with zero mean and variance σ 2 z /N 1 .
The noise scaling is chosen so the singular values of the teacher W and the noise Z are both O (1) , leading to non-trivial generalization effects.
Furthermore, we chose training inputs to be close to unit-norm, and make the input covariance matrix Σ 11 as white as possible (whitening is a common pre-processing step for inputs).
When P > N 1 , this can be done by choosing the rows of ˆ X to be orthonormal and then scaling up by (cid:113) P/N 1 , so the columns are approximately unit norm.
Then Σ 11 = P/N
I is proportional to the identity.
On the otherhand, if P < N 1 , we choose the columns of ˆX to be orthonormal, so that Σ 11 =
P || , where P || is a projection operator onto the P dimensional column space of ˆX spanned by the input examples.
Both these choices are intended to approximate the situation in which the columns of ˆX are chosen to be iid unit-norm vectors.
Finally, as generalization performance will depend on the ratio of teacher singular values to the noise variance parameter σ 2 z , we simply set σ z = 1 as in the main text.
Thus, given the unit-norm inputs, we can think of the teacher singular values as signal to noise ratios (SNRs).
We now examine how the dynamics of the test error evolves as we vary the number of training examples P .
We split our analyses into two distinct regimes: (1) the oversampled regime in which the data density
D ≡ P/N 1 > 1 , and (2) the undersampled regime in which D < 1 .
G. T HE OVERSAMPLED REGIME
The oversampled regime ( D > 1 ) is relatively simple.
First Σ 11 is scaled up by a factor of D .
And in the input-output covariance matrix, Σ 31 = W ˆXˆX T
+ ZˆX T , the signal component, W ˆXˆX T is scaled up by a factor of D while the noise component ZˆX T has the same singular value spectrum as the D = 1 case, up to an overall scaling by √D (since the rows of ˆX are orthogonal and all its singular values are equal to √D ).
This leads to an increase in the effective SNR by a factor of √D .
Thus overall, the test error curves for the case of D > 1 can be simply obtained from the theory of the test error curves for D = 1 through two modiﬁcations: (1) a boost in the SNR for the D = 1 case by a multiplicative factor of √D , and (2) and an overall speed up in the learning time by a multiplicative factor of D .
G. T HE UNDERSAMPLED REGIME
For the undersampled regime ( D < 1 ), we must account for the fact that the P training inputs do not span the full N 1 dimensional space of all inputs.
Thus the projection operator P || onto the
P dimensional column space of ˆX plays a crucial role.
Indeed the input-correlation Σ 11 = P || .
And Σ 31 = W P || + ZˆX T .
This implies that the learning dynamics only transforms the composite student map W from the P dimensional subspace spanned by the inputs to the N 3 dimensional output space.
In contrast, the student map from the N 1 − P dimensional subspace orthogonal to the image of P || remains frozen.
Tracing through the equations of the main paper and accounting for the projection operator P || , we ﬁnd the effective aspect ratio for this undersampled learning problem (when N 3 = N 2 = N 1 ) is no longer A
= N 3 /N
but rather D = P/N 1 .
Furthermore, in the limit
N 3 , N 1 → ∞ while N 2 remains O (1) , the singular values of the signal component W P || of Σare attenuated by a factor of √D , making the associated singular vectors more susceptible to noise.
Again tracing through the equations of the main paper, with all of these modiﬁcations, we ﬁnd the ﬁnal formula for test error curves in the undersampled measurement regime:
[Equation]
[Table]
This equation has several modiﬁcations compared to the case P = N 1 in (15).
First the term in the numerator involving N 3 − P reﬂects generalization error due to the N 3 − P dimensional frozen subspace, and the initial weight variance (cid:15) 2 contributes to this generalization error.
The second term in the numerator involves all the P − N 2 training modes which cannot be correlated with the teacher, and the average (cid:104)·(cid:105) is over a Marcenko-Pasteur distribution of singular values (see (13)) except with the aspect ratio A replaced by D .
The third term accounts for learned correlations between the student and teacher.
It involves the transformation from teacher singular values s to training data singular values ˆ s through the formula (11) except with the aspect ratio replacement
A → D , and the effective teacher singular value attenuation s → √D s .
Similarly, the computation of the singular vector overlap is done through (12) also with the replacements
A → D and s → √D s .
G. C OMPARISON OF THEORY AND EXPERIMENT FOR UNDER AND OVER SAMPLED MEASUREMENT REGIMES
In Fig. 13, we show an excellent match between our theory and empirical simulations for varying values of P , both in the oversampled and undersampled measurement regimes.
There are a number of interesting features to note.
First, although the minimum generalization error improves monotonically with P , the asymptotic ( t → ∞ ) generalization error does not, because of a frozen subspace (Advani & Saxe, 2017) of the modes that are not overﬁt when P < N 1 , because the training data rank is ≤ P .
Second, when P ≥ N 1 , the minimum generalization error is simply determined by SNR (cid:112) P/N 1 , so all curves converge to a single asymptotic line as P increases.
When P < N 1 , however, the curves for different SNRs separate because the projection and noise effects depend on initial SNR.
Finally, in Fig.
13D we show that approximately unit norm i.i.d.
gaussian inputs yield similar results to the orthogonalized data matrices we employed in the theory, although the gaussian inputs do result in slightly higher optimal stopping error.
H L ESS THAN FULL RANK STUDENTS
Although we generally assumed students were full rank in the main text to simplify the calculations, our theory remains exact for TA networks of any rank.
Furthermore, as shown in Fig. 14, the TA and random networks again show very similar optimal stopping generalization error, but with the optimal stopping time of the random networks lagging behind that of the TA networks.
Furthermore, this lag increases as the rank of the random network decreases (because a low rank network will have less initial projection onto the random modes, there is is more alignment to be done).
However, reducing the student rank does not change the optimal stopping error (as long as it is still greater than the teacher rank).
[Figure]