An analytic theory of generalization dynamics and transfer learning in deep linear networks

1 INTRODUCTION . 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 can not 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 ) . 2 THEORETICAL 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 define the training and test errors whose dynamics we wish to understand . 2.1 AN ENSEMBLE OF LOW-RANK NOISY TEACHERS . 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 . 2.2 STUDENT TRAINING AND TEST ERROR . 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 . 3 SINGLE 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.1 STUDENT TRAINING DYNAMICS AND TRAINING-ALIGNED ( TA ) NETWORKS . 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 .

Summary: 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.