Year: 2021

Abstract:
While a popular limit of infinite-width neural networks, the Neural Tangent Kernel (NTK) often exhibits performance gaps from finite-width neural networks on standard datasets, due to lack of feature learning. Although the feature learning maximal update limit, or μ-limit (Yang and Hu, 2020) of wide networks has closed the gap for 1-hidden-layer linear models, no one has been able to demonstrate this for deep nonlinear multi-layer perceptrons (MLP) because of μ-limit’s computational difficulty in this setting. Here, we solve this problem by proposing a novel feature learning limit, the π-limit, that bypasses the computational issues. The π-limit, in short, is the limit of a form of projected gradient descent, and the π-limit of an MLP is roughly another MLP where gradients are appended to weights during training. We prove its almost sure convergence with width using the Tensor Programs technique. We evaluate it on CIFAR10 and Omniglot against NTK as well as finite networks, finding the π-limit outperform finite-width models trained normally (without projection) in both settings, closing the performance gap between finiteand infinite-width neural networks previously left by NTK. NTK Finite-Width Network -Width -Limit (Ours) Figure 1: PCA of representations of images from 5 classes (= 5 colors) in Omniglot test set. We compare the representations from our best performing finite-width μ-parametrized network, our proposed∞-width π-limit, and NTK. The representations from former two form neat clusters according to their classes while those from the latter are mixed up together. See Appendix B.5.3.

heading: 1 INTRODUCTION
text: 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.

heading: 3 SOLUTION: PROJECT THE GRADIENT
text: 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.

heading: 3.1 CONTINUING THE 1-HIDDEN-LAYER EXAMPLE
text: 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.

heading: 4 EXPERIMENTS
text: Here we compare the performance of the relu π-limit on CIFAR10 (Krizhevsky, 2009) and Omniglot (Lake et al., 2015) against that of NTK, NNGP, and finite-width π- and µ-nets.12 As we will see, 1) π-limits of large enough r beat finite µ-nets; 2) finite π-nets underperform the above on CIFAR10 but, interestingly, outperform them on Omniglot; 3) all of the above beat NTK and NNGP.

heading: 4.1 CLASSIFICATION ON CIFAR10
text: Experimental Setup For π- and µ-networks, infinite or finite, we train for 50 epochs. We adopt a step learning rate schedule, with a learning rate drop of 0.15 at a certain milestone, which is a hyperparameter. We sweep over a variety of hyperparameters such as the learning rate, gradient clipping, weight decay, the LR drop milestone, etc, as well as width, r, and depth. For NTK and NNGP, we perform kernel regression following Lee et al. (2018) using centered labels. For them, we sweep over the initialization variances and learning rate multipliers, along with ridge coefficient. See Appendix B for more details.
Results In the literature, the performance gap between CNN and its NTK (Arora et al., 2019) is often cited for the deficiency of the NTK theory for explaining neural networks in practice. However, in fact, on MLPs we already see a nontrivial gap, as seen in Table 1 (compare µ-Net with NTK, NNGP).13 The π-limit closes this gap, having almost a 2-point advantage over NTK. The finite-width π-net outperforms NTK and underperforms π-limit both by about 1 point.
Feature Learning in π-Limit and Finite Networks We show the advantage of π-limit over NTK and NNGP is due to feature learning. To do so, we track the kernel regression performance of its (last-layer) feature kernel over the course of training. As seen in Fig. 4(Left), while the feature kernel at initialization underperforms both NNGP and NTK, it improves consistently over time to eventually exceed them.14 Similarly, the kernel performance of the best π- and µ-nets improves over time as well, though the accuracy of the feature kernel regression is slightly less than the network itself, likely due to the low rank property of the feature kernel. On the contrary, the π-limit benefits from feature kernel regression, improving the test accuracy from 61.5% to 61.85%; see Table 7.
Effects of Width, r, and Depth As shown in Fig. 4(Middle), accuracy of π-net increases monotonically with width and with r, approaching the accuracy of π-limit and µ-net, respectively, from below. This can also be seen in feature kernel regression across training time, Fig. 5. In contrast, performance is not monotonic in depth for any of µ-net, π-net, or π-limit, beyond a consistent improvement from 1 to 2 hidden layers; see Fig. 6.
12Because 1) standard parametrization differs from µ-parametrization only in factors depending on width, 2) we sweep such factors in hyperparameter optimization, and 3) our finite networks have width at most 2048 (so these factors are in practice constants), our best accuracies for µ-nets will also be best accuracies for standard-parametrized MLPs.
13Note that, contrary to Lee et al. (2018; 2020), which claimed that finite-width neural networks underperform the kernels, here we find the former outperform the latter. This is primarily due to the single learning rate drop we adopted, while Lee et al. (2018; 2020) used a constant learning rate. We believe this provides a fairer comparison to the kernels since 1) the step LR schedule is more common than constant learning rate in practice, and 2) the kernels would obtain the same performance if we did kernel gradient descent for infinite-time (which is equivalent to kernel regression) with learning rate drop, since this optimization is convex.
14In both NNGP and NTK limits, the feature kernel stays fixed throughout training (Yang and Hu, 2020), so the dotted line in Fig. 4(Left) shows the result if we do the same experiments for NNGP and NTK.
CIFAR10 Feature Kernel Regression
CIFAR10 Val Acc
Omniglot Val Acc

heading: 4.2 TRANSFER LEARNING FROM IMAGENET TO CIFAR10
text: Pretraining-and-transfer-learning is an important setting where feature learning is crucial. As pointed out in Yang and Hu (2020), the NTK limit trivializes pretraining, while the µ-limit both theoretical and empirically benefit from it. Here we investigate pretraining for the π-limit in the image domain by pretraining on ImageNet32 (Chrabaszcz et al., 2017)15 and transferring to CIFAR10.
Experimental Setup We pretrain the π-limit with r = 400 for 30 epochs on a fixed subset of ImageNet32 with 250 (out of 1000) randomly subsampled classes. To evaluate on CIFAR10, we compute the kernel induced by the pretrained final-layer-features and perform kernel regression.
Results As shown in Table 1, ImageNet32 pretraining nontrivially raises the downstream CIFAR10 performance over without pretraining, altogether creating a 5% gap compared to NTK (which has the same test accuracy whether or not it is pretrained on ImageNet32 because of the disparate classes, as shown in Yang and Hu (2020)). This again demonstrates the feature learning capability of the π-limit.

heading: 4.3 FEW-SHOT LEARNING VIA METALEARNING ON OMNIGLOT
text: Following Yang and Hu (2020), we also evaluate few-shot learning on Omniglot. Compared to traditional classification settings like CIFAR10, doing well on Omniglot requires learning features that can rapidly be adapted to new unseen data (Raghu et al., 2019). We will adopt a metalearning approach to this, following Finn et al. (2017). Unlike Yang and Hu (2020), we will train our models using ANIL (Almost No Inner Loop) (Raghu et al., 2019) rather than first-order MAML (Model Agnostic Meta-Learning). Briefly, ANIL is a variant of second-order MAML, where in the inner loop, we only adapt the output layer, while in the outer loop, we only train the network body. We adopt ANIL because: 1) We would like to train deep MLPs with SGD and without Adam or batchnorm, but empirically this makes optimization of first-order MAML difficult for both finite and∞-width networks compared to second-order MAML. 2) While a π-parametrized network trained by second-order MAML has an∞-width limit calculable using Tensor Programs, it does not stay in the fP form of Theorem 3.4 because of the second-order gradient, rendering the limit computation inefficient. 3) Fortunately, π-networks trained by ANIL do not have this issue because the inner and outer loop gradients are on different parameters. In addition, ANIL performs on par with MAML on standard finite-width networks (Raghu et al., 2019). 4) Furthermore, ANIL training more clearly delineates the role of the network body for learning reusable features and the role of the head for adapting to new data, which is more fitting for our goals in this paper.
Experimental Setup We focus on the 5-way, 1-shot task16, with only 1 step of ANIL adaption. For π- and µ-networks, infinite or finite, we train for 50 epochs, 1000 batches per epoch, and 8 tasks per batch. Following Antoniou et al. (2019), we use cosine annealing schedule on the meta-learning rate. For NTK and NNGP, ANIL meta-training has no effect, and meta-testing amounts to just taking 1
15i.e. ImageNet (Deng et al., 2009; Russakovsky et al., 2015) downsampled to 32× 32 resolution 16i.e. each task consists of 5 different classes, with 1 example provided for each class in training.
kernel gradient descent step on each task.17 We sweep over a variety of hyperparameters such as the outer and inner learning rates, gradient clipping, etc, as well as width, r, and depth. See Appendix B for more details.
Results As seen in Table 1, µ-net, π-net, and π-limit all outperform NNGP and NTK by about 40 points. Interestingly, while π-limit is slightly better than µ-net, they are both outperformed by π-net. This is related to the width nonmonotonicity seen on Omniglot, as we describe next.
Effect of Width, r, and Depth While π-net performance with r is still roughly monotonic, in contrast to in CIFAR10, here it can decrease with width, as seen in Fig. 4(Right). In addition, µ-net seems to underperform π-net of the same width for sufficiently large r. We find this is primarily due to optimization and not generalization because Omniglot is hard to overfit; see Fig. 7. Counterintuitively, π-net of large r can optimize faster than µ-net despite the projection. This is likely a side effect of π-initialization, as we see this advantage persist when r > width, where projection is a no-op. We do not see width nonmonotonicity on CIFAR10 because overfitting is much easier there, so the results depend much more on the implicit bias of each model. On the other hand, performance seems to monotonically increase along diagonals of fixed r/width ratio, peaking around r/width ≈ 1/2. This suggests a different limit of r ∝ width → ∞, which warrants investigation in future works. Unlike CIFAR10, test accuracy is more monotonic in depth here (Fig. 6). Again, this is probably because the extra weights help with optimization, which is the bottleneck on Omniglot.

heading: 5 RELATED WORKS
text: The π-limit bears some superficial similarities to hierarchical kernel processes such as deep GPs (Damianou and Lawrence, 2013; Salimbeni and Deisenroth, 2017) or deep kernel processes (Aitchison et al., 2021). The crucial difference here is that 1) during inference, the π-limit is deterministic and not a hierarchically sampled process, and 2) the π-limit is not trained variationally. Relatedly, deep GPs or kernel processes can be construed as a kind of randomly initialized infinitewidth neural network with finite bottleneck layers (Agrawal et al., 2020; Aitchison, 2020). However, their variational inference procedures do not correspond to the limit of SGD. In contrast, Littwin et al. (2021) derived such an SGD limit where the network between consecutive bottlenecks is in NTK parametrization, and the limit takes the form of a kind of intertwined system of kernel gradient descent. However, the computation of this limit is not scalable to datasets like CIFAR10. During training, the π-limit’s Al and Bl grow in size, reminiscent of works on growing neural networks (Liu et al., 2019; Wu et al., 2021; Gu et al., 2021; Gong et al., 2019). Of course, these works focus on improving training efficiency of practical models, which is not our goal here. Nevertheless, it will be interesting to explore whether insights from π-limit can contribute back to this literature. Our π-initialization (given fixed P) can be construed as a special form of the initialization in deep mean field limits of MLPs (Araújo et al., 2019; Sirignano and Spiliopoulos, 2020; Fang et al., 2020; Nguyen, 2019; Nguyen and Pham, 2020). However, without π-projection, these limits’ analytical computation suffers exponential runtime blowup like the µ-limit.

heading: 6 CONCLUSION
text: A good model for studying wide neural networks should 1) capture the desired behavior of neural networks in practice, including feature learning and performance on real datasets, and 2) come with a set of theoretical tools for working researchers. But an attractive model must also balance these properties with 3) computational simplicity for empirical investigations, and 4) mathematical simplicity for theoretical investigations. Previously, NTK and NNGP satisfy 2), 3), 4) but not 1), as shown in this and prior works. The µ-limit satisfies 1) and 2) but, arguably, not 3) and 4). In this work, we presented the π-limit which, we believe, satisfies all 4 properties and that can prove fruitful for future researches. While our work closed the∞-width performance gap, it also opens many new questions, e.g., Can we derive the π-limit for a fixed r/width ratio? When precisely does the width nonmonotonicity of π-net occur? What about modern architectures, beyond MLP? We leave these questions to future work.
17NTK and NNGP can be trained under second-order MAML, but we found their performance strictly decreases with metatraining, as the randomization of labels across tasks confuse the readout layer.

heading: A.1 BIASES
text: Biases are straightforward to add. Consider an L-hidden-layer µ-net f : Rd → Rdout with biases: such a neural network on input ξ ∈ Rd is given by h1(ξ) = √ nw1ξ + √ nb1 ∈ Rn, and
xl(ξ) = φ(hl(ξ)) ∈ Rn, hl+1(ξ) = wl+1xl(ξ) + √ nbl ∈ Rn, for l = 1, . . . , L− 1, (11)
and the network output is f(ξ) = n−1/2wL+1xL(ξ) + bL+1 for wL+1 ∈ Rdout×n. In µparametrization, the biases are initialized as blα ∼ N (0, σ2b/n) for l = 1, . . . , L, and b L+1 i ∼ N (0, σ2b ), for a hyperparameter σb. However, in π-parametrization, we will do this a bit differently.
π-Initialization Suppose we are given a collection P of matrices and vectors
P def= (RHS of Eq. (8)) ∪ {βl ∈ Rr}Ll=1 ∪ {βL+1 ∈ Rdout}. (12)
Then we can initialize weights and biases by first sampling a standard random Gaussian matrix Ω ∈ Rn×r,Ωαi ∼ N (0, 1), and 1) set wl as in Eq. (9), and 2) set biases as
bL+1 ← βL+1 and bl ← 1√ n Ωβl for all l = 1, . . . , L. (13)
This constitutes the π-initialization of biases.
π-Projection For π-projection, we also project the bias gradients by ΠΩ.
π-Limit Calculation Theorem A.1 (π-Limit Forward Pass with Bias). Let P be some collection as in Eq. (12). As n→∞,
fP(ξ)→ f̊P(ξ), for any ξ ∈ Rd,
where convergence is almost sure, and f̊P(ξ) is defined as follows. Write g1 def= A1>ξ + β1 ∈ Rr,
gl def= A l>Vφ(Blgl−1, Bl ◦Bl, ‖gl−1‖21) + βl ∈ { 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. Theorem A.2 (π-Limit Backward Pass). For the same setting as in Theorem 3.5 but with P as in Eq. (12), Al, Bl are updated exactly as in Theorem 3.5, and βl is updated by gradient accumulation like A1:
βlt+1 def = β l t − η∇βlL(f̊Pt(ξt), yt), for all l = 1, . . . , L+ 1

heading: A.2 PARAMETER MULTIPLIERS
text: We can insert to Eq. (11) constant parameter multipliers αw for each parameter w like so
hl+1(ξ) = αwl+1w l+1xl(ξ) + √ nαblb l ∈ Rn, for l = 1, . . . , L− 1, (14)
These multipliers are tuneable hyperparameters. They affect both the forward and backward passes of the network. In the π-limit, the forward pass is the same as in Theorem A.1 except we replace Al with αwlAl and βl with αblβl. The backward pass is the same as in Theorem A.2, but we just have to make sure that we backprop through the multipliers (in contrast, if we instead have absorbed the multipliers into the initialization, then we would not backprop through the multilpliers). In our experiments, we only consider the input weight multiplier, output weight multiplier, and a single multiplier for all biases.

heading: A.3 LEARNING RATE MULTIPLIERS
text: We may have custom learning rates for specific weights or biases. In our experiments, we implement this with learning rate multipliers relative to the global learning rate η, e.g., the learning rate of a parameter w becomes γwη if the multiplier is denoted γw. Then in the limit, we just need to replace η in Theorem 3.5 or Theorem A.2 with γwlη, i.e.
A1t+1 def = A 1 t − γw1η∇A1L(f̊Pt(ξt), yt)
Alt+1 def = append(A l t,−γwlη∇gltL(f̊ Pt(ξt), yt)), for all l = 2, . . . , L+ 1
βlt+1 def = β l t − γblη∇βlL(f̊Pt(ξt), yt), for all l = 1, . . . , L+ 1.
But note we do not modify the Bl update. In our experiments, we will only sometimes use a LR multiplier on the input layer (“Input Layer LR Mult”), one on the output layer (“Output Layer LR Mult”), and/or a single multiplier for all biases (“Bias LR Mult”).

heading: A.4 LARGE BATCH SIZE
text: For batch size S > 1, we still accumulate gradients into A1 and βl as if they are regular parameters, and for the hidden weights, we just append the S gradient vectors as if they are from S unit-sized batches.

heading: A.5 GRADIENT CLIPPING
text: The Frobenius norms of the projected gradients of the input weight w1 and all biases bl converge to exactly the Frobenius norms of the gradients of A1 and βl, i.e.
‖ΠΩ∇w1L‖F → ‖∇A1L‖F ‖ΠΩ∇blL‖ → ‖∇βlL‖
where the convergence is almost sure. For any hidden weights wl, suppose its gradient over an S-sized batch in the π-limit is given by
Alt+1 def = append(A l t,−ηÃ), Blt+1 def= append(Blt, B̃).
where Ã, B̃ ∈ RS×r, and append(B, B̃) means appending all S rows of B̃ into B, increasing the latter’s column length by S. Then
‖ΠΩ∇wlL‖F → √ tr ( Ã>φ̃(B̃B̃>)Ã ) where the convergence is almost sure. Then clipping the gradient in the π-limit just means clipping
the gradients of A1 and βl, and, for hidden weights, rescale Ã such that √ tr ( Ã>φ̃(B̃B̃>)Ã ) is at
most the clipping threshold.

heading: A.6 WEIGHT DECAY
text: The limit of decaying wl ← wl(1− lr · wd) is just Al ← Al(1− lr · wd).

heading: A.7 DECOUPLING LAYERS
text: We can actually use different rs and Ms for every layer, and a similar theorem can be proved. We can also use different, independently sampled Ω for each layer. However, the limit would be exactly the same as before. Nevertheless, this could make finite-width networks a less biased estimator of the ∞-width limit.

heading: B EXPERIMENTAL DETAILS
text: All of our experiments are done on V100 GPUs. All of our networks use relu activation. The V-transform of relu is (Cho and Saul, 2009)
Vrelu(r1r2c, r21, r22) = 1
2π
(√ 1− c2 + (π − arccos(c))c ) r1r2
for any r1, r2 > 0 and c ∈ [−1, 1].

heading: B.1 NETWORK INITIALIZATION AND PARAMETRIZATION
text: At initialization, 1) biases are always 0, 2) M is set to r so that Al, Bl are square matrices for l = 2, . . . , L, and 3) we sample Al, Bl as standard Gaussians and then scale them as follows:
A1 = A1/A1.norm(dim = 0, keepdim = True)
Bl = Bl/Bl.norm(dim = 1, keepdim = True) for all l = 2, . . . , L Al = Al/ √ A1.shape[0] for all l = 2, . . . , L, and
AL+1 = 0

heading: B.2 FEATURE KERNEL
text: When we talk about the “feature kernel” of a π-limit f̊P , we always mean the n → ∞ limit of the feature kernel of fP , and not the kernel induced by gL in f̊P . This feature kernel K on inputs {ξ1, . . . , ξk} is calculated as
Kij = Vφ(〈gLi , gLj 〉, ‖gLi ‖2, ‖gLj ‖2)
where gLi and g L j are the g L in Theorem 3.4 evaluated for two inputs ξi and ξj .

heading: B.3 CIFAR10
text: B.3.1 µ-NET, π-NET, AND π-LIMIT For π- and µ-networks, infinite or finite, we train for 50 epochs. We adopt a step learning rate schedule, with a learning rate drop of 0.15 at a milestone hyperparameter. We also clip gradients with a hyperparameter threshold, where the clip is triggered individually for each parameter by the parameter’s norm, rather than total parameter norm.
Hyperparameter Optimization We first perform random search on hyperparameters listed in Table 2 for 2-hidden-layer π-net, π-limit, and µ-net: We sample at least 512 hyperparameter combinations 1) for each (width, r) ∈ {128, 256, 512, 1024, 2048} × {50, 100, 200, 400} for π-net, 2) for each width ∈ {128, 256, 512, 1024, 2048} for µ-net, and 3) for each r ∈ {50, 100, 200, 400} for π-limit. The best accuracies per (width, r) are shown in Fig. 4(Middle). Then we take the best hyperparameter combinations overall for π-net, π-limit, and µ-net, and for depth ∈ {1, 2, 3, 4} (where depth denotes number of hidden layers), we resweep only learning rate and weight decay in a grid search, where the grid is (lr∗ · {2−2, 2−1.5, . . . , 21.5}) × (wd∗ · {2−2, 2−1.5, . . . , 21.5}) and lr∗ and wd∗ are the optimal hyperparameters from the 2-hidden-layer sweep. The best test accuracies per depth are shown in Fig. 6, while the overall best test accuracies over all depth are shown in Table 1.

heading: B.3.2 NNGP AND NTK
text: For NNGP and NTK, we perform kernel regression following Lee et al. (2018) using centered labels. For each depth ∈ {1, 2, 3, 4}, we sweep over hyperparameters listed in Table 3 (which change the kernels) using (complete) grid search, with weight initialization variance fixed at 1 in the NTK
parametrization.18 For each of NNGP and NTK, the best test accuracy over all depths is listed in Table 1.
B.4 IMAGENET TRANSFER
We pretrain the π-limit with r = 400 and 2 hidden layers for 30 epochs on a fixed subset of ImageNet32 with 250 (out of 1000) randomly subsampled classes. To evaluate on CIFAR10, we compute the kernel induced by the pretrained final-layer-features and perform kernel ridge regression. We optimize the hyperparameters in Table 4 via random search.

heading: B.5 OMNIGLOT
text: We focus on the 5-way, 1-shot task, with only 1 step of ANIL adaption.
B.5.1 µ-NET, π-NET, AND π-LIMIT
For π- and µ-networks, infinite or finite, we train for 50 epochs, 1000 batches per epoch, and 8 tasks per batch. In each epoch, we validate on 500 batches from the validation set. Following Antoniou et al. (2019), we use cosine annealing learning rate schedule. We also clip gradients with a hyperparameter threshold, where the clip is triggered by the total parameter norm.
Hyperparameter Optimization We first perform random search on hyperparameters listed in Table 5 for 2-hidden-layer π-net, π-limit, and µ-net: We sample at least 512 hyperparameter combinations 1) for each (width, r) ∈ {128, 256, 512, 1024, 2048} × {50, 100, 200, 400} for π-net, 2) for each width ∈ {128, 256, 512, 1024, 2048} for µ-net, and 3) for each r ∈ {50, 100, 200, 400} for π-limit. Then we take the best hyperparameter combinations overall (based on validation accuracy) for π-net, π-limit, and µ-net, and for depth ∈ {1, 2, 3, 4} (where depth denotes number of hidden layers), we resweep only the meta learning rate (i.e. outer learning rate) and step size (i.e. inner learning rate) in a grid search, where the grid is (ilr∗ · {2−2, 2−1.5, . . . , 21.5}) × (olr∗ · {2−2, 2−1.5, . . . , 21.5}) and ilr∗ and olr∗ are the optimal inner and outer learning rates from the 2-hidden-layer sweep. The best validation accuracies per depth are shown in Fig. 6. Then we take the models with overall best validation accuracies over all depth and evaluate them on the test set using 10000 batches. These test results are shown in Table 1. To investigate the width non-monotonicity more thoroughly, we further perform random search on hyperparameters listed in Table 5 for 2-hidden-layer π-net and µ-net. We sample at least 512 hyperparameter combinations 1) for each (width, r) ∈ {128, 256, 512, . . . , 8192} × {50, 100, 200, . . . , 32000}
18This is without loss of generality because relu is homogeneous and we are sweeping the bias variance.
for π-net, and 2) for each width ∈ {128, 256, 512, . . . , 8192} for µ-net. The best validation accuracies per (width, r) are shown in Fig. 4(Right).

heading: B.5.2 NNGP AND NTK
text: As discussed in Section 4.3, for NTK and NNGP, ANIL meta-training has no effect, and meta-testing amounts to just taking 1 kernel gradient descent step on each task. We fix the step size (i.e. inner learning rate) at 0.5. For each depth ∈ {1, 2, 3, 4}, we sweep over hyperparameters listed in Table 6 using (complete) grid search, with weight initialization variance fixed at 1 in the NTK parametrization. For each of NNGP and NTK, the best test accuracy over all depths is listed in Table 1.
B.5.3 VISUALIZATION OF IMAGE REPRESENTATIONS We sample 5 random classes and 10 random images from each class from the Omniglot test set, for a total of 50 images. We take the best performing NTK, µ-net, and π-limit and evaluate their feature kernels (c.f. Appendix B.2) on the 50 images. We then do PCA on these kernels to produce Fig. 1. In Fig. 8, we also do the same for our best performing π-net.

heading: C DETAILED CALCULATIONS OF 1-STEP SGD
text: Suppose we present an input ξ0 to the network and perform a step of gradient descent with learning rate η and loss L. Then simple calculations show that the updates ∆u,∆v to u, v are
√ n∆v = cφ(ξ0 √ nu), ∆u = cv φ′(ξ0 √ nu) (15)
where c = −ηL′. Then, via Tensor Programs, f(ξ) for any ξ now has a limit of the form
lim n→∞
f(ξ) = E(Z √ nv + Z √ n∆v)φ(Z √ nuξ + Z √ n∆uξ)
= E(Z √ nv + c̊φ(Z √ nuξ0))φ(Z √ nuξ + c̊Z √ nvφ′(Z √ nuξ0)),
where c̊ is the deterministic limit of c that is shown to exist by Tensor Programs.
D REMARKS ON THE π-LIMIT THEOREMS
1. π-projection actually ensures that, even for finite n, ft = fP for some P . However, this P is random, with fluctuation coming from the sampling of Ω (which is fixed at initialization). Taking n→∞ reduces this fluctuation to 0.
2. Again, even with the projection, we are optimizing in an infinite-dimensional space, though “linearly infinite” r ×∞ instead of “quadratically infinite”∞×∞ like in the µ-limit.
3. The forward pass of f̊P can be thought of as that of another MLP with nonlinearity Vφ, as illustrated in Fig. 3. In this view, M becomes the width of this MLP.
Table 7: Feature Kernel Regression (FKR) on CIFAR10. We take the best performing µ-net, π-net, and π-limit, and evaluate their learned feature kernels on CIFAR10 via kernel regression. We list their test accuracies in the middle block. For reference, we also include the NNGP and NTK numbers in the left block, as well as the ImageNet Transfer results in the right block.
NNGP NTK µ-Net π-Net π-Limit π-Limit ImageNet Transfer
FKR 58.92 59.63 59.12 59.72 61.85 64.39 vs Table 1 - - -2.19 -0.92 +0.35 -
0 10 20 30 40 50 Train Epochs
0.40
0.45
0.50
0.55 0.60 Te st K er ne l A cc ur ac y
500 1000 5000 10000 20000 30000 40000
Figure 5: Feature Kernel Regression of π-Nets vs Training Time, for Varying Widths. We took our best performing π-limit and trained finite-width versions of it, for width from 500 to 40000. Throughout training, we measure their feature kernel regression accuracy. Altogether, we see consistent increase in performance across width at any moment in time. Note that the visible gap between π-limit and the widest π-net (even at initialization) is to a large extent due to the dependence of kernel regression accuracy on the smallest eigenvalues of the kernel. See Fig. 11.
E NUMERICAL VERIFICATION OF THE π-LIMIT THEOREMS
Here we numerically show that wide π-nets have nearly identical loss curves as the π-limit. See Figs. 9, 10 and 12. In addition, Fig. 11 verifies the convergence of the feature kernel to its infinitewidth limit.

heading: F PROOFS
text: We will just prove Theorems 3.4 and 3.5, as Theorem 3.2 is a special case of them. At a high level, we need to do two things: 1) Show that fT converges almost surely to something, and 2) this something is f̊PT . Here we assume the reader is familiar with Tensor Programs, in particular the techniques used in Yang and Hu (2020).

heading: F.1 ALMOST SURE CONVERGENCE
text: Showing almost sure convergence is straightforward using the Tensor Programs technique, i.e. express π-initialization and the entire π-SGD training trajectory inside a Tensor Program (a NETSOR>+ program in particular) and apply the Master Theorem (c.f. Yang (2020b) in general and Yang and Hu (2020, Sec H.3, H.4) in particular). Below we sketch the program construction.
Initial Vectors and Matrices Unlike the program for the µ-limit, our program does not have initial matrix variables because in π-parametrization we do not initialize any n×nmatrices as iid Gaussians. This means that we do not use MatMul instructions in our program. The initial vector variables in our program are just the r columns Ω:1, . . . ,Ω:r of Ω.
Network Preactivations All preactivations of the network will be of the form ΩC for someC ∈ Rr whose entries are scalar variables in the program, so that ΩC can be expressed as a vector variable using Nonlin+.
Weight Matrices We will sketch the constructions surrounding hidden weight matrices; the input and output weight matrices are similar and easier. Like in Eq. (4), each hidden weight matrix in deep π-nets can be written mathematically as a sum of vector outer products.
wl = 1
n M+t∑ s=1 (ΩAls)⊗ φ(ΩBls), at time t (16)
where Als, B l s ∈ Rr are the sth rows of Al and Bl. Here φ(ΩBls) is the activation going into wl at time s, and ΩAls is the projected gradient ΠΩ∇hlL at time s. Note the sum here is from 1 to M + t, where M comes from the initialization and t is from t steps of training. All entries of Al and Bl will be constructed as scalar variables in the program, so ΩAls and φ(ΩB l s) are both vector variables. In the program, we do not expresswl directly, but rather through its matrix-vector product with vectors such as the incoming activation xl−1, which would be expressed via a combination of Moment and Nonlin+ instructions like so:
wlxl−1 = M+t∑ s=1 θs(ΩA l s) ∈ Rn Nonlin+
where θs = 1
n 〈φ(ΩBls), xl−1〉 ∈ R Moment
(17)
We express (wl)> indirectly through its matrix-vector products likewise, just with the roles of (ΩAls), φ(ΩB l s) reversed.
Gradient Projection In the program, we do not express ΠΩ directly, but rather indirectly by expressing the matrix-vector product ΠΩv for vector variables v, such as v = ∇hlL. The projection matrix ΠΩ is mathematically equal to Ω(Ω>Ω)+Ω> (for any width), where ()+ denotes pseudoinverse. In the program, we would first express (the entries of) 1nΩ >Ω ∈ Rr×r using many Moment
instructions
( 1
n Ω>Ω)ij =
1 n 〈Ω:i,Ω:j〉 ∈ R Moment
Then its pseudo-inverse can be expressed as another Moment instruction (with purely scalar arguments).
( 1
n Ω>Ω)+ij =
1
n n∑ α=1 fij(; {( 1 n Ω>Ω)ij}ij) ∈ R Moment
= fij(; {( 1
n Ω>Ω)ij}ij) ∈ R Moment
where fij takes a matrix to the ijth entry of its pseudoinverse. Note the above expressions depend on Ω only, and not on v. Finally, we express ΠΩv = Ωγ ∈ Rn, γ = ( 1nΩ >Ω)+θ ∈ Rr, θ = 1nΩ >v ∈ Rr in the program like so
ΠΩv = Ωγ = r∑ i=1 γiΩ:i ∈ Rn Nonlin+ (18)
where γi = r∑ j=1 ( 1 n Ω>Ω)+ijθj ∈ R Moment (19)
where θj = 1
n 〈Ω:j , v〉 ∈ R Moment (20)
Wrapping Up Other than what is discussed above, the unrolling of π-SGD follows identical to the unrolling of SGD in Yang and Hu (2020, Sec H.3, H.4). In particular, ft(ξ) for any input ξ and time t is a scalar variable in the program.
Getting Almost Sure Convergence We apply the Master Theorem (Yang and Hu (2020, Thm 7.4) or Yang (2020b, Thm E.15)) to the program to get almost sure convergence to some limit. We just need to check the conditions of the theorem. They are all straightforward except that we need to check the pseudoinverse operation we took is almost surely continuous (in order to satisfy Yang and Hu (2020, Assm F.4(1))).19 However, the only pseudoinverse we took was ( 1nΩ >Ω)+, and Ω has
19Yang and Hu (2020, Assm F.4(1)) actually requires Moment nonlinearities with only parameter arguments to be continuous eveywhere, but because our result is almost sure anyway, we can ignore any measure zero event.
1/
√
width to that of the π-limit.
rank r (full rank) almost surely for any n > r. Therefore, our pseudoinverse operation is almost surely continuous as pseudoinverse is continuous on matrices of constant rank.

heading: F.2 FORM OF THE LIMIT
text: Forward Pass (Theorem 3.4) In the large-n limit, by the Master Theorem, for hidden weights, Eq. (17) becomes
Zw lxl−1 = M+t∑ s=1 θs r∑ i=1 AlsiZ Ω:i = r∑ i=1 ZΩ:i M+t∑ s=1 θsA l si
where θs = E [ φ(Zh l−1 )φ ( r∑ i=1 BlsiZ Ω:i )]
Inductively, if gl ∈ Rr represents the coefficients of Zhl = Zwlxl−1 in terms of ZΩ def= (ZΩ:1 , . . . , ZΩ:r ) (which is distributed as a standard isotropic Gaussian vector), then
gli = M+t∑ s=1 Alsi Eφ(〈gl−1, ZΩ〉)φ(〈Bls, ZΩ〉),
which can be rearranged straightforwardly into the equation of Theorem 3.4. The equations for input and output layers can be derived similarly.
Backward Pass (Theorem 3.5) For hidden weightwl, since we maintainwl in the form of Eq. (16), the gradient update
wl ← wl − ηΠΩ∇wlL = wl + (−ηΠΩ∇hlL)⊗ φ(hl−1) in the limit corresponds to
Al ← append(Al, coef(Z−nηΠΩ∇hlL)), Bl ← append(Bl, coef(Zh l−1 ))
where coef(Zv) is the coefficient of Zv in terms of ZΩ = (ZΩ:1 , . . . , ZΩ:r ). (Note the factor of n in Z−nηΠΩ∇hlL) is just there so that nηΠΩ∇hlL has Θ(1)-sized coordinates). Of course, coef(Zh l−1 ) is just gl−1 ∈ Rr by definition. It remains to show that coef(Z−nηΠΩ∇hlL) = ∇gl limL, where limL = limn→∞ L(f(ξ), y) is the loss at the limit (which is deterministic). Using the Master Theorem, it is not hard to see that, for each preactivation hl, Zn∇hlL is the Frechet derivative ∂ limL
∂Zhl with respect to the random variable Zh
l
, where the Frechet derivative is defined
with respect to the Hilbert spaceH of square integrable random variables in the σ-algebra generated by ZΩ def= (ZΩ:1 , . . . , ZΩ:r ). Furthermore, using the limits of Eqs. (18) to (20) obtained from the Master Theorem, it is easy to see that
ZnΠΩ∇hlL = ΠZΩZ n∇ hl L
where ΠZΩ is the orthogonal projection to the linear subspace ofH spanned by the random variables ZΩ:1 , . . . , ZΩ:r . Then Lemma F.1 applies and we get that
ZnΠΩ∇hlL = r∑ i=1 ∂ limL ∂gli ZΩ:i
coef(ZnΠΩ∇hlL) = ∇gl limL,
where the derivative is now an ordinary partial derivative. Lemma F.1. Let H be a Hilbert space and V be a k-dimensional subspace of V , where k is finite. Let L : H → R be a Frechet differentiable function. Suppose Γ have for columns an orthonormal set of basis for V , and w = Γb ∈ V for some b ∈ Rk. Then
ΠV∇wL(w) = ∇bL(Γb),
where ΠV : H → V is the orthogonal projection to V .
Proof. By the theory of proximal gradients,
ΠV∇wL(w) = w −min v∈V
[ 〈v,∇wL(w)〉+ 1
2 ‖v − w‖2
] .
Changing coordinates via Γ, we have
ΠV∇wL(w) = w − min c∈Rk
[ 〈Γc,Γ∇bL(Γb)〉+ 1
2 ‖Γc− Γb‖2 ] = w − min
c∈Rk
[ 〈c,∇bL(Γb)〉+ 1
2 ‖c− b‖2 ] = ∇bL(Γb)

Review summaries:
The paper proposes an infinite-width parameterization ($\pi$-parameterization) that is similar to $\mu$-parameterization of Yang and Hu, with two important tweaks: the initialization involves a random projection of a certain form and the gradient update invokes the same projection. This results in an infinite-with description that is more amenable to computation and analysis. The paper demonstrates experimentally that this method achieves feature learning and better performance than NTK.

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.

In this work, the authors propose a new feature learning limit for infinite-width neural networks that resolves some of the computational issues with the previous $\mu$-limit. This new limit, called the $\pi$-limit, is based on projected gradient descent, and crucially allows the authors to evaluate the model on deep networks, which could not be done before. The authors show that it outperforms both finite-width and previous infinite-width networks on CIFAR10 and Omniglot.

The paper introduces an approach (named pi-limit) to compute explicitly the infinite-width limit performance of a multilayer perceptron (MLP). The approach is based on projecting the gradient of the loss at each epoch in a fixed direction, chosen upon initialization. If one does so,  the expectation value of the function learned by the network at infinite training time can be computed explicitly by the iterative procedure described in Theorem 3.5. The asymptotic accuracy on a network trained on CIFAR10 is 61.5 %, slightly (0.2 %) better than the accuracy obtained by the nu-net procedure, which does not allow estimating explicitly the limit, and almost 2 % better than the the accuracy of a Neural Tangent Kernel, which, as well known, is feature agnostic. Based on these tests, it is concluded that the pi-limit is the first procedure which allows computing exact expectation values in a MLP and that, at the same time, allows learning features.