Efficient Computation of Deep Nonlinear Infinite-Width Neural Networks that Learn Features
Authors: Greg Yang
Michael Santacroce
Edward J Hu
Venue: ICLR 2022 Poster (2022)
Keywords: infinite-width neural network
feature learning
maximal update parametrization
NTK

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 finite- and infinite-width neural networks previously left by NTK. Code for this work is available at github.com/santacml/pilim.

Published as a conference paper at ICLR 2022
EFFICIENT COMPUTATION OF DEEP NONLINEAR QO- WIDTH NEURAL NETWORKS THAT LEARN FEATURES
Greg Yang Microsoft
Michael Santacroce Microsoft
Edward J. Hu Microsoft
ABSTRACT
While popular limit of infinite- width neural networks, the Neural Tangent Ker- nel (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 j--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 ju-limit's com- putational difficulty in this setting. Here, we solve this problem by proposing a novel feature learning limit, the x-limit, that bypasses the computational issues. The n-limit, in short, is the limit of a form of projected gradient descent, and the r-limit of an MLP is roughly another MLP where gradients are appended to weights during training. We evaluate it on CIFARIÓ and Omniglot against NTK as well as finite networks, finding the -immm outperform finite-width models trained normally (without projection) in both settings, closing the performance gap between finite- and infinite-width neural networks previously left by NTK. Code for this work is available at g thub. com/ s santa /ppppplll im.
NTK
Finite-Wiss dth Network
w-Width r-Limit (Ours)
Figure PCA representations of images from 5 classes (= 5 colors) in Omniglot test set. We compare the representations from our best performing finite-width p-parametrized network, our proposed oo-width -immit, 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.
INTRODUCTION
The theory of Neural Tangent Kernel (NTK) (Jacot al., 2018) led to many theoretical discoveries of neural networks, but they ultimately apply only to limited, unrealistic regime where the network does not learn features. Recently, Yang and Hu (2020) discovered the maximal update, or p. parametrization that induces a feature learning oc-width limit, called the j-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-Eijnded 2018). They explicitly trained such wo-width limits Word2 and MAML, tasks relying crucially on feature learning. These ji-limits outperformed both NTK baselines and finite-widie width networks. Thus, the p-limit seems to capture key properties of neural networks in practice and is promising direction to further our theoretical understanding of them.
However, significant roadblock to this end is that the w-limii 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 MIPss where exact calculation is fast 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 the difficulty of theoretical analysis, especially in deep networks.
"oorresonndence to (gregyang, Mi chael. Santacroce @mi crosoft.com

---

Published as a conference paper at ICLR 2022
Contributions In this work, we bypass these issues by considering a form of projected gradient descent instead. Roughly speaking, with Va denoting the V-transform (aka dual) of á (Definition 3.1).
MLP with nonlincarity
width-20
MLP with nonlineurity Vo trained by gradient concotemation:
trained by projected gradient accumulation
Here gradient accumulation is just the normal training process where gradient is added to a parameter, whereas in gradient concatenation, gradient is appended a parameter. The RHS retains the feature learning capability while making the computation and the mathematics much simpler compared to the j-limmt. We explicitly evaluate new limit, which call the π T-limit, CIFARIO 10 and Omniglot, and compare against finite-width neural networks as well as NTK and NNGP. We find the т r-limit always outperform wide standard- Waanardd-paammttrizee or (4-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 CIFAR1O. To our knowledge, this is the first time deep nonlinear feature learning oc-width limits are evaluated on nontrivial natural datasets.
WHY IS THE LL-MMIT HARD TO COMPUTE? AN EXAMPLE
While the y-limit is the "optimal" feature learning limit of a neural network as described Yang and Hu (2020), it is difficult to calculate analytically. Let's demonstrate this difficulty with a simple 1 1-hidden-layer (1- u-parametrized network (abbreviated y-net, whose co-width limit the ę-limit) f with I-dimensional input & E IR and output ((E) E IR:
((E) = n 1/2 ((Vrun) E IR,
(I)
where u, V E R column vectors and Φ R - R is the nonlinearity (e.g, relu). Note here is the width of f. According to u-parametrization, we initialize uas Va ~ (0, 2))
At initialization, one easily sees lim (() = 0 because
n f({)
n Σ (VTv)a
is the average of a large number of iid random variables language of Tensor Programs (Yang, 2019b;a, 2020a;b), we have an almost sure convergence
with zero mean. In the
((E) - Te (EZZVTH) = 0, at initialization,
where Zv and VVT represent the independent standard Gaussian random variables describing the coordinate distributions of /nu and /nu.. However, one can see that f({) has limit of the form
S(E) + E(ZvTv + +... ).... + (...))((....))). after step of SGD,
(2)
where and still represent values at initialization, and represent terms coming from the gradient update.' Even though the ... appearing inside o() turn out to be all Gaussians, this expectation is hard to evaluate analytically due to the nesting of ø' inside 6, unless é is polynomial. Further- more, as more SGD steps are taken, this nesting quickly becomes more complicated, so even if 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 neural network under projected form of gradient descent which we call TSGG i.e., we update 0 t nIIV,L where II is a linear projection and 0 is the vector of all parameters, We shall first describe this projection II for the simple motivating example above before doing for the general case.
3.1 CONTINUING THE I-HIDDEN-LAYES EXAMPLE
In the example of Eq. (1), II acts on the gradients as follows: 1) leaves the output layer gradient v.C untouched, but 2) projects the input layer gradient V.C= = (... ( ) to the linear span of
up. the initial value of input weights u.
II,V.C, V.C) CV.C, Han vul) = (V cuo) for some E R.
See Appendix C for detailed calculation.
2 "The precise form is v.L- = L' d'vo'( (Vreufo).
'This projection may look excessively reductive, but this is just artifact of the 1-dimensional input and the shallow depth. See Scction 3.2 for the general case.
2

---

Published as a conference paper at ICLR 2022
Here Hls denotes the orthogonal projection to the span of un. In particular, after projection, VuL now has roughly iid Gaussian coordinates (proportional to un). Then unlike Eq. (2), we avoid the nesting of o' inside Φ:
(()) + B(ZVno + )o(ZZvm)) after step of TSGG
(3)
for some deterministic scalar ē. This expectation can then be evaluated routinely using the V- transform 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 the oc-width limit of f trained under TSGG which we call the 7-limit of f. See Theorem 3.2.
Definition 3.1. Let Vo : R + R denote the V-transform of Φ, such that (Y, ) E (X))))( for any centered jointly Gaussian random variables (X, R We will treat an activation function and automatically vectorize when V is applied to vectors of the same length.
V& like
3.1.1 COMPUTING THE т --LMMIT
Here we sketch how to compute the r-limit of Eq. (1). Let ue, Ut, ft denote corresponding objects at time I, with I = denoting initialization.
Memory and Time Requirements What information do we need store in order to calculate the u-limit? Because of the projection, u, Cg for some C2 E R. So to calculate the limit of ft, we need to track et (or rather its oo-width limit). This memory requirement is (1) in training time t. To see the form of vt, it helps to simplify our setup bit more by 1) initializing uo 0 and by 2) assuming each T-SG step involves minibatch containing a sole input & Correspondingly, because the gradient of is always proportional to ((t,
p(5:01/na there exist coefficients
{a, E IR. (B],==) (b. E ( (which are random, but deterministic conditioned on no; vo) such that
Vnu = Σ (.b(bbsnu))) a=0
(4)
Here by = c, and 02 is formed from the learning rate and loss derivative. So to calculate the limit of ft. we need {a, E [ {b. E [R] (or rather their n + oc limits). Note that Ct Ogz b, all have nontrivial fluctuations for finite n, but become deterministic as + 00.
This implies a memory requirement of (t)). which is also the total requirement for computing the limit of fe. As we will see, each forward and backward pass has runtime O well, the total runtime for calculating lim fe is ((2), compared to the exponential runtime of general p-limit.
Forward Pass
Using Eq. (4), we can intuit (and formalize using Tensor Programs) that
fe(E) = Σ
n
Σ
a.olb, a.olb. mmmaaa o(EcrV/noo) +
Σ à E Eo vin nwo ==0
1-1
) (EE Zvin)) (5)
where åz be, & denote the deterministic limits of the corresponding quantities. The expectation in the RHS can be evaluated using V-transforms. The terms in the summation implics a runtime of (t).
Backward Pass The gradient update for the output weights v, is clearly represented by setting (0(+1, b(+1) + (-nL', ; (cz), where L' denotes the loss derivative (f. lab However, a priori, seems unclear how to calculate the limit of the projected gradient IL, Vu,L of the input weights. Fortunately, because we can express the entire unrolled R-SG training inside a Tensor Program, the Master Theorem (Yang, 2020b) can automatically give the exact formulas for the gradient limit; see Appendix F.
But in fact, there is more intuitive way of obtaining the backward pass limit by recognizing that projected gradient is just the maximal ascent direction in the projected space. In our case, Il Vu, C
"This is because, while has fluctuations for finite n, it is roughly constant for large "Assummin minibatch size is constant, and we train for a constant numher of epochs, this translates to ( to train lim ft, where N is dataset size, and e ((N) to do inference on a single input. This compares favorably with Gaussian Processes, which requires ((V") to "train" (i.e. inverting the kernel,, and O(N) for inference. However, in our experiments here, the constant in our ((N²) in practice will make training lim ft slower than the corresponding NTK or NNGP kernel regression. "See Lemma F.1 for more details.
3

---

Published as a conference paper at ICLR 2022
(())
//vl
47
-98
( @ vfuua)
1
VyL
lim
1/-
Appeno
C3
R"
Vit
M
Cuo
V.C
77
Accurnllate
Accumulate
Figure 2: Summary of the r-limit for the 1-hidden-layer network of Eq. (1), (Left) Representation of width-n p-nett with nonlinearity & and its 5-SGD update. Here is depicted as / o(B &
vr which is equivalent to Eq. (4). In 7-SGD, the first layer gradient is projected to the space spanned by wo before being accumulated the weights. (Right) This network's oo-width limit can be roughly thought of as another I-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 uo. i.e. Iuo V., (V L) where V L the derivative w.r.t. the coefficient Gr of u, in terms of ug. In the oc-width limit, this derivative can be obtained by just auto-differentiaaing Eq. (5) when the V-transform has an explicit form, like for relu (Cho and Saul, 2009). So --GG step u, is equivalent to taking C++1
c v.a
Summary Below, for brevity, we say training routine to mean the package of learning rate n. training sequence of singleton minibatches {(E,, ::20 (where & is the input and ys is the label), and loss function ((f(E)) y) that continuously differentiable in the prediction of the model f())
Theorem 3.2. Consider the simple motivating example of --hhdden-layee network f Eq. (1) with dimensional input and outputs, initialized by ua ~ (0, 1/n) and va + O. Suppose its nonlinearity has a polynomially bounded 2nd derivative. For any training routine, trained by T-GG for T steps has an oo-width limit fr, in the sense that,
as xo. Fr(E) + (E),,, for any & E R
where the limit is almost sure. fr is given as follows:
Initialize An, By as empty column vectors, and initialize Co E R as Recall V& is the V-transform of b. For each t = 0, l,..,, define the function (where dot product of empty vectors is 0)
iA(E) th ACCEE,, Br Br, B, 1) for any & E IR,
(6)
where o denotes coordinatewise product and 1 denotes the all-Is vector of appropriate shape, and As, Br, Ot are inductively given by C++1 di Ct ((ee.(f((E)). and
A dd "append((A -n JAL.)) C()(())) u:))
B del
append( B., CrE.)
where append((v, (v. p) appends element p to the end of vector increasing the dimension of v by J.
Here A₁, By correspond to the column vectors formed from the oo-width limits of {a, E 1 (b, e B)++ and Ct corresponds to the limit of c. As discussed above, computing fr requires ((T) memory and @(T²) time. Theorem 3.2 is summarized by Fig. 2.
3.2 T PAPAAEETRIZATTOO FOR DEEP NETWORKS
We can straightforwardly generalize --GG and the 7-limit theorem (Theorem 3.2) to deep MLPs. However, due the x Gaussian matrix initialization in the middle of the network, the memory requirement will be (T?) and runtime will be ((T3) ) for training T steps, for the same reason as discussed in Yang and Hu (2020, Sec 8). This is not scalable to datasets like CIFARIO. Therefore, we propose a different initialization that brings down the memory requirement to ( and runtime OTT2), just like the --hidden- layer case. Consider an L-hidden-layer μ-net Rd + Ido:; For
7 For simplicity, we only consider batch size 1: it's straightforward to generalize larger batch sizes. Technically, should have written A,, B., Ć₁ to maintain the convention that Ö denotes limit of D. but for the sake of brevity, we drop this convention in Theorem 3.2.
4

---

Published as a conference paper at ICLR 2022
weight matrices w E Raxd and w², ..... ......" E Rax", and nonlinearity 0 R IR, such neural network on input { E Rd given by (E) = Vnwlt IR", and
=E( = (E)( R"; (E) = w E ER, for 1 -...=,.,,,,-
(7)
and the network output is f(g) m/ ( for wh+1 E Rdout x In p--arametriiation, we would initialize w ~ (0, 1/n) for any a, B,I. Next, we describe our alternative proposal.
T- m--nitializaaton 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 der (A', RMx 2 {A' E Rxxr U {BL+1 E Rxxx , E MM }.
(8)
Then we can initialize weights w' by first sampling a standard random Gaussian matrix Ω2 e Raxr Sai (00 1), before setting
w' - QA'T ((B'2T) E R=x" for all hidden = =2...., n
whtl 1 AA ((B++1
w QAIT ERdd, Vn
Rdoux
Vni
As example, the initialization of the 1-hidden-layer case in Theorem 3.2 corresponds d dout =-= 1, M = 0 and = 1. See Fig. 3(Lef) for an illustration.
Let's digest this initialization scheme: 1) Of course, A', B' need to be initialized themselves, and in this paper we will just do so with Gaussian random initialization; see Appendix B.1 for the pscudo-algorithm. 2) A', B², example above, whereas A 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 R=x" to m=x" below. 4) A B' will grow 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 saying w are generated by l-hidden-layer MLPs FA B R" + (R" or Rdont) ) with weights A', B', like so: wh. 1/a(Oa, FA (Qg)), , for all hidden
L). will play the same roles as A, B the l-hidden-layer
2..... L;; w
11 (Q0))) (10)
This is reminiscent hypernetworks (Ha et al., 2016), but the crucial difference is that the "batch dimension" (indexed by a, 3) of the generator FA becomes the width dimension of the generated network f, that the same generators can generate networks of arbitrary width.
T- m-Projection We leave the output weights u alone, but for any w E {w", w² ..... w²), where w has shape by H2g the projection matrix to the r-dimensional space spanned by the columns of f. This means that preactivations h always in this space. However, the input side of w is projected, so we are optimizing the hidden weights in a (r x n -i)dimensional space, which still becomes infinite- dimensional
or n x d, we project the gradient V.L of the same shape by left multiplying
- We refer to SGD with r-rojjected gradients as -SG..
This brings us to the following
Definition 3.3 (л- Parametrizzaiio and r-Limit). We define n-parametrization the package of r-initialization and TSGG,, where r and the initial M are clear from context. network in r-parametrization is abbreviated n-network or n-net, and we define x-limit as its wo-width limit.
3.2.1 COMPUTING THE T-LIMIT OF DEEP MIPP
Here show the --limit forward and backward passes can be efficiently computed.
T-Limit Forward Pass Let's write f for the (random) neural network r-initialized by PP as in Eqs. (8) and (9). Then it's straightforward to see f has a deterministic limit fe expressible composition of matrix multiplications and V-transforms:
Theorem 3,4 (imti Forward Pass). Suppose é has a polynomially bounded 2nd derivative. Then PP( (P(E), for any & Rd.
as n. cc, where convergence is almost sure, and fP((E) is defined as follows. Write g² de A E RR.
R
for = 2, .....I for = L+ 1,
g d A Vo
B¹ 11
and fP((E) g 5+1 g where B yields size-M column vector of squared norms of B's rows. with learning rate independent of width (but may vary with training time)
5

---

Published as a conference paper at ICLR 2022
Rloor
A /vm
n
771
((="u))
Φ
R*
9¹
8 R2"
DLA
Append
-
Pl
<((BnnT)
lim Teeio
14
Φ
8a
9²
g²
QA77
APT
78
- W
((BPnT)
g²
Accumulate
Φ
vi
g¹ D2"
w1
QATT/VR /vfi
Accumulate
Figure 3: Summary of the 7-limit for deep MLPs Eq. (7), in the same style as Fig. 2. take the example of 3 3-hidden-layee MLPs. See Theorems 3.4 and 3.5 and compare to Eq. (x).
Here we
See Fig. 3(ight for an illustration. Here, g represents the coefficients of preactivation hi (c.f. Eq. (7)) in the columns of Ω. As we will see next, x-projection ensures that at any point during training, fe has an wc-width limit of the form fP for some P.
T-Limit Backward Pass Just like the 1 1-hidden-layer case, we can interpret the projected gradient as the maximal ascent direction in the column space Ω2. Formalizing the intuition using Tensor Programs yields the following theorem that we also empirically verify in Appendix E.
Theorem 3.5 (i-Limit Backward Pass). Let Po denote the matrices (c.f. Eq. (8) used the w- initialization of f as in Eq. (9). Suppose its nonlinearity & has polynomially bounded 2nd derivative. Then for any training routine, trained by T-SG for T steps has an co-width limit which is equal to fPr for some PT given below, i.e.
as + DO, FT( - (P)(E), for any E IRd,
where the limit is almost sure. Pr is given inductively through its elements Ale. By, as follows: Ap41 def Al v A C( C(p (E:), vt), for, I L.
Alt1 def append - ((P))))))))))
B'u de append
Here df corresponds to g' in Theorem 3.4 evaluated for input & and function fo, and append B, .)) means appending 9 as a new row vector of B, increasing by the dimension represented by M.
Theorems 3.4 and 3.5 are summarized by Fig. 3. 10 As conveyed by Eq. (+), the -llimit can be thought of as another MIP with activation Vo and trained by gradient concatenation." Because p-parametrization has similar scaling with width as p-parametrization, former admits feature learning in the same way the latter does, i.c. the feature kernel of every layer evolves nontrivially during training. Using Tensor Programs, 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 the --limit theorems.
easy show that the
straightforward
extend the
T-Limit p-Limit While the projection means optimization is slower in the 7-limit, believe with sufficiently large r, the --limit should perform similarly to the p-limit. Indeed, prior works such as Li et al. (2018) has shown that optimizing neural network in much smaller, random parameter subspace can recover most the performance of training all parameters. This is also evidenced by our experimental results below, where the 7-limit generally only slightly outperform wide ne-networks.
O'e may wonder how would Theorem compare with just directly accumulating gradients on P. fact. this would train very badly because V4, very smooth around 0, j. would look very linear for long time. Empirically, direct SGD P would yield 53% test and 80% training accuracy on CIFAR1O.
Note that once a vector is appended to A' or B', it is not touched again. Therefore, despite the similarity hetween the -llimit and MI.P, A' and B' should not be thought of "parameters" in the usual sense.

---

Published as a conference paper at ICLR 2022
Table 1: Best Test Accuracies on CIFARIO and Omniglot, best of MLPs up hidden layers. width 2048, r 400, as well as random search over a host of hyperparameters: see Appendix B. Note the je-eet numbers are also the optimal numbers for standard parametrization finite networks, as discussed in Footnote 12. x- x-Limit ImageNet Transfer means pretraining a 7-llimit with r = 200 on ImageNet32 and perform kernel regression with its feature kernel (i.e. kernel of last layer activations) on CIFARIO, I0; see Section 4.2 and Appendix B.4 for details and finite network results, Also compare
with feature kernel regression without pretraining (Table 8).
NNGP
NTK
NIK perl ROP
Jt-Net
T-Net
T- T-Limit
IT Limit ImageNet Transfer
CIFARIO
58.92
59.63
61.31
60.64
61.50
64.39
Omniglot
43.80
51.72
+
91.22
92.21
91.46
4 EXPERIMENTS
Here we compare the performance of the relu x-limit on CIFARIO (Krizhevsky, 2009) and Omniglot (Lake et al., 2015) against that of NTK, NNGP, and finite-width л and j--ets. 12 As we will see, 1) r-limmiss of large enough r beat finite j--nets; 2) finite --nets underperform the above on CIFARIO but, interestingly, outperform them on Omniglot; 3) all of the above beat NTK and NNGP.
4.1 CLASSIFICATION ON CIFARIO
Experimental Setup For π- and p-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, 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 j-Ner with NTK, NNGP). 13 The -imim closes this gap, having almost 2- 2-point advantage over NTK. The width T--et outperforms NTK and underperforms 7-limit both by about point.
Feature Learning in --mmm and Finite Networks We show the advantage of -imit 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 j-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 r-limit benefits from feature kernel regression, improving the test accuracy from 61 1.5% to 61.85%; see Table
Effects of Width, r, and Depth As shown in Fig. 4(Middle), accuracy of n-net increases monoton- ically with width and with r, approaching the accuracy of -imit and p-net, respectively, from below. This can also be seen in feature kernel regression across training time, Fig. 6. In contrast, performance is not monotonic in depth for any of j--net xenet, or -imimt, beyond a consistent improvement from 1 to 2 hidden layers; see Fig. 7.
"eeaass standard parametrization differs from j-parametrization only in factors depending on width, 2) we sweep such factors in hyperparameter optimization, and 3) our finite networks have width most 2048 (so these factors are in practice constants), our best accuracies for μ-nets will also best accuracies for standard- parametrized MLPs. We also train --nets with untied Ωs; see Appendix A.7.
NNoe that, contrary to Lec 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 we did kernel gradient descent for infinite-time (which is equivalent to kernel regression) with learning rate drop, since this optimization is convex.
"In both NNGP and NTK limits, the feature kernel stays fixed throughout training (Yang and Hu, 2020), so the dotted line in Fig. 4(Teft) shows the result if do the same experiments for NNGP and NTK.
7

---

Published as a conference paper at ICLR 2022
Regression
Deniigitt
Test Aecurery
50 100 300 400 11TT
3200
n 11im nost
200
50
100
20 Train Epachs
40
120 256 532 1024 2048
220
210
912
2040
4090
8197
Width
Figure 4: (Left) Feature kernels of x r-limit, n-net, and fu-net all improve in quality with training, as measured by kernel regression on CIFARIO. All models are the best from our hyperparameter sweeps, but note that feature kernel regression causes accuracy decrease in finite models. (Middle) CIFARIO I0 validation accuracy is monotonic across width and r for æ-net, æ r-limit, and j--et. (Right) Omniglot validation accuracy, contrast, is not monotonic in width, and μ--net can underperform n-net of large r. (Note we did not run x-limit with r 2 800 in consideration of computational costs). All numbers in the heatmaps the best from our random hyperparameter searches for 2-hidden-layer networks.
4.2 TRANSFER LEARNING FROM IMAGENET TO CIFARIO
Pretraining-an cea----nsfer--anmiin is an important setting where feature learning is crucial. As pointed out in Yang and Hu (2020), the NTK limit trivializes pretraining, while the j--limit both theoretical and empirically benefit from it. Here we investigate pretraining for the --lmmit in the image domain by pretraining on ImageNet32 (Chrabaszez et al., 2017)15 and transferring to CIFARIO.
Experimental Setup We pretrain the --limit with r = 200 (as well as ju-Net, 7-Net with r = 200, and r-Net with 400) for 30 epochs on a fixed subset of ImageNet32 with 250 (out of 1000) randomly subsampled classes. To evaluate on CIFARIO, we compute the kernel induced by the pretrained final- final-layer-feattee layerrfeatures and perform kernel regression.
Results As shown in Table 1, ImageNet32 pretraining nontrivially raises the downstream CIFARIO
performance over without pretraining, altogether creating 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 r-limit. Table 7 shows the benefit of pretraining seems to be directly related to the capacity of the model.
4.3 FEW-SHOT LEARNING VIA METALEARNING ON OMNIGLOT
Following Yang and Hu (2020), we also evaluate few-shot learning on Omniglot. Compared to traditional classification settings like CIFARIO 10, 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 models using ANIL (Almost No Inner Loop) (Raghu et al., 2019) rather than first-order MAML (Model Agnostic Meta--Lernning. Briefly, ANIL is a variant of second-order MAMI,, 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 oo-width networks compared to second-order MAML. 2) While a n-parametrized network trained by second-order MAMI. has an co-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, -entwwrrks 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 new data, which is more fitting for our goals in this paper.
We focus on the 5-way, 1 1-shot task 16 with only step of ANIL adaption. For π- and μ-networks, infinite or finite, 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.
Experimental Setup
15: i.c ImageNet (Deng et al., 2009; Russakovsky et al., 2015) downsampled to 32 x 32 resolution
Mi.c. each task consists of 5 different classes, with example provided for each class in training.
8

---

Published as a conference paper at ICLR 2022
For NTK and NNGP, ANIL meta-training has no effect, and meta-testing amounts just taking 1 kernel gradient descent step on each task." 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 I, ju-net, nenet, and x-limit all outperform NNGP and NTK by about 40 points. Interestingly, while x-limit is slightly better than μ-net, they are both outperformed by π-net. describe next.
This is related to the width nonmonotonicity seen Omniglot,
Effect of Width, T, and Depth While n-net performance with is still roughly monotonic, in contrast to in CIFARIO, here it can decrease with width, as seen in Fig. 4(ight,) In addition, j--net seems to underperform æ-net of the same width for sufficiently large We find this is primarily due to optimization and not generalization because Omniglot is hard to overfit; see Fig. 8. Counterintuitively; g-net of large r can optimize faster than μ-net despite the projection. This is likely a side effect of T-initialization, as we see this advantage persist when r >> width, where projection do not see width nonmonotonicity on CIFAR1O because overfitting is much easier there, so the results depend much more the implicit bias of each model.
no-op. We
On the other hand, performance seems to monotonically increase along diagonals of fixed r/ r/width ratio, peaking around r/ /width = 1/2. This suggests different limit of x width + oc, which warrants investigation in future works.
Unlike CIFARIO I0, test accuracy is more monotonic in depth here (Fig. 7). Again, this is probably because the extra weights help with optimization, which is the bottleneck on Omniglot.
RELATED WORKS
The r-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 al., 2021). The crucial difference here is that 1) during inference, the --limit is deterministic and not a hierarchically sampled process, and 2) the r-limit is not trained variationally.
Relatedly, deep GPs or kernel processes can be construed as a kind randomly initialized infinite- width 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 kind of intertwined system of kernel gradient descent. However, the computation of this limit is not scalable to datasets like CIFAR10)
During training, the r-limit's A and B' 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 7-limit can contribute back to this literature.
Our m-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 p-limit.
6 CONCLUSION
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 thcoretical investigations.
Previously, NTK and NNGP satisfy 2), 3). 4) but not 1), as shown in this and prior works. The j-limit satisfies 1) and but, arguably, not 3) and 4). In this work, we presented the -llimit which, we believe, satisfies all 4 properties and that can prove fruitful for future researches.
While our work closed the oo-width performance gap, also opens many new questions, e.g., Can we derive the π -lmim for fixed r /width ratio? When precisely does the width nonmonotonicity of n-net occur? What about modern architectures, beyond MLP? We leave these questions future work.
NTK and NNGP can be trained under second-order MAML, but we found their performance strictly decreases with metatraining, the randomization of labels across tasks confuse the readout layer.

---

Published as a conference paper at ICLR 2022
REFERENCES
Devanshu Agrawal, Theodore Papamarkou, and Jacob Hinkle, Wide Neural Networks with Bot- tlenecks Deep Gaussian Processes. arXiv:2001.0092/ [cs, stat], January 2020. URL http:// http//axx.co aawiv. org/ abs/ .00921..
Laurence Aitchison. Why bigger is not always better: on finite and infinite neural networks. In Pro- ceedings of the 37th International Conference on Machine Learning, pages 156-164. PMLR. November 2020. URL. https:/ /ttpss/////cceeeeeeeeooo werroceedingg..... .m.r........ www..m........l1.m popress/vi1//aat
html.
Laurence Aitchison, Adam Yang, and Sebastian Ober. Deep kernel processes. In Proceedings of the 38th International Conference on Machine Learming, pages 130-140. PMLR, July 2021. URL https: https://ppoceedinng wwt.smmrr...... press /vcsss/v/33 v9////1119//// htt.////tttccinnna1... html.
Antreas Antoniou, Harrison Edwards, and Amos Storkey. How to train your MAML. arXiv:1810.09502 [cs, stat], March 2019. URL
09502..
Dyego Araújo, Roberto I. Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks, crXiv:1906.000 [cond-mat, stat], June 2019. URL 1906.000 00193.
org. wwwg//aaooo /bs/
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net, 2019.
Lenaic Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. arXiv:1805.09545 [cs, math, stat], May 2018. URL wwththtp://. //t///axiivvv// org/ waaaaaa19. 09545.
Youngmin Cho and Lawrence K. Saul. Kernel methods for deep learning. In Advances in neural information processing systems, pages 342-350, 2009. URL http:/ /ttp://ppprrrseee nips. cc/ paper/3628-ket -er222-kkeeelll www.eeeemmmmehedee for- deep- earning.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. downsampled variant of imagenet as an alternative to the cifar datasets, 2017.
Andreas C. Damianou and Neil D. Lawrence. Deep gaussian processes, 2013.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: large-scals scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009. doi: 10. 200.E.....009526666
Cong Fang, Jason D. Lee, Pengkun Yang, and Tong Zhang. Modeling from Features: a Mean-field Framework for Dverrparametterized Deep Neural Networks. arXiv:2007.01452 [cs, math, stat), July 2020. URL http: /ttt:://arxivv org/ abs 2007 2007.0145.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks. arXiv:1703.03440 [cs], July 2017. URL 03400.
org/abs //////bbb//3./. /1703.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient Training of BERT by Progressively Stacking. In Proceedings the 36th International Conference Machine Learning, pages 2337-2346. PMLR, May 2019. URL https: /tppp//pceeengg..eeeee mlr.
ppess/v97 vtpps///v/ggggllll wwwggongggggooo... html.
Xiaotao Gu, Liyuan Liu, Hongkun Yu, Jing Li, Chen Chen, and Jiawei Han. On the Trans- former Growth for Progressive BERT Training. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5174-5180, Online, June 2021. Association for Computational Linguis- tics. doi: 18653/ 1/2021 1aaccl..aaccmmaai maacl-main.406 naacl macc-mmain 406.
URL htt.aass//aeaaatoooooooooo//////
David Ha, Andrew Dai, and Quoe V. Le. HyperNetworks. arXiv:1609.09100 [cs], December 2016. URL hwtp: /awiiv waw..ivvoo. wwwp:vvvoorr/oooooooooo.o abs/ 1699.660..00 09106.
10

---

Published as a conference paper at ICLR 2022
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kemel: Convergence and Generalization in Neural Networks, arXiv:1806007577 [cs, math, stat], June 2018. URL http: //arxiv.. org/org////..ooo abs 1000. 07572.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Brenden M. Lake, Ruslan Salakhutdinov, and Joshua B. Tenenbaum, Human-level concept learning through probabilistic program induction. Science, 350062666 2015. ISSN 0036-8075. 1oi:10.1126//.. 1tt.1///ieence.aabb000 35C/ 6266/3226/332.
URL https: /ttpas:///ciceeeee/e/ sciencemag org/content/
Jaehoon Lee, Yasaman Bahri, Roman Novak, Sam Schoenholz, Jeffrey Pennington, and Jascha Sohl- dickstein. Deep Neural Networks as Gaussian Processes. In International Conference on Learning Representations, 2018. URL https: httpt:://peeeeeeeeeeeeeee /ttp///pppppeeeeeeee/ee/e.eoeeep net ftret/oorumiiii 2otp:/rrwwwwieeeeooooe...mmaeeee
Jachoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite Versus Infinite Neural Networks: Empirical Study. arXiv:2007.1580) [cs, stat], September 2020. URL 15801.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the Intrinsic Dimension of Objective Landscapes. arXiv:1804.088 [cs, stat), April 2018. URL htttt:://axxiv.rr/ abs/ 1804. 08838.
Etai Littwin, Omid Saremi, Shuangfei Zhai, Vimal Thilak, Hanlin Goh, Joshua M. Susskind, and Greg Yang. Implicit acceleration and feature learning in infinitely wide neural networks with bottlenecks, 2021.
Qiang Liu, Lemeng Wu, and Dilin Wang. Splitting steepest descent for growing neural architectures, 2019.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):77665- E7671, August 2018. ISSN 0027-8424, 1091-6490. 10.:10.11 10t.100.0//pass11800000 1806579115, URL https: www.pnase org/ wonnenttnnntt. 33/665555
Phan-Minh Nguyen. Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks. arXiv:1902.02880 [cond-mat, stat], February 2019, URL http: /ttxxixxxvvv////a///aaa org. arg/aba /abs//b//// /1902. 02880.
Phan-Minh Nguyen and Huy Tuan Pham. A Rigorous Framework for the Mean Field Limit of Multilayer Neural Networks. arXiv:2001.11443 [cond-mat, stat], January 2020. URL http:
// //arxiv.org/aaaatt....aaaa http:/aab/200111100
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid Learning Feature Reuse? Towards Understanding the Effectiveness of MAML. arXiv:1909.09157 [cs, stat], September 2019. URL http://///. www.i.vvvooooo..... parg/abs// 19999 C9157.
Grant M. Rotskoff and Eric Vanden-Fijnden. Neural Networks as Interacting Particle Systems: Asymptotic Convexity of the Loss Landscape and Universal Scaling of the Approximation Er- ror. 1rXiv:1805.00913 [cond-mat, stat], May 2018. URL http: /ttp///axvv.vvv/ http://axxiv..g//aaa/a. hatp/bb//110 org//bs/1800 00915.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander Berg, and Li Fei-Fei. Imagenet large scale visual recognition challenge, 2015.
Hugh Salimbeni and Mare Deisenroth. Doubly Stochastic Variational Inference for Deep Gaussian Processes. arXiv: /705.08933 [stat], November 2017. htt.ltttpp///xxi..orggg/ 08933.
Justin Sirignano and Konstantinos Spiliopoulos. Mean Field Analysis of Neural Networks. arXiv: 1oXiv:805.001e 010500 [math], May 2018. URL http:://arxiw..rrrbbb////// 01053.
11

---

Published as a conference paper at ICLR 2022
Justin Sirignano and Konstantinos Spiliopoulos. Mean Field Analysis of Deep Neural Networks. orXiv:/903 04440 [math, stat], February 2020. URL 04440.
/903.
Lemeng Wu, Bo Liu, Peter Stone, and Qiang Liu. Firefly neural architecture descent: a general approach for growing neural networks, 2021.
Greg Yang. Tensor Programs Wide Feedforward or Recurrent Neural Networks of Any Architecture are Gaussian Processes. arXiv:1910.1244 12478 [cond-mat, physies:math-phl, December 2019a. URL http:://rxixs //ppvv// org/ abs 19910. 12478.
Greg Yang. Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation. (00.iv:/902204400 [cond- mat, physics.mathh-ple star], February 2019b. htLpttp: httptppp//iiiii./t////aa/
abs / 1902. 04760.
Greg Yang. Tensor Programs II: Neural Tangent Kernel for Any Architecture. wxXi:2200.115588 [cond-mat, stat], August 2020a. URL http: http:/wwwww.r..a///////ba/a/a 145488,.
Greg Yang. Tensor Programs III: Neural Matrix Laws. arXivv2009.10688 [cs, math], May 2020b. 10685.
URL http:://aaxiiiii org/
Greg Yang and Edward J. Hu. Feature Learning in Infinite-Width Neural Networks. arXiv:201 2. 14522 [cond-mat), November 2020. URL
Greg Yang and Etai Littwin. Tensor Programs Ilb: Architectural Universality of Neural Tangent Kernel Training Dynamics. arXiv:2105.03003 [cs, math), May 2021. URL http: org/ ab///b///// 2205.115... 03703.
Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor. Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. arXiv:2203.02 03466 [cond-mat), March 2022. URL http: //rrivv..iivvv//aa wrg./bs///2......//// 03466.3666.....
12

---

Published as a conference paper at ICLR 2022
A EXTENSIONS OF --LMMI
A.1 BIASES
Biases are straightforward to add. Consider an L-hidden-layer p-net f IRd + Rdom with biases: such a neural network on input E Rd is given by (E) Vnwot & + v / l RR, and
x²( - (E) IR, l+) w I (E) + vnbl E ER" for Ly..., (11)
++ (E) + be+1 for wht1 E Rdouxn, In p parametrization, the biases are initialized as bba / for = 1,...,4, and b (0,, of), for hyperparameter dg However, in ग r--aaammtrizationn we will do this bit differently.
and the network output is f(E) n n11
n-Initialization Suppose we are given a collection P of matrices and vectors
p de (RHS of Eq. UB E E' U ( E ER
(12)
Then we can initialize weights and biases by first sampling a standard random Gaussian matrix Ω Rexr Oai NO, 1), and set w in Eq. (9), and 2) set biases
6 t gh+1 and b'+ -
28 for all / = L.
(13)
This constitutes the r-initialization of biases.
n-Projection For и projeecction, we also project the bias gradients by Igg
T-Limit Calculation
Theorem A.1 (æ-Limit Forward Pass with Bias). Let be some collection as in Eq. (12), Asn + co.
(P(E) + fP()) (E), for any & E IRd,
where convergence is almost sure, and (P) is defined as follows. Write n² 1f ALT B R
IR" Rdon
for 2..... T.
g de A (B B' B¹, 11g 11 + B¹
for
and fP((E) 1 L+1, where B yields size-M column vector of squared norms of B's rows.
Theorem A.2 (m-Li Backward Pass). For the same setting as in Theorem 3.5 but with P as in Eq. (12), A B are updated exactly as in Theorem 3.5, and 3 is updated by gradient accumulation like A):
844 def 8 - 1 (E), for all l 1,,,,,,
A.2 PARAMETER MULTIPLIERS
We can insert to Eg. (1I) constant parameter multipliers aw for each parameter w like so
h ( - 1/ (E) + Vnagib e R"" for 1 = 1...,, L 1-1,
(14)
These multipliers are tuneable hyperparameters. They affect both the forward and backward passes of the network.
In the -lmmmt, the forward pass is the same as in Theorem A.1 except we replace A' with a A' and 3! with 3'. 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, 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 single multiplier for all biases.
13

---

Published as a conference paper at ICLR 2022
A.3 LEARNING RATE MULTIPLIERS
We may have custom learning rates for specific weights or biases. In our experiments, we implement this with learning rate multipliers relative the global learning rate 7) e.g.. the learning rate of parameter w becomes " if the multiplier is denoted Yw-- Then in the limit, we just need to replace y in Theorem 3.5 or Theorem A.2 with Tw i.e.
A 1 del A
A (6)).))
1 de append( (A), Mon of (E),
for all = 2,.... ., for all | 1,.., I.+ 1....,,+++11
B 2 3 wVnV L( (E:). (E:).
But note we do not modify the B' update. In our experiments, we will only sometimes use LR multiplier on the input layer ("Input Layer LR Mult"), one on the output layer ("Output Layer LR Mult"), and/or single multiplier for all biases ("Bias LR Mult").
A.4 LARGE BATCH SIZE
For batch size 1, we still accumulate gradients into A and B 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.
A.5 GRADIENT CLIPPING
The Frobenius norms of the projected gradients of the input weight w and all biases t/ converge to exactly the Frobenius norms of the gradients of I and 3¹, i.e.
[L V. CI] V All V + Ll
where the convergence is almost sure. For any hidden weights w'. suppose its gradient over an S-sized batch in the --lmmi is given by
1 del append(AA,, -nA),
B append B,, B).
where A, Sxr, and append( B) means appending all S rows of B into increasing the latter's column length by S. Then
[L fII Vir VI (ATŌ(BBT)A)
where the convergence is almost sure. Then clipping the gradient in the 7-limit just means clipping the gradients of A and 3 and, for hidden weights, rescale À such that v (AT3(BBT)) is at most the clipping threshold.
A.6 WEIGHT DECAY
The limit decaying w +- w0'(1 Ir wd) just A' +- A'(1 br wd).
A.7
DECOUPLING LAYERS
We can actually use different rs and Ms for every layer, and similar limit theorem can be proved. We can also use different, independently sampled for each layer ("untied Qs*). However, the limit would be exactly the same as before, as is apparent in our proof. In addition, we verify that even in finite -entss there is difference in performance between tying across layers or not (Fig. 5). All of our experiments are actually done with untied Ωs.
B EXPERIMENTAL DETAILS
All of our experiments are done V100 GPUs. All of our networks use relu activation. The V-transform of relu is (Cho and Saul, 2009)
1 (V1 c2 + (πr arccos(e))c) rirz
Deeturrizz 12 (. 2m
for any rı, 72.2 and [-1,1].
14

---

Published as a conference paper at ICLR 2022
13
Teed Oreegas Unted Origgas
Thed Orregts
Unraed Omeegas
2.93
1.
Teat Loge
225
1.0
2.00
1.75
1.50
02
10
20
30
40
20
30
40
50
Train Epoche
Town Epechs
Figure Tying Untying Ω Across Layers Make No Difference in x-Nets. Using the same procedure as in Section 4.1, we train n-net with 400 using the best hyperparameter combination found (whose result shown in Table 50 times, each with different independently sampled Ωss either with Ω tied across layers (blue curve) or not (orange curve). We plot their mean training loss and test loss curves here, with shade indicating 95% confidence interval.
B. NETWORK INITIALIZATION AND PARAMETRIZATION
At initialization, 1) biases are always O, 2) M is set to r so that A', B' are square matrices for 1 ....., L, and 3) we sample A', B' standard Gaussians and then scale them follows:
A A /A' norm(dim = 0, keepdim True)
B' B' B normm 1, keeplim T True) for all 1 2,...,!!
A' = A / A [0] for all 1 = 2,..., L, and
A = 0
B.2 FEATURE KERNEL.
When we talk about the "feature kernel" of m-limit jP, we always mean the n -> 00 limit of the feature kernel of fP, and not the kernel induced by g² in fp. This feature kernel K on inputs tE....., } is calculated as
K₁ Vollot 11offl 12 11
where of and of are the gh in Theorem 3.4 evaluated for two inputs & and E
B.3 CIFARI0
B.3.1 NE-NET, TEN, AND L-MMMIT
For and je-etworks, infinite or finite, we train for 50 epochs. We adopt step learning rate schedule, with a learning rate drop of 0.15 at milestone hyperparameter. We also clip gradients with hyperparameter threshold, where the clip is triggered individually for each parameter by the parameter's 's norm, rather than total parameter norm.
Hyperparameter Optimization We first perform random search on hyperparameters listed in Table 2 for 2-hidden-layer nenet, --limit, and μ-net: We sample at least 512 hyperparameter com- binations 1) for each (width, r) E {128, 256, 512, 1024, 2048} (50, 100, 200, 400) for n-eet, 2) for each width E {128, 256, 512, 1024, 2048} for μ-net, and 3) for each r E {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 r-net, --limit, and μ-net, and for depth E (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 (Ir" (2-2,2 2-13. 21.51) x (wd" {2 2,2 ..... .1.51) and Ir* and wd" are the optimal hyperparameters from the 2-hidden-layer sweep. The best test accuracies per depth are shown in Fig. 7, while the overall best test accuracies over all depth are shown in Table 1.
15

---

Published as a conference paper at ICLR 2022
Table 2: CIFARIO Hyperparameter Grid for j--Net, -eNett and r-L-imit
Hyperparameter
Grid
Gradient Clip
0.4 [2-3 2-2 24
Learning Rate
0.5 (5-3,222222 24
Weight Decay
2 2.10 (2-3,
2*}
Bias Mult
0.5 (2- 2-2 ..... 2+}
LR Drop Milestone
(30,35, 40)
Input Weight LR Mult
0.1 (2-3,2-,2-- 24)
Output Weight LR Mult
16 (2 (2-8,22
24)
Input Weight Mult
{2 2- ....
Output Weight Mult
(2-3.2-2 3-2--2 24
Batch Size
(4, 16,
Table 3: CIFARIO Hyperparameter Grid for NNGP and NTK
Hyperparameter
GP
NTK
Bias Variance
(2
2-3.5
22
0.5 (2-1 2 2
22
Bias LR Multiplier
n.a.
{2 3.5
Input Weight LR Multiplier
n.a.
0.5 (2 2
2
Output Weight LR Multiplier Ridge
n.a.
{1, 11,20.25 2
{10
10
7
10
B.3.2 NNGP AND NTK
For NNGP and NTK, perform kernel regression following Lee et al. (2018) using centered labels. For each depth 31,2.3, 4). we sweep over hyperparameters listed 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 all depths is listed in Table
B.4 IMAGENET TRANSFER
We pretrain the -Lmii with 200 and hidden layers for 30 epochs on fixed subset of ImageNet32 with 250 (out of 1000) randomly subsampled classes, To evaluate on CIFARIO, we compute the kernel induced by the pretrained fira--reyesseaatures and perform kernel ridge regression. We optimize the hyperparameters Table 4 via random search.
B.5 OMNIGLOT
We focus on the 5-way, 1-shot task, with only step of ANIL adaption.
B.5.1 NENET, TENET, AND L-LIMIT
For π- and j-networks, infinite finite, train for 50 epochs, 1000 batches per epoch, and 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 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 7-net, -llimit, and μ-net: We sample at least 512 hyperparameter com- binations 1) for each (width, r) E {128, 256, 512, 1024, 20483 x (50, 100, 200, 400) for n-net, 2) for each width E {128, 256, 512, 1024, 2048} for penet, and 3) for each r E {50, 100, 200, 400} for -llimi..
Then take the best hyperparameter combinations overall (based validation accuracy) for n-net, -llimit, and μ-net, and for depth E {1, 3, 3,4 (where depth denotes number hidden layers), we resweep only the meta learning rate (i.e. outer learning rate) and step size (i.e. inner learning rate) in grid search, where the grid is (ilr* (2-2, 115 21.51) x (oir* (2-2, 2--3, ....,,..... 21.5))
This without loss of generality hecause relu homogeneous and are sweeping the hias variance.
16

---

Published as a conference paper at ICLR 2022
Table 4: ImageNet Transfer Hyperparameter Grid for j--Net, N-eett and x-Limit
Hyperparameter
x x-Limit Transfer
Bias Mult
0.5 (2-3 (2-3,2-2 23}
Batch Size
{6.8, 16)
Learning Rate
0.01
.... 26
Input Weight Mult
0.5 {1.5" 1.5
1.52.73
Output Weight Mult
(2-0.5 0.5 20 2
Weight Decay
(2-7 2-
20
Gradient Clip
0.4,2.04,0006 0.6, 0.8, 0.9,0 {10
Ridge
10
Table 5: Omniglot Hyperparameter Grid for p-Net, -Net and π --mimi
Hyperparameter
Grid
Step Size
0.5 0 2-7.5
2
Meta Learning Rate
16 (2-3.22.77
23
Gradient Clip
0.1 (2- 2.7.,. 2.75
.2.
Bias Mult
1. 1. 2---2,2-3.75
...., 22
Input Weight Mult
2. 2-
2
Input Weight LR Mult
0.2 0.2 (2 (2- 1.75
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. 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 shown in Table 1.
To investigate the width non -on-monotyn more thoroughly, we further perform random search on hy- perparameters listed in Table 5 for 2-hidden-layer n-net and μenet. We sample least 512 hyperparam- eter combinations 1) for each (width, r) E { [128, 256, 512, 8192] x {50, 100, 200,..., 32000) for æ-net, and 2) for each width E {128, 256, 512, ... 8192] for p-net. The best validation accura- cies per (width, r) are shown in Fig. 4(Right).
B.5.2 NNGP AND NTK
As discussed in Section 4.3, for NTK and NNGP, ANIL. meta-training has no effect, and meta-testing amounts to just taking kernel gradient descent step on each task. We fix the step size (i.e. inner learning rate) at 0.5. For each depth E {1. 2, 3, 4]. we sweep over hyperparameters listed in Table 6 using (complete) grid search, with weight initialization variance fixed 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 total 50 images. We take the best performing NTK, j-net, and w-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. we also do the same for our best performing neet
C DETAILED CALCULATIONS OF 1-STEP SGD
Suppose we present an input &o to the network and perform step of gradient descent with learning rate 17 and loss L. Then simple calculations show that the updates Au, Au to u, v
VnAv = mu), Au co @ ø' (Eovnu)
(15)
where c = -n Then, via Tensor Programs, f({) for any & now has limit of the form
lim f(E) E(Zvne + ( n-000
zvndug)
= -(Z + 80 /nu (00l:"""))(("VV"""" + /n o' (z go))
where ė is the deterministic limit of that is shown to exist by Tensor Programs.
17

---

Published as a conference paper at ICLR 2022
Table 6: Omniglot Hyperparameter Grid for NNGP and NTK
Hyperparameter
NNGP
NTK
Bias Variance
0.1 (2² 22.5
2
[2-1 2-0.5 25)
Bias LR Mult
n.a.
{2 2.5 ,22]
Input Layer LR Mult
n.a.
0.1 (2.-224,221
Output Layer LR Mult
0.1 0.
24)
Table 7: Pretraining ImageNet32 and Evaluating on CIFARIO, Full Results. We pretrained je-et, netet with r 200, n-net with 400, and --limit with r 400 on ImageNet32 and evaluated the result model on CIFARIO. Here, the π r-limit number 64.39 is the same as in Table under T-Limit ImageNet Transfer. For reference, we also include the NNGP and NTK numbers in the left block. The indicates we are comparing the r-limit transfer performance with r = 200 vs T r-limit CIFARIO number with 400, so the +2.79 is underestimate of the improvement due to pretraining. The benefit of pretraining seems to be directly related to the capacity of the model, as ---eet with r 200 < =NNet with 400 < j--Net < T- T-Lim
NNGP
NTK
J--Net
N-Net r=200
Net r=
T- Limit r=200
Transfer
58.92
59.63
61.84 +0.53
58.02
59.36
64.39
Table
+0
+0
-1.28
+2.79*
D REMARKS ON THE -LIMIT THEOREMS
1. n-projection actually ensures that, even for finite n, /: = P for some P. However, this P is random, with fluctuation coming from the sampling of f2 (which is fixed at initialization). Taking - 30 reduces this fluctuation O.
2. Again, even with the projection, "linearly infinite" r 20 instead of "quadratically infinite" oo oo like in the p-limit.
optimizing in an infinite-dimensional space, though
3. The forward pass of jP can thought of as that another ML.P with nonlinearity Va. as illustrated in Fig. 3. In this view, M becomes the width of this MIP.
E NUMERICAL VERIFICATION OF THE -LIMIT THEOREMS
Here we numerically show that wide --nets have nearly identical loss curves as the т T-limit. See Figs. 10, 11 and 13. In addition, Fig. 12 verifies the convergence of the feature kernel to its infinite- width limit.
F PROOFS
We will just prove Theorems 3.4 and 3.5, as Theorem 3.2 is special case of them. At high level, need to do two things: 1) Show that fr converges almost surely to something, and 2) this something fPr_ Here we assume the reader is familiar with Tensor Programs (Yang, 2019b;a; 2020a;b; Yang and Hu, 2020; Yang and Littwin, 2021; Yang et al., 2022), in particular the techniques used in Yang and Hu (2020).
F.1 ALMOST SURE CONVERGENCE
Showing almost sure convergence is straightforward using the Tensor Programs technique, i.e. express r-initialization and the entire 5-GG training trajectory inside Tensor Program NETSOR T 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 sketch the program construction.
Initial Vectors and Matrices Unlike the program for the p-limmi, our program does not have initial matrix variables because in p-parametrizztionn we do not initialize any a matrices as iid Gaussians. This means that we do not use MatMul instructions in our program. The initial vector variables in program are just the columns 2,,,..,, P2 of Ω
18

---

Published as a conference paper at ICLR 2022
Table 8: Feature Kernel Regression (FKR) on CIFAR1O. We take the best performing μi-net n--et, and æ-limii, and evaluate their learned feature kernels on CIFAR 10 via kernel regression. We list their test accuracies in the middle block. For reference, also include the NNGP and NTK numbers in the left block, as well as the ImageNet transfer results in the right block.
NNGP
NTK
J--Net
r-Net
--Liimi LiLiimi ImageNet Transfer
FKR
58.92
59.63
59.12
59.72
61.85
64.39
vs Table
+0
+0
-2.19
-0.92
+0.35
0.60
Test Kernel Accuracy
0.55
0.50
1000
5000
0.45
20000
40000
0.40
Train Epochs
Figure 6: Feature Kernel Regression of T-Nets vs Training Time, for Varying Widths. We took best performing n-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 any moment in time. Note that the visible gap between -llimit and the widest nenet (even at initialization) to large extent due to the dependence of kernel regression accuracy the smallest eigenvalues of the kernel. See Fig. 12.
Network Preactivations All preactivations of the network will be of the form QC for some C E R" whose entries are scalar variables in the program, so that 9C can be expressed vector variable using Nonlint.
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 n-nets can be written mathematically as a sum of vector outer products.
w Σ (PA)) @ ((IB)) time
(16)
where A'e, By, IR' are the sth of A' and B'. Here ((LB)) is the activation going into w' at time s, and QAA', is the projected gradient 1l. V at time Note the sum here is from to M + t, where M comes from the initialization and is from steps of training. All entries of A' and B' will be constructed scalar variables in the program, so QA and ((..B are both vector variables.
In the program, we do express w' directly, but rather through its matrix-vees -vetor product with vectors such as the incoming activation z!-1, which would be expressed via combination of Moment and Nonlint instructions like so:
w d Σ 0, ( E RD
Nonlin
(17)
where 0 (B(B + e R
Moment
We express (w")T indirectly through its matrix-vector products likewise, just with the roles of (DAA), ((DB)) reversed.
Gradient Projection In the program, we do not express Hg directly, but rather indirectly by expressing the matrix-vector product πg for vector variables v, such V,,C. The projection
19

---

Published as a conference paper at ICLR 2022
QFAR10
Ornniglot
60
85
menet 80
Jenet 75
58
depth
depth
Figure 7:
Performance vs Depth on CIFARIO and Omniglot. We take best performing p-net, T-het, and n-limit from thorough sweep of 2-hidden-layer networks, and resweep the learning rate and weight decay (for CIFAR10) or outer and inner learning rates (for Omniglot) for {1, 2,3, 4) hidden layers. We plot the best test accuracies of each depth here. See Appendix B for more details.
CIFARIO Train Acc
Omniglot Train Acc
1600
-90
96
400
100
90
80
128 256
512 1024 2048
:28
256
532
2048
4006
Width
Figure 8: Best Training Accuracy vs Width
CIFAR10 and Omniglot, taken all of our random hyperparameter searches. While networks with moderately large r and width can overfit CIFAR10 10 completely, no jenet, T-limit, or æ-net with width up to 8192 and up to 3200 is able to do so on Omniglot. See Appendix B.3.1 for experimental details.
matrix 1l mathematically equal to 22) (for any width), where (+ denotes pseudo- inverse. In the program, would first express (the entries of) } T Rexx using many Moment instructions
1 a 2) = (Q,, e R
Moment
Then its pseudo-inverse be expressed another Moment instruction (with purely scalar argu- ments).
L (T n
n Σ fols (T) e
Moment
n
n
fol'
ER
Moment
where fiy takes a matrix to the ijth entry of its pseudoinverse. Note the above expressions depend on Ω only, and not on Finally, we express Hg Qy E R" 7 (-22 2 E R", R",0 122 in the program like so
1l = Sh
N& E R
Nonlin
(18)
i=1
where
== = Σ
ER
Moment
(19)
72
/=1
where 0, = (1 E
Moment
(20)
Wrapping Up Other than what is discussed above, the unrolling of 7-GG follows identical the unrolling of SGD Yang and Hu (2020, Sec H.3, H.4). In particular, (((E) for any input & and time I. scalar variable in the program.
20

---

Published as a conference paper at ICLR 2022
NTK
Je Net
R-Limit
Figure 9: PCA of representations of images from 5 classes in Omniglot set. Same setting as in Fig. I, but here including our best performing n-net well.
Table 9: CIFARIO Compute Time (in Seconds) Comparison. We measure the average training time (in seconds) per epoch for 50 epochs of CIFARIO using half precision on a NVIDIA V100 GPU. We evaluate a y-Net, T-Net, and the L-Limi,, each of depth 1, 2. 3, and 4; the r-Nets and the r--immts have 400. Because the --Limi has linearly increasing compute time per epoch, also give an estimate expression for the compute time of the --Limmt in terms of epochs.
Hidden Layers
μ--Net
T-Net
T--imit Average T-Limit Epoch Estimate
6.78
6.83
83.01
30.97 + 2.23t
2
7.72
8,29
160.99
40.70 + 4.99t
3
9.03
9.80
200.31
48.413 7.756
4
10.17
11.29
263.11
61.20 + 10.36/
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 all straightforward except that need to check the pseudoinverse operation we took is almost surely continuous (in order to satisfy Yang and Hu (2020, Assm F.(())). 19 However, the only pseudoinverse we took was ( 22) and Q has rank r (full rank) almost surely for any n r. Therefore, our pseudoinverse operation almost surely continuous as pseudoinverse is continuous on matrices of constant rank.
F.2 FORM OF THE LIMIT
Forward Pass (Theorem 3.4) In the large-n limit, by the Master Theorem, for hidden weights, Eq. (17) becomes
M+1
M
2 zu'd
0
A
0 A,A
s=1
1=1
i=1
s=1
where 0, = E
6(2h''1
Bd S.
Inductively, if g' E R" represents the coefficients of ZM
zw'z'1 in terms of z def
( 20.)) (which is distributed as standard isotropic Gaussian vector), then
M++ Σ Bollg , 22)),
d.
which can 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 weight w', since we maintain wi in the form of Eq. (16), gradient update
" t w/ V w/ + (- (1h--1)
Yang and Hu (2020, Assm F.(()) actually requires Moment nonlinearities with only parameter arguments be continuous eveywhere, but because our result almost anyway, we ignore any measure event.
21

---

Published as a conference paper at ICLR 2022
Table 10: Omniglot Compute Time (in Seconds) Comparison. We measure average training time (in seconds) per epoch for 50 epochs of Omniglot using half precision on a NVIDIA V100 GPU. We evaluate ju-Net, æ--ett and the r-Limmt, each of depth 1, 2, and 4; the r-Nets and the r-Limits have r = 400. Because the r--iimi has a linearly increasing compute time per epoch due to gradient concatenation, also give an estimate expression for the compute time of the π- T-Limit in terms of t epochs. This is not necessary for the l-hidden-layer case, as there ANIL only trains the first layer,
which does gradient accumulation.
Hidden Layers
μ-Nett
T-Net
TT- T-Limit Average
T T--immt Epoch Estimate
1
28.68
29.33
36,48
N.A.
2
32.38
35.99
313.4
58.65 + 10.19t
3
38.47
43.06
596.98
73.48 + 20.94t
4
42.24
49.37
872.83
86.83 + 31.44t
Myg Mww
Figure 10: Wide 1-hidden-layer n-nets with r 2 have nearly identical loss curves their r-limmi. (Left) We train T-nets of r = 2 and width 23, 29, 215 as well as their common -imim on a 128- 28-image subset of CIFAR1O 110 over 200 steps, with batch size 32 per step. We plot the training loss vs steps on the left. While the width-8 -et fluctuates a bit around the --limit curve, width-512 and -32768 T-nets have nearly identical loss curves as the -liimit. (Right) With the same dataset and training procedure, we sweep widths 233.4....155 and 100 random seeds (which affect only the random initialization). For each width and seed, we calculate the median loss deviation of the n-net of that width from the r-limit, where the median is calculated over the 200 steps of training. Finally, for each width, we plot the median of those medians over the 100 seeds as the blue curve (with 95% confidence interval in shade). This curve shows that the loss curve of a n-net converges roughly as 1/ wiidh that of the -imitt.
in the limit corresponds to
A' + append((A' (A', n )), B' + append( coef(zz'))
where coef(Z" [")) ) is the coefficient of Z" in terms of z = (Z zorr). (Note the factor of n in --milla n L) is just there so that mella v C has ((1)-sized coordinates). Of course, coef( NN )) is just g² g- 1 Iº" by definition. It remains to show that coef( - V,, C) Vg lim C, where lim L. = limm (S)(S), y) is the loss at the limit (which is deterministicc.
Using the Master Theorem, it is not hard to see that, for each preactivation n', vt f is the Frechet derivative a2% L with respect the random variable Z²!, where the Frechet derivative is defined with respect to the Hilbert space H of square integrable random variables in the o-algebra generated by z def (ZR. .g....... 2P.r). Furthermore, using the limits of Eqs. (18) to (20) obtained from the Master Theorem, it is easy to see that Zalla I"
where In is the orthogonal projection to the linear subspace of H. spanned by the random variables z0n Z.. Then Lemma F.I applies and we get that
alim C 0..
Vnlla =
Og!
1=1
coef Vallaa
V. lim L,
22

---

Published as a conference paper at ICLR 2022
wmn 2374
sperd
04
92
Figure 11: Wide 2-hidden-layer g-nets with r = 2 have nearly identical loss curves as their r-limit. We repeat the procedure in Fig. 10 for 2-hidden-layed T-nets. (Left) Width- 32766 n-net has almost identical loss curve as the r-limit. Compared to the 1-hidden-layer case, the width 512 n-net is not close to the limit, but this is expected as depth slows down convergence with width. (Right) Nevertheless, we still see 1/ 1/ width convergence to the --limii in terms of training loss.
ompéiial deviadon
Fros.
ellml
2
2
10000 20000
25000
20020
Wide
Figure 12: Convergence of feature kernel at initialization, as measured by (left) Frobenius distance and (right) kernel regression accuracy. We perform all experiments here on subset of CIFARIO with 400 training and 400 testing examples. (Left) We empirically verify that, at initialization, the feature kernels of т-nets (with 2 hidden layers, kernel of the --mimit in Frobenius norm at a 1/Vwidth width rate. Here in blue we plot the Frobenius distance of the æ-net feature kernel to the limit kernel, normalized by the Frobenius norm of the limit kernel. The shade represents 95% interval of the mean, taken over 10 random seeds. (Right) We compute the feature kernel regression accuracy of randomly initialized i-nets of different widths (blue solid curve) and their common π r-limit (orange dashed curve). The shade represents 95% confidence interval of the mean, taken over 10 random seeds. We see convergence of this accuracy as one would expect from the theory. However, note that because the stability of kernel regression depends crucially on small eigenvalues of the kernel, the width needs to quite large compared to the data size (= kernel size) in order to visibly see convergence of accuracy; for data size beyond 400 training samples, we cannot see such convergence for width < 40, 000. This is why in Fig. 6, even at large gap in accuracy between æ --nett and the -lmimi..
400) converge to the feature
initialization
where the derivative is now an ordinary partial derivative.
Lemma F.1. Let H be Hilbert space and V be k-dimensional subspace of V, where k is finite. Let L: H IR be a Frechet differentiable function. Suppose T. have for columns an orthonormal set of basis for V, and w = Tb E V for some IR.. Then
V.. = VoL(Tb),
where Ny : H - V is the orthogonal projection to V.
Proof. By the theory of proximal gradients.
My V C(w) -min (v, Vw 110 w/|?
VEV
23

---

Published as a conference paper at ICLR 2022
2010
1/
200
Figure 13: Wide 2-hidden-layer g-nets with 400 have nearly identical loss curves their r-limit. We repeat the procedure in Fig. 10 for 2-hidden-layer æ-nets, but now with = 400 and widths 128, 2048, and 32768. (Left) Training loss. (Right) Training accuracy. In both subplots, width-32768 x-net has almost identical curves as the æ-limit, and the width-2044 curves follow them closely.
Changing coordinates via r, we have
(c.IVVC(TT)) + 1 [|IT - Tbll
Il V - ((w) min CER
(c, Vo.((()) 1 -1cc b11²
min
cER
V.L(T))
24


Reviews:
[{"review_id": "bWHg-Z5HAl2", "reviewer": "Reviewer_JW6k", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, this paper reads well and is clear despite the dense mathematical content, and solves a relevant problem.\nThe solution is technically sound, and works well empirically.\nI would be very happy to see this paper accepted.\n\nI still have some minor comments and concerns, however, that I hope the authors would consider addressing:\n\n- The proposed model ends up being quite different to a standard MLP, both in the way the weights are computed and the way training is performed.\nThis is unlike the previously proposed $\\mu$-net, which is much closer to the way standard MLPs are trained.\nClearly, the choices made are necessary for the computational speedups that this paper is about.\nHowever, I was hoping that the authors could comment more on the intuitive differences between these models?\nThe authors briefly discuss a similarity to the hypernetworks from Ha et al. (2016), but it would be great if this discussion could be expanded, particularly with respect to $r$.\nFor instance, would the authors expect the inductive biases to be much different between the two models?\nThis is likely a difficult question to answer, so I would simply ask the authors for some additional discussion along these lines.\nHowever, this question is quite relevant for determining what we can learn from these infinite width limits, which is one of the main motivations of studying infinite width limits.\n\n- I think that the results from Table 7 in the Appendix should be promoted to the main text.\nWhile not explicitly stated, I believe I am correct that the training loss in Table 1 is the standard softmax for the feature learning methods?\nIn this case, to ensure a fair comparison between methods, it is important that the same loss be used, as in my experience FKR is very different to training with a proper classification loss.\nAlong these lines, do the authors have any intuition as to why the $\\mu$-net suffers much more than the $\\pi$-net when used for FKR?\n\n- Would the authors be able to provide a comparison to $\\mu$-net and $\\pi$-net under transfer learning?\nIt would be understandable if not due to the computational requirements of doing so, but I think it would strengthen the paper.", "summary_of_the_paper": "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.\nThis 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.\nThe authors show that it outperforms both finite-width and previous infinite-width networks on CIFAR10 and Omniglot.", "summary_of_the_review": "In summary, I think this is a very good paper that I would personally like to see at the conference.\nI have some minor reservations which I hope the authors might address, but unless the other reviewers find something critical that I have missed I am unlikely to change my score.\nNevertheless, I look forward to reading the other reviewers' thoughts and the authors' responses. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "WMloWi4Sg9", "reviewer": "Reviewer_emnb", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Thanks a lot for the replies to many questions I asked!\n\nThrough the rebuttal, some central claims and contributions are clearer to me. But truthfully it turns out the paper now appears less promising and less exciting than I originally thought. Two major points that I learned from the rebuttal:\n\n- The $\\pi$-limit appears to give a qualitative view of the $\\pi$-net, but not the more precise details of the $\\pi$-net unless the width is 10 or 100 times larger than $r$. For example, the second figure in figure 12 shows that at width 2048 or even 32786, for $r=400$, the training accuracy could deviate from the $\\pi$-limit as much as 10% at various points of time. Figures 9 and 10 show better matches, but it is for $r=2$ which is atypical and very low.\n\n- The $\\pi$-limit can take a very long time to compute; its runtime is multiple times longer than training a $\\pi$-net (Table 8). Rather than “efficient computation”, what is demonstrated here is basically computability, against the fact the $\\mu$-limit is practically uncomputable.\n\nWithout further demonstration of what use the $\\pi$-limit can offer, this is quite hard to appreciate the $\\pi$-limit, against the empirical performances which not all reviewers here are excited about. Note, throughout the infinite-width literature, be it NTK or mean field, it is quite standard that the identification of the infinite-width limit comes with an additional analysis or a demonstration of use, such as proving theoretical global convergences.\n\nI’m also empathetic with reviewer MBzR’s concern, which is that the $\\pi$-parameterization simply goes further away from the usual initialization practice. The fact that the $\\mu$-parameterization bears similarity to the usual initialization practice (even though strictly speaking, it actually differs from the practice in a very important way) is probably one major reason why people pay much attention to it. By differing even more, the $\\pi$-limit disregards entirely this premise. As such, I would expect to see from the $\\pi$-limit something more interesting than a “fix” to the $\\mu$-limit.\n\nIn various ways, I currently find less reasons to be convinced it’s a promising direction.\n\nI will give the benefit of the doubt and raise my score, not because I want to align with other reviewers, but mainly because I simply like the idea. But I wish the paper were more carefully fleshed out than just this.\n\n=================================================\n\nThe proposal is interesting and novel to my knowledge. The fact that it allows for computation of the infinite-width limit, unlike $\\mu$-parameterization, while still achieving some nontrivial performance, is valuable in my opinion. I find that very encouraging, despite that the performance is of course nowhere close to strong modern neural nets yet. I also like that the paper provides a short description of the difficulty with $\\mu$-limit and takes a detour to the 1-hidden-layer example.\n\nI however find that the paper incomplete in several crucial aspects. The following should be considered in the finalized version:\n\n- The paper does not demonstrate convincingly if a finite-width net under $\\pi$-parameterization approaches the $\\pi$-limit as widths go to infinity; that is, a verification of Theorem 3.5. Now that the limit can be computed (as compared to $\\mu$-limit), this should have been done. In particular, it should demonstrate the closing gap to the $\\pi$-limit at every training steps, as indicated by Theorem 3.5. How large should the width be to be visually close to the limit? Is it close in just the classification accuracy metric? These questions should have been the first items to be answered when an infinite-width limit is claimed to appear more practical. \n\n    Figures 4 and 5 are insufficient in this aspect. Figure 5 looks concerning: as it shows, a width of 40,000 (!) still has a big gap to the limit on CIFAR10. I think this warrants deeper investigation.\n\n- The hyperparameter $r$ appears very important, but there is little discussion. From the experiments, $r$ should be large. But the need for large $r$ complicates the studies.\n\n    - Firstly this questions the relevance of the $\\pi$-limit itself, if one takes it as a theoretical limit of $\\pi$-net. In particular, $\\pi$-limit is derived with $\\text{width}\\to\\infty$ while keeping $r$ constant. Once $r$ is a nontrivial fraction of the width, one should expect some nontrivial random matrix effect to kick in and the $\\pi$-nets to have nontrivial deviation from the $\\pi$-limit, casting doubt on the validity of $\\pi$-limit. This is indeed the case in the experiments.\n\n    - Secondly if the proposal is to consider $\\pi$-limit as a new neural architecture with a specific training algorithm rather than a theoretical infinite-width limit of $\\pi$-nets, then the comparison should have been done differently. That is, it should not be compared against the simple NTK of feedforward nets, its expressiveness property should be investigated, its runtime should be studied, etc.\n\n- The claim (in the paper title) is that $\\pi$-limit can be efficiently computed, but again there is little discussion. In particular, can the paper provide a snapshot of the runtime for a reasonable, easily reproducible setup? There is a very short remark in page 4 that the complexity scales as $O(T^2)$, but this is insufficient to claim efficient computation. Is it running faster than or comparably to SGD on $\\pi$-nets, in terms of compute hours? Is it better / faster to just run SGD on $\\pi$-nets, and for what range of $r$?\n\nOther issues:\n\n- The paper claims outperformance of $\\pi$-limit and $\\pi$-nets over others in Table 1, but some of these numbers look too close to make a call. An increase of 2% for CIFAR10 could be borderline (given the nontrivial randomness in the generation of the matrices $A_l$ and $B_l$).\n\n- Footnote 12 is quite concerning: it is well-known that standard-parameterized MLPs should behave differently from $\\pi$-limit (and are close to NTK at best) as widths tend to infinity, but this footnote seems to suggest that the experimented $\\pi$-nets may operate in a regime far from the $\\pi$-limit?!\n\n- Footnote 13 seems unfair. Firstly the design of the learning rate schedule should be considered as a knob in hyperparameter tuning. Secondly the outcome of convex optimization would differ if one switches from a constant learning rate schedule to one with learning rate drops.\n\n- The hyperparameter tuning is rather intense for theoretically inclined audiences. It would be of great service if the paper can report the experimental results of a simple (not the best) hyperparameter setup for easy repetition/reproduction, in conjunction with these more tuned results. The paper should also specify the data processing step.\n\nQuestions:\n\n- Does the $\\pi$-parameterization inherit most properties of the $\\mu$-parameterization, other than being different in its limit’s analytical form and training procedure? \n\n    For example, the $\\mu$-parameterization has a somewhat paradoxical property (despite its name \"maximal\") that $\\Delta w(t) / w(0) \\to 0$ as width tends to infinity in most layers, for $\\Delta w(t)$ the gradient update at time $t$ and $w(0)$ the initial weight. That is, each weight coordinate moves very little from its initialization. Does the $\\pi$-parameterization share the same property?\n\n- How would one initialize the weights in $\\pi$-nets if the widths are not all equal? Here the same projection $\\Omega$ is used, restricting the widths to be all equal.\n\n- Footnote 10 seems to say that doing gradient accumulation on $\\pi$-nets could be difficult due to smoothness of $\\mathcal{V}_\\phi$. Is there such difficulty with the computation of the $\\pi$-limit?\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The proposal in the paper is interesting and novel, but the paper has deficiencies in its research methodology. It would be in a much better shape if several more fundamental issues are addressed. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LwvPdbhLJqb", "reviewer": "Reviewer_pNHj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The formal development presented in the paper is interesting and insightful. The idea is first presented in a simple example, in which the input and the output are a single real number, and this allows digesting the pretty complicated math. The idea itself, restricting the parameter dynamics in a few directions, is not new, as also stated by the authors  (it is also exploited in  Direct Feedback Alignment https://arxiv.org/abs/1609.01596, which in my opinion should be cited), but the explicit expression of the limit in the two theorems is, to the best of my knowledge, novel. \n\nMy main concern on the manuscript is on the meaningfulness of the numerical tests. The performance gap with respect to NTK on the CIFAR10 network is only 2 %, making the improvement brought by feature learning and by taking the explicit infinite-width limit difficult to appreciate. I suggest to repeat the test on a fully-connected architecture designed for image recognition, which achieves a test accuracy of 78 %, much closer to the state of the art on this data set: https://arxiv.org/abs/1511.02580\n\nIn this more realistic condition it should be possible appreciating better the improvement with respect to NTK and the nu-net.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The theoretical developments presented in the paper are (to my knowledge) novel, and can trigger other ideas for further improvements. The numerical tests presented to support the usefulness of the theoretical framework are not fully convincing, but can be rather easily improved. My score would be a 7, if 7 was available. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AcJ2hUXK824", "reviewer": "Reviewer_MBzR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\nInfinite-width models in their various flavours (NNGP, NTK, ...) form a useful set of tools for theoretically understanding the behaviour of neural networks in terms of classical kernel methods or GPs. Most (all?) of these types of limits usually mean that the limiting model keeps some notion of feature space constant during training. Motivated by the fact that practical neural networks actually update the feature space during training, the authors introduce a new limiting model (the pi-limit) that allows the features of the kernel to be updated during training. This is an important direction for the progression of infinite-width models.\n\nMain criticism:\nAs pointed out by the authors, there is a gap between these models and modern neural networks in terms of the architectures, training schemes, and training regime (i.e. is the network \"wide enough\" and trained with a \"small enough\" learning rate?). This means that when performing this style of analysis, one usually has to make some modification to what is typically done in practice. One runs the risk of making a modification so severe that the analysed model does not reflect a neural network that people usually compute with.\n\nIn this work, in order for the infinite width model to be tractable, several modifications are required for 1 layer case:\n-Initialise readout parameters at 0\n-Projected gradients (but only for hidden layers)\n-Minibatch size 1 (?)\n-Gradients are appended to the network instead of added (!)\n\nIn the multi-layer case, an additional modification is required:\n-pi-Initialisation. This initialisation appears to be substantially different to the typical scaling by sqrt n that is typical in infinite width works and in vanilla finite-width initialisation schemes (e.g. Glorot init).\n\nThese modifications, especially the fact that gradients are appended instead of added/subtracted in the usual way gradient updates are applied, mean that the model studied is quite different to one that would typically be used. This makes me wonder whether the studied model bears any resemblance to \"real\" neural networks. The model also does not achieve state of the art (it is not capable of expressing convolutional or transformer architectures). Since the model is not reflective of any simple practical neural network, nor does it achieve good performance, I am questioning the interest of this model in the broader ML community.\n\nQuestions for the authors:\n1. Can you say a bit more about this? \"For simplicity, we only consider batch size 1; it’s straightforward to generalize to larger batch sizes\". There is a comment in appendix A (A.4), but I do not see any statement or how to modify the result toa accommodate larger batch sizes.\n\n2. In equation (6), we use the V-transform, which is essentially the neural network kernel. The neural network kernel is symmetric in its second and third argument (i.e. swap X and Y and you get the same thing). In equation (6), the second and third argument don't look similar at all. Can you explain this asymmetry?\n\n3. There are number of fiddly modifications performed in the experimental evaluation. For example, the learning rate drop mentioned in footnote 13, gradient clipping, ANIL, cosine annealing, ... . This seems slightly counter to the overall objective of the paper, which was (I think) primarily to study an infinite width network that is able to perform feature learning for its theoretical niceness. When moving to the empirical evaluation, I am not convinced that the same amount of hyperparameter and model tuning was afforded to the NTK/NNGP/finite width models as was given to the pi-limit. I think that the fairest comparison would be to not include any \"tricks\" in any models.", "summary_of_the_paper": "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.", "summary_of_the_review": "After the reviewer response, I have upgraded my score to 6. My main concern remains that the model analysed is different to what is typically used in practice.\n\nThe paper appears to be technically sound, novel and written with the intention of solving an important problem. But it studies a model so far removed from a typical neural network (so it's theory is not widely applicable), and the model does not achieve state of the art performance (so it is empirically not widely applicable).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies deep non-linear infinite-width neural networks that go beyond the NTK and learn features. This paper extends the prior result on shallow neural networks to deep neural networks and empirically evaluates the deep inf-wide nn. The reviewers find the contributions in the paper valuable. The meta reviewer agrees and thus recommends acceptance.
Decision: Accept (Poster)