Efficient Computation of Deep Nonlinear Infinite-Width Neural Networks that Learn Features

1 INTRODUCTION . The theory of Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) led to many theoretical discoveries of neural networks , but they ultimately apply only to a limited , unrealistic regime where the network does not learn features . Recently , Yang and Hu ( 2020 ) discovered the maximal update , or µ , parametrization that induces a feature learning ∞-width limit , called the µ-limit of deep neural networks , which generalizes the mean field limit of shallow networks ( Sirignano and Spiliopoulos , 2018 ; Mei et al. , 2018 ; Chizat and Bach , 2018 ; Rotskoff and Vanden-Eijnden , 2018 ) . They explicitly trained such∞-width limits on Word2Vec and MAML , tasks relying crucially on feature learning . These µ-limits outperformed both NTK baselines and finite-width networks . Thus , the µ-limit seems to capture key properties of neural networks in practice and is a promising direction to further our theoretical understanding of them . However , a significant roadblock to this end is that the µ-limit for general neural networks is both hard to compute empirically and hard to analyze theoretically . Indeed , Yang and Hu ( 2020 ) restricted main experiments to linear MLPs where exact calculation is fast as a special case . According to them , the main complications come from the interleaving of forward and backward propagations that leads to the need to evaluate Gaussian integrals involving nested nonlinearities . Likewise , these complicated integrals also contribute to the difficulty of theoretical analysis , especially in deep networks . Contributions In this work , we bypass these issues by considering a form of projected gradient descent instead . Roughly speaking , with Vφ denoting the V-transform ( aka dual ) of φ ( Definition 3.1 ) . MLP with nonlinearity φ trained by projected gradient accumulation width→∞−−−−−−−−−−−−−−→ MLP with nonlinearity Vφtrained by gradient concatenation ( ? ) Here gradient accumulation is just the normal training process where gradient is added to a parameter , whereas in gradient concatenation , gradient is appended to a parameter . The RHS retains the feature learning capability while making the computation and the mathematics much simpler compared to the µ-limit . We explicitly evaluate our new limit , which we call the π-limit , on CIFAR10 and Omniglot , and compare against finite-width neural networks as well as NTK and NNGP . We find the π-limit always outperform wide standard-parametrized or µ-parametrized networks , and they all significantly outperform the kernels . For example , on Omniglot , this is evidenced by Fig . 1 . We also demonstrate transfer learning from ImageNet to CIFAR10 . To our knowledge , this is the first time deep nonlinear feature learning∞-width limits are evaluated on nontrivial natural datasets . 2 WHY IS THE µ-LIMIT HARD TO COMPUTE ? AN EXAMPLE While the µ-limit is the “ optimal ” feature learning limit of a neural network as described in Yang and Hu ( 2020 ) , it is difficult to calculate analytically . Let ’ s demonstrate this difficulty with a simple 1-hidden-layer µ-parametrized network ( abbreviated µ-net , whose∞-width limit is the µ-limit ) f with 1-dimensional input ξ ∈ R and output f ( ξ ) ∈ R : f ( ξ ) = n−1/2v > φ ( √ nuξ ) ∈ R , ( 1 ) where u , v ∈ Rn are column vectors and φ : R→ R is the nonlinearity ( e.g . relu ) . Note here n is the width of f . According to µ-parametrization , we initialize uα , vα ∼ N ( 0 , 1n ) . At initialization , one easily sees limn→∞ f ( ξ ) = 0 because f ( ξ ) = 1 n n∑ α=1 ( √ nv ) αφ ( ξ √ nuα ) is the average of a large number of iid random variables ( √ nv ) αφ ( ξ √ nuα ) with zero mean . In the language of Tensor Programs ( Yang , 2019b ; a ; 2020a ; b ) , we have an almost sure convergence f ( ξ ) → EZ √ nvφ ( ξZ √ nu ) = 0 , at initialization , where Z √ nv and Z √ nu represent the independent standard Gaussian random variables describing the coordinate distributions of √ nv and √ nu . However , one can see that f ( ξ ) has a limit of the form f ( ξ ) → E ( Z √ nv + · · · ) φ ( ξZ √ nu + ( · · · ) φ′ ( · · · ) ) , after 1 step of SGD , ( 2 ) where u and v still represent values at initialization , and · · · represent terms coming from the gradient update.1 Even though the · · · appearing inside φ ( ) turn out to be all Gaussians , this expectation is hard to evaluate analytically due to the nesting of φ′ inside φ , unless φ is polynomial . Furthermore , as more SGD steps are taken , this nesting quickly becomes more complicated , so even if φ is polynomial , the time needed for evaluation compounds exponentially . 3 SOLUTION : PROJECT THE GRADIENT . To alleviate this difficulty , we propose to study the limit of training a neural network under a projected form of gradient descent which we call π-SGD , i.e. , we update θ ← θ − ηΠ∇θL where Π is a linear projection and θ is the vector of all parameters . We shall first describe this projection Π for the simple motivating example above before doing so for the general case . 3.1 CONTINUING THE 1-HIDDEN-LAYER EXAMPLE . In the example of Eq . ( 1 ) , Π acts on the gradients as follows : 1 ) Π leaves the output layer gradient ∇vL untouched , but 2 ) projects the input layer gradient2 ∇uL = ( · · · ) φ′ ( · · · ) to the linear span of u0 , the initial value of input weights u:3 Π ( ∇vL , ∇uL ) = ( ∇vL , Πu0∇uL ) = ( ∇vL , cu0 ) for some c ∈ R. 1See Appendix C for a detailed calculation . 2The precise form is∇uL = L′vφ′ ( √ nuξ0 ) . 3This projection may look excessively reductive , but this is just an artifact of the 1-dimensional input and the shallow depth . See Section 3.2 for the general case . Here Πu0 denotes the orthogonal projection to the span of u0 . In particular , after projection , ∇uL now has roughly iid Gaussian coordinates ( proportional to u0 ) .4 Then unlike Eq . ( 2 ) , we avoid the nesting of φ′ inside φ : f ( ξ ) → E ( Z √ nv + · · · ) φ ( c̃Z √ nu ) , after 1 step of π-SGD , ( 3 ) for some deterministic scalar c̃ . This expectation can then be evaluated routinely using the Vtransform of φ ( e.g . Cho and Saul ( 2009 ) for relu ; see Definition 3.1 below ) , and likewise , so can the expectations involved in all later steps . After formalization inside a Tensor Program , this rigorously gives rise to the ∞-width limit of f trained under π-SGD , which we call the π-limit of f . See Theorem 3.2 . Definition 3.1 . Let Vφ : R3 → R denote the V-transform of φ , such that Vφ ( EXY , EX2 , EY 2 ) = Eφ ( X ) φ ( Y ) for any centered jointly Gaussian random variables ( X , Y ) ∈ R2 . We will treat Vφ like an activation function and automatically vectorize when Vφ is applied to 3 vectors of the same length . 3.1.1 COMPUTING THE π-LIMIT Here we sketch how to compute the π-limit of Eq . ( 1 ) . Let ut , vt , ft denote corresponding objects at time t , with t = 0 denoting initialization . Memory and Time Requirements What information do we need to store in order to calculate the π-limit ? Because of the projection , ut = ctu0 for some ct ∈ R. So to calculate the limit of ft , we need to track ct ( or rather its∞-width limit ) . This memory requirement is Θ ( 1 ) in training time t. To see the form of vt , it helps to simplify our setup a bit more by 1 ) initializing v0 = 0 and by 2 ) assuming each π-SGD step involves a minibatch containing a sole input ξt . Correspondingly , because the gradient of vt is always proportional to φ ( ξt √ nut ) = φ ( ξtct √ nu0 ) , there exist coefficients { as ∈ R } t−1s=0 , { bs ∈ R } t−1 s=0 ( which are random , but deterministic conditioned on u0 , v0 ) such that √ nvt = t−1∑ s=0 asφ ( bs √ nu0 ) . ( 4 ) Here bs = ξscs and as is formed from the learning rate and loss derivative . So to calculate the limit of ft , we need { as ∈ R } t−1s=0 , { bs ∈ R } t−1 s=0 ( or rather their n → ∞ limits ) . Note that ct , as , bs all have nontrivial fluctuations for finite n , but become deterministic as n→∞ . This implies a memory requirement of Θ ( t ) , which is also the total requirement for computing the limit of ft. As we will see , each forward and backward pass has runtime Θ ( t ) as well , so the total runtime for calculating limn→∞ ft is Θ ( t2 ) ,5 compared to the exponential runtime of general µ-limit . Forward Pass Using Eq . ( 4 ) , we can intuit ( and formalize using Tensor Programs ) that ft ( ξ ) = 1 n n∑ α=1 ( t−1∑ s=0 asφ ( bs √ nu0α ) ) φ ( ξct √ nu0α ) → t−1∑ s=0 ås Eφ ( ̊bsZ √ nu0 ) φ ( ξc̊tZ √ nu0 ) ( 5 ) where ås , b̊s , c̊t denote the deterministic limits of the corresponding quantities . The expectation in the RHS can be evaluated using V-transforms . The t terms in the summation implies a runtime of Θ ( t ) . Backward Pass The gradient update for the output weights vt is clearly represented by setting ( at+1 , bt+1 ) ← ( −ηL′ , ξct ) , whereL′ denotes the loss derivative ∂L ( ft ( ξ ) , label ) /∂ft ( ξ ) . However , a priori , it seems unclear how to calculate the limit of the projected gradient Πu0∇utL of the input weights . Fortunately , because we can express the entire unrolled π-SGD training inside a Tensor Program , the Master Theorem ( Yang , 2020b ) can automatically give us the exact formulas for the gradient limit ; see Appendix F. But in fact , there is a more intuitive way of obtaining the backward pass limit by recognizing that a projected gradient is just the maximal ascent direction in the projected space.6 In our case , Πu0∇utL 4This is because , while c has fluctuations for finite n , it is roughly constant for large n. 5Assuming minibatch size is constant , and we train for a constant number of epochs , this translates to Θ ( N2 ) to train lim ft , where N is dataset size , and Θ ( N ) to do inference on a single input . This compares favorably with Gaussian Processes , which requires Θ ( N3 ) to “ train ” ( i.e . inverting the kernel ) , and Θ ( N ) for inference . However , in our experiments here , the constant in our Θ ( N2 ) in practice will make training lim ft slower than the corresponding NTK or NNGP kernel regression . 6See Lemma F.1 for more details . √ nu0 ) , which is equivalent to Eq . ( 4 ) . In π-SGD , the first layer gradient is projected to the space spanned by u0 before being accumulated to the weights . ( Right ) This network ’ s∞-width limit can be roughly thought of as another 1-hidden-layer neural network with nonlinearity Vφ ( Eq . ( 6 ) ) . But updates are appended to A , B instead of added . See Theorem 3.2 and compare to Eq . ( ? ) . is just the maximal ascent direction of L in the span of u0 , i.e . Πu0∇utL = ( ∇ctL ) u0 , where∇ctL is the derivative w.r.t . the coefficient ct of ut in terms of u0 . In the∞-width limit , this derivative can be obtained by just auto-differentiating Eq . ( 5 ) when the V-transform has an explicit form , like for φ = relu ( Cho and Saul , 2009 ) . So a π-SGD step on ut is equivalent to taking ct+1 ← ct − η∇ctL . Summary Below , for brevity , we say training routine to mean the package of learning rate η , training sequence of singleton minibatches { ( ξt , yt ) } t≥0 ( where ξt is the input and yt is the label ) ,7 and a loss function L ( f ( ξ ) , y ) that is continuously differentiable in the prediction of the model f ( ξ ) . Theorem 3.2 . Consider the simple motivating example of 1-hidden-layer network f in Eq . ( 1 ) with 1- dimensional input and outputs , initialized by uα ∼ N ( 0 , 1/n ) and vα ← 0 . Suppose its nonlinearity φ has a polynomially bounded 2nd derivative . For any training routine , f trained by π-SGD for T steps has an∞-width limit f̊T , in the sense that , as n→∞ , fT ( ξ ) → f̊T ( ξ ) , for any ξ ∈ R , where the limit is almost sure . f̊T is given as follows : Initialize A0 , B0 as empty column vectors , and initialize C0 ∈ R as 1 . Recall Vφ is the V-transform of φ . For each t = 0 , 1 , . . . , define the function ( where dot product of empty vectors is 0 ) f̊t ( ξ ) def = A > t Vφ ( CtξBt , Bt ◦Bt , C2t ξ21 ) , for any ξ ∈ R , ( 6 ) where ◦ denotes coordinatewise product and 1 denotes the all-1s vector of appropriate shape , and At , Bt , Ct are inductively given by Ct+1 def = Ct − η∇CtL ( f̊t ( ξt ) , yt ) and At+1 def = append ( At , −η∇f̊t ( ξt ) L ( f̊t ( ξt ) , yt ) ) , Bt+1 def = append ( Bt , Ctξt ) where append ( v , p ) appends element p to the end of vector v , increasing the dimension of v by 1 . HereAt , Bt correspond to the column vectors formed from the∞-width limits of { as ∈ R } t−1s=0 , { bs ∈ R } t−1s=0 , and Ct corresponds to the limit of ct.8 As discussed above , computing f̊T requires Θ ( T ) memory and Θ ( T 2 ) time . Theorem 3.2 is summarized by Fig . 2 . 3.2 π-PARAMETRIZATION FOR DEEP NETWORKS We can straightforwardly generalize π-SGD and the π-limit theorem ( Theorem 3.2 ) to deep MLPs . However , due to the n× n Gaussian matrix initialization in the middle of the network , the memory requirement will be Θ ( T 2 ) and runtime will be Θ ( T 3 ) for training T steps , for the same reason as discussed in Yang and Hu ( 2020 , Sec 8 ) . This is not scalable to datasets like CIFAR10 . Therefore , we propose a different initialization that brings down the memory requirement to Θ ( T ) and runtime to Θ ( T 2 ) , just like the 1-hidden-layer case . Consider an L-hidden-layer µ-net f : Rd → Rdout : For 7For simplicity , we only consider batch size 1 ; it ’ s straightforward to generalize to larger batch sizes . 8Technically , we should have written Åt , B̊t , C̊t to maintain the convention that ̊ denotes limit of , but for the sake of brevity , we drop this convention in Theorem 3.2. weight matrices w1 ∈ Rn×d and w2 , . . . , wL ∈ Rn×n , and nonlinearity φ : R→ R , such a neural network on input ξ ∈ Rd is given by h1 ( ξ ) = √ nw1ξ ∈ Rn , and xl ( ξ ) = φ ( hl ( ξ ) ) ∈ Rn , hl+1 ( ξ ) = wl+1xl ( ξ ) ∈ Rn , for l = 1 , . . . , L− 1 , ( 7 ) and the network output is f ( ξ ) = n−1/2wL+1xL ( ξ ) for wL+1 ∈ Rdout×n . In µ-parametrization , we would initialize wlαβ ∼ N ( 0 , 1/n ) for any α , β , l. Next , we describe our alternative proposal . π-Initialization Choose integers r , M , to be explained shortly ; these numbers should all be thought of as constant in width n. Suppose we are given a collection P of matrices P def= { Al , Bl ∈ RM×r } Ll=2 ∪ { A1 ∈ Rd×r } ∪ { BL+1 ∈ RM×r , AL+1 ∈ RM×dout } . ( 8 ) Then we can initialize weights wl by first sampling a standard random Gaussian matrix Ω ∈ Rn×r , Ωαi ∼ N ( 0 , 1 ) , before setting wl ← 1 n ΩAl > φ ( BlΩ > ) ∈ Rn×n , for all hidden l = 2 , . . . , L ; w1 ← 1√ n ΩA1 > ∈ Rn×d ; wL+1 ← 1√ n AL+1 > φ ( BL+1Ω > ) ∈ Rdout×n ( 9 ) As an example , the initialization of the 1-hidden-layer case in Theorem 3.2 corresponds to d = dout = r = 1 , M = 0 and A1 = 1 . See Fig . 3 ( Left ) for an illustration . Let ’ s digest this initialization scheme : 1 ) Of course , Al , Bl need to be initialized themselves , and in this paper we will just do so with Gaussian random initialization ; see Appendix B.1 for the pseudo-algorithm . 2 ) Al , Bl , l ∈ [ 2 , L ] , will play the same roles as A , B in the 1-hidden-layer example above , whereas A1 plays the same role as C there . 3 ) r is a measure of dimension of the the projection . More precisely , we will project hidden weight gradients from Rn×n to Rn×r ; see below . 4 ) Al , Bl will grow in the M dimension with training time , just like how the sizes of A , B grow in the 1-hidden-layer example . 5 ) Eq . ( 9 ) can be interpreted as saying wl are generated by 1-hidden-layer MLPs FAl , Bl : Rr → ( Rr or Rdout ) with weights Al , Bl , like so : wlαβ ← 1/n〈Ωα , FAl , Bl ( Ωβ ) 〉 , for all hidden l = 2 , . . . , L ; wL+1 : β ← 1/nFAL+1 , BL+1 ( Ωβ ) . ( 10 ) This is reminiscent of hypernetworks ( Ha et al. , 2016 ) , but the crucial difference is that the “ batch dimension ” ( indexed by α , β ) of the generator FAl , Bl becomes the width dimension of the generated network f , so that the same generators can generate networks of arbitrary width . π-Projection We leave the output weights wL+1 alone , but for any w ∈ { w1 , w2 , . . . , wL } , where w has shape n × n or n × d , we project the gradient ∇wL of the same shape by left multiplying by ΠΩ , the projection matrix to the r-dimensional space spanned by the columns of Ω . This means that preactivations hl are always in this space . However , the input side of w is not projected , so we are optimizing the hidden weights in a ( r × n ) -dimensional space , which still becomes infinitedimensional as n→∞ . We refer to SGD with π-projected gradients as π-SGD . This brings us to the following Definition 3.3 ( π-Parametrization and π-Limit ) . We define π-parametrization as the package of π-initialization and π-SGD,9 where r and the initial M are clear from context . A network in π-parametrization is abbreviated π-network or π-net , and we define π-limit as its∞-width limit . 3.2.1 COMPUTING THE π-LIMIT OF DEEP MLP Here we show the π-limit forward and backward passes can be efficiently computed . π-Limit Forward Pass Let ’ s write fP for the ( random ) neural network π-initialized by P as in Eqs . ( 8 ) and ( 9 ) . Then it ’ s straightforward to see fP has a deterministic limit f̊P expressible as a composition of matrix multiplications and V-transforms : Theorem 3.4 ( π-Limit Forward Pass ) . Suppose φ has a polynomially bounded 2nd derivative . Then as n→∞ , fP ( ξ ) → f̊P ( ξ ) , for any ξ ∈ Rd , where convergence is almost sure , and f̊P ( ξ ) is defined as follows . Write g1 def= A1 > ξ ∈ Rr , gl def= A l > Vφ ( Blgl−1 , Bl ◦Bl , ‖gl−1‖21 ) ∈ { Rr for l = 2 , . . . , L Rdout for l = L+ 1 , and f̊P ( ξ ) def= gL+1 , where B ◦B yields a size-M column vector of squared norms of B ’ s rows . 9with learning rate independent of width ( but may vary with training time ) See Fig . 3 ( Right ) for an illustration . Here , gl represents the coefficients of preactivation hl ( c.f . Eq . ( 7 ) ) in the columns of Ω . As we will see next , π-projection ensures that at any point during training , ft has an∞-width limit of the form f̊P for some P . π-Limit Backward Pass Just like the 1-hidden-layer case , we can interpret the projected gradient as the maximal ascent direction in the column space of Ω. Formalizing the intuition using Tensor Programs yields the following theorem that we also empirically verify in Appendix E. Theorem 3.5 ( π-Limit Backward Pass ) . Let P0 denote the matrices ( c.f . Eq . ( 8 ) ) used in the πinitialization of f as in Eq . ( 9 ) . Suppose its nonlinearity φ has a polynomially bounded 2nd derivative . Then for any training routine , f trained by π-SGD for T steps has an∞-width limit which is equal to f̊PT for some PT given below , i.e . as n→∞ , fT ( ξ ) → f̊PT ( ξ ) , for any ξ ∈ Rd , where the limit is almost sure . PT is given inductively through its elements AlT , BlT as follows : A1t+1 def = A1t − η∇A1L ( f̊Pt ( ξt ) , yt ) , and , for l = 2 , . . . , L+ 1 , Alt+1 def = append ( A l t , −η∇gltL ( f̊ Pt ( ξt ) , yt ) ) , B l t+1 def = append ( B l t , g l−1 t ) . Here glt corresponds to g l in Theorem 3.4 evaluated for input ξt and function f̊Pt , and append ( B , g ) means appending g as a new row vector of B , increasing by 1 the dimension represented by M . Theorems 3.4 and 3.5 are summarized by Fig . 3.10 As conveyed by Eq . ( ? ) , the π-limit can be thought of as another MLP with activation Vφ and trained by gradient concatenation.11 Because π-parametrization has similar scaling with width as µ-parametrization , it is easy to show that the former admits feature learning in the same way the latter does , i.e . the feature kernel of every layer evolves nontrivially during training . Using Tensor Programs , it is straightforward to extend the above theorems to more general settings such as biases , large batch size , per layer learning rate , or metalearning ( Appendix A ) . Appendix D also makes several observations on the π-limit theorems . π-Limit vs µ-Limit While the projection means optimization is slower in the π-limit , we believe with sufficiently large r , the π-limit should perform similarly to the µ-limit . Indeed , prior works such as Li et al . ( 2018 ) has shown that optimizing a neural network in a much smaller , random parameter subspace can recover most of the performance of training all parameters . This is also evidenced by our experimental results below , where the π-limit generally only slightly outperform wide µ-networks . 10One may wonder how would Theorem 3.5 compare with just directly accumulating gradients on P . In fact , this would train very badly because Vφ is very smooth around 0 , so f̊P would look very linear for a long time . Empirically , direct SGD on P would yield ≤ 53 % test and ≤ 80 % training accuracy on CIFAR10 . 11Note that once a vector is appended to Al or Bl , it is not touched again . Therefore , despite the similarity between the π-limit and an MLP , Al and Bl should not be thought of “ parameters ” in the usual sense .

Summary: The authors study a certain variant of an MLP trained using a projected gradient descent-inspired update rule in the infinite width limit. This infinite width model admits closed-form rules for computing the predictor at any training iteration (as long as the V-transform or neural network kernel being known, which is the case for e.g. the arc-cosine kernel). The authors evaluate their model on Omiglot and CIFAR10 against NTK and finite width baselines.