AN ANALYTIC THEORY OF GENERALIZATION DYNAM- ICS AND TRANSFER LEARNING IN DEEP LINEAR NET- WORKS
Year: 2019

Abstract:
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 first, 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.

heading: 1 INTRODUCTION
text: 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 overfitting 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 first. 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 find 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).

heading: 2 THEORETICAL FRAMEWORK
text: 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 define the training and test errors whose dynamics we wish to understand.

heading: 2.1 AN ENSEMBLE OF LOW-RANK NOISY TEACHERS
text: We first consider an ensemble of 3-layer linear teacher networks with N i units in layer i, and weight matrices W
21 ∈ RN2×N1 and W32 ∈ RN3×N2 between the input to hidden, and hidden to output layers, respectively. The teacher network thus computes the composite map y = Wx, where W ≡W32W21. Of critical importance is the singular value decomposition (SVD) of W:
W = U S V T = N2∑ α=1 sαuαvαT , (1)
Where U ∈ RN3×N2 and V ∈ RN1×N2 are both matrices with orthonormal columns and S is an N2×N2 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 N1, N3 →∞ with an O(1) aspect ratio A = N3/N1 ∈ (0, 1] so that the teacher has fewer outputs than inputs. Also, we hold N2 ∼ O(1), so the teacher has a low, finite rank, and we study generalization performance as a function of the N2 teacher singular values.
We further assume the teacher generates noisy outputs from a set of N1 orthonormal inputs:
ŷµ = Wx̂µ + zµ for µ = 1, . . . ,N1. (2)
This training set yields important second-order training statistics that will guide student learning:
Σ11 ≡ N1∑ µ=1 x̂µx̂µT = I, Σ31 ≡ N1∑ µ=1 ŷµx̂µT = W + ZX̂T . (3)
Here the input covariance Σ11 is assumed to be white (a common pre-processing step), the inputoutput covariance Σ31 is simplified using (2), and Z ∈ RN3×N1 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 σ2z/N1. 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 σ2z , 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 = N1 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 6= N1, obtaining a good match between theory and experiment in both cases.

heading: 2.2 STUDENT TRAINING AND TEST ERROR
text: Now consider a student network with Ni units in each layer. We assume the first and last layers match the teacher (i.e. N1 = N1 and N3 = N3) but N2 ≥ N2, 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 = W32W21 ≡Wx. Its training error on the teacher dataset in (2) and its test error over a distribution of new inputs are given by
εtrain ≡ ∑N1 µ=1 ||Wx̂
µ − ŷµ||22∑N1 µ=1 ||ŷµ||22
, εtest ≡ 〈 ||Wx− y||22 〉 〈||y||22〉 , (4)
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 〈·〉 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
εtrain = Tr WTW − 2Tr WTΣ31 + Tr Σ31TΣ31
Tr Σ31TΣ31 , εtest = Tr WTW − 2Tr WTW + Tr WTW Tr W T W .
(5) Both εtrain and εtest can be further expressed in terms of the student, training data and teacher SVDs, which we denote by W = USVT , Σ31 = ÛŜV̂T , and W = U S V T respectively. Specifically,
εtrain =  N3∑ β=1 ŝ2β −1  N2∑ α=1 s2α + N3∑ β=1 ŝ2β − 2 N2∑ α=1 N3∑ β=1 sαŝβ ( uα · ûβ ) ( vα · v̂β ) , (6) εtest =
 N2∑ β=1 s2β −1  N2∑ α=1 s2α + N2∑ β=1 s2β − 2 N2∑ α=1 N2∑ β=1 sαsβ ( uα · uβ ) ( vα · vβ ) . (7) 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 fixed training data {ŝα, ûα, v̂α} and teacher {sα,uα,vα} singular modes respectively.

heading: 3 SINGLE TASK GENERALIZATION DYNAMICS: THEORY AND EXPERIMENT
text: 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.

heading: 3.1 STUDENT TRAINING DYNAMICS AND TRAINING-ALIGNED (TA) NETWORKS
text: We assume the student weights undergo batch gradient descent with learning rate λ on the training error ∑ µ ||ŷµ−W32W21x̂µ||22, which for small λ is well approximated by the differential equations:
τ d
dt W21 = W32
T ( Σ31 −W32W21Σ11 ) , τ d
dt W32 =
( Σ31 −W32W21Σ11 ) W21 T ,
(8) (where τ ≡ 1/λ), which must be solved from an initial set of student weights at time t = 0 (Saxe et al., 2013a). We consider two classes of student initializations. The first initialization corresponds to a random student where the weights W21 and W32 are chosen such that the composite map W = W32W21 has an SVD W = UVT, where U and V are random singular vector matrices and all student singular values are . As such a random student learns, the composite map undergoes a time dependent evolution W(t) = U(t)S(t)V(t)T = ∑N2 α=1 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 difficult to obtain analytically from random initial conditions.
Thus we also consider a special class of training aligned (TA) initial conditions in which W21 and W32 are chosen such that the composite map W = W32W21 has an SVD W = 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 . 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
s(t, ŝ) = ŝe2ŝt/τ
e2ŝt/τ − 1 + ŝ/ , t(s, ŝ) =
τ 2ŝ ln ŝ/ − 1 ŝ/s− 1 . (9)
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 . As shown in Fig. 1A, for each ŝ, s(t, ŝ) is a sigmoidal learning curve that undergoes a sharp transition around time t/τ = 12ŝ ln (ŝ/ − 1), at which it rises from its small initial value of 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.

heading: 3.2 HOW THE TEACHER IS BURIED IN THE TRAINING DATA: A RANDOM MATRIX ANALYSIS
text: 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
sα(t) = s(t, ŝα), u α(t) = ûα, vα(t) = v̂α. (10)
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 N1, N3 →∞ with the aspect ratio A = N3/N1 ∈ (0, 1], and N2 ∼ O(1).
In this limit, the top N2 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
ŝ(s) =
{ (s)−1 √ (1 + s2)(A+ s2) if s > A1/4
1 + √ A otherwise.
(11)
The associated top N2 singular vectors of Σ31 can also acquire a nontrivial overlap with the N2 modes of the teacher through the relation |ûα · uα| |v̂α · vα| = O(sα), where the singular vector overlap function is given by
O(s) =  [ 1− A(1+s 2) s2(A+s2) ]1/2 [ 1− (A+s 2) s2(1+s2) ]1/2 if s > A1/4
0 otherwise (12)
The rest of the N3−N2 singular vectors of Σ31 are orthogonal to the top N2 ones, and their singular values are distributed according to the the Marchenko-Pastur (MP) distribution:
P (ŝ) =
{√ 4A−(ŝ2−(1+A))2
πAŝ ŝ ∈ [1− √ A, 1 + √ A]
0 otherwise. (13)
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 = A1/4. For s ≤ A1/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 > A1/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 inflated by the noise, though the inflation 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 > A1/4, and the alignment approaches unity as s increases (Fig. 2C).

heading: 3.3 PUTTING IT ALL TOGETHER: AN ANALYTIC THEORY OF GENERALIZATION DYNAMICS
text: 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
εtrain(t) =  N3∑ α=1 ŝ2α −1(N3−N2)〈ŝ2〉Rout+(N2−N2)〈(s(ŝ, t)− ŝ)2〉Rin + N2∑ α=1 [sα(t)− ŝα]2 
(14) Here, sα(t) = s(ŝα, t) as defined in (9) are the TA singular values, and ŝα = ŝ(sα) as defined in (11) are the training data singular values associated with the teacher singular values sα. Also 〈·〉R denotes an average with respect to the MP distribution in (13) over a regionR. Two distinct regions contribute to training error. First Rin contains those top N2 − N2 training data singular values that do not correspond to the N2 singular values of the teacher but will be learned by a rank N2 student. Second,Rout corresponds to the remaining N3 −N2 lowest training data singular values that cannot be learned by a rank N2 student. In terms of the MP distribution,Rout = [1− √ A, f ]
andRin = [f, 1 + √ A], where f is the point at which the MP density has 1−N2/N3 of its mass to
the left and N2/N3 of its mass to the right. In the simple case of a full rank student, f = 1− √ A, and one need only integrate overRin 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 find
εtest(t) =  N2∑ α=1 s2α −1(N2−N2)〈s(ŝ, t)2〉Rin+ N2∑ α=1 [ (sα(t)− sα)2 + 2sα(t)sα(1−O(sα)) ] (15)
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 N2, teacher SNRs {sα}, and the teacher aspect ratio A = N3/N1), the architectural complexity of the student (i.e. its rank N2, its number of layers Nl, and the norm of its initialization), and the training time t. They yield considerable insight into the dynamics of good generalization early in learning and overfitting later, as we show below.

heading: 3.4 NUMERICAL TESTS OF THE THEORY OF NEURAL NETWORK GENERALIZATION DYNAMICS
text: 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 N2. 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 sufficiently 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 overfitting.
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 toptgradient ≡ argmintεtest(t) and the minimal test error achieved at this time ε opt gradient ≡ mintε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 Nl = 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 finally followed by overfitting, are replicated. However, there are some discrepancies, in particular nonlinear networks (especially deeper ones) begin overfitting 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 benefits of transfer.

heading: 3.5 RANDOMIZED DATA VS. REAL DATA: A LEARNING TIME PUZZLE
text: An intriguing observation that resurrected the generalization puzzle in deep learning was the observation by Zhang et al. (2016) 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 σ2r = 1 N3 [∑N2 i=α s̄ 2 α ] + 1 N1 σ2z , 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 reflecting the teacher, and distributes them across all randomized data
N A
2 = N B 2 = 1, N2 = 50, N3 = 50.)
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 data (Fig. 6B) since the singular mode detection wave encounters the first signal singular values in structured data earlier than it encounters the edge of the stretched MP sea in randomized data.

heading: 3.6 OUT-PERFORMING OPTIMAL EARLY STOPPING THROUGH A NON-GRADIENT ALGORITHM
text: For the case of a rank 1 teacher, it is straightforward to derive a good analytic approximation to the important quantities εoptgradient and t opt gradient. We assume the teacher SNR is beyond the phase transition point so its unique singular value s1 > A1/4, 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 first term in (15). Then optimizing the second term yields the optimal student singular value s1 = s1O(s1). Inserting this value into (15) yields εoptgradient = 1−O(s1)2, and inserting it into (9) yields toptgradient. 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 opt gradient 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, s1(t)→ ŝ1 which is greater than the optimal s1 = s1O(s1) for mode 1. Thus as learning progresses, the student overfits on the first 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:
εoptnon-gradient =  N2∑ α=1 s2α −1  N2∑ α=1 s2α(1−O(sα)2)  . (16) 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 sa(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 ≤ ε opt gradient.

heading: 4 A THEORY FOR THE TRANSFER OF KNOWLEDGE ACROSS MULTIPLE TASKS
text: Consider two tasks A and B, described by N3 by N1 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 εoptA 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 N3 neurons per head. Optimal early stopping on each head of the student yields test errors ε opt A←B and εoptB←A. We define the transfer benefit that task B confers on task A to be T A←B ≡ ε opt A − ε opt A←B . A postive (negative) transfer benefit 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 benefit T A←B depends on the two tasks defined 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 εoptA depends on W A only through S A . In App. D we show that the transfer benefit 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 N1 dimensional input space that is important for each task, then Q reflects the input feature subspace similarity matrix. Interestingly, the transfer benefit 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 benefit for the simple case of two rank one teachers. Then S A , S B
, and Q are simply scalars sA, sB and q, and we explore the function T A←B(sA, sB , 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, taskB must be very well aligned to it for transfer to be beneficial – otherwise there will be interference. The degree of alignment required increases as the task A SNR increases, but the quantity of benefit 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.

heading: 5 DISCUSSION
text: 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 first, 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.

heading: A LEARNING DYNAMICS FOR DEEPER NETWORKS
text: 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 Nl layers (i.e. Nl − 2 hidden layers), the strength of the mode obeys the differential equation:
τ d
dt u = (Nl − 1)u2−2/(Nl−1)(s− u)
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:
t(s, ŝ) = τ
2
[ tanh−1 (√ u ŝ ) ŝ3/2 − 1 ŝ √ u ]s
This expression cannot be analytically inverted to find s(t, ŝ), so we numerically invert it where necessary.

heading: B ALIGNMENT LAG IN RANDOMLY INITIALIZED NETWORKS
text: 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 fig. 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.

heading: C TRAIN AND TEST ERRORS AFTER A PROJECTION
text: 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:
εtrain ≡ ∑N1 µ=1 ||PWx̂
µ −Pŷµ||22∑N1 µ=1 ||Pŷµ||22
, εtest ≡ ∑N1 µ=1 ||PWx
µ −Pyµ||22∑N1 µ=1 ||Py µ||22 , (17)
respectively. As in the main text, we can rexpress these as:
εtrain = Tr WTPTPW − 2Tr WTPTPΣ31 + Tr Σ31TPTPΣ31
Tr Σ31TPTPΣ31 , (18)
εtest = Tr WTPTPW − 2Tr WTPTPW + Tr WTPTPW
Tr W T PTPW
. (19)
Using the cyclic property of the trace, we can modify these to get:
εtrain = Tr PWWTPT − 2Tr PΣ31WTPT + Tr PΣ31Σ31TPT
Tr PΣ31Σ31TPT , (20)
εtest = Tr PWWTPT − 2Tr PWWTPT + Tr PWWTPT
Tr PWW T PT
. (21)
As before, we express these in terms of the student, training data and teacher SVDs, W = USVT , Σ31 = ÛŜV̂T , and W = U S V T respectively. Specifically, εtrain = [∑N3 β=1 ŝ 2 β ||Pû α||22 ]−1[∑N2 α=1 s 2 α||Pu α||22+ ∑N3 β=1 ŝ 2 β ||Pû α||22−2 ∑N2 α=1 ∑N3 β=1 sαŝβ(Pu α·Pûβ)(vα·v̂β) ] ,
(22)
εtest = [∑N2 β=1 s 2 β ||Pu α||22 ]−1[∑N2 α=1 s 2 α||Pu α||22+ ∑N2 β=1 s 2 β ||Pu α||22−2 ∑N2 α=1 ∑N2 β=1 sαsβ(Pu α·Puβ)(vα·vβ) ] .
(23)

heading: D TRANSFER LEARNING DERIVATIONS & DETAILS
text: Thm 1 (Transfer theorem) The transfer benefit T A←B:
• Is unaffected by the UA and UB .
• Is completely determined by only σ2z , S A , S B , and the N A 2 by N B
2 similarity matrix Q = V AT V B .
Proof: We define
U AB =
N A 2 N B 2[ ] N3 U A 0
N3 0 U B
S AB =
N A 2 N B 2[ ] N A 2 S A 0
N B 2 0 S B
V AB = N A 2 N B 2[ ] N1 V A V B
W A+B =
[ U A
0
0 U B
][ S A 0
0 S B ] VAT V BT  (24) 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 define 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 N1 dimensional input space that is important for each task, then Q reflects 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
W A+BT W A+B = V AB S AB U ABT U AB S AB V ABT = V AB S AB2 V ABT
Now if ~c is an eigenvector of this matrix:
V AB S AB2 V ABT ~c = λ~c
This implies that
V ABT V AB S AB2 V ABT ~c = λV ABT ~c
Hence eigenvalues of V AB S AB2 V ABT are also eigenvalues of V ABT V AB S AB2
, 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 AB2 is full rank. From this, it is clear that
rankV ABT V AB S AB2 = rankV ABT V AB = rankV AB .
Furthermore, S AB2 is positive definite, so
rankV AB S AB2 V ABT = rankV AB .
Now that we know the eigenvectors of these matrices are in bijection, note that:
V ABT V AB S AB2 =  VAT V BT [ VA VB ]  SA2 0 0 S B2 =[ I Q QT I ] SA2 0 0 S B2  Because the output modes don’t matter (as noted above), the alignment between the eigenvectors of
V AB S AB2 V ABT and V A , weighted by their respective eigenvalues, gives the transfer benefit.
For any given tasks, the transfer benefit 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 (sA = sB = s), the matrix
[ I Q
QT I ] SA2 0 0 S B2  reduces to [
1 q q 1
] s2
with eigenvalues s √
1± q and eigenvectors[ 1 1 1 −1 ] 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.1 MISALIGNMENT 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 YA 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 sA), 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 first 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 beneficial but sometimes detrimental.

heading: E NON-GRADIENT TRAINING ALGORITHM
text: 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.

heading: F TRANSFER RESULTS GENERALIZE TO NON-LINEAR NETWORKS
text: 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 N1 = 100 N3 = 50, N2 = 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 N2 = N3 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 benefits in more complicated architectures.
G VARYING 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 modified as the number of training examples P is varied. For simplicity, when P 6= N1, we focus on the case of a full rank student with aspect ratio A = 1 (so that N1 = N2 = N3). 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:
ŷµ = Wx̂µ + zµ for µ = 1, . . . ,P. (25)
This training set yields important second-order training statistics that will guide student learning:
Σ11 ≡ X̂X̂T Σ31 ≡ ŶX̂T = WX̂X̂T + ZX̂T . (26) Here X̂, Ŷ, and Z are each N1 by P , N3 by P , and N3 by P matrices respectively, whose µ’th columns are x̂µ, ŷµ, and ẑµ, respectively. Σ11 is an N1 by N1 input correlation matrix, and Σ31 is an N3 by N1 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 σ2z/N1. 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 > N1, this can be done by choosing the rows of X̂ to be orthonormal and
then scaling up by √ P/N1, so the columns are approximately unit norm. Then Σ11 = P/N1I
is proportional to the identity. On the otherhand, if P < N1, 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 σ2z , 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/N1 > 1, and (2) the undersampled regime in which D < 1.
G.1 THE 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 = WX̂X̂T + ZX̂T , the signal component, WX̂X̂T is
scaled up by a factor of D while the noise component ZX̂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 modifications: (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.2 THE UNDERSAMPLED REGIME
For the undersampled regime (D < 1), we must account for the fact that the P training inputs do not span the full N1 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 = WP || + ZX̂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 N3 dimensional output space. In contrast, the student map from the N1 − 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 find the effective aspect ratio for this undersampled learning problem (when N3 = N2 = N1) is no longer A = N3/N1 but rather D = P/N1. Furthermore, in the limit N3, N1 →∞ while N2 remains O(1), the singular values of the signal component WP || of Σ31 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 modifications, we find the final formula for test error curves in the undersampled measurement regime:
εtest(t) =
[ (N3 − P ) 2 + (P−N2)〈s(ŝ, t)2〉+ ∑N2 α=1 [ (sα(t)− sα)2 + 2sα(t)sα(1−O( √ Dsα)) ]] [∑N2
α=1 s 2 α ] (27) This equation has several modifications compared to the case P = N1 in (15). First the term in the numerator involving N3 − P reflects generalization error due to the N3 − P dimensional frozen subspace, and the initial weight variance 2 contributes to this generalization error. The second term in the numerator involves all the P −N2 training modes which cannot be correlated with the teacher, and the average 〈·〉 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 → √ Ds. Similarly, the computation of the singular vector overlap is done through (12) also with the replacements A → D and s→ √ Ds.
G.3 COMPARISON 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 overfit when P < N1, because the training data rank is ≤ P . Second, when P ≥ N1, the minimum generalization error is simply determined by SNR √ P/N1, so all curves converge to a single asymptotic line as P increases. When P < N1, 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.

heading: H LESS THAN FULL RANK STUDENTS
text: 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).
(a) Best generalization error is quite similar between aligned and random initializations (b) Optimal stopping time is quite similar between aligned and random initializations
(c) Optimal generalization error vs. optimal stopping time for randomly initialized networks (d) Optimal generalization error vs. optimal stopping time for initially aligned networks
Figure 14: Empirical verification that the simplifying assumptions of our theory are approximately valid in the regime we are considering at different student ranks. Initializations with random initial weights (random init.) and initializations with initial weight aligned to the noisy data SVD (aligned init.) are compared across varying student ranks. (a) The minimum generalization errors are almost identical between the different initializations and different student ranks. (b) The optimal stopping time in the randomly initialized networks consistently lags behind the aligned networks, because it takes time for the alignment to occur. This lag increases as the students rank decreases. (c) Randomly initialized networks of varying ranks obey qualitatively similar trends of increase in optimal stopping error and optimal stopping time as SNR decreases. (d) The theory predicts the aligned networks trends of increase in optimal stopping error and optimal stopping time with decreasing SNR almost perfectly. (All plots are made with a rank 1 teacher and N1 = N3 = 100)

Review summaries:
Pros: The paper tackles an interesting problem of generalization and transfer learning in deep networks. They start from a linear network to derive the theory, identifying phases in the learning, and relating learning rates to task structure and SNR. The theory is thorough and backed up by numerical simulations, including qualitative comparisons to nonlinear networks.