A DAPTIVE O PTIMIZATION IN THE ∞ -W IDTH L IMIT
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent works have developed detailed understanding of large neural networks’ behaviors via their inﬁnite-width limits, e.g., the neural tangent kernel (NTK) and the feature learning ( µ ) 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.
Speciﬁcally, for deep
MLPs, we derive the NTK and µ parametrizations as well as their inﬁnite-width limits.
We ﬁnd 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 ﬁxed by the µ 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.
I NTRODUCTION
Inﬁnite 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.
Speciﬁcally, 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 it cannot account for data dependent feature learning, a staple of deep learning in practice.
In contrast, the µ limit (of which the well-known mean ﬁeld limit is a speciﬁc case in 1-hidden-layer perceptrons) refers to a regime of training where the weights adapt to the data during training in a 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 to initialization scale and learning rate determine the type of inﬁnite-width limit one can associate with a trained neural network.
Notably, the µ parameterization was identiﬁed as the unique parameterization which gives rise to "maximal" feature learning dynamics in the inﬁnite-width 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 identiﬁcation and prescription of two types of inﬁnite-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 a fundamentally different dynamics for adaptive optimization, referred to as the adaptive neural tangent kernel (ANTK) regime.
In this limit, the training dynamics can no longer be described by kernel gradient descent, since the kernel function itself depends non-linearly on the loss derivative.
Our results lay a clear path to theoretically analyze the implicit biases of AGO in the inﬁnite-width limit.
Key Technical Contribution Analyzing the dynamics of adaptive optimization of arbitrary neural network architectures in the inﬁnite-width limit presents a 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 down the relevant neural network computation (e.g. the ﬁrst forward pass in the NNGP case) as a 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 ﬂexible, 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 functionities, and proving a new master theorem which enables our analysis.
While we present a simple application of our new framework on MLPs in Theorem 4.
and Theorem 4.2, it is applicaple in a much wider setting, including most practical architectures and algorithms.
As an additional technical contribution, we prove a O ( n − 1 / 2 ) (where n represents the width) convergence rate guarantee for all variables produced by the program, which might be of independent interest.
Our Contributions
This paper presents the following major contributions:
We present the ﬁrst rigorous inﬁnite-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.
We develop a new tensor program framework along convergence rate guarantees, unlocking the inﬁnite-width analysis of adaptive optimization in an architecturally universal sense.
Paper Organization
This paper is organized as follows: We survey related work in Section 2.
In Section 3 we set up preliminaries and notations used extensively in Section 4.
In Section 4 we illustrate ANTK and µ 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.
R ELATED W ORK
A large body of literature exists on both the kernel (NTK) limit Arora et al.
; Jacot et al.
Lee et al.
; Yang (2020c); Yang & Littwin (2021) and the mean ﬁeld limit forlayer neural network Chizat & Bach (2018); Mei et al. (2018b); Nguyen & Pham (2020); Rotskoff & Vanden-Eijnden (2018); Sirignano & Spiliopoulos (2020).
Various papers describe the kernel and feature learning regimes more generally without taking an inﬁnite-width limit.
Chizat et al.
(2019) describes the "lazy training" regime in arbitrary differentiable programs, and is controlled by a single parameter α which scales the output.
It is shown that when α is large, the weight need only move slightly to ﬁt the training data, and network essentially performs kernel learning.
Many papers Allen-Zhu et al.
; 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 inﬁnite-width limit, and others derive ﬁnite width corrections to the NTK for ﬁnite 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 inﬁnity.
This way, kernel and feature learning behaviour are separated by the parameterization employed at initialization, and not width or training time.
TPs, ﬁrst introduced in Yang (2019) and expanded upon in Yang (2020a;b), were developed as a theoretical framework to analyze the inﬁnite-width limits of any architecture expressible in the TP language, in an attempt to rid the per architecture analysis prevalent in the literature Alemohammad et al.
; Du et al.
Hron et al.
(2020); Littwin et al.
(2020b).
Yang & Hu (2020) deﬁned a natural space of neural network parametrizations (abc-parametrizations), and classiﬁed all resulting inﬁnite-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 identiﬁed as the "optimal" parameterization for arbitrary architectures in which all layers learn features, and was later heuristically extended to AGOs Yang et al.
Unrelated, AGOs Duchi et al.
; Kingma & Ba (2015); Zhou et al.
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.
Liu et al.
; Zhang et al.
Crucially, no previous work has yet developed a theory for inﬁnite-width neural network trained with AGOs.
P RELIMINARIES
Adaptive Optimizers Generically, if g 0 , g 1 , ..., g t ∈ R denote the gradients of some scalar parameter w ∈ R at steps 0 , 1 , ..., t , an adaptive update ∆ w t = w t +1 − w t at step t takes the form ∆ w t = − η m √ v + (cid:15) where η is the learning rate and m and v are both functions of the past gradients g 0 , . . .
, g t .
For example, in Adam, m and v are the exponential moving averages of g ( i ) and g 2( i ) .
Here, we consider an even more general notion of adaptive updates, encompassing all modern AGOs.
Deﬁnition 3.1.
We say an update ∆ w t ∝ Q t ( g 0 , g 1 , ..., g t ; (cid:15) ) to a weight w at time t is adaptive if it is proportional (up to a constant factor) to a function Q t : R t +
→ R such that ∀ c (cid:54) =0 , Q t ( cg 0 , cg 1 , ..., cg t ; c(cid:15) )
=
Q t ( g 0 , g 1 , ..., g t ; (cid:15) ) .
Moreover, if Q t ( g 0 , g 2 , ..., g t ; (cid:15) ) =
Q ( g t ; (cid:15) ) (only depends on g t ), then we furthermore say ∆ w t is memoryless .
To maximize clarity, we focus on the simpler case of memoryless adaptive updates.
For example, in
Adam this implies setting β 1 , β 2 = 0 .
This simpliﬁcation will already highlight the key differences between the adaptive and non-adaptive case, and we maintain that our analysis is applicable in the more general setting without introducing any additional arguments.
MLPs and ABC(D)
Parameterization We use a 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 ξ ∈ R d in , weight matrices W L +1 ∈ R n , { W l } Ll =2 ∈ R n × n , W 1 ∈ R n × d in and an activation function φ which we assume has a pseudo lipschitz ﬁrst derivative, the output f ( ξ ) ∈ R is given by:
[Equation]
We adopt the ABC parameterization convention from Yang & Hu (2020).
Namely, for any layer l , each weight matrix is parameterized using W = n − a l w l where w l are the learnable weights, which are initially sampled iid from a normal distribution N (0 , n − 2 b l ) .
Finally, the learning rate is parameterized using ηn − c l where we plug η = 1 for simplicity.
In this paper, we assign speciﬁc values to { a l } l , { b l } l , { c l } l for the ANTK and µ parameterizations.
Additionally, we will parameterize the (cid:15) parameter in the AGO with (cid:15) l = n − d l (cid:15) , where (cid:15) > 0 .
The per-layer scaling for (cid:15) l will turn out to be crucial to prevent the adaptive gradients from collapsing to either 0 or a step function as n → ∞ .
We summarize the two parameterizations in the following table:
Parameterization a
[Equation]
ANTK (cid:26)l > else (cid:26)
L + 1 > l > else (cid:26)
L + 1 > l > else µ   − 1l =
l = L + else (cid:26)
L + 1 > l > else (cid:26)
L + 1 > l > else and µ parameterizations.
Representing (pre)activation vectors via random variables As we 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 to see).
Hence a vector’s behavior can be tracked via a random variable that reﬂects the distribution of its entries.
Concretely, if x ∈ R n is one such vector, then we write Z x for such a random variable, such that x ’s entries look like iid samples from Z x .
When x is scaled to have typical entry size independent of n , 1 then Z x can be taken to be a random variable independent of n as well.
In general, given two such vectors x, y ∈ R n , their random variables Z x and Z y will be correlated, in such a way that lim n →∞ x yn = E
Z x Z y .
Generally, inferring with initialized networks entail computing expectations with gaussian Z variables, which take a relatively simple form.
However, a fundamental question is how the Z variables evolve during training, which we address next.
A DAPTIVE O PTIMIZATION OF AN MLP
In the following section we illustrate the inﬁnite-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 a cleaner presentation, we assume the ﬁrst and last layers are ﬁxed, however our results are easily extended to the general case.
In our setup we assume the network is trained using an AGO according to Deﬁnition 3.1,and a batchsize of 1.
notations Slightly abusing notation, we use subscripts to denote both the step index t , and coordinates of vectors with α, β .
We assume ξ t is a training sample fed to the neural network at step t (starting from ξ 0 ).
and we use y t ( ξ ) 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 y t =
y t ( ξ t ) and y =
y 0 ( ξ 0 ) ).
We use ˜ y t
= y t ( ˜ ξ ) to express the dependency of y on an arbitrary input ˜ ξ at step t .
We will also denote ∆ y t ( ξ ) =
y t +1 ( ξ ) − y t ( ξ ) and δy t ( ξ ) = √ n ∆ y t ( ξ ) .
We assume the network is trained using a generic loss function L , with a loss derivative L (cid:48) t = ∇ f t L t .
We use the notation dh l based on context : for ANTK parameterization, we set dh l ( ξ ) def = √ n ∂f∂h l ( ξ ) ∈ R n , whereas for µ parameterization, we set dh l ( ξ ) def = n ∂f∂h l ( ξ ) ∈ R n .
This context dependent notation is convenient since it insures that the components of dh l ( ξ ) are roughly in the order of Θ(1) for both parameterizations.
Finally, we use ˚ • to denote the inﬁnite-width limit of a (possibly random) scalar • (i.e lim n →∞ ˜ f t = ˚˜ f t ).
Using the above notation, we can express the gradient of any intermediate layer w l at step t for both parameterizations by 1 n dh lt x l − 1 (cid:62) t
L t .
Using
Deﬁnition 3.1 the adaptive weight update ∆ w l for both parameterizations is given by:
[Equation]
where the function Q is applied elementwise on the matrix dh lt x l − 1 (cid:62) t
L
[Equation]
∈ R n × n .
For the remainder of the paper we will suppress the explicit the dependency on (cid:15) and simply absorb it into Q .
4. T HE ANTK L IMIT
In the NTK limit, intuitively, the weights of the neural network move by a negligible amount during training, such that the network function may be approximated by its ﬁrst 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 t : ∆ ˜ h lt for any layer l will be of order Θ( n − 12 ) .
By deﬁnition ˜
h lt +1 =
˜
h lt + ∆˜ h lt , hence the coordinates of ˜
h lt for any layer l do not change in the limit, and, for any input, the coordinate distributions remain constant ∀ l ∈ [1 ,l ] Z ˜ h
lt = Z ˜ h l , Z d ˜ h lt = Z d ˜ h l .
Instead of training f , we consider training the ﬁrst order linearization of f , denoted by f lin , around its initial parameters.
The function updates ∆ ˜ f lin t are given by
[Equation]
Under the ANTK parameterization, as with SGD, training f lin and f using AGO is equivalent.
The following theorem describes the evolution of ˚˜ f t exactly:
Theorem 4.1.
Let f ( ξ ) ∈ R denote an MLP as in Eq.
(1) parameterized using the ANTK parameterization described in Section 4.1, where φ (cid:48) is pseudo-Lipschitz.
Assume layers { w l } Ll =2 are trained using a memoryless AGO with a pseudo-Lipschitz function Q according to Deﬁnition 3.1 and a batchsize of 1, using a loss function L with a pseudo-Lipschitz ﬁrst derivative.
Then, at any step t and for any sample ˜ ξ , it holds that ˜
f
t a.s → ˚˜ f t where ∆˚˜ f t = −K Adp ( ξ t , ˜ ξ
| ˚ L (cid:48) t ) , where:
[Equation]
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 ˜ f t , and by extension the loss derivatives L (cid:48) t 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 2 .
Second, it is straightforward to show that by setting Q ( x ) = x we would get the SGD equivalent (when setting the learning rate to be 1 ) of Eq.
(162), which takes the form ˜ f sgd t ≈ − (cid:80)
Ll =2 dh l (cid:62) t d ˜ h lt n
[Equation]
For SGD, one may naively apply the the law of large numbers argument (LLN) and derive the inﬁnite-width limit under plain SGD: ∆˚˜ f sgd t a.s → −K ( ξ t , ˜ ξ )˚ L (cid:48) t where K is the NTK function deﬁned as:
[Equation] [Equation]
Hence Theorem 4.1 is a generalization of the well known NTK limit.
At a glance, the transition from Eq.
(3) to its inﬁnite-width counterpart seems like a 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 ﬁrst 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 a matrix (more generally a 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 a 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 K Adp is nonlinear in the loss derivative L (cid:
) t , 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 inﬁnite-width limit, which we analyze next.
4. F EATURE L EARNING WITH µ PARAMETERIZATION
We now turn to analyzing the inﬁnite-width training dynamics under µ parameterization.
Fundamentally, each weight update ∆ w lt will cause each preactivation vector ˜
h
lt to change entrywise by something of order Θ(1) , and the coordinate distributions at the limit will evolve non-trivially.
Generally, the dynamical equations equivalent of an inﬁnite-width neural network in the feature learning regime (using µ or otherwise) is much more complex to derive for deep networks.
Although our new
TP formulation discussed in Section 5 provides a complete set of tools to tackle most architectures, we will be content with illustrating the main points using a 2 hidden layer MLP where only the middle layer weights w 2 are trained.
Using Eq.
(2), we can express ˜ h
2 t +1 , ˜ x 2 t +1 , ˜ f
t +1 using:
[Equation]
Note that under µ the coordinates of √ nw 3 are randomly distributed according as N (0 , 1) , hence we expect ∆ ˜ f t to be Θ(1) .
Due to the Q function applied to the gradient, the components of the
R n × n matrix Q ( dh 2 t x t L (cid:48) t ) are not vanishing as n → ∞ , and the update 1 n Q ( dh 2 t x 1 (cid:62) t L (cid:48) t )˜ x 1 is by consequence generally not vanishing as well.
Since the updates ∆˜ h 2 t are non vanishing, the feature vector ˜ x 2 t 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 where Q = Identity, which results in the update equation ˜ h
2 t +1 =
˜
h
2 t − ˜ x 1 (cid:62) x 1 ( ξ t ) n dh 2 t L (cid:48) t , and can be expressed purely using operations between vectors.
For general nonlinear Q functions, we must deal with a matrix - vector multiplication as in Eq.
This implies that unlike with SGD where we must reason about how a ﬁnite collection of R n vectors behave in the limit, we must now reason about the behaviour of R n × n matrices (see Section 6 for further discussion on this).
The following theorem describes the evolution of ˚˜ f t under µ exactly:
Theorem 4.2.
Let f ( ξ ) ∈ R denote an MLP as in Eq.
(1) with L = 2 parameterized using the µ parameterization described in Section 4.1, where φ (cid:48) is pseudo-Lipschitz.
Assume layers w 2 is trained using an AGO with a pseudo-Lipschitz function Q function according to Deﬁnition 3.1 and a batchsize of 1, using a loss function L with a pseudo-Lipschitz ﬁrst derivative.
Then at any step t and for any sample ˜ ξ , it holds that ˜
f
t a.s → ˚˜ f t where ˚˜ f t can be computed as follows:
[Equation]
where the expectations are taken over all Z variables (including ζ d = N (0 , 1) ).
From Theorem 4.2 it is clear in what sense µ parameterization enables feature learning in the limit: from Eq.
(8) the random variable encoding the updates to the hidden representation Z ˜ h 2 t is of order Θ(1) for non degenerate Q and φ functions, and is generally no longer gaussian at steps t > 0 , allowing the neural network to learn data dependent features.
Once again, substituting Q ( x ) = x would default the equations in Eq.
(8) back to those of SGD with an appropriate step size, hence our
Theorem 4.2 generalizes feature learning with SGD.
A T ENSOR P ROGRAM FOR A DAPTIVE O PTIMIZERS
In Section 4 we have derived two types of limits in a 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 in a 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 inﬁnitewidth limits of trained neural networks by explicitly expressing the updates to the network internal representation and output at step t , 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 a framework designed to automate these computations, while providing theoretical guarantees.
Generally, computations involving adaptive optimization of most architectures contain a few repeating operations (i,e 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 R n vectors or scalars), and edges represent operations performed on the variables.
As a 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 afﬁne transformation in each layer represent an edge between nodes.
We give a more formal description of our framework in the following.
5. N E ⊗ ORT P ROGRAMS
Deﬁnition 5.1.
A N E ⊗ ORT program is a sequence of R n -vectors and R -scalars inductively generated via one of the following ways from an initial set C of random scalars, V of random R n vectors, and a set W of random R n × n matrices (which will be sampled with iid Gaussian entries in Setup 5.2).
Concretely, using weights W and some pseudo-lipschitz function ψ : R k ( r +1)+ l → R for k, l, r ∈ N , the program generates new vectors and scalars from previously generated vectors x = { x 1 , ..., x k } ∈ R n and scalars Θ = { θ 1 , θ 2 , ..., θ l } ∈ R by the following instructions (using the notation x i = { x 1 i , ..., x ki } ):
T ENSOR Generates a vector x ∈ R n by x α = 1 n r (cid:80) nβ 1 ,...,β r
[Equation]
1 ψ ( x α , x β 1 , ..., x β r ; Θ) .
T ENSOR M OMENT Generates a scalar θ ∈ R by θ =
[Equation]
M ATMUL Generates a vector x ∈ R n by x = W ¯ x or x = W (cid:62) ¯ x where ¯ x ∈ x , W ∈ W .
Let us unpack Deﬁnition 5.1.
We can think of the T ENSOR operation as a generalized version of the standard pointwise nonlinearity which acts on vectors (or tensors where only one dimension increases to inﬁnity, akin to the N ONLIN instruction in Yang (2020b).
Instead, the T ENSOR instruction applies a pointwise nonlinearity ψ to a tensor of rank r + 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 to vectors by setting r = 0 .
The T ENSOR M OMENT operation allows us to fully contract a tensor of rank r + 1 to a scalar.
Finally, the M ATMUL operation is copied over from Yang (2020b), and implements a standard linear layer.
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 W αβ ∼ N (0 , σ 2 W /n ) for some variance σ 2 W associated to W , independent of other W (cid:48) ∈ W ; 2) for some multivariate Gaussian Z V =
{ Z x : x ∈ V} ∈ R V , we sample the initial set of vectors V like { x α : x ∈ V} ∼ Z V iid for each α ∈ [ n ] .
For each initial scalar θ ∈ C , we require ∀ p> 0 , n 1 − p ( θ − ˚ θ ) 2 a . s .
−−→ 0 for some deterministic ˚ θ ∈ R .
Note that the initial set of vectors V are assumed to be gaussian in R n .
In a typical neural network training scenario, the initial vectors correspond to input/output layer weights and biases at initialization, and the initial matrices correspond to hidden layer weights at initialization, so their Gaussian sampling reﬂects their Gaussian initialization.
Example A program encoding the ﬁrst forward,backward and adaptive update (Eq.
(7)) using µ parameterization is provided in Table
Expression
Op type
Implementation
[Equation]
∆˜
[Equation]
[Table]
In the above, a, b, c, θ ∈ R represent inputs to some function ψ implementing a T ENSOR or a
T ENSOR M OMENT instruction.
5. T HE M ASTER T HEOREM
We can guarantee certain properties hold for vectors and scalars generated by a tensor program in the inﬁnite-width limit.
In short, any generated scalar θ will almost surely converge to a deterministic limit ˚ θ as n → ∞ , at a rate of O ( 1 √ n ) .
For any generated vector x ∈ R n , the coordinates of x will approach iid samples from some distribution.
Adopting the notation from Section 3, we denote by
Z x a random variable distributed like the coordinates of x as n → ∞ .
The following constructs the random variable Z x for every vector x and a deterministic scalar ˚ θ for every scalar θ in the program, where we assume x = { x 1 , ..., x k } , Θ = { θ 1 , ..., θ l } are previously generated vectors and scalars, and we use the abbreviated Z x to denote the set of k random variables { Z x i } ki =1 for all x i ∈ x .
Deﬁnition 5.3 ( Z x and ˚ θ ) .
We recursively deﬁne Z x def =
ˆ Z x + ˙ Z x for each vector x and ˚ θ for each scalar θ as follows:
ZI NIT If x ∈ V , then Z x is deﬁned as in Setup 5.2.
We also set ˆ Z x def =
Z x and ˙ Z x def = 0 .
ZT
ENSOR If x is generated by T ENSOR (see Deﬁnition
1), then Z x def = f ( Z x ) where f ( ζ ) def = E Z x 1 ,...,
Z x r [ ψ ( ζ, Z x 1 , Z x 2 , Z x r ; ˚Θ)] with Z x i being iid copies of Z x .
ZT ENSOR
M
OMENT If θ is generated by T ENSOR M OMENT (see Deﬁnition 5.1), then ˚ θ def =
E Z x ,Z x 1 ,...,Z x r [ ψ ( Z x , Z x 1 , ..., Z x r ; ˚Θ)] .
Here ˚Θ = ˚ θ 1 , . .
, ˚ θ l are deterministic, so the expectation is taken over Z x , Z x 1 , ..., Z x r , where { Z x j } rj =1 are r iid samples drawn from the same distribution as Z x .
ZM AT M UL
If x = W ¯ x for ¯ x ∈ x , W ∈ W then Z Wx def
= ˆ Z Wx + ˙
Z Wx , where:
ZH AT ˆ
Z
Wx is a Gaussian variable with zero
mean
Let
V
W denote the set of all vectors in the program of the form Wy for some y .
Then { ˆ Z Wy : Wy ∈ V W } is deﬁned to be jointly Gaussian with zero mean and covariance Cov (cid:16)
ˆ Z Wx , ˆ Z Wy (cid:17) def = σ 2 W E Z x Z y , for any Wx, Wy ∈ V W . .
Furthermore, { ˆ Z Wy : Wy ∈ V W } is mutually independent from { ˆ Z v : v ∈ V ∪ (cid:83)
¯ W (cid:54) =
W V ¯ W } , where ¯ W ranges over W ∪ { A (cid:62) :
A ∈
W
ZD OT
We can always unwind Z x
=
Φ( · · · )
, for some arguments (
=
{ ˆ Z W (cid:62)
y i } ki =1 ,
ˆ Z z i } ji
=
; { ˚
θ i }
li
=
) , z i
(cid:54)∈
V W (cid:62) (where V W (cid:62) is deﬁned in 5.3), and deterministic function Φ : R k + j + l → R .
Deﬁne ∂Z x /∂ ˆ Z W (cid:62)
y i def =
∂ i Φ( · · · ) .
Then we set ˙
Z Wx def = σ 2 W (cid:80) ki
=1 Z y i E ∂Z x ∂ ˆ
Z
W (cid:62) yi
There is some nuance in this deﬁnition, so see Remark A.1 and A.2.
The following theorem ties the symbolic nature of the Z s to the analytic nature of a Tensor Program.
Theorem 5.4 (N E ⊗ ORT Master Theorem) .
Fix a N E ⊗ ORT program initialized accordingly to
Setup 5.2.
Assuming all nonlinearities are pseudo
Lipschitz in all arguments, then
For any collection of vectors x = { x 1 , ..., x k } and scalars Θ = { θ 1 , ..., θ l } in the program, and for any pseudo-Lipschitz ψ :
R k ( r +1)
+ l → R , as n → ∞ :
[Equation]
where { Z x j } r +1 j =1 are r + 1 iid samples drawn from the same distribution as Z x , and
Z x = { Z z 1 , ..., Z x k } are deﬁned according to Deﬁnition 5.3.
Any scalar θ in the program tends to ˚ θ almost surely such that ∀ p> 0 , n 1 − p ( θ − ˚ θ ) 2 a .
s .
−−→ 0 , where ˚ θ is as deﬁned in Deﬁnition 5.3.
Theorem 5.4 along with Deﬁnition 5.3 provide a general tool set to analyze adaptive (and standard) training of neural networks in the inﬁnite-width limit, as long as the computational graph expressing the training process can be implemented in a N E ⊗ ORT program.
Moreover, Theorem 5.4 provides a universal O ( n − 1 / 2 ) asymptotic rate of convergence for all scalars produced by the program.
P ROOF S KETCH
N E ⊗ ORT programs equipped with Theorem 5.4 provide the main tool to proving Theorem 4.1 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 N E ⊗ ORT program, mechanically compute the Z variables according to Deﬁnition 5.3, and apply Theorem 5.4 to compute the limit (see proofs in appendix Appendix C).
What remains is to prove Theorem 5.4 using a strategy which we now outline.
In a program, all vectors can be collected into an n × M matrix V where n is the dimension of each vector, and M is the total number of vectors.
The Master Theorem can be interpreted as saying that each row of V (i.e., the slice for each α ∈ [ n ] ) is roughly an iid sample from some distribution D on R M (which can be derived via the calculus of Z random variables as in Deﬁnition 5.3).
Speciﬁcally,
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 (cid:48) of iid samples from D , as measured by applying arbitrary pseudo-Lipschitz “test function ψ ” to both sides and taking averages.
Core Insight Our core insight here is that V is in fact similar to V (cid:48) in a stronger sense without needing to refer to any test function ψ : There is a “small” matrix δV of the same size as V such that
V − δV is distributed exactly the same as V (cid:48) .
In general, if this happens, then we say V is equivalent to V (cid:48) .
The deﬁnition of “small” roughly means that each entry of δV has typical size O ( n − 1 / 2 ) .
Then, to recover Theorem 5.4, we just note that the test function ψ is (by assumption) smooth enough that δV contributes a vanishing amount to the LHS of Eq.
To prove this core insight, there are two parts.
Part 1 We show that, in any N ETSOR (cid:62) program (i.e., a program with no scalar variables and no T ENSOR operation), V is equivalent to V (cid:48) .
This can be done by re-analyzing the proof of the
N ETSOR (cid:62) Master Theorem in (Yang, 2020b) in a fairly straightforward way.
Part 2 For any N E ⊗ ORT program π (the subject of our work here), we construct a parallel
N ETSOR (cid:62) program and show, by induction,
that the vectors of the two programs are 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 Deﬁnition 5.3(ZT ENSOR M OMENT ), and 2) all T ENSOR operations by N ONLIN operations, as computed in Deﬁnition 5.3(ZT ENSOR ).
In this induction, we need to prove and use a lemma that, in the simplest case as an illustration, says the following: For any pseudo-Lipschitz function ψ : R k
→ R and random vector x ∈ R n with iid standard Gaussian entries, the following two tensors T and T (cid:48) are equivalent: 1) the tensor T with entries T β 1 ...β k = ψ ( x β 1 , . . .
, x β k ) , and 2) the tensor T (cid:48) with entries T (cid:48) β 1 ...β k
= ψ ( x 1 β 1 , . . .
, x kβ k ) where x 1 , . . .
, x k are iid copies of x .
The proof of this lemma interestingly requires Baranyai’s theorem, a classical theorem from the theory of combinatorial design.
C ONCLUSION
Adaptive optimizers are a staple in the modern deep learning toolkit, an are a necessary ingredient in most large scale neural network training.
In this work, we have derived adaptive counterparts to the
NTK and µ 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 a path to study the implicit bias of adaptive optimizers by studying their evolution equations in the inﬁnite-width limit.
A F ULL P ROOF OF T HEOREM 5.
In this section we provide the proof for Theorem 5.4, restated:
Theorem 5.4 (N E ⊗ ORT Master Theorem) .
Fix a N E ⊗ ORT program initialized accordingly to
Setup 5.2.
Assuming all nonlinearities are pseudo-Lipschitz in all arguments, then 1.
For any collection of vectors x = { x 1 , ..., x k } and scalars Θ = { θ 1 , ..., θ l } in the program, and for any pseudo-Lipschitz ψ :
R k ( r +1)+ l → R , as n → ∞ :
[Equation]
where { Z x j } r +1 j =1 are r + 1 iid samples drawn from the same distribution as Z x , and
Z x = { Z z 1 , ..., Z x k } are deﬁned according to Deﬁnition 5.3.
Any scalar θ in the program tends to ˚ θ almost surely such that ∀ p> 0 , n 1 − p ( θ − ˚ θ ) 2 a .
s .
−−→ 0 , where ˚ θ is as deﬁned in Deﬁnition 5.3.
Remark A.1 (Partial derivative) .
The partial derivative in 5.3 should be interpreted as follows.
By a simple inductive argument, Z x for every vector x in the program is deﬁned uniquely as a deterministic function ϕ
( ˆ Z x 1 , . .
, ˆ Z x k ) of some x 1 , . . .
, x k in V or introduced by M ATMUL (notationally, we are suppressing the possible dependence on limit scalars ˚ θ 1 , . .
, ˚ θ l ).
For instance, if in a program we have A ∈ W , v ∈ V , y =
Av, x
= A (cid:62) y , then Z x = ˆ Z x + ˆ Z v , so ϕ is given by ϕ ( a, b )
= a + b .
Then ∂Z x /∂ ˆ Z x i def =
∂ i ϕ
( ˆ Z x 1 , . . .
, ˆ Z x k ) , and ∂Z x /∂ ˆ Z z def = 0 for any z (cid:54)∈ { x 1 , . . .
, x k } .
Note this deﬁnition depends on the precise way the program is written, not just on the underlying mathematics.
For example, if y, z ∈ V and x = φ ( W ( y + z )) , then Z x = φ ( ˆ Z W ( y + z ) )
so that ∂Z x /∂ ˆ
Z Wy = ∂Z x /∂ ˆ Z Wz = 0 .
If instead, we have x = φ
( Wy + Wz ) , then Z x = φ ( ˆ Z Wy + ˆ Z Wz ) so that ∂Z x /∂
ˆ Z W ( x + y )
= 0 .
However, in both cases, ˙ Z W (cid:62) x
= ( Z y + Z z ) E φ (cid:48)
( ˆ Z W ( y + z ) ) .
Remark A.2 (Partial derivative expectation) .
The quantity E ∂Z x ∂ ˆ Z W (cid:62) y is well deﬁned if Z x is differentiable in ˆ Z W (cid:62) y .
However, even if this is not the case, e.g. if x = θ ( W (cid:62) y ) where θ is the Heavyside step function, we can still deﬁne this expectation by leveraging Stein’s lemma:
In 5.3, suppose { W (cid:62) y i } ki
=1 are all elements of V W (cid:62) introduced before x .
Deﬁne the matrix
C ∈ R k × k by C ij def
=
E Z y
i Z y j and deﬁne the vector b ∈ R k by b i def =
E ˆ Z W (cid:62)
y i Z x .
If a = C + b (where C + denotes the pseudoinverse of C ), then in 5.3 we may set
[Equation]
This deﬁnition agrees with the partial derivative expectation by Stein’s lemma when the latter is well deﬁned.
Pseudo-Lipschitz functions are, roughly speaking, functions whose weak derivatives are polynomially bounded.
Deﬁnition A.3.
A function f : R k
→ R is called pseudo-Lipschitz of degree d if | f ( x ) − f ( y ) | ≤
C (cid:107) x − y (cid:107) (1 + (cid:80) ki
=1 | x i | d + | y i | d ) for some C .
We say f is pseudo-Lipschitz if it is so for any degree.
Here are some basic properties of pseudo-Lipschitz functions:
• The norm (cid:107) · (cid:107) in Deﬁnition A.3 can be any norm equivalent to the (cid:96) 2 norm, e.g. (cid:96) p , p ≥ 1 , norms.
Similarly, (cid:80) ki =1 | x i | d + |
y i | d can be replaced by (cid:107) x (cid:107) dp
+ (cid:107) y (cid:107) dp , for any p ≥ 1 .
• A pseudo-Lipschitz function is polynomially bounded.
• A composition of pseudo-Lipschitz functions of degrees d 1 and d 2 is pseudo-Lipschitz of degree d 1 + d 2 .
• A pseudo-Lipschitz function is Lipschitz on any compact set.
Indexing notations We use superscripts to distinguish between different tensors in the program, and subscripts to index into coordinates of tensors.
(i.e x ij denotes the j ’th coordinate of vector x j ∈ R n .
We typically use β = { β 1 , β 2 , ..., β r } to denote a set of coordinates (typically containing r coordinates unless speciﬁed otherwise).
For any vector x ∈ R n , x β denotes the set { x β 1 , x β 2 , ..., x β r } .
For summation over all indices in β , we use the abbreviated notation (cid:80) n β =1 def = (cid:80) nβ 1 =1 ... (cid:80) nβ r =1 .
We typically use x to denote a set of vectors { x 1 , ..., x k } , where the size of the set is implied by the context.
The notation x β refers to the set of scalars { x 1 β 1 , ..., x kβ 1 , ..., x 1 β r , ..., x kβ r } (note that in this case | x β | = kr .
We additionally use α to index into tensors.
The notation x α, β for a vector x refers to the set { x α , x β 1 , ..., x β r } .
Similarly, the notation x α, β for x = { x 1 , ..., x k } refers to the set of scalars { x 1 α , ..., x kα , x 1 β 1 , ..., x kβ 1 , ..., x 1 β r , ..., x kβ r } (in this case | x α, β |
= k ( r + 1) .
We use c, C, ˜ C as arbitrary constant scalars throughout the appendix (Their value might change between in different lines).
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 deﬁnitions throughout the proof:
Deﬁnition A.4.
For any set of vectors x = { x 1 , x 2 , ..., x k } in the program, deﬁne the matrix Λ x ∈
R
[Equation] [Equation]
We say the set x has stable rank if lim n →∞ rank (Λ x ) a.s = rank (lim n →∞ Λ x ) .
Deﬁnition A.5.
We say a random vector x ∈ R n is vanishing
if ∀ p> 0 lim n →∞ (cid:107) x
(cid:107) 2 n p a.s = 0 .
Deﬁnition A.6.
We say a random vector x ∈ R n is regular with constants { ˇ c ( p ) , ˆ c ( p ) } if for all p ∈ (0 , ∞ ) there exists constants 0 < ˇ c ( p ) ≤ ˆ c ( p ) <
∞ such that ˇ c ( p ) ≤ lim n →∞ (cid:107) x (cid:107) pp n ≤ ˆ c ( p ) almost surely.
Deﬁnition A.7.
We say a random scalar θ ∈ R is vanishing if it converges almost surely to ˚ θ , and it holds that ∀ p> 0 , lim n →∞ ( θ − ˚ θ ) 2 n p − 1 a.s = 0 .
Equivalently, ( θ − ˚ θ ) 1 n is a vanishing vector.
Deﬁnition A.5 and Deﬁnition A.6 extend naturally to tensors, which will come in handy in our analysis.
We index a rank r tensor x ∈ ⊗ r R n with indices β = { β 1 , ..., β r } , such that x β = x [ β 1 , β 2 , ..., β r ] .
Deﬁnition A.8.
We say a rank r random tensor x ∈ ⊗ r R n is vanishing
if ∀ p> 0 lim n →∞ (cid:107) x
(cid:107) 2 n p a.s = 0 .
Deﬁnition A.9.
We say a rank r random tensor x ∈ ⊗ r R n is regular with constants { ˇ c ( p ) , ˆ c ( p ) } if for all p ∈ (0 , ∞ ) there exists constants 0 < ˇ c ( p ) ≤ ˆ c ( p ) <
∞ such that ˇ c ( p ) ≤ lim n →∞ (cid:107) x (cid:107) pp n r ≤ ˆ c ( p ) almost surely.
A. I NDUCTION H YPOTHESIS
Setup A.10 (Induction Setup) .
We will keep track of a set of vanishing scalars Θ (see Deﬁnition A.7), and two sets of vectors: The core set x which contains regular vectors produced by a M ATMUL operation (see Deﬁnition A.6 and Deﬁnition 5.1), and a vanishing set ˆ x of vanishing vectors (see
Deﬁnition A.5).
We will denote by x W the set of vectors y :
Wy ∈ x .
Given the sets of vanishing scalars Θ , corset x and vanishing set
ˆ x at some step m in the program, let h ∈ R n , θ ∈ R deﬁne a new vector and scalar via T ENSOR and T ENSOR M OMENT operations respectively.
Namely:
[Equation] [Equation]
where ψ : R ( | x | + | ˆ x | )( r +1)+ l → R is pseudo lipschitz in all of its arguments.
We further deﬁne:
[Equation]
Our induction hypothesis IH ( m ) asserts that the following hold simultaneously:
ReWrite ( m )
Any vector produced by M ATMUL can be written as a linear combination of vectors from x and ¯ x . 2.
StableRank ( m ) For any W ∈ R n × n , the set x W has a stable rank.
Dichotomy ( m ) It holds that: (a) ¯ h is either a regular or a vanishing vector.
(b) ∆ h, ∆¯ h are vanishing vectors.
TensorMoment
( m )
It holds that θ a.s → ˚ θ , where:
[Equation]
ConvRate ( m )
It holds that θ ∈ R is a vanishing scalar, or:
[Equation]
A. H ELPER L EMMAS
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 u, v ∈ R n are vanishing vectors, then ν = u
+ v is a vanishing vector.
Proof.
[Equation]
Lemma A.12.
If u, v ∈ R n are regular vectors, then ν = | u | + | v | is a regular vector.
Proof.
for p ∈ [1 , ∞ ) , using triangle inequality for p norms:
[Equation]
For p ∈ (0 , 1) :
[Equation] [Equation]
One the other hand the lower bound is trivially ∀ p ≥ 1 , (cid:107) ν
(cid:107) pp n
[Equation]
.s
[Equation]
(ˇ c u ( p
[Equation]
ˇ
[Equation]
v ( p
[Equation]
Lemma A.13.
If u ∈ R n is a vanishing vector, then for any p > 0 , it holds that 1 n (cid:107) u (cid:107) pp a.s → 0 (i.e u has vanishing moments).
Proof.
For p ≥ 2 , we have that:
[Equation] [Equation]
For 0 < p < 2 , using the fact that ∀ 0 <p<q ,
[Equation]
u
[Equation]
u
[Equation]
and assigning q = 2 :
[Equation] [Equation]
which proves the claim.
Lemma A.14.
If u ∈ R n is a vanishing vector, and v ∈ R n is a regular vector, then ν = u
+ v is a regular vector.
Proof.
This is immediate from Lemma A.13 and Lemma A.12, and setting the constants for u { ˇ c u ( p ) , ˆ c u ( p ) } = { 0 , 0 } .
Lemma A.15.
If u ∈ R n is a vanishing vector, then for any r ≥ 1 , ν = | u | r is a vanishing vector.
If u ∈ R n is a regular vector, then for any r ≥ 0 , ν = | u | r is a regular vector.
Proof.
For vanishing u , using elementary norm bounds, that ∀ 1 <p<q , (cid:107) u (cid:107) q ≤ (cid:107) u (cid:107) p and assigning q = 2 :
[Equation]
The proof for regular u follows immediately from the deﬁnition of a regular vector.
Lemma A.16.
If u ∈ R n is a vanishing vector, and v ∈ R n is a regular vector with constants { ˇ c ( p ) , ˆ c ( p ) } , then ν = u (cid:12) v is a vanishing vector.
Proof.
For any p > 0 , choose m, l ∈ (0 , 1) such that p > m , and l + m = 1 .
Using Holders inequality:
[Equation]
where we used Lemma A.15 to assert that (cid:107)| u
| 1 l (cid:107)
[Equation]
Lemma A.17.
If θ is a vanishing scalar, then f ( θ ) for f : R → R is a vanishing scalar if f is locally lipschitz at ˚ θ , and f (˚ θ ) = 0 .
Proof.
WLOG assume ∀ ( θ − ˚ θ ) 2 <(cid:15)
f
( θ ) <
A | θ − ˚ θ | for some (cid:15) ∈ R .
Deﬁne :
[Equation]
Since θ → ˚ θ almost surely, this implies that Prob (lim n →∞ ( θ − ˚ θ ) 2 < (cid:15) )
= 1 , hence ( θ − ˚ θ ) 2 < (cid:15) almost surely.
Therefore:
[Equation] [Equation]
Lemma A.18.
Let x , ˆ x denote sets of regular and vanishing vectors, and let ψ : R | x | + | ˆ x | → R be pseudo lipschitz.
Then:
[Equation]
Proof.
We have that:
[Equation]
Since ψ is pseudo lipschitz:
[Equation]
Deﬁne the vectors u α = (cid:0) 1 + (cid:107) ˆ x α (cid:107) dd + (cid:107) x α (cid:107) dd (cid:1) , v α = (cid:107) ˆ x α (cid:107) Note that u, v are regular and venishing vectors respectively, hence by Lemma A.16 u
(cid:12) v is a vanishing vector.
Then by Lemma A.13,n (cid:80) nα =1 ( u (cid:12) v ) α → 0 almost surely.
Plugging into Eq.
(30) and taking the limit, we have that:
[Equation]
proving the claim.
Note that Lemmas A.11 to A.16 and A.18 trivially extend to vanishing and regular tensors.
Lemma A.19.
If u ∈ ⊗ r +1 R n , r ≥ 1 is a vanishing tensor, then ν ∈ R n : ν α = 1 n n 2 (cid:80) n β =1 u α, β is a vanishing vector.
Proof.
Using the elementary inequality (cid:107) v (cid:107) 1 ≤ √ n (cid:107) v (cid:107) for any vector v ∈ R n , we have:
[Equation]
Lemma A.20.
Let { x n } n> 0 be a sequence of random variables.
If for some t ∈ N , and for all n it holds that E x 2 tn ≤ cn − 1 − λ for c, λ > 0 then x n → 0 almost surely.
Proof.
by markov’s inequality, for any (cid:15) > 0 :
[Equation]
By the Borel-Cantelli Lemma, x n → 0 almost surely.
Lemma A.21.
Let m, n, r ∈ N such that n is divisible by r .
Let { ν i } mi =1 , ∀ i , ν i ∈
R
nr denote random (possibly dependent), zero mean vectors with iid coordinates and ﬁnite moments of any order ∀ q ∈ N , E
[( ν iα ) 2 q ] = C 2 q .
Deﬁne S
i = 1 √ n (cid:80) nr α =1 ν iα .
Then, there exists a function f ( q ) ∈ [0 , ∞ ) independent on n, m such that:
[Equation] [Equation]
Proof.
Using Hölder’s inequality:
[Equation]
The inner expectations are given by:
[Equation]
Note that since the random variables ν iα have zero mean, the expectation E (cid:104) (cid:81) 2 qj
=1 ν iα j (cid:105) does not vanish only when the indices α 1 , ..., α 2 q do not contain an entry which appears in isolation.
In other words, the number of non-zero terms n (cid:63) in the sum is:
[Equation]
Note that the the only term which depends on n is
[Equation]
which is bounded by:
[Equation]
where ˜ f ( q ) < ∞ independent on n
[Equation]
m .
It follows:
[Equation]
where:
[Equation]
In addition, from Hölder’s inequality:
[Equation]
Hence it follows that E (cid:2) ( S i ) 2 q (cid:3) ≤ ( nr ) q f ( q ) n q
[Equation]
q
[Equation]
q ) r q C 2 q .
Inserting into Eq.
(38), we ﬁnally get:
[Equation] [Equation]
Before delving into the full proof of the induction hypothesis, we note the following fact that immediately holds at any arbitrary step in the induction.
Let h, h 0 , ∆ h, θ be deﬁned as in Setup A.10.
Then the following claim holds:
Claim A.1.
∆ h is a vanishing vector.
Proof.
Since ψ is psuedo lipshitz, there exists some d ≥ 0 such that:
[Equation]
where we have deﬁned the tensors τ
[Equation]
T ∈ ⊗ r +1 R n such that:
[Equation]
Note that by Lemmas A.11 to A.15, T is a regular tensor, while τ is a vanishing tensor.
Hence, T (cid:12) τ is a vanishing tensor by Lemma A.16, hence ∆ h is a vanishing vector by Lemma A.19.
Claim A.1 is useful due to the following general claim:
Claim A.2.
If Dichotomy ( m ) , TensorMoment ( m ) and ConvRate ( m ) apply for the vector h (as deﬁned in Setup A.10), then it applies for h
+ δ
if δ is a vanishing vector.
Proof.
This is true due to the function ψ being pseudo lipschitzness.
More speciﬁcally: 1.
If Dichotomy ( m ) holds for h , then it holds for h
+ δ
This is trivially true due since we can expand ˆ x := ˆ x ∪ δ , and invoke Lemmas A.11 and A.14.
If TensorMoment ( m ) holds for h , then it holds for h + δ .
This is since we may trivially assert that: θ = 1 n (cid:80) nα =1 ( h α + δ α ) → 1 n (cid:80) nα =1 h α , which stems from Lemma A.13.
If ConvRate ( m ) holds for h , then it holds for h + δ .
This is since:
[Equation]
Given Claim A.2 and Eq.
(53), we may prove Dichotomy ( m ) , TensorMoment ( m ) and ConvRate ( m ) for h 0 instead of h .
[Figure]
A.2. H YPERGRAPHS AND Baranyai’s theorem
Baranyai’s theorem in combinatorial mathematics deals with the number of ways one can partition a complete hypergraph into 1-factors.
A complete hypergraph G nr is a hypergraph containing n vertices in which every subset of r vertices forms a hyperedge.
A 1 factor of this graph is a partition of the hypergraph into nr hyperedges in which each vertex touches exactly one hyperedge.
Theorem A. gives an informal statement of Baranyai’s theorem.
Theorem A.22 (Baranyai’s theorem - Informal) .
The n vertices of a hypergraph G nr | n, r ∈ N such that r divides n can be partitioned into 1-factors in (cid:0) nr (cid:1) rn different ways such that each hyperedge in
G nr 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 as is stated in Theorem 5.4.
Concretely, we will need to reason about the moments of inﬁnite sums of random variables.
Speciﬁcally, let { z 1 , ..., z k } ∈ R n denote independent and normally distributed vectors, let ψ β = ψ ( z 1 β 1 , ..., z kβ r ) where ψ : R kr → R is polynomialy bounded and ∀ β , E [ ψ β ] = 0 .
Consider the expression:
[Equation]
Theorem A.23.
There exists C ( q ) ∈ [0 , ∞ ) such that lim n →∞ m q ≤ C ( q ) .
Proof.
We can express m q by breaking the sum (cid:80) n β =1 :
[Equation]
≤ A ( n ) + B ( n ) where:
[Equation]
where ˜ C ∈ [0 , ∞ ) does not depend on n .
We now prove the following:
lim n →∞ B → 0 .
To see this, notice that there are n r − n !
( n − r )!
non zero terms in the sum (cid:80) n β =1 ψ β (1 − 1 β 1 (cid:54) =
β 2 (cid:54)
= ...
(cid:54) = β r ) , hence:
[Equation]
Notice that lim n →∞ n
r
− n !
( n − r )!
n r − 0 .→ 0 , and max β E (cid:2)(cid:0) ψ β (cid:1) 2 q (cid:3)
is bounded and does not depend on n .
We can therefore conclude B vanishes as n → ∞ .
We use Lemma A.21 and Theorem A.22 to show that lim n →∞ A ≤ C ( q ) for some C ( q ) < ∞ .
Let Θ = { β 1 , β 2 , ..., β n !
( n − r )! } denote the set of all possible conﬁgurations of r indices β such that β 1 (cid:54)
=
β 2 (cid:54)
= ...
(cid:54) =
β r .
Let { Θ i } mi =1 denote m sets where each set Θ i contains nr conﬁgurations Θ i = { β 1 ,i , ..., β nr ,i } such that:
(a) ∀ i ; β , β (cid:48)
∈
Θ i | β (cid:54) =
β (cid:48) , β ∩ β (cid:48)
= ∅ (b) ∀ i (cid:54) =
j ; β ∈ Θ i ; β (cid:48) ∈
Θ j , β (cid:54) =
β (cid:48)
Finally, let R denote the remaining conﬁgurations that do not appear in any set { Θ i } (i.e
R = Θ / (Θ 1 ∪ Θ 2 ∪ ... ∪ Θ m ) .
We then have:
[Equation]
65) where S i = 1 √ n (cid:88) β
[Equation]
[Equation]
Note that by construction of the sets { Θ i } mi =1 , and since the vectors z 1 , ..., z k contain iid coordinates, we can conclude that for all
i , random variables { ψ β } β ∈ Θ i are independent.
Provided we can construct m = rn r − 1 − O ( n r − 2 ) sets, then | R |
= O ( n r − 1 ) .
In that case, then by Lemma A.21:
[Equation] [Equation]
We are then left with proving that we can in fact partition the set Θ into { Θ i } mi =1 ∪ R where m = rn r − 1 − O ( n r − 2 ) .
To show this, we deﬁne a complete hypergraph G nr with n vertices in which every vertex corresponds to an integer in { 1 , 2 , ..., n } .
We can think of a hyperedge in G nr as an edge connecting r integers without ordering, hence the set of all hyperedges in
G nr has cardinality 1 r !
| Θ | (this is since ordering matters in Θ ).
By Theorem A.22, we can partition the vertices in G nr into nr hyperedges (sets of r unique integers with no ordering) in (cid:0) nr (cid:1) rn different ways, where each hyperedge appears in a single partition.
For each partition i , we can assign Θ i where any β ∈ Θ i corresponds to a single hyperedge in partition i , with the ordering of the vertices decided arbitrarily.
Notice that for each partition Θ i , we can construct ( r ! − 1) additional partitions by reordering β in any β ∈ Θ i
Therefore, the total number of valid partition is given by (cid:0) nr (cid:1) rn r !
=
rn r − 1 − O ( n r − 2 ) , proving the theorem.
A. B ASE C ASE
WLOG, we start with an initial corset of Gaussian iid vectors x = { x 1 , ..., x k } (which are regular), an initial vanishing set of vanishing vectors ˆ x = { ˆ x 1 , ..., ˆ x k (cid:48) } and a set of vanishing scalars Θ = { θ 1 , ..., θ l } .
Note that ReWrite (1) and StableRank (1) trivially hold.
We proceed to prove
Dichotomy (1) , TensorMoment (1) and ConvRate (1) .
We deﬁne the functions ψ : R k
→ R , ¯ ψ : R ( r +1) k
→ R , ¯ ψ : R k
→ R using the pseudo lipschitz function ψ : R ( r +1) k
→ R and vectors x :
[Equation]
Note that ¯ ψ ( y ) is a random function that depends on the vectors x .
Theorem A.24.
∆¯ h is a vanishing vector.
Proof.
From Lemma A.20, it sufﬁces to show that for every p > 0 , there exists t ∈ N such that
E
[Equation]
n tp ≤ cn − 1 − λ for some c, λ > 0 (which may depend on p ).
Fix p > 0 , and choose t = 1+ λp , and let q = (cid:100) t (cid:101) .
Then by Jensen’s inequality :
[Equation]
We are now left with the task of proving that E (cid:104)(cid:16) (cid:80) n β =1 ¯ ψ
[Equation]
x
[Equation]
β
[Equation]
n r −
[Equation] [Equation]
(cid:17) 2 q (cid:105) is ﬁnite for any (ﬁxed) integer q and for n → ∞ .
Firstly, we may express the over all indices (cid:80) n β =1 as (cid:80) n β =2 + (cid:80) n β |∃ j : β j =1 .
That is, we ﬁrst 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:
[Equation]
where ˜ C ≤ ∞ does not depend on n .
we now make the following observations:
Claim A.3.
There exists a constant C < ∞ that is independent of n such that ∀ n,r , E (cid:104)(cid:16) (cid:80) n β |∃ j : β j =1 ¯ ψ ( x 1 , x β ) n r − 0 .
5 (cid:17) 2 q (cid:105) ≤ C .
To see this, note that the summation (i.e (cid:80) n β |∃ j : β j =1 ) effectively sums over n r − ( n − 1) r ∼ O ( n r − 1 ) terms.
Since ¯ ψ is a centered pseudo lipschitz function of normally distributed variables, we can bound the second expectation:
[Equation]
for some ˜ C < ∞ that do not depend on n .
Claim A.4.
There exists a constant C < ∞ that is independent of n such that ∀ n,
[Equation]
, E
[Equation]
≤ C .
Note that since the summation over the indices β do not include the ﬁrst index, the random variable x β can be treated as independent from x 1 .
WLOG we can then bound the following term instead:
[Equation]
For some random variable y independent of x β for all values of β , with the same dimensions as x 1 .
We can now condition on y , and apply Theorem A.23 to complete the proof.
Now, we can express h =
¯ h + ∆ h
+ ∆¯ h , where ∆ h, ∆¯ h are vanishing vectors.
Invoking Claim A.2, it is enough to prove the base case holds for ¯ h alone:
TensorMoment (1) is immediate from the law of large numbers given that x are iid gaussians.
Namely:
[Equation]
Dichotomy (1) holds since ψ is a smooth, pseudo lipschitz (given by gaussian averaging of ψ ) function.
From TensorMoment (1) , for any p > 0 :
[Equation]
where N ( Z x ) is the gaussian measure.
Therefore, if 1 n lim n →∞ (cid:107) ¯ h (cid:107) pp = 0 for some p > 0 , then ψ is identically zero.
In that case, we can write h = 0 +∆ h +∆¯ h , which is a vanishing vector.
If ψ is not identically zero, then h is a regular vector from TensorMoment (1) , proving Dichotomy (1) .
ConvRate (1) holds since for all p > 0 :
[Equation]
which holds from Theorem A.24 and assigning r = 0 .
We have therefor concluded the base case.
A. I NDUCTION S TEP
We prove the induction step assuming IH ( m ) holds.
Namely, we must show that IH ( m ) → IH ( m +1) .
Assume a new vector is introduced via M ATMUL , namely Wν where ν is given by T ENSOR operation of corset x , vanset ˆ x and vanishing scalars Θ .
By Dichotomy ( m ), we may express ν = ¯ ν +∆ ν +∆¯ ν where ∆ ν + ∆¯ ν is a vanishing vector, and ¯ ν is either regular or vanishing.
A.4. IH
[Equation]
RITE ( m + 1 ) + S
We can express Wν = W
¯ ν + W (∆ ν + ∆¯ ν ) , and point to the following fact:
Claim A.5.
If ¯ ν is a regular vector, then W ¯ ν is a regular vector.
Moreover, for any vanishing vector δ ∈ R n , Wδ is a vanishing vector.
Proof.
If δ is vanishing then Wδ is vanishing:
This is true since W is a gaussian matrix with a uniformally bounded (in n ) operator norm.
The ﬁrst part Claim A.5 holds since ¯ ν depends only on vectors from x , for which the set x W has a stable rank from StableRank ( m ) .
We can therefore use the gaussian conditioning trick (conditioning on all vectors in x and x W ).
We can now expand the vanset with ˆ x := ˆ x ∪ W (∆ ν + ∆¯ ν ) , and proceed by casework: 1.
If ¯ ν is vanishing then we expand ˆ x := ˆ x ∪ W ¯
ν .
In that case x remains unchanged and
StableRank ( m + 1) trivially holds.
If ¯ ν is regular then we expand x := x ∪ W ¯ ν , and we get StableRank ( m + 1) using
Theorem A.25.
Theorem A.25.
Let x 1 , ..., x k ∈ x W .
If Z x = α 1 Z x 1 + α 2 Z x
+
+ Z x k , then almost surely, for large enough n , x = α 1 x 1 + α 2 x 2 + ... + α k x k .
Proof.
The set x is constructed as a standard N ETSOR (cid:62) program (without scalars), and we may immediately apply theorem 6.3 in Yang (2020b).
A.4.
[Table]
We are left with proving Dichotomy ( m + 1) , TensorMoment ( m + 1) and ConvRate ( m + 1) .
Note that if ¯ ν is a vanishing vector, the set x remains unchanged, and we immediately get IH ( m + 1) by
Claim A.2.
Hence we proceed assuming ¯ ν
(hence W ¯ ν is a regular vector.
Getting Dichotomy ( m + 1) Assume a new vector is introduced in the program via a T ENSOR operation:
[Equation]
Where we made explicit the inclusion to x of the new vector W ¯ ν .
Let h 0 , ¯ h, ∆ h, ∆¯ h be deﬁned as in Eq.
Note that ∆ h is vanishing by Claim A.1, which holds generally.
We next prove ∆¯ h
= h
− ¯ h is a vanishing vector, where (using ψ ( − ) ≡ ψ ( − ; 0 ; ˚Θ ) to ease notational burden):
[Equation]
The key insight to proving ∆¯ h
= h
− ¯ h is indeed vanishing is that both h 0 , ¯ h can be written as a sum of a shared regular vector, and a vanishing vector (i.e we can express h 0 = µ + δ 1 , ¯ h = µ + δ 2 where µ is regular, and δ 1 , δ 2 are vanishing).
Their difference h 0 − ¯ h
= δ 1 − δ 2 is therefore the difference of two vanishing vectors, which is itself vanishing.
To show this, we can make explicit the distribution of
W ¯ ν by using the gaussian conditioning trick (see Appendix A.5).
Denote by X, Y ∈ R n × r , U, V ∈ R n × s the matrices with { x i } ∈ x , { y i } ∈ x W , { u i } ∈ x , { v i } ∈ x W (cid:62) as columns respectively, representing previously generated vectors in the program,
such that X =
WY, U = W
(cid:62) V .
Using the Gaussian conditioning trick (conditioning on all the vectors in x ), g =
W ¯ ν is distributed as:
[Equation]
where d → ˚ d, e → ˚ e, σ → ˚ σ , and z ∼ N (0 , I n ) .
Deﬁne:
[Equation]
We now note that:
• ∀ i , ( d i − ˚ d i ) x
i is a vanishing vector since x i is regular ( x ∈ x W ), and ( d i − ˚ d i ) is vanishing by the induction hypothesis, and
Lemma A.17.
• ∀ i , ( e i − ˚ e i )
v i is a vanishing vector since v i is regular ( v ∈ x W (cid:62) ), and ( e i − ˚ e i ) is vanishing by the induction hypothesis, and
Lemma A.17.
• ( σ Π ⊥ V − ˚ σ ) z is a vanishing vector.
To see this, note that:
[Equation]
(˚ σ − σ ) z is vanishing due to the induction hypothesis and Lemma A.17.
V ( V V (cid:62) ) † V (cid:62) z is vanishing as well.
To see this, note that:
[Equation]
By the induction hypothesis ( V V (cid:62) n ) † converges almost surely ( V has stable rank).
Then it is enough to show that ∀ i,p> 0 , lim n →∞ 1 n p 2 z (cid:62) v i √ n a.s → 0 .
From TensorMoment ( m ) and
ConvRate ( m ) :
[Equation]
We therefore conclude that b = g − a is vanishing.
We have:
[Equation] [Equation]
where A ( m ) = 1 n r n (cid:88) β
= ψ
[Equation] [Equation]
From Claim A.1, B ( m ) is a vanishing vector.
Furthermore, a is a deterministic function of previous vectors in x , and an iid gaussian noise vector z .
We can now recursively expand A ( m )
=
A ( m − 1) + B ( m − 1) where B ( m − 1) is a vanishing vector, until we are left with A (1) (i.e A ( m ) =
A (1) + (cid:80) m −
[Equation]
m
[Equation]
=1 B
[Equation]
, where { B ( m (cid:48) ) } m (cid:48) are all vanishing vectors).
Note that A (1) can be expressed as a pseudo lipschitz function of normally distributed vectors, with coordinate distributions given by Z g , Z x 1 , ..., Z x | x | .
We can apply the same decomposition to ¯ h and get ¯ h
= ¯ A (1) + (cid:80) m − 1 m (cid:48)
=1 ¯ B ( m (cid:48) ) .
We have that:
[Equation]
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 ∆¯ h are vanishing.
Getting TensorMoment ( m +1) and ConvRate ( m +1)
These are immediate since we can express h
= ¯ A (1) + δ where ¯ A (1) is a function of gaussian vectors, and δ is a vanishing vector, and by
Claim A.2 invoke the base case on ¯ A (1) .
A. G AUSSIAN C ONDITIONING
Let g =
Wh where g, h ∈ R n are vectors in a N ETSOR (cid:62) program.
Denote by X, Y ∈ R n × r , U, V ∈ R n × s the matrices with { x i } , { y i } , { u i } , { v i } as columns respectively, representing previously generated vectors in the program, such that X = WY, U = W
(cid:62) V .
Using the Gaussian conditioning trick (conditioning on all the vectors in x ), g =
Wh is distributed as:
[Equation]
where we have deﬁned A
= Eν 0 ,
[Equation]
=
Π ⊥ V ˜ W Π ⊥ Y h , Π V , Π Y are projection matrices, and ˜ W is a fresh iid sample of W , and:
[Equation]
Rewriting the conditional distribution of g , we get:
[Equation]
99) with Θ def = Eh ∈ R n , σ def = σ A (cid:114) (cid:107) Π ⊥ Y h (cid:107)n ∈ R (100)
Moreover, σ converges to a deterministic limit ˚ σ , and Θ can be written as:
[Equation]
101) where ˆ (cid:15) ∈
R
r , ˇ (cid:15) ∈
R s are vanishing vectors, and ˚ d ∈ R r , ˚ e ∈ R
s are deterministic vectors.
B P ROOFS OF T HEOREM 4.1 AND T HEOREM 4.
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 at each training step at ﬁnite width, implement it as a set of TP instructions, convert it to inﬁnite-width computation according to Deﬁnition 5.3 and apply the master theorem.
B. T HE T ENSOR P ROGRAM FOR T HEOREM 4.
In the next section, we construct the Tensor Program that encodes the training of an L-hidden layer
MLP as inEq.
(1) under the ANKT parametrization.
Here we ﬁrst describe the initial matrices, vectors, and scalars of the program, along with necessary notations.
Initial Matrices, Vectors and Scalars We ﬁrst deﬁne the initial set of matrices W , vectors V and scalars C :
• Initial matrices { w l } Ll =2 ∈ R n × n are all sampled iid from N (0 , 1) .
We set W =
W 2 ∪ W 3 ∪ ... ∪ W L +1 .
• The initial vectors V are given by the ﬁrst layer
h 1 ( ξ ) for all inputs, and the last weight vector w L +1 ∈ R n , all samples iid from N (0 , 1) .
• Initial scalar C
=
1 √ n } .
Notations We use := to more clearly denote assignment happening in the program, as opposed to mathematical equality.
We will use the notation
T ENSOR ( y 1 , .., y k ; Θ) , T ENSOR M OMENT ( y 1 , .., y k ; Θ) to denote an arbitrary implementation of these instructions give vectors y 1 , ..., y k and scalars Θ .
Initial Forward Pass Starting with our initial vectors
h 1 ( ξ ) :
= w 1 ξ , we compute all { x l ( ξ ) } Ll =1 , { h l ( ξ ) } Ll =2 using T ENSOR and M AT M UL instructions:
[Equation]
The initial outputs are given by f ( ξ ) =
[Equation]
since T ENSOR M OMENT only allows division by n r for integer r , rather than √ n .
However recall that in Theorem 4.1 we assume WLOG that the outputs for any input ξ is ﬁxed to f ( ξ )
=
g ( ξ ) .
Let X denote a matrix composed of all x L ( ξ ) as columns.
Denote by e the event that f ( ξ ) = 0 for all inputs.
Then, using gaussian conditioning, the conditional distribution of w L +1 e given e is:
[Equation]
where ˜ w L +1 is an independent sample of w L +1 , and Π =
I − 1 n X
( X (cid:62) Xn ) † X (cid:62) and • † is the pseudo-inverse of • .
Then, we can write:
[Equation]
n
[Equation]
By the master theorem X (cid:62) Xn a.s → γ , ( X (cid:62) Xn ) † a.s → γ † and X (cid:62) ˜ w L +1 n a.s → 0 .
Hence, after conditioning on f ( ξ ) = 0 , the distribution of w L +1 e is still identical to that of w L +1 at the limit.
At this point we can just implement w L +1 in the program with ˜ w L +1 −
[Equation]
which can be implemented with T ENSOR and T ENSOR M OMENT instructions.
We now have a program that encodes the initial forward pass of the MLP conditioned on f ( ξ ) = 0 for all ξ .
Initial Backward Pass and Loss Derivatives For any input sample ˜ ξ , we can implement d ˜ h l using
T ENSOR :
[Equation]
Then, for all 1 ≤ l < L , using M AT M UL and T ENSOR :
[Equation]
The initial loss derivatives L (cid:48) ( ξ ) are all deterministic scalars since we have conditioned on the initial outputs f ( ξ ) .
Forward and Backward Passes at Any t The forward and backward and loss computation for any t are given by:
[Equation]
with the weight updates given by:
[Equation]
which can all be implemented with T ENSOR and Matmul operations as before.
Adaptive Update at Time t Using Eq.
(2) and w lt = w l + (cid:80) tt
[Equation]
w lt
[Equation]
recall w l = √ nW l ), we have that:
[Equation]
which can be implemented using T ENSOR :
[Equation]
and similarly for any layer 2 < l ≤ L .
The (pre)activations at any step t can be implemented as follows using T ENSOR :
[Equation]
Output Updates The function updates can be implemented using T ENSOR M OMENT :
[Equation]
The loss derivatives can be implemented using using T ENSOR M OMENT :
[Equation]
B. P ROOF OF T HEOREM 4.
After writing the program using TP operations, we are ready to prove Theorem 4.1 by taking the inﬁnite-width limit.
First, note that from Eqs.
(119), (120) and (166) and ??
and applying
Deﬁnition 5.3, we have that:
[Equation] [Equation]
[Equation]
Applying Deﬁnition 5.3 to Eqs.
(117) and (118), we have that:
[Equation] [Equation]
where we have that Z W l δ ˜ x
l − 1 =
ˆ Z W l δ ˜ x
l − 1 + ˙ Z W l δ ˜ x
l − 1 according to Deﬁnition 5.3.
Then using
Theorem 5.4:
[Equation]
We now use lemma L.3 from Yang (2020b) restated:
Lemma B
[Equation]
For any x, y ∈ R n and W ∈ R n × n in the program, it holds that:
[Equation]
Applying Lemma B.1 to Eq.
[Equation]
[Equation]
Similarly expanding E (cid:2) Z d ˜ h
L
[Equation]
1 Z W L
[Equation]
1 δ ˜ x
L
[Equation]
2 (cid:3) we arrive at:
[Equation]
Finally
[Equation]
using Eq.
[Equation] [Equation]
B. T HE T ENSOR P ROGRAM FOR T HEOREM 4.
In the next section, we 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 deﬁne ¯ w 3 = √ nw 3 , so the output is given by f ( ξ )
= 1 n ¯ w 3 (cid:62) x 2 ( ξ ) .
Here we ﬁrst describe the initial matrices, vectors, and scalars of the program, along with necessary notations.
Initial Matrices, Vectors and Scalars We ﬁrst deﬁne the initial set of matrices W , vectors V and scalars C :
• Initial matrices w 2 ∈ R n × n sampled iid from N (0 , 1 n ) .
We set W = W 2 .
• The initial vectors V are given by the ﬁrst layer
h 1 ( ξ ) for all inputs, and the last weight vector ¯ w 3 ∈ R n .
Notice that ¯ w 3 is normally distributed.
• Initial scalar C
=
1 √ n } .
Notations As in the ANTK case, we use := to more clearly denote assignment happening in the program, as opposed to mathematical equality.
To clearly demonstrate the application of T ENSOR , we will also freely introduce function symbols ψ to put things into T ENSOR form.
Initial Forward and Backward Passes Starting with our initial vectors
h 1 ( ξ ) :
= w 1 ξ , we compute x 1 ( ξ ) , h 2 ( ξ ) , x 2 ( ξ ) , dh 2 ( ξ ) , f ( ξ ) , L (cid:48) ( ξ ) for all inputs using T ENSOR , T ENSOR M OMENT and M AT M UL instructions at step any t :
[Equation]
[Equation] [Equation]
[Equation]
ENSOR
[Equation]
) for T ENSOR
[Equation]
Note that with µ parameterization we can express the output f ( ξ ) directly without conditioning using a T ENSOR M OMENT .
Expressing ˜ h
2 t +
Using Eq.
(2), we have:
[Equation]
x 1 t
[Equation]
˜
[Equation] [Equation]
L
[Equation]
) for T ENSOR (
[Equation]
[Equation]
B. P ROOF OF T HEOREM 4.
After writing the program using TP operations, we are ready to prove Theorem 4.2 by taking the inﬁnite-width limit.
Applying Theorem 5.4 to Eqs.
(145) to (151), we have that:
[Equation]
[Equation] [Equation]
155) where Z ¯ w 3 ∼ N (0 , 1) .
C E XTENSIONS OF T HEOREM 4.1 AND T HEOREM 4.2 TO AGO S WITH M EMORY
So far we have dealt with the case of memoryless adaptive optimizers, and a batchsize of 1, however our results can be trivially extended to 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 deﬁnition Deﬁnition 3.1, if g 0 , g 1 , ..., g t ∈ R denote gradients of some scalar parameter w at times 0 , 1 , ..., t , a general adaptive update can be described by a function
Q t ∈ R t +1 → R such that ∆ w t = − µQ t ( g 0 , g 1 , ..., g t ; (cid:15) ) .
Concretely, in the case of Adam, Q takes the following form (replacing β 1 , β 2 with γ 1 , γ 2 to prevent confusion with other indices): [Equation]
[Equation]
In the context of optimizing an MLP, we can write the equivalent of Eq.
(2) for a general Q function, and a general batchsize for both parameterizations:
[Equation]
where (cid:80) β t denotes summation over samples in the minibatch β t at step t (i.e if β t
[Equation]
= { ξ [Equation] , ξ j , ξ k } then (cid:80) β t u β t = u t ( ξ i ) + u t ( ξ j ) + u t ( ξ k ) ), and Q t operates elementwise on the components of its inputs.
Note that we have used Deﬁnition 3.1 to conveniently remove the 1 n factors from inside the Q function, as in Eq.
Since (cid:15) is a constant, we will absorb it into the deﬁnition of Q from now onward.
Note that for any vector v , the matrix vector product ∆ w lt v can be implemented as a
TENSOR instruction (see Deﬁnition 5.
[Equation]
where { u β t } β t is a collection of all vectors u evaluated at time t on minibatch β t (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 a more general result. C. E XTENSION OF T HEOREM 4.1 (NTK)
We now state a general theorem for an AGO with memory and arbitrary batchsize.
Theorem C.1.
Let f ( ξ ) ∈ R denote an MLP as in Eq.
(1) parameterized using the ANTK parameterization described in Section 4.1, where φ (cid:48) is pseudo-Lipschitz.
Assume layers { w l } Ll = are trained using an AGO applied on minibatches of arbitrary size, where Q t is pseudoLipschitz deﬁned according to Deﬁnition 3.1, using a loss function L with a pseudo-Lipschitz ﬁrst derivative.
Then, at any step t and for any sample ˜ ξ , it holds that ˜ f
t a.s → ˚˜ f t where ∆˚˜ f t =
[Equation]
Adp ( {
[Equation]
β 0 } , ..., { ξ β t } , ˜ ξ |{ ˚ L (cid:48) β 0 } , ..., { ˚ L (cid:48) β t } ) , where:
[Equation]
where the expectation is taken over all Z variables at initialization
[Equation]
Proof.
The proof of Theorem C.1 is a straightforward extension of the proof of Theorem 4.1.
The forward and backward passes for any t are again given by:
[Equation]
only with weight updates that are given by:
[Equation]
[Equation]
As in the memoryless case, using Eq.
(2) and w lt = w l + (cid:80) tt (cid:48) =
∆ w lt (cid:48) (recall w l = √ nW l ), we have that:
[Equation]
For any vector v , the matrix vector product ∆ w lt v can be implemented as a TENSOR instruction:
[Equation]
where { u β t } β t is a collection of all vectors u evaluated at time t on minibatch β t (and likewise for scalars).
Hence, we may proceed exactly as in the base proof of Theorem 4.1 (i.e expressing the optimization process as a tensor program, applying Deﬁnition 5.3 to get the coordinate distributions in the limit, and applying Theorem 5.4).
Note that in the concrete case of Adam, we get the following function update:
[Equation]
C. E XTENSION OF T HEOREM 4.2 ( µ P)
Theorem C.
[Equation]
Let f ( ξ ) ∈ R denote an MLP as in Eq.
(1) with L = 2 parameterized using the µ parameterization described in Section 4.
[Equation]
where φ (cid:48) is pseudo-Lipschitz.
Assume layers w 2 is trained using an AGO with a pseudo-Lipschitz function Q function according to Deﬁnition 3.1 (for general batchsize)
[Equation]
using a loss function L with a pseudo-Lipschitz ﬁrst derivative.
Then at any step t and for any sample ˜ ξ
[Equation]
it holds that ˜ f t a
[Equation]
s → ˚˜ f t where ˚˜ f t can be computed as follows:
[Equation]
where the expectations are taken over all Z variables (including ζ d = N (0 , 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 µ P.
D N UMERICAL V ERIFICATION
We conduct numerical experiments to verify our results.
For both parameterizations, the exact network dynamics at the inﬁnite width limit is not tractable in the general case, since the expectations involved do not admit an analytical solution (unlike the standard NTK for ReLU networks).
Even for the ANTK parameyterization, the inﬁnite-width dynamics cannot be separated to a ﬁxed kernel and a loss derivative, as with the NTK dynamics for SGD.
We therefore must resort to MC simulations to approximate the expectations involved in evaluating the inﬁnite width dynamics in both regimes.
We verify Theorem C.1 and Theorem C.2 by training a ReLU MLP ( L = 4 for ANTK and L = 2 for µ ) on R 10 gasussian inputs and a unit output.
For a loss we use the standard L2 loss function, regressing to random targets.
We train networks with varying widths using Adam with β 1 = 0 .
9 , β 2 = 0 .in full batch mode, on 100 training samples, and run 10 trials per width.
We use a learning rate of 0 .n , and (cid:15) = 1 e − 4 n (where n is the width).
In order to account for different initial outputs and loss derivative per weight initialization, we subtract the initialized network output from the output for each sample, such that the output is identically zero at initialization for all inputs.
To approximate the inﬁnite-width training dynamics, we approximate the expectation in Eq.
(174) and Eq.
(176) using
MC simulations 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 inﬁnite width dynamics can be approximated without actually constructing a network.
To compare the evolution of the ﬁnite vs inﬁnite architectures, we evaluate the output at each iteration on random inputs.
Our results are summarized in Fig.
2 and Fig.
As expected, as the width increases the training dynamics converge to that of the inﬁnite dynamics.
[Figure] [Figure]