Adaptive Optimization in the $\infty$-Width Limit
Authors: Etai Littwin
Greg Yang
Venue: ICLR 2023 poster (2023)
Keywords: Infinite width
neural tangent kernels
feature learning
theory
adaptive optimization
tensor programs

Abstract: Recent works have developed detailed understanding of large neural networks' behaviors via their infinite-width limits, e.g., the neural tangent kernel (NTK) and the feature learning ($\mu$) limits. These theories were developed for stochastic gradient descent. Yet, in practice, all large NN are trained using Adam or other adaptive gradient optimizers (AGO), which are not covered by such previous works. Here, we close this gap via the Tensor Programs framework. Specifically, for deep MLPs, we derive the NTK and $\mu$ parametrizations as well as their infinite-width limits. We find 1) The NTK limit of AGO, in contrast to that of SGD, now depends nonlinearly on the loss derivative but nevertheless still fails to learn features; 2) this is fixed by the $\mu$ limit of AGO (as in the case of SGD). To obtain these results, we extend the Tensor Programs language with a new instruction that allows one to express the gradient processing done by AGOs.

Published as a conference paper at ICLR 2023
ADAPTIVE OPTIMIZATION IN THE OO- OW-WIDTH LIMIT
Etai Littwin
Greg Yang
Apple
Microsoft Research
elittwin@@pple.com
gregyang crosoft.com com
ABSTRACT
Recent works have developed detailed understanding of large neural networks' behaviors via their infinite-width limits, e.g., the neural tangent kernel (NTK) and the feature learning (1) limits. These theories were developed for stochastic gradient descent. Yet, in practice, all large NN are trained using Adam or other adaptive gradient optimizers (AGO), which are not covered by such previous works. Here, we close this gap via the Tensor Programs framework. Specifically, for deep MLPs, we derive the NTK and j parametrizations as well as their infinite- wiith limits. We find 1) The NTK limit of AGO, in contrast to that of SGD, now depends nonlinearly
the loss derivative but nevertheless still fails learn features; 2) this is fixed by the μ limit of AGO (as in the case of SGD). To obtain these results, we extend the Tensor Programs language with new instruction that allows one express the gradient processing done by AGOs.
INTRODUCTION
Infinite width limits of neural networks have been a major focus of study in the last several years, underlying some of the most profound recent breakthroughs in our theoretical understanding of deep learning. Specifically, two types of limits have garnered the lions share of attention from the research community. The kernel limit, popularized by the seminal work of Jacot et al. (2018) refers to a regime of training where weights remain roughly in their initialized values, and training may be entirely characterized in function space by a constant kernel of a particular form, which depends on the network architecture. While easier to analyze. this limit does not permit updates to the internal representation of the network, hence cannot account for data dependent feature learning, a staple of deep learning in practice. In contrast, the μ limit (of which the well- mean field limit a specific case in 1 l-hidden-layer perceptrons) refers to a regime of training where the weights adapt the data during training nonlinear fashion, facilitating representation learning. It was recently shown in Yang & Hu (2020) that, under vanilla gradient based training, the precise setting of various hyperparameters relating initialization scale and learning rate determine the type of infinite-width limit one can associate with trained neural network. Notably, the i parameterization was identified as the unique parameterization which gives rise to "maximal" feature learning dynamics in the infinite-widde limit, where maximal refers to the fact that every layer learns features. However, quite remarkably, no such limits have yet been formally established for adaptive gradient based optimization of neural networks, which we make the focus of the present paper. Our main results in the paper are the identification and prescription two types of infinite-width limits relating to popular AGO, the counterparts of the kernel and feature learning limits for vanilla GD. For the kernel limit counterpart, we uncover fundamentally different dynamics for adaptive optimization, referred to as the adaptive neural tangent kernel (ANTK) regime. In this limit, the training dynamics can longer be described by kernel gradient descent, since the kernel function itself depends non-linearly on the loss derivative. Our results lay clear path to theoretically analyze the implicit biases of AGO in the infinite-width limit.
Key Technical Contribution: Analyzing the dynamics of adaptive optimization of arbitrary neural network architectures in the infinite-widie limit presents major technical challenge. As a main technical tool, we build upon the TP framework introduced and developed in a series of recent papers Yang (2019; 2020a;b). At a high level, the mechanics of the TP technique involves 1) write
Please see vrorg for the full, updated version of this paper

---

Published as a conference paper at ICLR 2023
down the relevant neural network computation (e.g. the first forward pass in the NNGP case) principled composition of matrix multiplication and coordinatewise nonlinearities, called a Tensor Program, and 2) recursively calculate the distribution of coordinates of each vector via what's called the Master Theorem. However flexible, the "language" of TP is not expressive enough to represent the necessary computations involving adaptive optimization since it does not support the application of nonlinear functions to high order tensors, In the present paper, we solve this issue by expanding the TP framework with additional ffunctiontiess, and proving a new master theorem which enables our analysis. While we present simple application of our new framework on MLPs in Theorem 4. and Theorem 4.2, it applicaple algorithms. As an additional technical contribution, we prove On 1/2) width) convergence rate guarantee for all variables produced by the program, which might be of independent interest.
much wider setting, including most practical architectures and (nn11/2) (where represents the
Our Contributions: This paper presents the following major contributions:
1. We present the first rigorous infinite-width analysis of adaptive optimization of MLPs parameterized using the ANTK and μ parameterizations. Our results rigorously equate training of such networks to discrete time dynamical equations.
2. We develop new tensor program framework along convergence rate guarantees, unlocking the infinite- width analysis of adaptive optimization an architecturally universal sense.
This paper is organized as follows: We survey related work in Section 2. In Section we set up preliminaries and notations used extensively in Section 4. In Section 4 we illustrate ANTK and ya limits for MLPs. Section 5 is dedicated to a formal introduction to the new TP framework. Although it is used as a main tool to prove our results in Section 4, Section 5 is more general and can be read as a standalone.
Paper Organization:
RELATED WORK
A large body of literature exists on both the kernel (NTK) limit Arora et al. (2019); Jacot et al. (2018); Lee et al. (2019); Yang (2020c); Yang & Littwin (2021) and the mean field limit for 2 layer neural network Chizat & Bach (2018); Mei et al. (2018b); Nguyen & Pham (2020); Rotskoff & Vanden-Fiijnen (2018); Sirignano & Spiliopoulos (2020). Various papers describe the kernel and feature learning regimes more generally without taking an infinite-width limit. Chizat et al. (2019) describes the "lazy training" regime in arbitrary differentiable programs, and is controlled by single parameter a which scales the output. It is shown that when a is large, the weight need only move slightly to fit the training data, and network essentially performs kernel learning. Many papers Allen-Zhu et al. (2019); Huang & Yau (2020); Mei et al. (2018a) view the kernel and feature learning regimes as learning in different timescales, explicitly incorporating the time dependence in the affnite-widtt limit, and others derive finite width corrections to the NTK for finite width networks Hanin& & Nica (2020); Littwin et al. (2020a). In this paper, we consider training time to be constant, and take only the width to infinity. This way, kernel and feature learning behaviour separated by the parameterization employed at initialization, and not width or training time. TPs, first introduced in Yang (2019) and expanded upon in Yang (2020a;b), were developed analyze the infinite-width limits of any architecture expressible in the TP language, in an attempt to rid the per architecture analysis prevalent in the literature Alemohammad al. (2021), )); Du et al. (2019); Hron et al. (2020); Littwin et al. (2020b). Yang & Hu (2020) defined a natural space of neural pate-aaattttizaioonss) and classified all resulting infinite-widie width limits into two possible catagories: 1) the kernel limit, in which weights and activations remain roughly in their initialization state, and 2) feature learning limit, in which weights move substantially and adapt to data. The μ parameterization was then identified as the "optimal" parameterization for arbitrary architectures in which all layers learn features, and was later heuristically extended to AGOS Yang et al. (2021). Unrelated, AGOs Duchi et al. (2010); Kingma & Ba (2015); Zhou et al. (2018) and their variants were developed to accelerate learning by adapting the learning rate on a per parameter basis, and currently serve as a prerequisite for training large scale transformer models Huang et al. (2020); Liu et al. (2020); Zhang et al. (2019). Crucially, no previous work has yet developed a theory for infinite-wide neural network trained with AGOs.
theoretical framework to
network parametrizations
2

---

Published as a conference paper at ICLR 2023
3 PRELIMINARIES
Adaptive Optimizers: Generically, 90-91 91......... E IR denote the gradients of some scalar pa- rameter E at steps 0, 1,.,,,, L, an adaptive update Aw, Wt+11 W; at step t takes the form Aw; 11 Vote where v) is the learning rate and m and are both functions of the past gradients 90..... ......9. 9t- For example, in Adam, m. and are the exponential moving averages of 9()) and Poy- Here, we consider an even more general notion of adaptive updates, encompassing all modern AGOs.
Definition 3.1. We say update Aw, & Qe(yo, 91, .... Gti é) to weight at time L is adap- tive if it proportional (up to a constant factor) to a function Q. R++1 - such that Ve40, Qc(ego, cy1, .... c9ti cE) = Q:(g0.90090 11.00-911--0 91.... 9r; e). Moreover. if Qi 92. gric) Q(gie) (only depends on 9t), then we furthermore say Aw, is memoryless.
To maximize clarity, we focus on the simpler case of memoryless adaptive updates in the main text. For example, in Adam this implies setting 31, 32 = O. This simplification will already highlight the key differences between the adaptive and non-adaptive case, We provide an extension of these results to the case of AGOs with memory in Appendix C, and provide numerical verification of our results in Appendix D.
MLPs and ABC(D) Parameterization: We use standard scalar output MLP f with L hidden layers as a working example to illustrate the adaptive kernel and feature learning limits. Given an input sample & Rd,, weight matrices WL+1 E R", (W"),, E R= W R=xd, and an activation function & which we assume has a pseudo lipschitz first derivative, the output f({) E R is given by:
[a((E) - (E) for S0Is VISIEL ((())= (E) E
(I) We adopt the ABC parameterization convention from Yang & Hu (2020). Namely, for any layer I, each weight matrix is parameterized using W ====-=a:: -- which are initially sampled iid from normal distribution (N, n 2b)). Finally, the learning rate is parameterized using yyn - where we plug 1) = for simplicity. In this paper, we assign specific values to {arh1, (b {ct) for the ANTK and μ parameterizations. Additionally, we will parameterize the 6 parameter the AGO with e -n-ic,, where > The per-layer scaling for e will turn out to be crucial to prevent the adaptive gradients from collapsing to either O or step function as n co. We summarize the two parameterizations the following table:
f(E) = W (
where w the learnable weights,
Parameterization
at
b
Ct
di
1>
[ L 1>1
[ L+1>
ANTK
0
else
else
else
[=1
{1
-+1>1>1
[1 L+1>1>1
1
[+1+1
22
else
else
else
Table 1: ANTK and μ parameterizations.
Representing Activaaciittion Vectors via Random Variables: As will see, as width becomes large. the entries of the activation and preactivation vectors will become roughly iid (just like in the SGD case), both at initialization (which is easy to see) and training (which is harder see). Hence vector's behavior can be tracked via a random variable that reflects the distribution of its entries. Concretely, if E RR is one such vector, then we write 7" for such a random variable, such that a's entries look like iid samples from Z". When is scaled to have typical entry size independent of nt then - can be taken to be a random variable independent of n as well. In general, given two such vectors x, R", their random variables z and Z will be correlated, in such
lic. || /n- = @11) as
3

---

Published as a conference paper at ICLR 2023
way that -im z'v EE" Generally, inferring with initialized networks entail computing n expectations with gaussian Z variables, which take relatively simple form. However, a fundamental question is how the variables evolve during training, which we address next.
ADAPTIVE OPTIMIZATION OF AN MLP
In the following section we illustrate the Simnite-width limits of adaptive optimization for simple MLPs. For each parameterization, we begin by laying the basic intuition, culminating in Theorem 4. and Theorem 4.2. For cleaner presentation, our results are easily extended to the general case. In our setup we assume the network is trained using an AGO according Definition 3.1 l,and& batchsize of 1.
assume the first and last layers are fixed, however
Notations: Slightly abusing notation, we use subscripts to denote both the step index i, and coordinates of vectors with a, B. We assume & is training sample fed to the neural network at step (starting from 5o). and we use (E) for any input dependent vector/scalar y to denote its evaluation given & at step t. To reduce clutter we remove explicit dependency on input if it is implied by the step index (i.e 3/ ye(Ee) and y 3o(Eo)). We use ÿr vo(E) to express the dependency on an arbitrary input { at step f. We will also denote (g) (5) )e((E and 6y: (E) /nnyy(E) We assume the network is trained using generic loss function L, with loss derivative C = V Lt. We use the notation dh' based on context: for ANTK parameterization,
we set dh'(E) def T (5) E IK", whereas for μ parameterization, we set dh' (E) def af (5) E R". nI
(<)
V ग.
This context dependent notation is convenient since it insures that the components of dE) are roughly in the order of ((1) for both parameterizations. Finally, use denote the infinite-width
limit of (possibly random) scalar (i.e lim, J² f.). Using the above notation, can express the gradient of any intermediate layer w step for both parameterizations by n _dlhaa, L. Using Definition 3.1 the adaptive weight update Aw" for both parameterizations is given by:
VIcISh
Awi
IT ci; { 1 (diiat
1-1 Ciie)
(2)
where the function Q is applied element-wise on the matrix dhi/ C2 Raxx, For the remainder of the paper we will suppress the explicit the dependency and simply absorb it into Q.
4.1 THE ANTK LIMIT
In the NTK limit, intuitively, the weights of the neural network move by negligible amount during training, such that the network function may be approximated by its first order Taylor expansion around its initialization. This intuition carries over to the ANTK limit as well. At a high level, the following hold at any step Ah, for any layer will be of order O(n $). By definition h/41 i2 Ani, hence the coordinates of h for any layer I do not change in the limit, and, for any input, the coordinate distributions remain constant VIE(I.)) Zh! z zatt zzhi, Instead of training f. we consider training the first order linearization of f. denoted by f around its initial parameters. The function updates A A are given by
A
n² Σ an (ATTQaakaataa0 (,)
(3)
Under the ANTK parameterization, as with SGD, training flin and f using AGO equivalent. The following theorem describes the evolution of i. exactly:
Theorem 4.1. Let f(E) E I denote an MLP as in Eq. (1) parameterized using the ANTK parameter-
ization described in Section 4.1, where o' is Sreudo-Lipschite Assume layers L. are trained using a memoryless AGO with a greudo-Lipschit function Q according to Definition 3. 1 and a
li=2
batchsize of I, using loss function L with pseudo-I.pschhic first derivative. Then, any step t.
4

---

Published as a conference paper at ICLR 2023
and for any sample E, it holds that J. = j. where st. -K,uplEe, where:
KAdp (1C))
Σ >EEEE B[zuli Q(22a'(s)) 2
(4)
(((((((((EE))
(5)
where the expectation is taken over all Z variables at initialization.
Let us discuss Theorem 4.1 in a bit more detail. First, note that after the output values ft. and by extension the loss derivatives C are deterministic after conditioning on the outputs f at initialization, hence the only source of randomness in Eq. (162) is from the Z variables at initialization ?. Second, it straightforward to show that by setting Q(x)
we would get the SGD equivalent (when setting db, ab; - Cy-- For SGD, one may naively apply the the law of large numbers argument (LLÑ) and derive the
the learning rate to be 1) of Eq. (162), which takes the form Foed EE
indth limit under plain SGD: A " -K((,, where K. is the NTK function defined as:
(Ex) K(E,, E) - Σ Elzza zald " 2z0
(6)
Hence Theorem 4.1 is generalization of the well known NTK limit. At a glance, the transition from Eq. (3) to its infinite-width counterpart seems like straightforward application of LLN. However, the validity of Theorem 4.1 is not at all straightforward, and cannot be obtained by applying gaussian conditioning based arguments as in Yang & Littwin (2021), even for the first weight update. Technically, the complication arises from nonlinearity of the Q function: unlike SGD where nonlinear functions are only applied to vectors, in Eq. (3) we construct matrix (more generally tensor) using a nonlinear function Q. Operations of this type make even the simplest case, where all inputs are iid gaussian, tricky to analyze. Developing a general framework to handle such operations will be key to developing general framework to prove our main results later on. We discuss this technicality in more detail in Section 6.
For general adaptive updates, Theorem 4.1 implies that Kadp is nonlinear in the loss derivative Li- inducing a fundamentally different dynamics then the kernel gradient descent featured with SGD, and we leave the more in depth analysis of its nature for future work. Similar to NTK however, the ANTK limit does not allow data dependent feature learning, since the weights and activations remain roughly at their initialized values. This allows us to adopt a function space view of the training dynamics, where the output updates depend solely on the values of the outputs in the previous iteration, and without any dependence on the state of the internal representation computed by the network, which remain unchanged. In contrast, the μ parameterization allows data dependent feature learning in the infinite-wide winth limit, which we analyze next.
4.2 FEATURE LEARNING WITH PARAMETERIZATION
We now turn to analyzing the infinite-widhe training dynamics under ju parameterization. Funda- mentally, each weight update Auf will cause each preactivation vector hς change entrywise by something of order (11) and the coordinate distributions at the limit will evolve non-trivially. Gen- erally, the dynamical equations equivalent of an infinite-wide neural network the feature learning regime (using μ or otherwise) is much more complex to derive for deep networks. Although new TP formulation discussed in Section 5 provides a complete set of tools to tackle most architectures, will be content with illustrating the main points using a 2 hidden layer MLP where only the middle layer weights 2² are trained. Using Eq. (2),
express h2 2. fit1 using:
Viow 3T 222
R2 = b - 1-1Q(dh)))) x C))',, 27t1 = ((R
f
(1+1)
(7)
7
Note that under y the coordinates of Vnw" randomly distributed according (0, 1), hence we expect fe. to be (1)). Due to the Q function applied the gradient, the components of the Thhe Z variables in fact independent from the outputs f(E). This is made rigorous in the proof.
5

---

Published as a conference paper at ICLR 2023
co, and the update 1.2 Q(dhhf 2)) is by consequence generally not vanishing as well. Since the updates Ahς are non vanishing, the feature vector 27 evolves substantially (for non degenerate Φ), which enables feature learning. Taking the limit of Eq. (7) to derive dynamical equations is again not an easy task. Consider the case
R"*" matrix ((dh)x((i) are not vanishing
where Q Identity, which results in the update equation h2t1 - hī >')''( dh)
dh Ly and can be expressed purely using operations between vectors. For general nonlinear Q functions, we must deal with matrix- vector multiplication as in Eq. (7). This implies that unlike with SGD where we must reason about how a finite collection of R" vectors behave in the limit, we must now reason about the behaviour of R=x" matrices (see Section 6 for further discussion on this). The following theorem
describes the evolution of under μs exactly:
Theorem 4.2. Let f(E) E IR denote an MLP as in Eq. (1) with L = parameterized using the parameterization described in Section 4.1, where o' is pseudo-Lipschit Assume layers w² is trained using an AGO with a
batchsize of I, using loss function C with function Q function according to Definition 3.1 and a for any sample É, it holds that ft 4 fe where Se. can be computed as follows:
first derivative. Then any step t and
zhi zhi -Ez=ee,,2 2= [Q(Co' (z" (5.) E.)
(8)
z = (Z fo 0, ft = = C,(((E)
(9)
where the expectations are taken over all Z variables (including Ç d (V
From Theorem 4.2 it is clear what sense is parameterization enables feature learning in the limit:
from Eq. (8) the random variable encoding the updates to the hidden representation zdz is of order
((1) for non degenerate Q and é functions, and generally no longer gaussian at steps
allowing the neural network to learn data dependent features. Once again, substituting Q(z) x
would default the equations in Eq. (8) back those of SGD with an appropriate step size, hence our Theorem 4.2 generalizes feature learning with SGD.
5 A TENSOR PROGRAM FOR ADAPTIVE OPTIMIZERS
In Section 4 we have derived two types of limits in relatively restricted setting of training an MLP. In the following section, we go into more detail into the TP framework that allows such principled derivations much broader setting. While doing so, we will highlight the additional functionities introduced in the present paper that are key to unlocking the analysis of adaptive optimization, and removing certain assumptions preventing previous iterations from achieving full architectural generality. In the previous section we have provided intuitive calculations for computing the infinite- width limits of trained neural networks by explicitly expressing the updates to the network internal representation and output at step f, and then naively converting coordinates of vectors to iid samples from a known distribution (the Z variables). However, these computations become exceedingly complex with an arbitrary number of updates and complex architectures, and it is not clear whether these arguments can be applied in a general setting. Tensor programs is framework designed to automate these computations, while providing theoretical guarantees. Generally. computations involving adaptive optimization of most architectures contain a few repeating operations (i.c matrix multiplications, point wise non linearities...), applied to variables of different types (i.e matrices, vectors and scalars). This brings forth the notion of a program: A directed computational graph where nodes represent variables (typically " vectors or scalars), and edges represent operations performed on the variables. As concrete example, the forward pass of an MLP given some input can be expressed as a tensor program where the input represents the root node, and the affine transformation in each layer represent edge between nodes. We give a more formal description of our framework in the following.
5.1 NETORT PROGRAMS
Definition 5.1. A NESORT program is a sequence of "I"--vctors and RRscalars inductively generated via one of the following ways from an initial set C of random scalars, V of random R" vectors, and
Once again, the loss derivatives C deterministic in Eq. (8)

---

Published as a conference paper at ICLR 2023
set W of random R"** matrices (which will be sampled with iid Gaussian entries in Setup 5.2). Concretely, using weights W and some pseudo-lipschitz function tb R*(++1)+ + R for I, r E N.
the program generates new vectors and scalars from previously generated vectors 2: = {a,, x²) E
In and scalars ④ {0i, 02................ ...., 01} E R by the following instructions (using the notation æ; {a)....;)))
TENSOR Generates a vector I E R by To = 1 > 2. 8,= (T 1 B. Dz. f)
11
TENSORMOMENT Generates a scalar 0 E R by 0 = - _1 9, =1 1h ax F 3,5 O)
MATMUL. Generates a vector I E R by x = Wžorr= W I where I E σ W E W.
Let us unpack Definition 5. We can think of the TENSOR operation as a generalized version of the standard pointwise nonlinearity which acts vectors (or tensors where only one dimension increases to infinity, akin to the NONLIN instruction in Yang (2020b). Instead, the TENSOR instruction applies pointwise nonlinearity & to tensor of rank + 1 where all dimensions are of size n, and then contracts r dimensions to produce a vector. We note that the instruction subsumes the standard applications of (non)linear functions applied vectors by setting operation allows us to fully contract a tensor of rank + to a scalar. Finally, the MATMUL operation copied over from Yang (2020b), and implements standard linear layer.
O. The TEMSORMOMENT
The initial sets of vectors and scalars V C. and weights W are randomly sampled according to Setup 5.2:
Setup 5.2. 1) For each initial W W. we sample iid Was -N(0, /n) for some variance By associated W, independent of other W' W: 2) for some multivariate Gaussian Z E x V} RV. we sample the initial set of vectors V like {Ia : Z E V} z iid for each [n]. 3) For each initial scalar 0 E C, we require Vp>o,n' p0 0)2 a.s. 0.for some deterministic 0 e R.
Note that the initial set of vectors V are assumed to be gaussian in R". In typical neural network training scenario, the initial vectors correspond to input/output layer weights and biases at initializa- tion, and the initial matrices correspond to hidden layer weights at initialization, so their Gaussian sampling reflects their Gaussian initialization.
Example: A program encoding the first harkard and adaptive update (Eq. (7)) using parameterization is provided in Table 2
Expression
Op type
Implementation
2 = W
MATMUL. TENSOR
x² ((h2)
whh)) for vb(a) = x(a)
f Vnuu
TENSORMOMENT TENSORMOMENT
1 (Vnwa x²)) for (1 b) ab
L'(f) Th dh?
=== 41 f)) for vi( (1)
TENSOR
16 knww h22) for Va b) aδ' (b)
AT2
TENSOR
- bldhh ry. f3. L')
for b,,,,,;0000) (ab00)
Table 2: A NETORT Program encoding the forward/backward and adaptive update of an MLP. In the above, a, b, 6 E IR represent inputs to some function to implementing a TENSOR or a TENSORMOMENT instruction.
5.2 THE MASTER THEOREM
We can guarantee certain properties hold for vectors and scalars generated by a tensor program in the infinite- width limit. In short, any generated scalar 0 will almost surely converge to a deterministic limit i as + co, at a rate of O ).)) For any generated vector E RR,, the coordinates of will approach iid samples from some distribution. Adopting the notation from Section 3, we denote by
7

---

Published as a conference paper at ICLR 2023
Z" a random variable distributed like the coordinates of 2: - The following constructs the random variable 7 for every vector and deterministic scalar 0 for every scalar in the program, where we assume x = {a ..... x² θ = (01,-..,0) 0t) are previously generated vectors and scalars, and we the abbreviated n to denote the set of lo random variables {Z" Vi for all x' æ.
Definition 5.3 (Z" and 0). We recursively define z def 20 + 20 for each vector x and for each scalar follows:
ZINIT I E V. then Z" is defined as in Setup 5.2. We also set ~0 def z" and 2 def 0.
ZTENSOR x is generated by TENSOR (see Definition 5.1), then Z def f(Z") where f(c) def
Ezz "z......z (s, 2T 23 2,;; Ô)] with Z," being iid copies of z
ZTENSORMOMENT If is generated by TENSORMOMENT (see Definition 5.1), then 0 def 2= 2 .F [ ZT 1... ..., ZT {)]. Here Ô = 01...., 02 are deterministic, so the ex- pectation is taken over 7 7 3 T,, where ( J iid samples drawn from the same distribution as Z
ZMATMUL I Wa for 2 E æ, W E W then z def zw + WW, where:
ZHAT 2 is Gaussian variable with mean. Let Vw denote the set of all vec- tors in the program of the form Wy for some y. Then (2W', :y E Vw} is
defined to be jointly Gaussian with zero mean and covariance Cov (2 2ww)
def
w E. z" 7" for any Wx Wy E Vw.. Furthermore, { Wy E Vw} is mutually independent from (2 : u VU Vw where W ranges over WOfA A W).
ZDOT We can always unwind 7" :(...), for some arguments (...) 1/= P0 z' & VwT (where VwT is defined in 5.3), and / /02W/ y' def 0, Then There some nuance this definition, so T
(12w s'
deterministic function Φ Rb+j++ + R. Define 02 we set ZWa def Ow Σk= Zd oz"
see Remark A. and A.2.
The following theorem ties the symbolic nature of the Zs to the analytic nature of a Tensor Program.
Theorem 5.4 (NESORT Master Theorem). Fix NETORT program initialized accordingly to in all arguments, then
Setup 5.2. Assuming all nonlinearities are
..... 1. For any collection of vectors = "=.'...."") and scalars θ = =f....., 0.} in the program, and for any pseslo-lipschitt v/ R R((++))+ + R. as co:
T1 Σ
(Ta Dg, x3. 0) 8.8. Egg 25
[b(27 Z3 2 é)
nt+1
(10)
where { are r + 1 iid samples drawn from the same distribution as z and
7 (7. ..... 7:² are defined according to Definition 5.3.
2. Any scalar in the program tends to 0 almost surely such that Vp n2 -2-pp(0- - 00² asa where is defined in Definition 5.3.
Theorem 5.4 along with Definition 5.3 provide general tool set analyze adaptive (and standard) training of neural networks in the infinite-width limit, as long as the computational graph expressing the training process can be implemented in a NETORT program. Moreover, Theorem 5.4 provides universal -1/2) asymptotic rate of convergence for all scalars produced by the program.
8

---

Published as a conference paper at ICLR 2023
6 PROOF SKETCH
NE@ORT programs equipped with Theorem 5.4 provide the main tool to proving Theorem 4. and Theorem 4.2, and indeed their generalization to most common architectures and adaptive (and non adaptive) optimizers. This is done by adopting the following strategy: express the optimization dynamics using a NE@ORT program, mechanically compute the Z variables according Defini- tion 5.3, and apply Theorem 5,4 to compute the limit (see proofs in appendix Appendix B). What remains to prove Theorem 5.4 using strategy which we now outline.
In program, all vectors can be collected into an X M matrix V where is the dimension of each vector, and M the total number of vectors. The Master Theorem can be interpreted as saying that each row of V (i.e., the slice for each E [nz) is roughly an iid sample from some distribution 'D on RM (which can be derived via the calculus Z random variables as in Definition 5.3). Specifically, Theorem 5.4 and all previous versions of the Master Theorem formalize this by saying: this matrix V of vectors looks similar to a matrix V of iid samples from D, as measured by applying arbitrary pseudo- Lipschitz "test function r to both sides and taking averages.
Core Insight: Our core insight here is that V is in fact similar to V in a stronger sense without needing refer any test function t:: There is "small" matrix 8V of the same size V such that V - 5V is distributed exactly the same as V In general, if this happens, then we say V is equivalent to V'. The definition of "small" roughly means that each entry of áV has typical size O(n-1/2) Then, recover Theorem 5.4, we just note that the test function &b is (by assumption) smooth enough that 8V contributes a vanishing amount to the LHS of Eq. (10).
To prove this core insight, there are two parts.
Part 1: We show that, in any NETSOR T program (i.e., program with no scalar variables and no TENSOR operation), V is equivalent to V". This can done by re-analyzing the proof of the NETSORT Master Theorem in (Yang, 2020b) in a fairly straightforward way.
For any NETORT program π (the subject of our work here), we construct a parallel NETSOR program and show, by induction, that the vectors of the programs equivalent (i.e., distributed exactly the same after subtracting "small" vectors). This parallel program essentially replaces 1) all scalar variables in the original program by their deterministic limits, as computed in Definition 5, 2.32TENSSRRMOENTT and 2) all TENSOR operations by NONLIN operations, as computed in Definition 5.3SRT)
Part 2:
In this induction, we need to prove and use a lemma that, in the simplest case an illustration, says the following: For any standard Gaussian entries, the following two tensors T and T are equivalent: 1) the tensor T with entries Ts..... ( xz,) and 2) the tensor T" with entries T Ty,.... = ( ....) where x² ...., thth are iid copies of z. The proof of this lemma interestingly requires Baranyai's theorem, classical theorem from the theory of combinatorial design.
function v E* - and random vector R with iid
CONCLUSION
Adaptive optimizers are staple in the modern deep learning toolkit, an are necessary ingredient in most large scale neural network training. In this work, we have derived adaptive counterparts to the NTK and yu limit in prior works, which had only been derived for SGD, More generally, we have extended the Tensor Programs framework to allow the expression of any computation graph involving adaptive optimizers and the calculation of their large width limits. Our work lays path to study the implicit bias of adaptive optimizers by studying their evolution equations in the infinite-wine limit.
9

---

Published as a conference paper at ICLR 2023
REFERENCES
Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk. The recurrent neural tangent kernel. ArXiv, Sis/2006.102455 2021.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. ArXiv, abs/1811 1t///811.03399 1.03962, 2019.
Sanjeev Arora, S. Du, Wei Hu, Zhiyuan Li, R. Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. NeurIPS, 2019.
Lénaïc Chizat and Francis R. Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In NeurIPS, 2018.
Lénaïc Chizat, Edouard Oyallon, and Francis R. Bach, On lazy training in differentiable programming. In NeurlPS, 2019.
Simon Shaolei Du, Kangcheng Hou, Barnaáás Póczos, Ruslan Salakhutdinov, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. ArXiv, abs/1905.131 13192, 2019.
John C. Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. In J. Mach. Learn. Res., 2010.
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. ArXiv, (bbs/1909.05988 2020.
Jiri Hron, Yasaman Bahri, Jascha Narain Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In ICML, 2020.
Jiaoyang Huang and H. T. Yau. Dynamics of deep neural networks and neural tangent hierarchy. ArXiv, abs/1909.08
2020.
Xiaoshan Huang, Felipe Pérez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In ICMI., 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kerne:: convergence and generalization in neural networks (invited paper). Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR. abs/1412.6980, 2015.
Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. ArXiv, abs/1902.06620, 2019.
Etai Litwin, T. Galanti, and L. Wolf. On random kernels of residual architectures. arXiv: Learning, 2020a.
Etai Littwin, Tomer Galanti, Lior Wolf, and Greg Yang. On infinite-width hypernetworks. arXiv: Learning, 2020b.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. ArXiv, abs/2004.08249, 2020.
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 of the United States of America, 115:E7665 E7671, 1,,77,, 2018a.
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 of the United States of America, 115:E7665 E7671, 2018b.
10

---

Published as a conference paper at ICLR 2023
Phan-Minh Nguyen and Huy-Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. ArXiv, abs/2001. 1443, 2020.
Grant M. Rotskoff and Eric Vanden-Eijnden. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. ArXiv, abs/I 105001005,0915, 2018.
Justin A. Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM J. Appl. Math., 80:725-752, 2020.
Greg Yang. Tensor programs i: Wide feedforward or recurrent neural networks of any architecture gaussian processes. In NeurlPS, 2019.
Greg Yang. Tensor programs ii: Neural tangent kemee for any architecture. ArXiv, abs/2006.14548, 2020a.
Greg Yang. Tensor programs ii:: Neural matrix laws. ArXiv, abs/2009. 2020b.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. ArXiv, abs/2006.15555 2020c.
Greg Yang and Edward J. Hu. Feature learning in infinite-width neural networks. ArXiv, abs/201 1.15222, 2020.
Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality neural tangent kernel training dynamics. ArXiv, abs/2105.03703, 2021.
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub W. Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In NeurlPS, 2021.
J. Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Surinder Kumar, and Suvrit Sra. Why adam beats sgd for attention models. ArXiv, abs/1912.03194, 2019.
Zhiming Zhou, Qingru Zhang, Guansong Lu, Hongwei Wang, Weinan Zhang, and Yong Yu. Adaptive learning rate methods. 2018.

---

In 5.3, suppose TW''') are all elements of Vwr introduced before x. Define the matrix def Published as a conference paper at ICLR 2023
Appendix organization The appendix is organized as follows: In Appendix A we prove Theorem 5.4, which serves as the main tool to prove Theorem 4.1 and Theorem 4.2.
We then proceed to prove Theorem 4.1 and Theorem 4.2 in Appendix B. Appendix C we extend the proofs of Theorem 4.1 and Theorem 4.2 the of AGOs with memory, and provide numerical verification to our results in Appendix D.
A FULL PROOF OF THEOREM 5.4
In this section we provide the proof for Theorem 5.4, restated:
Theorem 5,4 (NE@ORT Master Theorem). Fix a NE@ORT program initialized accordingly to Setup 5.2. Assuming all nonlinearities are pseudo-Lipschitg in all arguments, then
1. For any collection of vectors a {x² ,.., } and scalars = {01,.... 01} the program. and for any pseudo- pseulo-Lipschitzz v [b(++1)++ + R, as n 7 00:
(Ea LB T3, Q) 8.6 Egg 29 ...... 2 [2 2 22 Z 6)]
n
0,01
(10)
where (Z 1
j==1 are r + üdd samples drawn from the same distribution as ze, and 7 (7* ..... 7=² are defined according to Definition 5.3.
2. Any scalar 0 in the program tends to 0 almost surely such that Vp>a, - 0)2 a.s. where is as defined in Definition 5.3.
Remark A.I (Partial derivative). The partial derivative in 5.3 should be interpreted as follows. By simple inductive argument, 2 for every vector æ in the program is defined uniquely as deterministic function p(ž²' ..... ....,2"") of some x² ..... _x in V introduced by MATMUL (notationally, suppressing the possible dependence on limit scalars OI,..., 01).
For instance, if in a program b.
have A E W.U. V. y Au, x A y. then Z 2 + 2". so yp is given by plaa b) Then
Oz 1022 def 0:p((
2n0 and az /82= del 0 for any E (x"....,,""
Note this definition depends on the precise way the program written, not just on the underlying mathematics. For example, if y, = E V and = ( (y + z)), then z oi - ((++))) so that
07 /82 = =07 /2=//22= = 0. If insteed we have x = Φ( W y Wz), then 7" = 62W so that 0Z" /82W(=+)) 0. However, in both cases, 2w
nZ Z E o' 2'(22 (w(++)))
Remark A.2 (Partial derivative expectation). The quantity E 8Z
azw is well defined if 2 is differen-
tiable in 2wTT, y. However, even if this is not the case, e.g. if x = ((WT y) where 0 is the Heavyside step function, we can still define this expectation by leveraging Stein's lemma:
C' E R by Cij = E Zd zn' and define the vector b E R by b; def 2w 2z, If C+b (where C+ denotes the pseudoinverse of C"), then in 5.3 we may set
ow
2 E
aZWTy² of
(11)
This definition agrees with the partial derivative expectation by Stein's lemma when the latter is well defined.
Pseudo-Lipschit functions are, roughly speaking, functions whose weak derivatives are polyno- mially bounded.
Definition A.3. A function f R² R is called pseudo-Lipschitz of degree d [f(x) f(v) <
C||2= yl + Et |r,]d + ())ld) for some C. We say f is Lipschitt if it is so for any degree.
12

---

Published as a conference paper at ICLR 2023
Here some basic properties of pseudo-Lipschitz functions:
The norm || || in Definition A.3 can be any norm equivalent to the norm, e.g. Pp, 2 1, norms. Similarly, Σ |a;]|* d + lys| can be replaced by | + llvll i for any p 2 1.
A pseudo-Lipschizz function is polynomially bounded.
A composition of pseudo-L.pschitz functions of degrees di and d2 is pseudo-L.ppshhe of degree d1 + dg.
A pseudo-Lipschitz function is Lipschitz on any compact set.
Indexing notations We use superseripts to distinguish between different tensors in the program, and subscripts to index into coordinates of tensors. (i.e z" denotes the j'th coordinate of vector X E IR". We typically use B 82. ..., }.} to denote a set of coordinates (typically containing r coordi- nates unless specified otherwise). For any vector E IR",, zg denotes the set ( F8g: ..., ",,,,,,,, Hg For
summation over all indices in B3 we use the abbreviated notation ΣØ= del - >.. We
typically use æ to denote set of vectors {u'...., ...,.."., }},, where the size of the set is implied by the context. The notation Zg refers to the set of scalars ( .................. ..,,,,, ",..... ....,
(note that in this case |an = kr. We additionally use a to index into tensors. The notation Fa for a vector J. refers refers to the set of scalars {=....,, v B, ..... 26 B, ...,..... (in this case |za,al = k(r + 1). We use c, C, as arbitrary constant scalars throughout the appendix (Their value might change between in different lines).
23 ...,
to the set {.a FB,, ...; Fz.}. Similarly, the notation aa for x {a ...,.. x*}
Proof. As in previous versions of Tensor Programs, we prove Theorem 5.4 by inducting on the vectors and scalars in the program. We use the following definitions throughout the proof:
x², ..... x*) in the program, define the matrix A E Definition A.4. For any set of vectors æ = {
Rkxk A T" We say the set has stable rank if lim, rank(A*) as rank A A*).
Definition A.5. We say random vector R" is vanishing if pp>0 limn |||| as 0.
Definition A.6. We say random vector E IR is regular with constants {é(p), [(p)} if for all
E (0, oo) there exists constants 0 < é(p) < a(p) < oc such that ö(p) < lim,, ||= < e(p) almost surely.
Definition A.7. We say random scalar E R is vanishing if it converges almost surely to 0 and it (0-0)" no-
holds that Vp>0: limn-
0. Equivalently, (0 0)1,, ,, is a vanishing vector.
Definition A.5 and Definition A.6 extend naturally to tensors, which will come in handy in our analysis. We index a rank tensor x E @'Ø" with indices & {B1...., 8r such that zg x131, 32, ... 3r].
||||| 0. 12
Definition A.8. We say a rank r random tensor x E 8R is vanishing if Vyso lim
Definition A.9. We say a rank r random tensor x E @ R' is regular with constants (é(p), A(p)) if for
|| n < a(p)
all p E (0, oo) there exists constants < a(p) S c(p) < co such that o(p) lim almost surely.
A.I INDUCTION HYPOTIISSIS
Setup A.10 (Induction Setup). We will keep track of set of vanishing sealars (see Definition A.7), and two sets of vectors: The core set which contains regular vectors produced by MATMUL operation (see Definition A.6 and Definition 5.1), and a vanishing set & of vanishing vectors (see Definition A.5). We will denote by aw the of vectors y : W y æc Given the sets of vanishing
13

---

Published as a conference paper at ICLR 2023
sealars θ corset and vanishing set è at some step in the program, let h E R" , R,, E R define new vector and scalar via TENSOR and TENSORMOMENT operations respectively. Namely:
(eo θ) 0 0- Σ ha
No n
(12)
1t
where v 1) R is pseudo lipschitz in all of its arguments. We further define:
1
ha
36 ( 3: 0:6)
(13)
B=1
ha Ezz ...... 77 (1 29, 22 .... 2,; 0: θ)
(14)
Ah: h 1xx x.
(15)
Ah ht" h
(16)
Our induction hypothesis IH (m) asserts that the following hold simultaneously:
1. ReWrite(m) Any vector produced by MATMUL can be written as a linear combination of vectors from z and æ.
2. StableRank(m) For any W RRRx,, the set aw has stable rank.
3. Dichotomy(m) It holds that:
(a) h is either a regular or vanishing vector.
(b) Ah, Ah are vanishing vectors.
4. (ensorMoment((m) It holds that 0 # ô, where:
Ezz,... 27..
(b( (
1 6)]
(17)
5. ConvRate(m) It holds that 0 R is vanishing scalar,
(0- 0)2
V Vp>0,
ag 0
(18)
np-1
A.2 HELPER LEMMAS
We use the following Lemmas regularly throughout the proof. Note that some of the lemmas are stated for the vector, however their extension to tensors are immediate.
Lemma A.11. If uu U. E R are vanishing vectors, then = u + v is a vanishing vector.
Proof.
| n
|/]|2 + ||0||2 + 2u
||0113
pp>0,
(19)
n2
np
72
Lemma A.12. If fu,
Idd are regular vectors, then |ua| + |u| is regular vector
Proof. for E [1, oo), using triangle inequality for norms:
|
|| + lulll
|| 1
$ p as (c((p) + +2(())²)²
Yp21,
(20)
For prrpe 10;
|
||
+
Vo<p<l.
(.) + 2x(())
(21)
n
| 2 max (p), 2, (p)).
One the other hand the lower bound trivially Vpa
14

---

Published as a conference paper at ICLR 2023
Lemma A.13. If IR is vanishing vector, then for any p 0, it holds that I/ 3 0 (i.e has vanishing moments).
Proof. For p 2 2, have that:
|uu1"
||ul|
ag
0
(22)
12
For <p< 2, using the fact that Vo<pear
|ull S no
-luull and assigning g = 2:
n
$
|ull
(23)
n
which proves the claim.
Lemma A.14. If E IRth is vanishing vector, and E IR is a regular vector, then u + u is regular vector.
Proof. This is immediate from Lemma A.13 and Lemma A.12, and setting the constants for (Cu(p), (p()) = (0.
Lemma A.15. If E RR" is vanishing vector, then any 2 v = lu|" is vanishing vector. If E R" is regular vector, then for any r 2 0, |u]" is a regular vector.
Proof. For vanishing u, using elementary norm bounds, that Viepens
||lll I l!ull and assigning
q 2:
|M|
1 12: 12
|ull
2 20
p>0.
(24)
nP
np
n
The proof for regular u follows immediately from the definition regular vector.
Lemma A.16. If u I is a vanishing vector, and E R regular vector with constants {o(p), } then v = u O v is a vanishing vector.
Proof. For any p choose m, E (0, such that > m, and [ + m 1. Using Holders inequality:
||/1
14
1/
m
Vp20,
(25)
n2
12
m
n
n
||12|
|/0||
2 T 0 (26) S 0 (8(
)mT
n
n
||| 1
where we used Lemma A.15 to assert that
2g 0. P-m n
Lemma A.17. 0 is a vanishing scalar, then for f R + R is vanishing scalar if f is locally lipschitz at 0, and f (0) O.
Proof. WLOG assume "0o- f(0) < Al0 8) for some e R. Define
A 0 (())
8) (0 8)2 <c
g(0)
(27)
else
Since 0 - 0 almost surely, this implies that Prob(lim,,.) (loom 8)2 < t) 1, hence (0 - 002 < almost surely. Therefore:
lim S(0)
lim g
(0 01² 0
V
Vps0,
<
lim A
(28)
70+00 n1-p
000
1-pp
+0000
n n
15

---

Published as a conference paper at ICLR 2023
Lemma A.18. Let a, x denote sets of regular and vanishing vectors, and let R|
R be
pseudo lipschitz. Then:
lim 7-000 71
2 (T La)
lim 110000
v axi 0)
(29)
71
Proof. We have that:
(mra; 0)
(maa aid) vx 0)) I (:aaaaa S (a 0) + [b Da) v(xa:)))
(30)
Since v is pseudo lipschitz:
(Iaa - (b 0)
I ( (1+ | + ||
(31)
Define the vectors Uo = (1 + | @ + L lid) Va = | Note that u, are regular and venishing vectors respectively, hence by Lemma A.16
vanishing vector. Then by Lemma A.13,
-a= (u v)a + 0 almost surely. Plugging into Eq. (30) and taking the limit, we have that:
lim Σ
1
(w 0) 0 lim
1 7
(Dxi Ta)
lim
(x 0) +
(32)
a=1
c=1
proving the claim.
Note that Lemmas A.11 A.16 and A.18 trivially extend to vanishing and regular tensors.
Lemma A.19. Ifu E * r is a vanishing tensor, then R Vo n Ei un is a vanishing vector.
1
Proof. Using the elementary inequality ||0/]] /nllool for any vector R", we have:
Ea=1
(Vn" V EB= u
||ul12 n²
p230,
(33)
n"+
n²+p
n"
Ea. U Da,e=1 ua ag 0
(34)
n
Lemma A.20. Let {zn]nzo be a sequence of random variables. for some N, and for all holds that Et S en en-i- -1- for X > 0 then In - almost surely.
Proof. by markov's inequality, for any 0:
Ex
P([xn > c) P( > 22) <
(35)
21
P([zn >6)
2n <.00
(36)
By the Borel-Cantelly Lemma, Fn 0 almost surely.
Lemma A.21. Let m, n, r E such that n. is divisible by Let {v- Y,, v¹ E R denote random (possibly dependent), zero mean vectors with iid coordinates and finite moments of any order VyEN, Bl( )2] Czq Define S In Eãe Por Then, there exists a function f(4) E 10, x))
independent m such that:
S 20 )a
E E[
Σ i=1
S C2
(37)
16

---

Published as a conference paper at ICLR 2023
Proof. Using Hölder's inequality:
178
S
24
176
E
E
St::
(38)
m
24
1=1
172
2q
E (SS)))2
(39)
m2
The inner expectations are given by:
20
[[(())2*)= E
Σ
E
(40)
Note that since the random variables V/ have zero mean, the expectation E vanish only when the indices 01, 022g do not contain an entry which appears isolation. In other words, the number of non-zero terms n² in the sum is:
11 v does not
n
Σ
29 Ugg
2q
Ug
U1
..
29
Ei= ui
21
( x- 10
(41)
EN (ur 1
By 2g2uh 12..2a1
11 Ei 2.
E!
Note that the the only term which depends on n
which is bounded by:
( Di= 1 Tupo
"!
t 9)! $ (
(42)
e Σl S0)
where f(q) < oo independent n, m. It follows:
n¹ < ())((
Σ {u, [u Upg 122ug-12..22
22 2q ).
(24 -
E1=
(43)
Ug
Ug-1
11
1/,14,11
Σ1- :::29
I n1 f(g)
where:
(44)
f(1) def f(2)
Σ
24
24 Ugg
(2q Σi= Ui
(45)
Ug
Uq-1
"1
2 12..
Σ 24=29
In addition, from Hölder's inequality:
24
2g 26
S
(ut ay
Czn
(46)
05
Hence it follows that E
({()))
20]
(g))
C2g
f(9)
Cz
Inserting into Eq. (38), finally get:
no
St
24
2q
f(a)
Czal
(47)
24
4
1=1
B 11...0000 11
f(4)
<
Ozg S 52))
(48)
74
17

---

Published as a conference paper at ICLR 2023
Before delving into the full proof of the induction hypothesis, note the following fact that immediately holds at any arbitrary step in the induction. Let h, ho Ah, 0 be defined as in Setup A.10.
Then the following claim holds:
Claim A.1. Ah is a vanishing vector.
Proof. Since & is psuedo lipshitz, there exists some d 2 such that:
Aba =
Σ
(3; 0:
(49)
72
1=1
1
0
ô11 +
|| 2 | 13
ellll + ||ô + || + ||
(50)
+ || lå + ||
(51)
2
TT. a
(52)
n2
where we have defined the tensors T @*+1 12 such that:
Ta.3 1 |01l + ||ôl| [
li
||aa
(53)
| θ + الههة |I
|
T
5
(54)
Note that by Lemmas A.11 to A.15, T is regular tensor, while r is vanishing tensor. Hence, T ( a vanishing tensor by Lemma A.16, hence Ah is vanishing vector by Lemma A.19.
Claim A.I is useful due to the following general claim:
Claim A.2. If Dichotomy(m) TensorMmennt m) and ConvRate(m) apply for the vector h (as defined in Setup A.),), then it applies for h + 8 if 8 is a vanishing vector.
Proof. This is true due to the function ve being pseudo lipschitzness. More specifically:
1. If Dichotomy(m) holds for h, then it holds for This trivially true due since we can expand à : 2 J... 6, and invoke Lemmas A.11 and A.14.
2. If Temsor holds for h. then it holds for f + 8. This is since may trivially assert that: = 1 n. E (ha + 6.) + 1 > has which stems from Lemma A.13.
3. If ConyRate(m) holds for ht then it holds for h + 8. This since:
pp>c>0:
1
8a + 0)2
n
60
2 n2
(0
8)
(55)
np
1 71
nP-1
n
+ nP-
(0 0)2
(56)
||0|| )² + 21
nn0 8)2
|
1
(0 8)²
(57)
71
n 2
n
7p-1
||
(0 8)²
|
(0 8)
+ +21
V
+
25 0
(58)
np
p
n
nP
n
Given Claim A.2 and Eq. (53), we may prove (mate(m) for 1 instead of h.
and Con-
18

---

Published as a conference paper at ICLR 2023
Figure I: Baranyai's Theorem. A graphical illustration of Baranyai's Theorem for = 8, = 2. A partition of 8 vertices into factors, represented by different colors. Each factor a partition the vertices into hyperedges (in this case, since = 2. simply edges) where vertex is shared between two edges, and no edge shared between two I factors. Baranyai's Theorem states that there are ( such I factors.
A.2.1 HYPERGRAPHS AND Baranyai's theorem
Baranyai's theorem combinatorial mathematics deals with the number of ways one can partition complete hypergraph into 1-factors. complete hypergraph G. is hypergraph containing vertices in which every subset of r vertices forms hyperedge. factor of this graph is a partition of the hypergraph into z hyperedges in which each vertex touches exactly one hyperedge. Theorem A.22 gives an informal statement of Baranyai's 's theorem.
Theorem A.22 (Baranyai's 's theorem Informal). The vertices of a hypergraph G ln, N such that r divides can be partitioned into I-fuctors in (?) different ways such that each hyperedge in G. appears in exactly one of the partitions (see Fig. 1 for a graphical illustration).
Theorem A.22 turns out to be useful in getting the almost sure convergence is stated in The- orem 5.4. Concretely, we will need to reason about the moments of infinite sums of random variables. Specifically, let (=t,..,,ee ,.., 2*) E Rn denote independent and normally distributed vectors, let eig = v'2) 2f..) where ti Rkrr + is polynomialy bounded and V/g E[vig] O. Consider the expression:
my -E
dB n2 0.5
)20]
(59)
3=1
Theorem A.23. There exists C(q) E [0, oo) such that lim mg S C(g).
Proof. We can express my by breaking the sum Ep=1:
my
=1(
[( 11 -- 0.5
1
tha $14
1
1h I .1.........00..)) 29
(60)
0.5 0.5
n²
$ A(n) + B(n)
where:
(61)
A E1( 0.5
w 31 31 81 291
(62)
B 01(
d 1 4..
(63)
-05
19

---

Published as a conference paper at ICLR 2023
where C [0, oa) does not depend on n. We now prove the following:
1. lim - 0. To this, notice that there are n"
m! [n-r] non zero terms in the sum
2g vb( - 1 S. #8,) hence:
B < Ō
(n-rj)
2q
max
E (vag) [4]
(64)
n 0.5 0.5
B
Notice that limm-soc n _ O, and max EE
((a))
24
is bounded and does not depend
We can therefore conclude B vanishes
2. We use Lemma A.21 and Theorem A.22 to show that 1im < C(g) for some C(1) <
00. Let {,3,, B ..... ....BImme [i=r]]] denote the set of all possible configurations of indices
I such that B1 * 32 * +.. Br. Let [0])m ]==1 denote sets where each set O contains m
configurations Oi = (g ..... 3 #11 such that:
(a) VE 3 =
(b) Vegjseerr B + 3'
Finally, let R denote the remaining configurations that do not appear in any set (0*) (i.e
== 0/(@ u u...u U.T)). We then have:
S
Σ
3 81 81
(65)
12"
0.5
2
-..
1=1
BER
where S
V/
(66)
BEB
Note that by construction of the sets {0* J= and since the vectors 2 ... z contain iid coordinates, we can conclude that for all i, random variables h are independent. Provided we can construct m = rn"-1_ Onn sets, then [R] = On In that case, then by Lemma A.21:
m St 24
lim A lim Ö( n-000
2
lim ČE
Σ
(67)
n
m
110000
BE
|
< C(g) + lim Č(
22
max (((b))))
(68)
-.3
= C(q)) +0
(69)
Ji=1 RR where rrn"1_ - 2). To show this, we define complete hypergraph GH with vertices in which every vertex corresponds to an integer in (1, 2..,...., n). We can think of a hyperedge in G7
We are then left with proving that we can in fact partition the set θ into {0')
edge connecting integers without ordering, hence the of all hyperedges in G has cardinality | (this is since ordering matters in θ). By Theorem A.22, we can partition the vertices in G into 2 hyperedges (sets of unique integers with no ordering) in
y ( different ways, where each hyperedge appears single partition. For each partition i, we can assign @' where any [3 O corresponds to single hyperedge in partition i, with the ordering of the vertices decided arbitrarily. Notice that for each partition O' we construct (r! 1) additional partitions by reordering 33 in any / E e' Therefore, the total number of valid partition is given by (n) Er! rn"- O(n 2), proving the theorem.
20

---

Published as a conference paper at ICLR 2023
A.3 BASE CASE
WLOG, we start with an initial corset of Gaussian iid vectors a = fx² "ur"..."") (which are regu- lar). an initial vanishing set of vanishing vectors ø = L'',...,'"" "...,"" and a set of vanishing scalars .... 0t). Note that Re and StableRank(1) trivially hold. We proceed to prove (ensorMmeetm)) and ConvRate(1).
We define the functions & R + R, 3 1g( + R, v: R + R using the pseudo lipschitz function &b (r++) - and vectors x::
(b())) def Ezg .g.......
[b: Zi Z2 z
(70)
(3 xg del pyy Rzg) (1)
(71)
(vi x2))
(72)
1=1
ha def (Maa)
(73)
Al. del' (Tx))
(74)
Note that (vy) is a random function that depends on the vectors æ.
Theorem A.24. Ah is vanishing vector.
Proof. From Lemma A.20, it suffices to show that for every p > 0, there exists / E N such that
E 20 )((
< en
en-l-À 1 for some c, À > 0 (which may depend on p). Fix p > 0, and choose
nep
t 1+A and let q [t]. Then by Jensen's 's inequality: p
Σ
12 D==1
(x
12
E
E B[(
(75)
np
n++1
(z TD) )
(76)
n
1+
n"
a=1
B=1
x33)
20
nitx
B
-0.
(77)
We are now left with the task of proving that E
(w ms) Eاا
24
is finite for any (fixed) integer
n
q and for n -> op. Firstly, we may express the over all indices Σ3= as ΣB + Σn E That
is, we first sum over all the indices where each one is bigger than 1, and then sum over all indices where at least one of them is 1. Then we have:
(x (3))
29
(x1
( mp) 24
E
E
(78)
n--
0.5
B
B=2
# 1y:8
< ČB
( mp)
24
ČE
Σ
n
( 233)
24
(79)
n 0.5
--5.
where Ō < oo does not depend on n. we now make the following observations:
oo that is independent of such that S. To see this, note that the summation (i.e E effectively sums over n" -(n 1y O(n"-1) terms. Since T is a centered pseudo lipschira.unntion of normally distributed variables, can bound the second expectation:
Claim A.3. There exists constant
C
Varrs
E
Vlay nr-a.b [
24
Vn,r, E
Σ
((01,
24
Č nd
max B. (2 SC
(80)
21

---

Published as a conference paper at ICLR 2023
for some Ö < 3o that do not depend
Claim A.4. There exists a constant C < o that is independent of such that Vngr E
((m) 207 } - C. Note that since the summation the indices B do not include the first index, the random variable a can be treated as independent from Et. WLOG we can then bound the following term instead:
(E
(2 3=1
(1, a3))
207 SC
Vnn, E
(81)
For some random variable y independent of ag for all values of B3, with the same dimensions as æ-- We can now condition on and apply Theorem A.23 to complete the proof.
Now, we can express h = h Ah + Ah, where Ah, Ah are vanishing vectors. Invoking Claim A.2, it is enough prove the base case holds for h alone:
1. TemsorMoment(() is immediate from the law of large numbers given that are iid gaus- sians. Namely:
3 Taa
0.5
Eg= (Z
(82)
71
2. Dichotomy(1) holds since 1 is smooth, pseudo lipschitz (given by gaussian averaging of v function. From TensorMoment(1), for any p > 0:
1 lim |11
Eas1
hha 1
/ MP
lim 22000
Ez | |Z
ddz
1 40000
(83)
where ((*)) is the gaussian measure. Therefore, limm | = for some p > then &e is identically zero. that case, we can write = 0 + Ah + Ah, which is a vanishing vector. If d is not identically zero, then h is a regular vector from TensorMoment(1) proving Dichotomy(1)
3. ConvRate(1) holds since for all > 0:
Eo a= (ha - 0) lim
Ex= (ha - -0) -0)}² vn
lim <00 np-1
(84)
n- n2
which holds from Theorem A.24 and assigning r =
We have therefor concluded the base case.
A.4 INDUCTION STEP
We prove the induction step assuming IH(m) holds. Namely, we must show that IH IH IH (mm++)), Assume a new vector is introduced via MATMUL, namely Wv where v is given by TENSOR operation of corset ær vanset and vanishing scalars Ô By Dichotomy(m), we may express v = - Av Ad
where Av + AD is vanishing vector, and 7 is either regular or vanishing.
A.4.1 IH (m) + REWRITE(M + 1) + STABLERANK( + 1)
We can express Wy Wi + (( + AD), and point to the following fact:
Claim A.S. Df is a regular vector, then W is a regular vector. Moreover, for any vanishing vector 8 E R., W8 vanishing vector.
22

---

Published as a conference paper at ICLR 2023
Proof. If ā is vanishing then W6 is vanishing: This is true since W is gaussian matrix with uniformally bounded (in n) operator norm. The first part Claim A.5 holds since E depends only on vectors from æ, for which the set æw has a stable rank from StableRank(m)) (m)). We can therefore use the gaussian conditioning trick (conditioning on all vectors in x and aw)
IT
We can now expand the vanset with é = ( + Az). and proceed by casework:
1. If y is vanishing then we expand क := ② U Wz. In that case a remains unchanged and StableRank (m + 1) trivially holds.
2. If regular then we expand == x W7, and we get StableRank(m + 1) using Theorem A.25.
Theorem A.25. Let x: y.... Itth Eaw. 7f on zu' + 02770 + +... + za then almost surely, for large enough
a/æ + aag2 + + ag.x
Proof. The set æ is constructed as a standard NETSOR T program (without scalars), and we may immediately apply theorem 6.3 in Yang (2020b).
A.4.2 IH (m) + REWRITE(M + 1) STABLERANK(M IH +1
We are left with proving Dichotomy( (m + 1), TensorMoment((m + 1) and ConvRate(m + 1). Note that if D a vanishing vector, the set z remains unchanged, and we immediately get IH + +) by Claim A.2. Hence proceed assuming & (hence Wi is a regular vector.
Getting Dichotomy (m + 1) Assume new vector is introduced in the program via a TENSOR operation:
ha
- 11" Σ (Fa.S (W 3; 9)
Σ
(85)
Where we made explicit the inclusion to of the vector W7. Let h h, Ah, Ah be defined as in Eq. (13). Note that Ah is vanishing by Claim A.1, which holds generally. We next prove
Al. = 11" h is a vanishing vector, where (using x-(-) = (1=-;; θ) to ease notational burden):
Ahg 1
Σ
( (WD)a.g) - -B2q.zy "....2 Z5... [tb (WD).. Z Z ZT ZWP)
n
8=1
(86)
The key insight to proving Ah = 1 h indeed vanishing is that both n., 7 can be written as a sum of shared regular vector, and a vanishing vector (i.c μ is regular, and 81, 82 are vanishing). Their difference to" - = 8² 82 is therefore the difference of two vanishing vectors, which is itself vanishing. To show this,
express Λ = μ + 6², h μ + 82 where
make explicit the distribution of We by using the gaussian conditioning trick (see Appendix A.5). Denote by X.Y Y E Rx U, V e Ra** the matrices with {a') E x, {y'} E Iw, {u') x, {u²) æw as columns respectively, representing previously generated vectors in the program, such that X WY = WTV. Using the Gaussian conditioning trick (conditioning all the vectors in x). g WD is distributed as:
WE Σ daa +
Σ Eev +oll
(87)
where d + d. + ē, +Ø G, and & ~ NN0, Define:
a = Σ der' + Σ E, +
(88)
b = g a = Σ (14 diyxt + E(e; 2:)0' + + - -
(89)
We now note that:
23

---

Published as a conference paper at ICLR 2023
v. (d, - -d,)z* is vanishing vector since x² is regular (æ zw), and (d; d,) is vanishing by the induction hypothesis, and Lemma A.17.
v/s (e=-,,Ju e:je* is vanishing vector since ef is regular (v E IwT), and (e; - é,) is vanishing by the induction hypothesis, and Lemma A.17.
(olly - =) a vanishing vector. To see this, note that:
(oll) -)2 (6 - o)= V(VVT)ivT=
(90)
(a - o) is vanishing due to the induction hypothesis and Lemma A.17. V(VVT)tV is vanishing as well. To see this, note that:
IV(VV ) L
V VVV TT2l
P220
n
(91)
n2
n V Vv VV = V VVT VVT vn
(92)
n+t1
np
By the induction hypothesis ( (VY converges almost surely (V has stable rank). Then = ng 0. From TensorMoment(m)) and
it is enough to show that Vepso, limn-xc n² vn ConyRate(m):
1 TT o²
2 2 1
V
lim n n00 n Vn a 400
lim
n
)2
(93)
72
)2 2$ 200
lim n
p-p(
(94)
n2
We therefore conclude that b = a is vanishing. We have:
11
1b (En. Sa.gg
Σ
A(m) + B(m)
(95)
n2"
n"
B=1
B=1
where
A(m)
wbl (Warp Edia'
drat Σ êvv" + @z)),;))
(96)
n1"
B=1
B(m)
(Fa
Eej 200 +8z)a.o)
(97)
n2
From Claim A.I, B(m) is vanishing vector. Furthermore, is deterministic function of previous vectors æ, and an gaussian noise vector z.
We can recursively expand A(m) = A(m 1) + Bmm 1) where B(m 1) is vanishing vector, until we are left with A(1) (i.e A(m) A(1) + Σ Σm=1 B(m'), where (mm'm))m are all vanishing vectors). Note that A(1) can be expressed as pseudo lipschitz function of normally
distributed vectors, with coordinate distributions given by Z -=a ..... 2=|m| We can apply the same decomposition and get h A(1) + m B(m'). We have that:
Ah- = h2 h= A(1) - A(1) Σ (B(m') B(m'))
(98)
Finally, it is easy to see that A(1) = A(1), and hence we may invoke the base case (in particular Theorem A.24) and conclude that A(m) A(m). and by extension Ah are vanishing.
Getting Tensor mm1) and ConyRate These are immediate since express = A(1) + where A(1) is a function of gaussian vectors, and a vanishing vector, and by Claim A.2 invoke the base case A(1).
24

---

Published as a conference paper at ICLR 2023
A.5 GAUSSIAN CONDITIONING
Let Wh where h E R are vectors in a NETSOR T program.
Denote by X, S.Y. _axr U.V Rax* the matrices with (x'), {y"), {u'), {u²} as columns respec- tively, representing previously generated vectors in the program, such that X = WY, = W
Using the Gaussian conditioning trick (conditioning on all the vectors in æ), g Wh is distributed as:
g 4 (E + + B
where we have defined A Ev., B 11 Wl1 Ily, Ily are projection matrices, and W is fresh iid sample of W, and:
Rewriting the conditional distribution of g. we get:
g P 0
(99)
||
with del ChER", del oAo
ER
(100)
Moreover, converges to a deterministic limit o and è can be written as:
+ x + V +
(101)
where R,,", E R are vanishing vectors, and d E R,, ė R are deterministic vectors.
B PROOFS OF THEOREM 4.1 AND THEOREM 4.2
Now that we have proven Theorem 5.4, we prove Theorem 4.1 and Theorem 4.2 by writing out the TPs which implements the training process, and apply the master theorem. To accomplish this, we start by expressing the explicit computation done each training step at finite width, implement it as set of TP instructions, convert it to infinite-width computation according to Definition 5.3 and apply the master theorem.
B. THE TENSOR PROGRAM FOR THEOREM 4.1
In the next section, we construct the Tensor Program that encodes the training of an L-hidden layer MLP inEq. (1) under the ANKT parametrization. Here first describe the initial matrices, vectors, and scalars of the program, along with necessary notations.
Initial Matrices, Vectors and Scalars scalars C:
We first define the initial set of matrices W, vectors V and
Initial matrices {w/ Me E R are all sampled iid from A/ 1). We W W U W3
The initial vectors V are given by the first layer he (E) for all inputs, and the last weight vector w++1 E IR", all samples iid from.N (0,
Initial scalar C = }.
Notations We use = to more clearly denote assignment happening in the program, as opposed to mathematical equality. We will use the notation TENSOR(y's y O), TENSORMOMENTIg ... ::,)) to denote an arbitrary implementation of these instructions give vectors y ...,," and scalars O.
25

---

and is the Published as a conference paper at ICLR 2023
Initial Forward Pass Starting with our initial vectors h (() (x)( (2))2 using TENSOR and MATMUL instructions:
w'g
compute all
VISISL, d'(E) == (hh) for v(y) tht o(v)
(102)
VICISL, (²(()) = w 2')
(103)
The initial outputs are given by f(() 2 since TENSOR only allows division by n" for integer r, rather than nÝ. However recall that in Theorem 4.1 assume WLOG that the outputs for any input & is fixed to f(E) = g(E). .) Let X denote a matrix composed all x x²(() as columns. Denote by the event that f({) 0 for all inputs. Then, using gaussian conditioning, the conditional distribution of w given e is:
Vn
w W
11
14
(104)
where w is an independent sample of wht1 and II = 1 ---4XX ( X pseudo-inverse of Then, we can write:
X X
XT whi+1
W J++ D
T+1
X(
(105)
72
n
" 0. Hence, after conditioning on By the master theorem X 33 7, X xx y and x m111
f(E) 0, the distribution of w is still identical to that of w ++ at the limit. At this point we can just implement w in the program with x/+1 - X n with TENSOR and TENSOR MOMENT instructions. We now have program that encodes the initial forward pass of the MLP conditioned on f()) for all g.
xx x g which can be implemented n
Initial Backward Pass and Loss Derivatives TENSOR:
For any input sample &, we can implement di' using
dh L = TENSOR (E), wh+1) for TENSOR . del ('(y') Oy 2?
(106)
Then, for all 1 51 L, using MATMUL and TENSOR:
di's = TENSOR(W'*IT I R') for TENSOR(y ./) def 06 0
(107)
The initial loss derivatives (() are all deterministic scalars since we have conditioned on the initial outputs f().
Forward and Backward Passes at Any t any are given by:
The forward and backward and loss computation for
h1 ( 1
Aupe
(108)
Vn #=
dhf = ((())) O w
T++1
(109)
Vigiet dit = [
Σ Autt1) T
ditt1]
o (k()
(110)
==
with the weight updates given by:
Aw! Qudh E n
(111)
which can all be implemented with TENSOR and Matmul operations before.
26

---

Published as a conference paper at ICLR 2023
Adaptive Update at Time have that:
Using Eq. (2) and w/ = =w' + Ei= Awp, (recall w' vnW"),
t Au
(112)
σ Autat
+ vn
(w' Σ Swo ) 60611
Auf
(113)
vn
Q(antat L.)at + (wi
nvT Σ
Qddhh.ar
1
(114)
d==0
1
(115)
nyn
_ = (o 8h4 (ht)
(116)
which can be implemented using TENSOR:
5hद == TENSOR (dhi, xr, #t f.) for TENSOR(J'.''.'' def Qu''' 0)y²
(117)
n
d = TENSOR (R2 oh?; 1 for TENSOR(y',. :)::)) der 500(1* - -0y²) :
> 0
(0'(y') on v²
0 0=00 0
(118)
and similarly for any layer 2 < 1 S L.
The (pre)activations at any step t can be implemented follows using TENSOR:
VIsiSL, B141 == TENSOR(hL (hi, ohe; Vn for TENSOR (3', .2:00) def y + By²
(119)
def VIKISL, #4 = TENSOR G. dāg; for TENSOR(y',O ,22:0) = =y' + Oy²
(120)
Output Updates The function updates can be implemented using TENSORMOMENT:
Ai. == TENSORMOMENT(( (z))) for TENSORMOMENT(//y>) def Σ Valla (121) a=1
The loss derivatives can be implemented using using TENSORMOMENT:
C = TENSORMOMENT(;;1 ft) for TENSORMOMENT( def 1 n
Σ C()) n
(122)
B.2 PROOF OF THEOREM 4.1
After writing the program using TP operations, we are ready to prove Theorem 4.1 by taking the infinite-win winth limit. First, note that from Eqs. (119), (120) and (166) and ?? and applying Definition 5.3, we have that:
VISISL, zol =z'
(123)
VISISI, ze1 7.5¹
(124)
VISKL, zahi zah Z dilt1 8((zh)
(125)
zdilt zahl zwttt (2222)
(126)
Applying Definition 5.3 to Eqs. (117) and (118), we have that:
o = F.z.' zz5 (Q(a" zn' LT) 2.04
(127)
Z = 6'(z²²) )
(128)
(129)
27

---

Published as a conference paper at ICLR 2023
And similarly for Eqs. (113) and (116):
V2c1gr. 2 Sh - E z,1-1 200-1
1,
( (2(h za'-1
21)20""1 + 2
(130)
Yasist,288 = (Z )
(131)
(132)
where we have that 2 2w' Theorem 5.4:
+ z according to Definition 5.3. Then using
Af El Zwttt 2024
(133)
El 2m
0' h 20h
(134)
E Zut
8 (Zh
Bzah-
[Q(2ah) zot
[() z0t-1
(135)
E [Z Φ 8'(Zh', >2 12W"8z"'1
(136)
- El QZZar" zot-11 [i) 2²)
(137)
+ E [ =
(138)
We now use lemma L.3 from Yang (2020b) restated:
Lemma B.1. For any x, E R and W E RRxn in the program, holds that:
E|Z*2 y = Bl x =
(139)
Applying Lemma B.1 Eq. (138):
ah² zwL684-1 11 = F.. F2 W d 7624-11
(140)
z (E (Ezor-2 2222- 1000=220((00 0lth 1--1 20- (1'
Q(
L) L2) 20""² -2 3
(141)
--EEzah Qz 1-11 200 -00-0220 L')2z" 1--2
E2 Zahh zW -W,-1185" 586/12
(142)
Similarly expanding E 1ith- 1 zwL-1854 we arrive at:
1. Aj
Σ
E zahi Q(zai' zp' (1) -=)==="==---- =-x(= El1")
(143)
1=2
Finally, using Eq. (152):
Afr = (E, EIC)
(144)
B.3 THE TENSOR PROGRAM FOR THEOREM 4.2
In the next section, construct the Tensor Program that encodes the training of an 2-hidden layer MLP as inEq. (1) under the μ parameterization. Since the last layer is not trained, we define 3 Vnw3, the output is given by f({) 3 2²) Here first describe the initial matrices, vectors, and scalars of the program, along with necessary notations.
Initial Matrices, Vectors and Scalars We first define the initial set of matrices W, vectors V and scalars C:
Initial matrices 2 E R sampled iid from NiO, 4). We set W = W
The initial vectors are given by the first layer 1e² (S) for all inputs, and the last weight vector 3 E IR" Notice that 3 is normally distributed.
Initial scalar C {
28

---

Published as a conference paper at ICLR 2023
Notations As in the ANTK case, to more clearly denote assignment happening in the program, as opposed to mathematical equality. To clearly demonstrate the application of TENSOR, will also freely introduce function symbols & put things into TENSOR form.
Initial Forward and Backward Passes Starting with our initial vectors / (²( == w we compute x² (E). (e) (S), dh and MATMUL instructions at step any t:
f(E), (E) for all inputs using TENSOR, TENSORMOMENT
z' (E) = TENSOR( (E)) for TENSOR(U) del ((y)
(145)
(E) = w²z*(E) (E)
(146)
(() == (E)) for TENSOR(y) d ((v)
(147)
f(() = TENSORMOMENT( 2E) for TENSORMOMENT(y" 222 def n
yaya
(148)
0=1
C'(E) : C TENSORMOMENT(. ft) for TENSORMOMENT(:0) del
C'
(149)
11
dh (E) == TENSOR(W' h² for TENSOR(y' ,3²) def y o o')())
(150)
Note that with l parameterization we can express the output f({ directly without conditioning using a MOMENT.
Expressing h.41 Using Eq. (2), we have:
1 2 O)y² h2 = TENSOR die u|. ali () for TENSOR(U) '3?,, :0;;) y
(151)
B.4 PROOF OF THEOREM 4.2
After writing the program using TP operations, we are ready to prove Theorem 4.2 by taking the infthite-width limit. Applying Theorem 5.4 to Eqs. (145) (151), we have that:
z z - Ezz [Q(Z (t) Zz)
(152)
z Φ((²i)
(153)
z z01 0b (z²²)
(154)
fc = 00 z
(155)
where ~ 1).
C EXTENSIONS OF THEOREM 4.1 AND THEOREM 4.2 TO AGOS WITH MEMORY
So far we have dealt with the case of memoryless adaptive optimizers, and a batchsize of 1, however our results can be trivially extended the more general case. To illustrate this, we now show how the proofs of Theorem 4.1 and Theorem 4.2 can be easily adapted to general AGOs with memory, and a general batchsize. Recall from definition Definition 3.1. 9o-91 91..... 9t E R denote gradients of some scalar parameter w at times 0, 1, ...,, l, general adaptive update can be described by function Q E I + IR such that Awy x (90- 91, ... e)) Concretely, in the case of Adam, Q takes the following form (replacing 31, B2 with 7/, 72 prevent confusion with other indices):
(1-1 Σl=
7 4 gi
Qe(gv. 91........................ def
(156)
11 11-21 Ei
1/2 o2 +
1-73
29

---

Published as a conference paper at ICLR 2023
In the context of optimizing MLP, we write the equivalent of Eq. (2) for general Q function, and a general batchsize for both parameterizations:
VIISL AwA
(157)
Σ 130
d
1-1T
Lip.
Es,
dli 1-1T ", L/g, 1
Es, dli 1-1T
B. B, Lip, {
Q1(
(158)
Q0 Q.
dhi a
1-1T
Σ
dh'g, J3z L'g B,
Σ B.
dhis 3, L33,it)
1-1T
(159)
Bo
Bo
where Eg, denotes summation samples in the minibatch Gr step (i.e if 3. (E,, Ey+ then Eg. "13, = (Cs) ue(E,) + (Ex))) and Q: operates element-wise on the components of its inputs. Note that we have used Definition 3.1 to conveniently remove the factors from inside the Q function, as in Eq. (2). Since is a constant, we will absorb it into the definition of Q from now onward. Note that for any vector the matrix vector product Awie can be implemented as TENSOR instruction (see Definition 5.1):
Aw/u TENSOR( ( (1 , ( ), (1/s,
1 "s (L/,,)
(160)
1 Q. 2 dile 1 C'p.... Edhhy 1 )
(161)
B.
72
3.
3.
where {us, }a, is collection of all vectors evaluated at time on minibatch 3 (and likewise for scalars). We can now conveniently plug in Eq. (160) into the tensor programs in Theorem 4. 1. and Theorem 4.2 to prove more general result.
C.I EXTENSION OF THEOREM 4.1 (NTK)
We now state a general theorem for an AGO with memory and arbitrary batchsize.
Theorem C.1. Let f(g) E R denote an MLP as in Eq. (1) parameterized using the ANTK pa- rameterization described in Section 4.1, where o' is pseudo-Lipschirz Assume layers {w'}/= L trained using an AGO applied on minibatches of arbitrary size, where Qt is pseudo- Lipschitz defined according to Definition 3.1, using a loss function L with a pseudo-Lipschit first derivative. Then, at any step and for any sample E, holds that f. =3 where where:
Ai. KAdp (LE,, (60,). ,,.... ()),,)) (..
KAdp CAdp((EAnp (LEA (EA, } ( ... (6/,,))
(162)
Σ
E a
dh Qe( Σ
zahion 2o 1-1 L'p.... Σ zello, zolai 22a,)2*'11
(163)
9o
c((((E)))
(164)
where the expectation taken over all Z variables at initialization.
Proof. The proof of Theorem C.1 is straightforward extension of the proof of Theorem 4.1. The forward and backward passes for any L are again given by:
72 = (
Au x
(165)
K=0
d ('(ttt O w
(166)
Vigiet di. = [(w1t1
Σ
Awuty)
1+1
dilt1] k'((h)
(167)
==
only with weight updates that are given by:
Aw! n
Q
1 1-17 dhla, #B
C/p, Σ
Lah dh/g
1-1T L's,)
(168)
9
9.
30

---

Published as a conference paper at ICLR 2023
As in the memoryless case, using Eq. (2) and w/ w/ + Ee Awy, (recall wt /nlW), we have that:
3 Auξh!
(169)
V2cisL:
dig Aufat
+
ww Σ
Aufy
53
1 Autdat
(170)
vn
vnn
x=0
5h4
5 Vno(h/+
(h/)
(171)
vn
For any vector v, the matrix vector product Awic can be implemented as a TENSOR instruction:
Awçu = TENSOR (fdii/s, ), (=/)0"), ,,.,, {dhigs, ), (x/2_ 1B, v; (C/,
(172)
-.2000 1 dhlg, 1-1T L'a,... Σ dl =/ 1-1T C/
(173)
P.
P.
where {ug, }a, is collection of all vectors u evaluated at time on minibatch 3, (and likewise for scalars). Hence, may proceed exactly as in the base proof of Theorem 4.1 (i.e expressing the optimization process as a tensor program, applying Definition 5.3 to get the coordinate distributions in the limit, and applying Theorem 5.4). Note that in the concrete case of Adam, get the following function update:
(1Sp, I..,, (Ea, J.a.C ...,, (C',,)) (.',,,,,,)))1 =
(174)
7 Σ zah
(1-1l Ei 11 1-1 Es, 2 2 p. L's,
z
(175)
(1-72) 1-72
Eizo - (E,, 2 14 Lt 1-1
22
C.2 EXTENSION OF THEOREM 4.2 (nP)
Theorem C.2. Let f(() E R denote an MLP as in Eq. (1) with = 2 parameterized using the μ parameterization described in Section 4.1, where ő' is pseudo- -iischhi, Assume layers w² is trained using an AGO with a pseudo-Lipsele Lipschitz function batchsize), using loss function L. with pseudo-Lipschitz first derivative. Then at any step I and for any sample E, it holds that Je - Se where Se can be computed as follows:
Lincciitz,unnctin Q function according to Definition 3.1 (for general
zde+1 z²
Eo' yz2 (Eao) Ls,.... A.
Σ (z (z=)))))"(())))))))))))))))))( Lg.
z
Q:
B,
(176)
z =zz =
f. ) 4
-(c((((())))
(177)
where the expectations are taken over all Z variables (including Ç _ N/ 1)).
Proof. A similarly straightforward application of the Master Theorem Theorem 5.4 to the tensor program in described in Theorem 4.2 together with Eq. (160) in j.P.
D NUMERICAL VERIFICATION
We conduct numerical experiments to verify our results. For both parameterizations, the exact network dynamics at the infinite width limit not tractable in the general case, since the expectations involved do not admit analytical solution (unlike the standard NTK for ReLU networks). Even for
Once again, the loss derivatives C deterministic in Eq. (8)
31

---

Published as a conference paper at ICLR 2023
the ANTK parameterization, the infinite-widee width dynamics cannot be separated a fixed kernel and loss derivative, as with the NTK dynamics for SGD. We therefore must resort to MC simulations to approximate the expectations involved in evaluating the infinite width dynamics in both regimes. We verify Theorem C. and Theorem C.2 by training a ReLU MI.P (I. = 4 for ANTK and 7, for (1) the standard 1.2 loss function, regressing
R.10 gasussian inputs and unit output. For loss to random targets. We train networks with varying widths using Adam with 31 = 0.9, 82 0.99 in full batch mode, on 100 training samples, and run 10 trials per width. We use learning rate of 0.2 and le-4 (where n is the width). In order to account for different initial outputs and loss a
n derivative per weight initialization, we subtract the initialized network output from the output for each sample, such that the output is identically zero initialization for all inputs. To approximate the infinite-widd training dynamics, we approximate the expectation Eq. (174) and Eq. (176) using MC simu lations where we sample the Z random variables from gaussian processes corresponding to the network architecture at initialization. Since the initial loss derivatives are deterministic (given that the outputs are zero), the infinite width dynamics can be approximated without actually constructing a network. To compare the evolution of the finite vs infinite architectures, we evaluate the output at each iteration random inputs. Our results summarized in Fig. 2 and Fig. 3. As expected, the width increases the training dynamics converge to that of the infinite dynamics.
5
20
25
15
26
2
25
30
35
15
10
0.9
380
3.
1.5
1.0
1.2
-31
-22
-2.
40
-24
-02
-35
-05
-25
4ril
p/
25
25
28
154007
(a)
(b)
(c)
Figure 2: Training dynamics of finite and infinite-width networks in the ANTK parameterization. We train networks of widths 64 (a), 512 (b), 7000 (c). and track the outputs for 4 random inputs (one per row) at each iteration as the network trains. We compute the output distribution over 10 independent runs for each network, and compare with the infinite-widd dynamics (black curve). As the width grows, the network function converges to that of the infinite-wide width dynamics captured in Eq. (174)
32

---

Published as a conference paper at ICLR 2023
odeds
15.0011
with 10
vet 7800
-02
-3
-04
-04
-35
-05
12
4.2
-14
-14
35
89
35
150
412
41-
1.5
1.0
40
35
36
-2
2
25
022
15
50
09
-3
-02
-04
-05
(a)
(b)
(c)
Figure 3: Training dynamics of finite and infinite-width networks in the μ parameterization. We train networks of widths 64 (a), 512 (b), 7000 (c) and track the outputs 4 random inputs (one per row) at each iteration as the network trains. We compute the output distribution over 10 independent runs for each network, and compare with the infinite wiith dynamics (black curve). As the width grows, the network function converges to that of the infinite-wiet width dynamics captured in Eq. (176)
33


Reviews:
[{"review_id": "xP_d7SQdntd", "reviewer": "Reviewer_1g2Y", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "This paper studies two infinite-width limits (kernel and feature learning limits) of fully-connected neural networks (MLPs) under adaptive gradient-based optimization. Both results generalize that of MLPs trained by a non-adaptive way. In addition, the authors modify a framework (called Tensor Program) that allows to express the adaptive gradient processing as well as the convergence guarantee.\n", "summary_of_the_review": "This paper extends infinite-width limits to the adaptive optimization setting, which has never been discovered before. But, it seems that a more in-depth comparison/analysis of the non-adaptive results would make the contributions much stronger. The writing quality needs to be improved for better readability.", "strength_and_weaknesses": "Strength:\n\n- This paper extends two infinite-width limits of MLPs for adaptive gradient-based optimization, which have not been studied before. The proposed results can generalize both kernel and feature learning limits of non-adaptive settings. It is interesting that the adaptive NTK regime is still data-independent like the non-adaptive NTK.  \n\n- They provide a general Tensor Program framework for the adaptive optimization setting with O(1/sqrt(n)) convergence rate guarantees. \n\nWeakness:\n\n- Adaptive optimizations have various benefits over the SGD (e.g., fast convergence, better generalization, etc). Can this point of view be captured by the proposed adaptive infinite-width limits? It would be great if more clear motivations and meaningful results of adaptive infinite-width limits are provided. \n\n- The paper focuses on simple settings (e.g., memoryless adaptive gradient descents, update of parameters in a single layer). But, it is unclear whether the proposed results hold for more general settings. \n", "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": "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": "QEyOx16ag9", "reviewer": "Reviewer_k9hR", "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": "", "summary_of_the_paper": "In this work the authors tackle the question of the infinite-width limit of neural networks trained using adaptive optimizers (like ADAM). They show that if the step taken in the optimizer depends on a non-linear, scale-invariant function of the previous step, then the NTK can be computed by applying said function to $\\frac{\\partial f_{train}}{\\partial\\theta}\\nabla_{\\theta}\\mathcal{L}_{train} $, and then multiplying by $\\frac{\\partial f_{out}}{\\partial\\theta}$. An analogous result holds for the $\\mu$-parameterization (the unique family of infinite-width feature learning initializations).", "summary_of_the_review": "The paper does an excellent job of building intuitive and rigorous arguments for memoryless adaptive optimizers. However, I feel that this is not a significant enough contribution on its own, and the case of non-memoryless optimizers like ADAM should be explicitly covered if they are in fact trivial extensions of the framework as suggested by the authors. I did not have the time nor expertise to rigorously check the proofs. I also suggest that numerical evidence of convergence to the theoretical limits would be very helpful to readers, especially those without strong theoretical backgrounds.", "strength_and_weaknesses": "The arguments for the memoryless case are very clean. In addition, these additions to the TP framework make it more flexible and can give insight in many of the more complex practical cases which arise when training large models. The authors make a good mix of intuition-building arguments and links to the more rigorous ones.\n\nOne point of clarification: in the mean-field limit, does $\\epsilon$ have to go to $0$ to get all the scaling arguments right? Is there an intuitive reason that this doesn't mess up the Gaussianity arguments via divisors which are random variables which can have significant probability mass around $0$? The final limits seem to be $\\epsilon$-independent.\n\nOne of the big weaknesses is that the promise of extending the analysis to the case of an adaptive optimizer with memory (depending on gradients $g$ further back in time) is not clearly made. I believe without that, the result in the memoryless case is interesting but not sufficiently significant for an acceptance rating. It would also be good to have a discussion on how the total time $t$ must scale with $n$ for the arguments to hold.\n\nI think the paper would be significantly strengthened if there was an explicit formulation and proof of at least the NTK for the ADAM algorithm. From the comments in the paper these seem like results that are already completed, or easily derivable for the framework. This paper is positioned to become a standard reference for NTK/$\\mu$-parameterization for adaptive algorithms, and it would be strange for the most popular adaptive algorithm to not be explicitly covered.\n\nDue to the limits of my technical ability and the review timescale, I have not been able to rigorously check the main proofs in the TP framework. I followed along with the arguments in the main text and they seem sound. I believe that there needs to be numerical evidence that the proposed scaling limits hold - not to replace the proofs, but to supplement them and decrease the possibility that some key point in the complex machinery did not go awry.\n\nAs a small note, on page 3 it is unclear why the labeling of the gradients $g$ has even time indices only.\n\nUpdate: most of the concerns above have been addressed; my updated review score reflects those changes.", "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": "Not applicable"}, "extra_text": {}}, {"review_id": "DBjMlgCEH8", "reviewer": "Reviewer_8D7G", "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": "", "summary_of_the_paper": "This paper studies the infinite width limit of Multi-layer perceptrons, in particular the existence of feature learning. \nIn contrast to previous work, it studies adaptive optimizers instead of Gradient Descent, but results are roughly the same: NTK limit still results in lazy training (and no feature learning), while mu-parameterization supports feature learning.\nUnfortunately, no experiments on real (or even synthetic) data are provided in support of theoretical claims (It has been addressed in the rebuttal).\n", "summary_of_the_review": "I would be much more convinced of the significance of this work if empirical validation of the theoretical results was provided (It has been addressed in the rebuttal).\n", "strength_and_weaknesses": "Strengths:\n\n- Infinite width limit of MLP studied for the first time with adaptive optimizers.\n\n\nWeaknesses:\n\n- No empirical data is provided in support of theoretical findings.\n\n- NTK limit is provided (theorem 4.1) in memoryless case only, which is of very limited interest.\nThe authors state: \"we maintain that our analysis is applicable in the more general setting without introducing any additional arguments\", but no proof is provided as far as I understand.\n\n(It has been addressed in the rebuttal)\n", "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": "42xPUvJoxo", "reviewer": "Reviewer_7U4M", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "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": "", "summary_of_the_paper": "This submission analyzes the dynamics of adaptive optimization of arbitrary neural network architectures in the infinite-width limit, and as a by-product,  a new tensor program framework is proposed.", "summary_of_the_review": "The submission adds some new tools in the previously developed framework (Tensor programs) for helping analyze the dynamic of $\\infty$-width neural network trained by adaptive optimizers. The only disadvantage is that the theory does not give more insight into the difference between SGD and adaptive optimizers (e.g., Adam).\n\nOverall, I still hold a positive score for this work since it is a necessary expansion of Tensor programs.", "strength_and_weaknesses": "#### Strength\n- Close the gap between adaptive optimizer and NTK analysis.\n- Extend the guarantee in NTK and feature learning to the non-linear cases in the sense of the Tensor program language.\n\n\n#### Weaknesses\n- The adaptive optimizer considered in the analysis is memoryless and one-dimensional (Thm.4.1-4.2) given batch size is one, which is far from the practical case.\n- The theoretical results seem to be a generalized version of SGD, and we cannot see the difference between AGO and SGD from the theory. In fact, AGO and SGD behaviors differently during training, and we cannot treat SGD as a particular case.", "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": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster