Tensor Programs IVb: Adaptive Optimization in the Infinite-Width Limit

Chapter 1 Introduction

While historically the deep learning theory literature has by-and-large (carelessly in hindsight) identified “infinite-width neural networks” with “neural tangent kernel” or “gaussian process” , by now we understand these are just particular kinds of infinite-width limit with simple mathematics. Indeed, there also exists a “feature learning regime” with much more complex mathematics but also all the actually desired properties of a neural network . Yang & Hu precisely characterized this so-called Dynamical Dichotomy: there is no other regime that can happen for MLPs trained by SGD in finite time.

In practice, of course, most neural networks of importance are trained by adaptive optimizers like Adam . Can new phenomenon arise in the infinite-width limit of adaptive optimizers? For example, if one invokes the (not-quite-correct) intuition that neural network trained with small learning rate exhibits kernel behavior, then one might suppose that these optimizers may adaptively enforce large effective learning rates. This prevents kernel behavior and may even “supercharge” feature learning.

But it turns out there’s nothing special about adaptive optimizers in this regard. Essentially the exact dichotomy of feature learning vs kernel regime plays out for any adaptive optimizer. This means there is also a “neural tangent kernel” limit for any adaptive optimizer, where the network evolution can be captured solely by some kind of evolution equation in the function space — albeit no longer a linear equation. This also means that “maximal update” parametrization can be defined for e.g., Adam, that maximizes feature learning in a suitable sense. Indeed, already contained an intuitive derivation of this $\mu$P for Adam, and showed it preserves optimal hyperparameters as ones scales the width of a model (e.g., Transformer ). Part of this work serves to fill the gap in ’s theoretical foundations.

To achieve this result, we leverage the Tensor Programs framework: express the adaptive optimization of a network in any parametrization in a Tensor Program, and invoke the Master Theorem to take the infinite-width limit of the whole computation, obtaining in particular the limit at the end of training.

Yet, there is one problem: no previous Tensor Program language can express adaptive optimization! The main issue is the expression of the entrywise “normalization” of the gradient done by the optimizer: For example, in the first step of training, Adam essentially just takes the sign of the gradient; while previous Tensor Program languages like $\textsc{Netsor}\top$ can express this operation for the input and output weights, they cannot do so for the hidden weights.More generally, $\textsc{Netsor}\top$ can only express this for vector-like (1 dimension tends to infinity) parameters but not matrix-like (2 dimensions).

We prove the Master Theorem for $\textsc{Ne}{\otimes}\textsc{or}\top$ (i.e., how to take the infinite width limit for any $\textsc{Ne}{\otimes}\textsc{or}\top$ program), for both the classical Gaussian case (where matrices are sampled from Gaussians) as well as the general non-Gaussian case, following the strategy of .

The Tensor Programs series is known for architecturally universal results — results that hold for all “natural deep learning architectures”, past and future. For example, established the architectural universality of Neural Network-Gaussian Process (NNGP) Correspondence, established the same for the Neural Network-Tangent Kernel Correspondence, and likewise for Free Independence Principle. Yet, ’s theoretical development of maximal update parametrization only focused on multi-layer perceptrons.with a brief discussion of defining $\mu$P for any architecture in the appendix, but nothing about its limit; did a thorough empirical investigation of muP for a variety of architectures but nothing theoretical

Here we write down the $\mu$-limit (as well as the neural tangent limit) for any architecture and any adaptive optimizer. The key innovation here is the definition of what “any architecture” means, while the proofs follow essentially the MLP examples.

Prior works did not write down the $\mu$-limits for all “natural” architectures mainly because the Tensor Programs notation was not efficient enough to deal with this arbitrary complexity. The mundane looking but vital innovation of this work is a new set of Tensor Programs notations that enables concise expression of all of the above: the bra-ket (aka Dirac) notation, borrowed from quantum physics. For readers familiar with prior Tensor Programs papers, in short:

The expectation inner product becomes notably succinct in the new notation, which also enables much more efficient expressions of the nonlinear outer product that is at the center of adaptive optimization.

We formalize a general notion of adaptive optimization in deep learning, called entrywise updates — the property that gradients are processed entrywise — satisfied by common optimizers like SGD and Adam .

We define the neural tangent and maximal update parametrizations for entrywise optimizers and derive their infinite-width limits. While we focus on MLPs in most of this paper for pedagogical purposes, we eventually write down the limits for any “reasonable” architecture.

More generally, like , we identify all “natural” infinite-width limits in this setting and dichotomize them into feature learning vs a nonlinear version of kernel regime. The maximal update limit remains the “optimal” feature learning limit for all entrywise optimizers.

All of the above results are made possible by a new version of Tensor Program, called $\textsc{Ne}{\otimes}\textsc{or}\top$, that can express the adaptive updates using the new instruction of nonlinear outer product. This forms the bulk of the technical advances made in this work.

Even so, the most vital contribution of this work is perhaps introducing the bra-ket notation to drastically simplify expressions and calculations common in Tensor Programs.

The infinite-width limits of adaptive optimizers take the headline for this paper and may be what piques readers’ interest in the near term. But the new Tensor Program and the new notation will likely have longer lasting impact, pushing forward our fundamental knowledge of large neural network behavior and lowering the translation tax between how this knowledge is stored on paper and in our heads.

In a way, this work gathers and generalizes all previous results in the Tensor Programs series about infinite-width limits: 2.9.19 generalizes the architecturally universal Neural Network-Gaussian Process correspondence and Neural Network-Tangent Kernel correspondence ; the new Tensor Program, $\textsc{Ne}{\otimes}\textsc{or}\top$ (2.6.1), generalizes the $\textsc{Netsor}\top^{+}$ of ; its Master Theorem (2.6.10) holds for both Gaussian and non-Gaussian matrices, generalizing ; the Dynamical Dichotomy for entrywise updates (2.8.20) generalizes that for SGD ; our $\mu$-limit equations generalize those of for SGD and, for the first time, we even write down the general $\mu$-limit equations for any architecture (2.9.25).

1 Related Work

Here we briefly overview past works on infinite-width neural networks, but we recommend the reader to refer to [45, Sec 2] for a more comprehensive review. A large body of literature exists on both the kernel (NTK) limit and the mean field limit for 2-layer neural network . Various papers describe the kernel and feature learning regimes more generally without taking an infinite-width limit. describes the “lazy training” regime in arbitrary differentiable programs, and is controlled by a single parameter $\alpha$ which scales the output. It is shown that when $\alpha$ is large, the weight need only move slightly to fit the training data, and network essentially performs kernel learning. Many papers view the kernel and feature learning regimes as learning in different timescales, explicitly incorporating the time dependence in the infinite-width limit, and others derive finite width corrections to the NTK for finite width networks . In this paper, as in , we consider training time to be constant, and take only the width to infinity.

Tensor Programs, first introduced in and expanded upon in , were developed as a theoretical framework to analyze the infinite-width limits of any architecture expressible in a simple formal language, in an attempt to unify the per-architecture analysis prevalent in the literature . defined a natural space of neural network parametrizations (abc-parametrizations), and classified all resulting infinite-width limits into two possible catagories: 1) the kernel regime, in which the neural network function evolves as a linear model, and 2) the feature learning regime, in which the representations change and adapt to data over the course of training. The $\mu$ parametrization was then identified as the “optimal” parametrization for arbitrary architectures in which all layers learn features, and was later heuristically extended to adaptive optimizers .

Adaptive optimizers and their variants were developed to accelerate learning by adapting the learning rate on a per parameter basis, and currently is a critical component of large scale pretraining of transformer models . No previous work has developed their theory for infinite-width neural network, but a concurrent work has derived the infinite-width NTK for SignSGD in the batch-size 1 setting (which is not equivalent to the general batch-size setting).

2 Notations

One of the the key innovations of this work is a set of much cleaner notations to express ideas in Tensor Programs. While this sizeable section may be off-putting to some readers, it’s better to explain the notation sooner than later. We recommend skimming until the end of the Outer Product subsection and then move on, coming back to read other parts when necessary.

2.1 The Tensor Program Ansatz: Representing Vectors via Random Variables

As we will see, as width becomes large, the entries of the (pre-)activation vectors and their gradients 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 any such vector’s behavior can be tracked via a random variable that reflects the distribution of its entries. While we call this the “Tensor Program Ansatz”, it is a completely rigorous calculus as seen below in Section 2.6 (as well as in previous papers in the Tensor Programs series for SGD).

Furthermore, this notation cleanly handles the multi-vector case when $\bm{x}=(x^{1},\ldots,x^{k})$ is an $n\times k$ matrix where $k$ is fixed as $n\to\infty$:

$\langle\bm{x}\hstretch{0.7}{\talloblong}\bm{y}\rangle$ represents the limit of $\bm{x}^{\top}\bm{y}/n$.Note that later, we will consider $\bm{x}$ of shape $n\times k_{1}\times\cdots\times k_{r}$, in which case $\hstretch{0.7}{\talloblong}\bm{x}\rangle$ and $\langle\bm{x}\hstretch{0.7}{\talloblong}$ both have shape $k_{1}\times\cdots\times k_{r}$, and $\langle\bm{x}\hstretch{0.7}{\talloblong}\bm{x}\rangle$ has shape $k_{1}\times\cdots\times k_{r}\times k_{1}\times\cdots\times k_{r}$.

Because we will often need to multiply a ket with a diagonal matrix, we introduce a shorthand:

if $\bm{x}$ is $n\times k$ and $\bm{\chi}$ is a $k$-dimensional vector.

Likewise, if both $\bm{x}$ and $\bm{y}$ have shape $n\times k$, the expression

which corresponds to the limit of $\frac{1}{n}\bm{w}^{\top}\bm{x}\bm{y}^{\top}$. This definition of $\hstretch{0.7}{\talloblong}\bm{x}\rangle\langle\bm{y}\hstretch{0.7}{\talloblong}$ makes the expressions

unambiguous (since any way of ordering the operations give the same answer).

where \left(\hstretch{0.7}{\talloblong}y^{1}\rangle^{\mathord{\ThisStyle{\LMpt\kern 1.0pt\LMpt\framebox{\SavedStyle 1}\kern 1.0pt\LMpt}}},\ldots,\hstretch{0.7}{\talloblong}y^{k}\rangle^{\mathord{\ThisStyle{\LMpt\kern 1.0pt\LMpt\framebox{\SavedStyle 1}\kern 1.0pt\LMpt}}},\hstretch{0.7}{\talloblong}\bm{z}\rangle^{\mathord{\ThisStyle{\LMpt\kern 1.0pt\LMpt\framebox{\SavedStyle 1}\kern 1.0pt\LMpt}}}\right) is an iid copy of $\left(\hstretch{0.7}{\talloblong}y^{1}\rangle,\ldots,\hstretch{0.7}{\talloblong}y^{k}\rangle,\hstretch{0.7}{\talloblong}\bm{z}\rangle\right)$ independent from $\hstretch{0.7}{\talloblong}\bm{x}\rangle$ and the expectation is taken only over the former. This is just like, in the finite $n$ case,

where $\otimes$ denotes outer product of vectors and expectation is taken over everything.

Note $\phi(\hstretch{0.7}{\talloblong}\bm{x}\rangle\langle\bm{y}\hstretch{0.7}{\talloblong})$ is not the image of the operator $\hstretch{0.7}{\talloblong}\bm{x}\rangle\langle\bm{y}\hstretch{0.7}{\talloblong}$ under $\phi$ in the continuous function calculus of operators, but rather a “coordinatewise application” of $\phi$. For example, if $\phi(t)=t^{2}$, then $\phi(\hstretch{0.7}{\talloblong}x\rangle\langle y\hstretch{0.7}{\talloblong})$ is not $\hstretch{0.7}{\talloblong}x\rangle\langle y\hstretch{0.7}{\talloblong}x\rangle\langle y\hstretch{0.7}{\talloblong}$, the latter being what typically “squaring an operator” means, but rather $\hstretch{0.7}{\talloblong}x\rangle^{2}\langle y\hstretch{0.7}{\talloblong}^{2}=\hstretch{0.7}{\talloblong}x\odot x\rangle\langle y\odot y\hstretch{0.7}{\talloblong}$.

In later applications, when $\phi$ is an update function (such as $Q_{t}$ in Eq. 2.1), this will be clear from context. For example, the bar notation in $\overline{\hstretch{0.7}{\talloblong}d\bm{h}^{l}_{t}\rangle{}_{\bm{\chi}_{t}}\langle\bm{x}^{l-1}_{t}\hstretch{0.7}{\talloblong}}$ abbreviates $Q^{l}_{t}$ where $l,t$ are the same as in $d\bm{h}^{l}_{t}$ inside. Then we use the much lighter “bar” notation

In contrast to the cases above, when an expression involves only kets (or bras), then the usual calculus of kets as random variables or vectors apply, e.g., $\hstretch{0.7}{\talloblong}\bm{x}\rangle\odot\hstretch{0.7}{\talloblong}\bm{y}\rangle$ is just the random vector formed from entrywise product of $\hstretch{0.7}{\talloblong}\bm{x}\rangle$ and $\hstretch{0.7}{\talloblong}\bm{y}\rangle$.From readers with quantum mechanics background, beware that $\hstretch{0.7}{\talloblong}x\rangle\hstretch{0.7}{\talloblong}y\rangle$ in our context is the product of random variables $\hstretch{0.7}{\talloblong}x\rangle$ and $\hstretch{0.7}{\talloblong}y\rangle$, which is not equal to their “tensor product” (which would be written \hstretch{0.7}{\talloblong}x\rangle\hstretch{0.7}{\talloblong}y\rangle^{\mathord{\ThisStyle{\LMpt\kern 1.0pt\LMpt\framebox{\SavedStyle 1}\kern 1.0pt\LMpt}}}).

For readers familiar with the Tensor Programs papers, this new “bra-ket” notation (aka Dirac notation) relates to the old $Z^{\bullet}$ notation by

2.2 IID Copies

2.3 Big-O Notation

Following previous papers of the Tensor Programs series, we adopt the following semantics of big-O notation, which concerns the “typical size” of entries of a tensor rather than the norm of the tensor (as is the more common usage of big-O notation). Therefore, the reader must internalize this notation sooner rather than later to avoid confusion.

Given a sequence $\bm{x}=\{\bm{x}(n)\}_{n=1}^{\infty}$ of random tensors, where $\bm{x}(n)$ can have different shapes for different $n$, we write $\bm{x}=\Theta(n^{-a})$ and say $\bm{x}$ has coordinates (or entries) of size $\Theta(n^{-a})$ if there exist constants $A,B>0$ such that almost surely,Here “almost surely” is with respect to the probability of the entire sequence $\bm{x}$. for sufficiently large $n$,

where $\#\bm{x}(n)$ is the number of entries in $\bm{x}(n)$. We make similar definitions for $O(n^{-a})$ and $\Omega(n^{-a})$.

Note the constants $A,B$ can depend on everything except $n$; in concrete contexts below, such constants can, for example, depend on neural network architecture, training time, optimizer, etc, but just not width.

Most often, $\bm{x}$ will have “approximately iid” coordinates, so the notation $\bm{x}=\Theta(n^{-a})$ can be interpreted intuitively to say $\bm{x}$ has coordinates of “empirical standard deviation” $\Theta(n^{-a})$, which justifies the name.

Chapter 2 Exposition of Main Results

Here we explain our main results while later chapters will prove them. We begin by isolating a concept that captures most of adaptive optimizers, namely entrywise optimizers (Section 2.1), which forms the focus of this work. By considering how to scale the learning rate, initialization, and multipliers (a so-called abcd-parametrization), we catalogue all natural ways of taking infinite-width limits (Section 2.2). We study the archetypical examples, the (canonical generalizations of) neural tangent (NT) (Section 2.4) and the maximal update ($\mu$) (Sections 2.5 and LABEL:{sec:deepmaximalupdate}) parametrizations, and describe their infinite-width limits. More generally, we classify all possible limits of abcd-parametrizations (Section 2.8): while most parametrizations are degenerate in one way or another, the rest can be divided into the feature learning and the operator regimes, the latter being the nonlinear counterpart of kernel regime. The $\mu$ and NT limits are respectively the “maximal” elements of each regime in that all parameters contribute to the function evolution. Nevertheless, like in the SGD case, all operator regime limits, including the NT limit, do not learn features and trivialize transfer learning. While all of the above stars the MLP as the instructional architecture, finally we write down the NT and $\mu$ limits for any architecture (Section 2.9).

Underlying these results is the new Tensor Program language, $\textsc{Ne}{\otimes}\textsc{or}\top$, that expresses the so-called nonlinear outer products (Section 2.6). We formulate the algorithm, aka the Master Theorem, to compute the infinite-width limit of any $\textsc{Ne}{\otimes}\textsc{or}\top$ program. New techniques, such as new notions of equivalence of random vectors, are needed to prove the Master Theorem; we overview the proof in Section 2.11 before giving it in full in Chapter 4.

What does one mean by adaptive optimization? Both SGD and Adam, prototypical optimizers in deep learning, have entrywise updates, where parameter updates take the form of a function of the current and/or past gradients. This turns out to be a concept that captures most of adaptive optimizers.

where $m$ and $v$ are both functions of the past gradients $g_{0},\ldots,g_{t}$ and $\epsilon>0$ is there for numerical stability. For example, in Adam , $m$ and $v$ are respectively the exponential moving averages of them and their squares, resulting in the following unwieldy expression:

where $\beta_{1},\beta_{2}$ are Adam’s momentum hyperameters. We can also consider a simpler “memoryless” version of this, namely SignSGD :

We will also write memoryful and nonstationary for the opposite of memoryless and stationary. In this sense, SGD and SignSGD are both memoryless and stationary but Adam is neither. We will always present the memoryless stationary versions of our theorems first, as they carry across the main ideas. The full version (i.e., memoryful nonstationary) will always be a straightforward modification, though often requiring more notations.

Many other optimizers are covered by this entrywise update framework Eq. 2.1, including RMSProp, Adagrad, Adadelta, NAdam, Adamax, etc . However, some other ingredients of “adaptive” optimization, such as gradient clipping, weight decay, or momentum factoring as in Adafactor , are not directly covered. Nevertheless, using our new extension of Tensor Programs discussed in Section 2.6, it is straightforward to derive and classify the infinite-width limits including such ingredients, and we do so in Section 2.10. Such theorems will be by-and-large the same as what we have here (e.g., Dynamical Dichotomy 2.8.20 still holds), but the definitions of neural tangent and maximal update parametrizations (2.4.1,2.5.1) can change, as well as, e.g., the equations characterizing feature learning. See [47, Sec B.3] for further intuitive discussions.

If $Q_{t}$ in Eq. 2.1 is nonlinear, then $\textsc{Netsor}\top^{+}$, the most advanced version of Tensor Programs before this work, is unable to express Eq. 2.1 for hidden weights of a network (weight matrices where both dimensions tend to infinity). In this work, we extend $\textsc{Netsor}\top$ to $\textsc{Ne}{\otimes}\textsc{or}\top$ that can express it and develop its Master Theorem.

2 abcd-Parametrizations

In this work, we consider how such optimization should be parametrized wrt the width of a neural network, generalizing the the abc-parametrization of .

We parametrize each weight as $W^{l}=n^{-a_{l}}w^{l}$ for actual trainable parameter $w^{l}$

We initialize each $w_{\alpha\beta}^{l}\sim\mathcal{N}(0,n^{-2b_{l}})$

The learning rate is $\eta n^{-c_{l}}$ for some width-independent $\eta$

The gradients of $w^{l}$ are multiplied by $n^{d_{l}}$ before being processed by $Q^{l}_{t}$: i.e., the update at time $t$ is

where $g_{s},s=0,\ldots,t$, are the gradients of $w^{l}$ at time $s$ and $Q^{l}_{t}$ is applied entrywise.

A simple example is the “standard parametrization” that is the default for, e.g., PyTorch, where nothing scales with width other than the initialization.

The standard abcd-parametrization (SP) is defined by

In 2.2.1, beyond the obvious addition of scaling exponent $d_{l}$, compared to abc-parametrization of , we also now have layer dependent $c_{l}$. This is without loss of generality, because of the redundancy in $a_{l},b_{l},c_{l}$, as shown in [45, Eq 5]. This takes the more general form as follows for abcd-parametrization:

In the case of SGD (i.e., when $Q^{l}_{t}(g_{0},\ldots,g_{t})=g_{t}$), an abcd-parametrization reduces to an abc-parametrization with the mapping

As in , our concern here is the correct way to scale with width. In practice, there should be tunable hyperparameters in front of the powers of $n$ in 2.2.1, as investigated in .

In the idealized case of Adam and similar adaptive optimizers where the $\epsilon$ in Eq. 2.2 is 0, $Q^{l}_{t}$ is degree-0 homogeneous and $d_{l}$ itself is redundant. When $\epsilon>0$, this is no longer true. But the almost homogeneity yields an alterative but equivalent way to define $d_{l}$: instead of $g_{s}$ being multiplied by $n^{d_{l}}$, we let $\epsilon$ be multiplied by $n^{-d_{l}}$.

If one sets the global learning rate for SP (2.2.2) to its largest stable value, then for SGD, SP is in the kernel regime . But for SignSGD and Adam, assuming perfect scale invariance, SP’s largest stable learning rate is $\Theta(1/n)$, so that with this setting (i.e., setting $c_{l}=1$ for every $l$), SP is actually in the feature learning regime. The reader is not expected to understand the underlying reasoning at this point, but the claim above can be derived by calculating $r=0$ from 2.8.5 and invoking 2.8.12 and 2.8.19. This difference in default scaling may be a contributing factor to the success of Adam compared to SGD.

3 Setup

Here we set up the notation and conventions regarding the data, (pre)activations, and training of the network as well as the main technical assumptions for our rigorous results.

While considered only the batch-size-1 case to simplify notation, here, because of our notational advance, we can afford to consider the following more general setting.

Furthermore, we can implement train-test split via $\varepsilon_{t}$: For example, $\bm{\xi}$ can be split into two parts, $\bm{\xi}=(\bm{\xi}^{train},\bm{\xi}^{test})$ such that $\varepsilon_{t}$ is always 0 on the $\bm{\xi}^{test}$. Then the evolution of $\bm{f}_{t}$ can track the evolution of function values on the test set due to changes from the training set.

The $\varepsilon_{t}$ framework more generally covers settings like reinforcement and online learning where the error signal is not obtained from just a simple loss function.

For all rigorous results in this work, we will consider the following smoothness assumption. This is sufficient for deriving the NTK and $\mu$P limits but more assumptions, stated later, are required for Dynamical Dichotomy, i.e., the classification of abcd-parametrizations.

This is a very weak assumption satisfied by typical loss functions (e.g., MSE or cross entropy), update functions (e.g., SGD or Adamspecifically, Adam in the form Eq. 2.2 with $\epsilon>0$), and nonlinearities (e.g., tanh or gelu). The notable exception here is that relu itself is not covered because its derivative has a discontinuity. But this is a common technicality not treated in the theoretical literature. Nevertheless, we expect all theorems in this work should apply to relu as well and can be proven rigorously in the future.

With this setup in mind, we next describe the two prototypical infinite-width limits, the neural tangent and maximal update limits, for nonlinear entrywise updates before completely classifying the space of abcd-parametrization.

4 Neural Tangent

The “classical” neural tangent (abc-)parametrization (NTP) can be generalized easily to an abcd-parametrization using the intuition that the input to any nonlinear update function should be $\Theta(1)$ (we won’t go through this calculation here but c.f. 2.8.8 and LABEL:{lemma:faithfulconditionsinit}). After defining this generalization next, we adapt the well-known continuous-time heuristic for deriving the NTK limit to the nonlinear updates case (Eq. 2.7), before writing down a succinct expression of the infinite-width limit made possible by our new bra-ket notation (Eqs. 2.8 and 2.9).

The neural tangent abcd-parametrization (NTP) is defined (modulo Eq. 2.5) by

Reducing to abc-parametrization in the SGD case via Eq. 2.6 (i.e., subtracting the last 2 rows), we recover exactly the classical NTP [45, Table 1].

4.1 Continuous Time Intuition

where $\mathcal{L}_{t}=\mathcal{L}(\bm{f}_{t})$ and $\mathcal{L}$ is the loss function. For a general (memoryless stationary) update function $Q$, this just becomesTechnically, we should include terms involving $d_{l}$ from 2.2.1 in Eq. 2.7, but for simplicity, let’s just assume that $Q$ is temporarily redefined to have already included them

Now we turn to make this intuition rigorous.

4.2 Infinite-Width Limit

where the “bar” notation abbreviates application of $Q$ as in Eq. 1.2.

As an example, when $Q$ is identity, the “bar” can be removed, and $\mathcal{K}_{Q}$ reduces to the linear operator represented by the NTK:Again, $\bm{\chi}$ is a row vector. In prior works, it’s usually treated as a column vector in which case one would write $K\bm{\chi}$ instead.

As another example, consider $Q=\operatorname{sign}$ (Eq. $\mathtt{SignSGD}$ with $\epsilon=0$). If the batch size is 1, i.e., $\bm{\chi}$ is nonzero on exactly one input, say $\xi^{b},b\in[\mathcal{M}]$, then Eq. 2.8 is linear in $\operatorname{sign}(\bm{\chi})$ because $\operatorname{sign}(xy)=\operatorname{sign}(x)\operatorname{sign}(y)$. Thus,

for each $a\in[\mathcal{M}]$. This expression was concurrently derived in . However, when batch size is larger than 1, $\mathcal{K}_{Q}(\bm{\chi})$ is no longer linear in $\operatorname{sign}(\bm{\chi})$ generally, and there is no simplification like this. In particular, in the continuous time gradient flow setting, SignSGD with a large batch size is not equivalent to that with batch size 1 but very small learning rate, in contrast to SGD.

We can now state the generalized Neural Tangent limit for memoryless stationary updates. In particular, this covers SGD and SignSGD (Eq. $\mathtt{SignSGD}$).

Consider any training routine (2.3.1) with memoryless stationary update function $Q$ (2.1.1). Adopt 2.3.2. Then

The proof can be found in Section 3.3. Note, as in the SGD case, $\mathring{f}_{t}$ is deterministic conditioned on $\mathring{f}_{0}$. We remind the reader that Eq. 2.9 simultaneously covers full batch, mini-batch, train-test split, and other schemes by changing $\varepsilon_{t}$, as discussed under 2.3.1. This will be the same for all theorems in this work.

The memoryless stationary condition allowed a clean mathematical formulation of the NT limit. But we can remove it easily at the cost of some more notation.

where $\overline{\hstretch{0.7}{\talloblong}d\bm{h}^{l}\rangle{}_{\bm{\chi}_{\leq t}}\langle\bm{x}^{l-1}\hstretch{0.7}{\talloblong}}$ is shorthand for $Q^{l}_{t}\left(\hstretch{0.7}{\talloblong}d\bm{h}^{l}\rangle{}_{\bm{\chi}_{0}}\langle\bm{x}^{l-1}\hstretch{0.7}{\talloblong},\ldots,\hstretch{0.7}{\talloblong}d\bm{h}^{l}\rangle{}_{\bm{\chi}_{t}}\langle\bm{x}^{l-1}\hstretch{0.7}{\talloblong}\right)$.

With this in mind, the following theorem yields the NT limit of Adam (Eq. $\mathtt{Adam}$) as a corollary.

If the update functions $\bm{Q}$ are memoryless but not necessarily stationary, then 2.4.7 holds with Eq. 2.9 replaced by

For general $\bm{Q}$, not necessarily memoryless, 2.4.7 holds with Eq. 2.9 replaced by

The proof is a straightforward adaptation of the proof of 2.4.7 in Section 3.3.

4.3 Lack of Feature Learning

Just as for the SGD case, the neural tangent limit cannot learn features, for example, in the sense that the feature kernel (the Gram matrix of the input representations) does not evolve during training (2.8.19). Likewise, pretraining is futile in this limit as finetuning it would be no different than finetuning a randomly initialized network (2.8.22). Nevertheless, the NT limit is “maximal” among all nondegenerate limits without feature learning in that all other limits are just given by a neural tangent operator that involves a subsum of Eq. 2.8 (2.8.21).

5 Maximal Update

As for NTP, the “classical” maximal update (abc-)parametrization can be generalized easily to an abcd-parametrization using the intuition that the input to any nonlinear update function should be $\Theta(1)$ (c.f. 2.8.8 and LABEL:{lemma:faithfulconditionsinit}). After defining this generalization next, we study its limit for shallow MLP. The deep case will have to wait until we develop the new Tensor Program theory in the next section.

The maximal update abcd-parametrization ($\mu$P) is defined (modulo Eq. 2.5) by

If we assume that the Adam update function (Eq. $\mathtt{Adam}$) is perfectly scale-invariant, then the $d_{l}$ row can be dropped, yielding [47, Table 8] regarding Adam LR scaling.

To recover the abc version of $\mu$P in [45, Table 1] for SGD, just apply Eq. 2.5 to the $l=1,L+1$ columns with $\theta=\nicefrac{{-1}}{{2}}$ and then apply Eq. 2.6 (i.e., subtracting the last 2 rows).

5.1 Shallow Infinite-Width Limit

We focus on the shallow case first because its $\mu$-limit is fairly easy to describe. We shall cover the general case after we describe the outer product tensor program $\textsc{Ne}{\otimes}\textsc{or}\top$ in Section 2.6. Adopt the following leaner notation:

Consider any training routine (2.3.1) with memoryless stationary update function $Q$ (2.1.1). Adopt 2.3.2. As $n\to\infty$, $\bm{f}_{t}$ for the network in Eq. 2.13 converges almost surely to some $\mathring{\bm{f}}_{t}$ for every $t$, which is recursively defined from $t=0$ by the following dynamics:

where the bar notation (Eq. 1.2) abbreviates application of $Q$ and $\cdot$ denotes dot product.

Again, one can note that if $Q$ is identity, then we recover the SGD equations from [45, Theorem 6.1].

The proof of 2.5.2 is given in Section 3.2. This can be adapted straightforwardly to cover the nonstationary or memoryful cases:

If $\bm{Q}$ is memoryless but not stationary, then 2.5.2 holds if the bar in Eq. 2.15 (resp. Eq. 2.14) is interpreted as $Q^{2}_{t}$ (resp. $Q^{1}_{t}$).

If $\bm{Q}$ is not memoryless, then 2.5.2 holds if Eqs. 2.15 and 2.14 are replaced with

2.5.2 and 2.5.3 also hold if the output dimension is greater than 1, in which case the equations should be interpreted slightly differently; see Eq. 2.35.

As mentioned above, the gradient processing done by entrywise optimizers in general cannot be expressed previously by even the most expressive Tensor Program language. Here we fix this issue by adding “nonlinear outer products” to Tensor Programs. Like in previous works, we algebraically construct such programs’ limit objects (“kets”, in our new notation) and link them to the analytic properties of vectors in the corresponding programs via a Master Theorem. Unlike previous works, we also construct the limits of matrices, which are operators on kets. This is purely a conceptual change, but this perspective helps express the deep $\mu$-limit much more efficiently than possible before.

choose a matrix $W\in\mathcal{W}$ and vector $x\in\bm{x}$, and append to $\bm{x}$ the vector

where $\bm{x}_{\gamma}$ is the $\gamma$th row in $\bm{x}$ as a matrix and $|\bm{x}|$ is the number of vectors in $\bm{x}$. We call $r+1$ the order of $\psi$ in this context.

Note that while 2.6.1 doesn’t directly give an instruction to transform scalars into scalars (in contrast to MatMul and OuterNonlin that transforms vectors into vectors), this can be done by combining instructions.

Think of $\psi$ as a function that ignores vector arguments and depend only on scalars, and use it in OuterNonlin to create the vector whose entries are identically equal to $\psi(\bm{c})$. Applying Avg to this vector gives the desired result. ∎

6.2 Setups

We are interested in the behavior of $\textsc{Ne}{\otimes}\textsc{or}\top$ programs in two typical settings:

AssumeCompared to , we have WLOG simplified the setup by assuming 1) $\sigma_{W}=1$ for every $W$, 2) $Z^{\bm{x}^{0}}=\hstretch{0.7}{\talloblong}\bm{x}^{0}\rangle=\mathcal{N}(0,I)$, and 3) $\mathring{\theta}=0$ for every $\theta\in\bm{c}^{0}$. This is WLOG because $\sigma_{W}$, $\mathring{\theta}$, and the mean and covariance of $\hstretch{0.7}{\talloblong}\bm{x}^{0}\rangle$ can all be absorbed into OuterNonlin via the appropriate linear functions.

Every entry of every $W\in\mathcal{W}$ is sampled iid from $\mathcal{N}(0,1/n)$.

Every entry of every initial vector $x\in\bm{x}^{0}$ is sampled iid from $\mathcal{N}(0,1)$.

The initial scalars $\bm{c}^{0}$ converge almost surely to 0.

All functions $\psi$ used in OuterNonlin are pseudo-Lipschitz.

Assume the same as 2.6.3 but replace 1) and 4) with

there exists a sequence $\nu_{3},\nu_{4},\ldots>0$ such that all matrices have independent entriesFor all of our results, it does not matter how the matrices for different $n$ are correlated, e.g., whether they are independent or the matrices are all upper left submatrix of fixed infinite iid matrix. This is because our proof only depends on how moments of vectors behave with $n$, which does not care about such inter-$n$ correlations. drawn from distributions with zero mean, variance $n^{-1}$, and all higher $k$th moment bounded by $\nu_{k}n^{-k/2}$; andInitial vectors are still sampled from $\mathcal{N}(0,1)$, as in .

We further require initial scalars $\bm{c}^{0}$ to have moments of all orders bounded in $n$.But the moments do not need to be bounded as a function of the order.

While 2.6.4 allows more general distributions for matrix entries, its nonlinearities need to have more smoothness than 2.6.3. See for more discussions on these setups.

6.3 Limit Objects

As before, when the width $n$ of the program goes to infinity, one can infer how the program behaves via a calculus of random variables. We define them below via the new ket notation instead of the earlier $Z$ notation.

We recursively define the random variable $\hstretch{0.7}{\talloblong}x\rangle$ (called a ket) for each vector $x$ and deterministic number $\mathring{\theta}$ for each scalar $\theta$ in the program. For a vector $Wx$ produced by MatMul, we also define random variables $\hstretch{0.7}{\talloblong}Wx\hat{\rangle\mkern-3.0mu}\mkern 3.0mu$ and $\hstretch{0.7}{\talloblong}Wx\dot{\rangle\mkern-3.0mu}\mkern 3.0mu$ (called hat-ket and dot-ket respectively) such that $\hstretch{0.7}{\talloblong}Wx\rangle=\hstretch{0.7}{\talloblong}Wx\hat{\rangle\mkern-3.0mu}\mkern 3.0mu+\hstretch{0.7}{\talloblong}Wx\dot{\rangle\mkern-3.0mu}\mkern 3.0mu$. Their recursive definitions are given below.

If $x$ is generated by OuterNonlin as in Eq. 2.17, then

Every dot-ket is a linear combination of previous kets, expressed by the following equation

Eq. 2.19 is the same equation [45, Zdot] but formulated much more succinctly in the bra-ket notation:

The bra $\langle\nabla_{W^{\top}\bm{x}}\hstretch{0.7}{\talloblong}$ is really the “dual” of $\mkern 3.0mu\hat{\mkern-3.0mu\langle}W^{\top}\bm{x}\hstretch{0.7}{\talloblong}$ in the sense thatBut note this identity only holds when $\bm{x}$ contains all vectors $z$ where $\hstretch{0.7}{\talloblong}y\rangle$ depends on $\hstretch{0.7}{\talloblong}W^{\top}z\hat{\rangle\mkern-3.0mu}\mkern 3.0mu$.

This follows from Stein’s lemma. In the language of Riemannian geometry, if we think of $\langle\bm{x}\hstretch{0.7}{\talloblong}\bm{x}\rangle$ as a metric tensor in a Riemannian manifold, then $\langle\nabla_{W^{\top}\bm{x}}\hstretch{0.7}{\talloblong}$ is obtained from $\mkern 3.0mu\hat{\mkern-3.0mu\langle}W^{\top}\bm{x}\hstretch{0.7}{\talloblong}$ by “lowering the index.” Thus, a more appropriate notation for $\langle\nabla_{W^{\top}\bm{x}}\hstretch{0.7}{\talloblong}$ is perhaps

However, this “duality” is not essential for understanding this paper, so we keep the more intuitive notation $\langle\nabla_{W^{\top}\bm{x}}\hstretch{0.7}{\talloblong}$ instead.

Let $W$ be an initial matrix in a $\textsc{Ne}{\otimes}\textsc{or}\top$ program. We define $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$ to be the linear operator on kets To be rigorous, we need to specify the “Hilbert space” of kets. This is somewhat pedantic and not crucial to the key points of this paper, but the Hilbert space can be constructed as follows: Let $\sigma(\pi)$ be the $\sigma$-algebra generated by the kets of the program $\pi$. Let \Sigma(\pi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{{def}}}{=}}}}\bigcup_{\pi^{\prime}\supseteq\pi}\sigma(\pi) be the union (more precisely, the direct limit) of $\sigma(\pi^{\prime})$ over all programs $\pi^{\prime}$ extending $\pi$. Then the Hilbert space in question is the $L^{2}$ space of random variables over the $\Sigma$ of our program. that acts by

Any linear operator that is equal to $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$ for some initial matrix $W$ is called an initial operator. A set of initial operators is called independent if their corresponding initial matrices are distinct.i.e., independently sampled in 2.6.3 or 2.6.4.

We have already seen an example of a linear operator on kets: expressions like $\overline{\hstretch{0.7}{\talloblong}y\rangle{}_{\chi}\langle z\hstretch{0.7}{\talloblong}}$. 2.6.7 puts $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$ in the same space as $\overline{\hstretch{0.7}{\talloblong}y\rangle{}_{\chi}\langle z\hstretch{0.7}{\talloblong}}$. This allows us to add them in the sequel, which simplifies the presentation of the $\mu$-limit.

We can immediately see a few properties of $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$ by considering the counterpart when $n$ is finite.

For any initial matrix $W$, the operator $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$ is bounded.i.e., there exists real number $L>0$ such that for any ket $\hstretch{0.7}{\talloblong}x\rangle$, $\langle Wx\hstretch{0.7}{\talloblong}Wx\rangle\leq L\langle x\hstretch{0.7}{\talloblong}x\rangle$.

This follows from the classical operator norm tail bounds on iid random matrices (see, e.g., ), which passes to the limit via 2.6.10 below. ∎

For any initial matrix, the operator $\hstretch{0.7}{\talloblong}W^{\top}\hstretch{0.7}{\talloblong}$ is the adjoint“adjoint” in the sense of Hilbert space operators; see Footnote 24. of the operator $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$.

i.e., $\hstretch{0.7}{\talloblong}W^{\top}\hstretch{0.7}{\talloblong}$ is the adjoint of $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$. ∎

6.4 The Master Theorem

Our key foundational result is that the Master Theorem of earlier Tensor Programs generalizes to $\textsc{Ne}{\otimes}\textsc{or}\top$ programs. This underlies all of our theorems about adaptive optimization.

Consider a $\textsc{Ne}{\otimes}\textsc{or}\top$ program with (Gaussian) 2.6.3 or (non-Gaussian) 2.6.4. Then, as $n\to\infty$, its scalars $\bm{c}$ satisfy

7 Maximal Update for Deep MLP

In this section we describe the infinite-width limit of $\mu$P for arbitrarily deep MLP. The main difference here compared to the shallow case (Section 2.5.1) is the presence of $n\times n$ iid matrices in the middle of the network, which behaves like initial operators (2.6.7) in the limit.

Consider any training routine (2.3.1) with memoryless stationary update function $Q$ (2.1.1). Adopt 2.3.2. As $n\to\infty$, $\bm{f}_{t}$ converges almost surely to some $\mathring{\bm{f}}_{t}$ for every $t$, which is recursively defined from $t=0$ by the following dynamics: We remind the reader that, in $\mu$P (2.5.1), $w_{t}^{L+1}$ is the output layer weights normalized so that $w_{t}^{L+1}$ has $\Theta(1)$-sized entries (whereas $W_{t}^{L+1}=\frac{1}{n}w_{t}^{L+1}$). The same point applies to $w_{t}^{1}$ (but $W_{t}^{1}=w_{t}^{1}$). We use lower case $w$ for input and output weights while upper case $W$ for other layers to emphasize that the former are vector-like parameters (one dimension going to $\infty$) while others are matrix-like (two dimensions going to $\infty$).

(Initialization) $\hstretch{0.7}{\talloblong}W^{2}_{0}\hstretch{0.7}{\talloblong},\ldots,\hstretch{0.7}{\talloblong}W^{L}_{0}\hstretch{0.7}{\talloblong}$ are independent initial operators (2.6.7), and

One can check that when $L=1$, we recover 2.5.2.

Here, Eq. 2.21 uses the operator semantics discussed in and below 2.6.7 to cleanly express the parameter update. When unwinded, Eq. 2.21 is equivalent to

The proof of 2.7.1 and 2.5.2 can be found in Section 3.2.

Observe that all of the equations above are essentially the “ket” versions of what one does in a finite network. This holds for general architectures: the $\mu$-limit can always be obtained straightforwardly transcribing the tensor operations in a finite network to their counterparts acting on kets in the infinite-width limit. See 2.9.25. In this sense, $\mu$P is the most natural parametrization.

As before, the general case without stationarity or memorylessness is straightforward given 2.7.1, albeit with some more notation.

If $\bm{Q}$ is memoryless but not stationary, then 2.7.1 holds if the bars in Eqs. 2.21, 2.20 and 2.22 are interpreted as $Q^{l}_{t}$ (where $l$ is the same as the layer index appearing in $W^{l}_{t+1}$ or $w^{l}_{t+1}$ on the LHS).

If $\bm{Q}$ is not memoryless, then 2.7.1 holds if Eqs. 2.21, 2.20 and 2.22 are replaced with

2.5.2 and 2.5.3 also hold if the output dimension is greater than 1, but the equations need to be interpreted slightly differently. See Eq. 2.35.

For SGD, showed $\mu$P is the unique stable parametrization where every weight matrix is updated maximally [45, Defn 5.2] and the output weight matrix is also initialized maximally [45, Defn 5.4]. These definitions still make sense here and this statement holds as well if “stable” is replaced with “stable and faithful” (2.8.7, 2.8.8).

8 Dynamical Dichotomy: The Classification of abcd-Parametrizations

In this section, we characterize the infinite-width limits of all possible abcd-parametrizations under reasonable assumptions of the optimizer and network nonlinearity, generalizing the work done for SGD in . First, we filter out the uninteresting limits: the unstable (training blows up), the trivial (training gets stuck at initialization), and the unfaithful (the update functions $\bm{Q}$ and/or nonlinearities $\phi$ are trivialized). We sort all other limits into a Dynamical Dichotomy (2.8.20) between feature learning and operator regimes (the latter being the nonlinear version of kernel regime in the SGD case). The $\mu$ and NT limits are respectively their archetypes (indeed, the maximal parametrizations in these regimes (LABEL:{rem:maximality})). Like in the SGD case, this dichotomy is not tautological: it implies certain network training dynamics cannot be the infinite-width limit of any abcd-parametrization (2.8.24). Likewise, pretraining is still futile in the operator regime even with adaptive optimizers (2.8.22).

Our results here hold for not only memoryless stationary but also memoryful nonstationary updates.

For proving the Dynamical Dichotomy Theorem for entrywise updates, we will make the following technical assumptions. Roughly speaking, we will focus only on “relu-like” nonlinearities and sign-preserving update functions with mild smoothness.

$\phi$ and $\phi^{\prime}$ are nonnegative and pseudo-Lispchitz.

$Q^{l}_{t}$ is pseudo-Lipschitz for all $l,t$, and $Q^{l}_{0}$ preserves sign for all $l$.

As in , the pseudo-Lipschitzness of $\phi,\phi^{\prime},Q^{l}_{t}$ are sufficient for letting us use our new Master Theorem (2.6.10) to take the limits of any parametrization, get the operator limits, and prove their properties. Any assumptions beyond such is required only for proving that $r=0$ implies feature learning (3.1.3). In particular, the reason we only require $Q^{l}_{0}$ to preserve signs (instead of for all $t$) is because we will only need to show that features evolve in the first step.

8.2 Size of Feature Learning

In , the number $r$ of an abc-parametrization measures how much the features change over training. We adapt its definition to abcd-parametrizations.

Morally, in any “reasonable” abcd-parametrization (in the sense of stable (2.8.7) and faithful (2.8.8) discussed below), we have $\Delta W_{t}^{l}\bm{x}_{t}^{l-1}=\Theta(n^{-r_{l}})$ and $\Delta\bm{x}_{t}^{L}=\Theta(n^{-r})$, where $\Delta\bullet_{t}=\bullet_{t}-\bullet_{0}$ is the cumulative change of $\bullet_{t}$. Concretely, for NTP and $\mu$P we have, for all $l\in[1,L]$,

The reader should sanity check that $r_{l}$ and $r$ are invariant to the symmetry in Eq. 2.5.

Ostensibly, 2.8.5 is different from and much simpler than [45, Defn 3.2] for abc-parametrizations. But, in fact, they are equivalent for stable and faithful abcd-parametrizations (under the reduction Eq. 2.6 to abc-parametrizations). The comparative simplicity of 2.8.5 is also due to the faithfulness.

8.3 Stability and Faithfulness

We will only care about any parametrization satisfying two basic properties: 1) does not blow up during at initialization or during training as width $\to\infty$ and 2) does not trivialize the update functions $Q_{t}^{l}$. The former property is known as stability and has already been studied in for SGD. The latter is specific to nonlinear entrywise updates, and we will call this the faithful property. By “trivialize,” we mean that the input to $Q_{t}^{l}$ is either i) too small and thus linearizing $Q_{t}^{l}$ around the origin or ii) too large and only depends on $Q_{t}^{l}$’s behavior “around infinity”. Both scenarios ignore the bulk of $Q_{t}^{l}$’s values as a function. If such behaviors are actually desired, then one can change $Q^{l}_{t}$ to such effects. For example, the linearizing behavior in case i) can be implemented by choosing a linear $Q_{t}^{l}$ and modifying $c_{l},d_{l}$ appropriately so that the input to $Q_{t}^{l}$ has constant typical size (wrt width).

Recall from Section 1.2.3 the entry-wise semantics of big-O notation. We formalize the stability and faithfulness properties below.

We say an abcd-parametrization of an $L$-hidden layer MLP is

stable during training if for any training routine, any time $t\geq 0$, $l\in[L]$, we have

We say the parametrization is stable if it is stable both at initialization and during training.

We say an abcd-parametrization is faithful at time $t$ if the input to $Q_{{}_{t}}^{l}$ is $\Theta(1)$ for every $l\in[L]$. We also say it is faithful at initialization if this is true at $t=0$.

The condition $\bm{h}_{0}^{l},\bm{x}_{0}^{l}=\Theta(1)$ in Eq. 2.24 is in truth more of a “faithfulness to $\phi$” condition than just stability (which would strictly speaking be more like $O(1)$ than $\Theta(1)$). But we never need to distinguish these notions, so following , we will keep the definition as is.

Adopt 2.8.2. An abcd-parametrization is stable at initialization iff

This condition is the same as [45, Thm H.6(1)]. For example, SP, NTP, and $\mu$P are all stable at initialization. In this situation, some easy calculation shows that, at initialization, the last layer gradients have entry size $\Theta(n^{-a_{L+1}})$ and while all other layers’ gradients have entry size $\Theta(n^{-a_{l}-a_{L+1}-b_{L+1}})$. Hence,

In Lemma 2.8.10, the abcd-parametrization is furthermore faithful at initialization iff

For example, NTP and $\mu$P are faithful at initialization but SP (2.2.2) is not (but if $\bm{Q}$ is scale invariant then SP is equivalent to a faithful parametrization).

Adopt 2.8.2. Suppose, at initialization, an abcd-parametrization is both faithful and stable. Then it remains so for all $t\geq 1$ iff

In words: 1) $r_{l}\geq 0$ for all $l\in[L]$ ensures the features do not blow up, while 2) $r_{L+1}\geq 0$ and $a_{L+1}+b_{L+1}+r\geq 1$ resp. ensure that $W_{t}^{L+1}\bm{x}_{t}^{L},W_{0}^{L+1}\Delta\bm{x}_{t}^{L}=O(1)$, so $\bm{f}_{t}$ does not blow up;Recall $\Delta\bullet_{t}=\bullet_{t}-\bullet_{0}$ is the cumulative change of $\bullet_{t}$. finally, 3) $b_{L+1}\leq c_{L+1}$ ensures that $W^{L+1}$ does not change scale after updates, since otherwise the $d_{l}$ in Eq. 2.26 is no longer faithful.

One can check that 2.8.12 reduces to [45, Thm 3.3] (the stability characterization of abc-parametrizations) plus the additional constraint that $W^{L+1}_{t}=O(W^{L+1}_{0})$ for all $t$ (which is imposed by $b_{L+1}\leq c_{L+1}$ in Eq. 2.27). As remarked above, this constraint is due to the faithfulness requirement. But in fact, if we allow $d_{l}$ to vary during training, then this constraint can be removed and the appropriate version of 2.8.12 would be equivalent to [45, Thm 3.3].

8.4 Nontriviality

Even among stable and faithful parametrizations, we are only interested in nontrivial parametrizations (as in ), where in the infinite width limit the network will not be stuck at initialization during training.

Nontriviality is characterized by a disjunction of equations in $a_{l},b_{l},c_{l}$, just as for SGD in .

Adopt 2.8.2. A stable and faithful abcd-parametrization is nontrivial iff

This is essentially equivalent to [45, Thm 3.4]. For example, NTP and $\mu$P are both nontrivial.

8.5 Feature Learning and Operator Regimes

Finally, we ask: what are the different possible behaviors among the nontrivial, stable, and faithful parametrizations? As in , we will see a dichotomy between feature learning and a nonlinear version of the kernel regime which we call the operator regime.

For memoryless, nonstationary updates, we allow $\mathcal{K}$ to depend on $t$. For general entrywise updates, we make the same definition if for all $t$,

Notice that the operator regime is defined solely in the function space, without talking about the internals of the network, in contrast to feature learning.

We say an abcd-parametrization of an $L$-hidden layer MLP admits feature learning in the $l$th layer if there exists some training routine such that

for some $t\geq 0$. We say the parametrization admits feature learning if it does so in any layer.

We say the parametrization fixes the $l$th layer features if for all training routine,

for all $t\geq 0$. We say the parametrization fixes all features if it does so in every layer.

We make similar definitions as above replacing feature with prefeature and $\bm{x}^{l}$ with $\bm{h}^{l}$.

We say an abcd-parametrization of an $L$-hidden layer MLP evolves the $l$th layer feature kernel if there exists some training routine such that

for some $t\geq 0$. We say the parametrization evolves feature kernels if it does so in any layer.

We say the parametrization fixes the $l$th layer feature kernel if for all training routine,

for all $t\geq 0$. We say the parametrization fixes all feature kernels if it does so in every layer.

We make similar definitions as above replacing feature with prefeature and $x^{l}$ with $h^{l}$.

8.6 Classification of abcd-Parametrizations

The classification of abcd-parametrizations is similar to that of abc-parametrizations [45, Thm H.13]. We remind the reader that this holds for not only memoryless stationary but more generally memoryful nonstationary updates.

Adopt 2.8.2. Consider a nontrivial, stable, and faithful abcd-parametrization of an $L$-hidden layer MLP. Then

feature kernel evolution in the $L$th layer

prefeature kernel evolution in the $L$th layer

fixes prefeature kernel in the $L$th layer

If there is feature learning or feature kernel evolution or prefeature learning or prefeature kernel evolution in layer $l$, then there is feature learning and feature kernel evolution and prefeature learning and prefeature kernel evolution in layers $l,\ldots,L$.

Consequently, we can generalize Dynamical Dichotomy to (nonlinear) entrywise updates.

A nontrivial, stable, and faithful abcd-parametrization either admits feature learning or is in the operator regime, but not both.

Of course, the canonical examples here are $\mu$P in the feature learning regime and NTP in the operator regime. In the SGD case, LABEL:{thm:abcdclassification} and 2.8.20 are equivalent to their counterparts [45, Thm H.13, Cor H.14] other than a minor technical difference as discussed in 2.8.13.

The $\mu$ and NT limits are resp. the “maximal” limits in the feature learning and operator regimes, in the sense that all parameter tensors contribute to the function update, and that any other limits in those regimes are just “downgrades” of $\mu$ and NT limits by zeroing out the initialization or learning rate of some parameters. See also 2.7.5.

[45, Thm H.17] holds almost verbatim in our case as well, after replacing “stable” with “stable and faithful” and “kernel regime” with “operator regime”: Finetuning any pretrained network in the operator regime would be equivalent to finetuning a randomly initialized network. Thus, pretraining in the operator regime is useless.

In the memoryless stationary case, an operator regime limit resides solely in the function space picture, i.e. $\bm{f}_{t+1}$ being solely determined by the function values $\bm{f}_{t}$ themselves (as opposed to the internal activations of $f$ as well) along with learning rate $\eta$ and error signals $\varepsilon_{t}$. However, as in [45, Remark 3.11], this is not true of any feature learning limit because one can construct counterexamples where $\bm{f}_{t}$ are close for two infinite-width limits but $\bm{f}_{t+1}$ are far.

Compared to the kernel regime in the SGD case, the operator regime now allows nonlinear evolution in the function space picture. Nevertheless, in such dynamics, $\bm{f}_{t+1}-\bm{f}_{t}$ must be linear in $\eta$ for every $t$. Thus, any function space evolution nonlinear in $\eta$ cannot be the infinite-width limit of any entrywise optimizer.The same holds for any adaptive optimizer with ingredients discussed under Optimizer Coverage of Section 2.1. For example, $\bm{f}_{t+1}-\bm{f}_{t}=-\eta\bm{f}_{t}-\eta^{2}\bm{f}_{t}^{2}$ is not a valid limit.

[45, Sec G] identified a subclass of abc-parametrizations, called uniform parametrizations, where all layers “learn the same amount of features” and the output layer is initialized and updated maximally. This is used in to give an alternative presentation of $\mu$P as well as discussion of joint width-depth limit. This notion also makes sense for abcd-parametrizations: For every $s\in[0,\nicefrac{{1}}{{2}}]$, there is a unique stable and faithful abcd-parametrization, called UPs such that $r_{l}=s$ for all $l=1,\ldots,L$ and $r_{L+1}=1$ and $a_{L+1}+b_{L+1}=1-s$. For example, UP0 is $\mu$P and UP${}_{\nicefrac{{1}}{{2}}}$ is NTP.

9 Infinite-Width Limits for Any Architecture

Having written down the infinite-width limits of adaptive optimizers for MLPs, we now turn to general architectures. An astute reader may have already absorbed the key insights from the previous sections and can use them to derive the NT or $\mu$ limits for each new architecture in an ad hoc fashion. In contrast, here we describe the algorithm to do this once and for all, uniformly for all “reasonable” architectures. This uniformity of course requires abstraction, which is not conducive to quick comprehension; on the flip side, this algorithm will always be here for someone to fall back to if the ad hoc approach does not work out.

The main work here is not the proofs (which can be easily adapted from the MLP case and so are omitted), but the definitions: What is an architecture? What architecture counts as “reasonable”? How to define abcd-parametrization for any such architecture? What are their infinite-width limits? Answering these questions in the most generality requires careful thought.

As demonstrated in , 2.9.1 covers essentially every architecture in practice: RNNs, residual network, transformers, etc.

In this definition, only the symbolic structure of the program matters; the random sampling of 2.6.3 and 2.6.4 plays no role.

For simplicity, we only considered the case when the notion of “width” is the same throughout the network. Nevertheless, this definition can be easily modified to cover the nonuniform case, but stating it would be much more complex.

In truth, we could have phrased 2.9.1 using $\textsc{Netsor}\top$ instead of $\textsc{Ne}{\otimes}\textsc{or}\top$, since we do not know of any neural network in the wild that is not $\textsc{Netsor}\top$-representable ($\textsc{Ne}{\otimes}\textsc{or}\top$ is only required for expressing adaptive optimization like Adam). But there is little cost in stating the more general version, which can potentially matter in the future.

also defined a notion of representable functions using $\textsc{Netsor}\top$. In comparison, our definition is much more general: Beyond the superficial difference of $\textsc{Ne}{\otimes}\textsc{or}\top$ here vs $\textsc{Netsor}\top$ there, dealt with input and output layer weights in special ways, whereas here we do not, instead opting to uniformly deal with scalar, vector, and matrix parameters. The input and output weights of are vector parameters in this view.

The $L$-hidden-layer MLP in Eq. 2.3 has $d+1$ vector parameters (where $d$ corresponds to input dimension and $1$ corresponds to output dimension) and $L-1$ matrix parameters. It is represented by the program that generates 1) $h^{1}$ using OuterNonlin and $h^{l},l\geq 2,$ using MatMul; 2) $x^{l}$ using OuterNonlin; and 3) generate function output by summing $W^{L+1}\odot x^{L}$ (so we can take $x^{1}$ in 2.9.1 to be $W^{L+1}\odot x^{L}$).

9.2 abcd-Parametrization for Any Architecture

We parametrize $W$ as $W=n^{-a_{W}}w$ for actual trainable parameter $w$;

We initialize each entry of $w$ iid from $\mathcal{N}(0,n^{-2b_{W}})$;

The learning rate is $\eta n^{-c_{W}}$ for some width-independent $\eta$;

The gradients of $w$ are multiplied by $n^{d_{W}}$ before being processed by $Q^{W}_{t}$: i.e., the update at time $t$ is

where $g_{s},s=0,\ldots,t$, are the gradients of $w$ at time $s$ and $Q^{W}_{t}$ is applied entrywise;

As always, we are only concerned with scaling with $n$ here, but there can be a tunable constant hyperparameter in front of every power of $n$ in 2.2.1.

The random initialization in 2.9.7 is always mean-zero. For some applications, such as layernorm/batchnorm weights $W$ (that is initialized as all 1s), this may seem insufficient. However, one can just refactor the parameter: For this particular example, we can refactor $W=1+W^{\prime}$ where $W^{\prime}$ is the initial vector of the program $\pi$. $W^{\prime}$ can then be initialized as $\mathcal{N}(0,\sigma n^{-2b_{W^{\prime}}})$ for some tunable constant hyperparameter $\sigma$ (which is set to 0 by practitioners typically).

NTP and $\mu$P naturally generalize to general representable architectures.

In comparison to their counterparts for MLP, the NTP and $\mu$P above consider scalar parameters (which are not present in the MLP in Eq. 2.3). Otherwise, 2.4.1 can be recovered from 2.9.11 by mapping the columns $[2,L]$ to “matrix,” $1$ to “vector,” and $L+1$ to “vector” but with the value of $a$ taken from the “out” column. 2.5.1 can be recovered likewise from 2.9.12.

Before we formulate their limits, we need to discuss how to construct the “backpropagation” of an arbtrary of $\textsc{Ne}{\otimes}\textsc{or}\top$ program.This has been constructed previously for $\textsc{Netsor}\top$ programs in .

9.3 Interlude: Backpropagation and Total Programs

Consider any $\textsc{Ne}{\otimes}\textsc{or}\top$ program $\pi$ and a vector $x$ in $\pi$. Then $\pi$’s backpropagation program wrt $x$ is an extension of $\pi$ defined by constructing the following objects on top of $\pi$: (Intuitively, one should interpret $d^{x}y=n\frac{\partial\langle x_{\alpha}\rangle_{\alpha}}{\partial y}$ if $y$ is a vector and $d^{x}c=\frac{\partial\langle x_{\alpha}\rangle_{\alpha}}{\partial c}$ if $c$ is a scalar.)

For any MatMul instruction $z:=Wy$ in $\pi$, we construct $d^{x|z}y:=W^{\top}d^{x}z$ (via another MatMul)

Suppose $y=\langle\psi(\bm{x};\bm{x}_{\beta_{1}};\cdots;\bm{x}_{\beta_{r}};\bm{c})\rangle_{\beta_{1}\cdots\beta_{r}}$. For each $i=0,\ldots,r$, let

Finally, for every vector or scalar $y$ in $\pi$ other than $x$,

where $u$ ranges over all vector or scalar in $\pi$ whose construction used $y$.If this sum is empty, then the RHS is set to 0.

The subprogramThe notion of subprogram is formally defined in [44, Defn I.1]. Roughly it means a contiguous subset of instructions in the program. constructing all of these new objects is denoted $d^{x}\pi$ (so that the backpropagation program is $\pi|d^{x}\pi$).

Recall ([44, Defn I.1]) that “$|$” (as in “$\pi|\pi^{\prime}$”) signifies the concatenation of programs.

If $r=0$ in OuterNonlin (i.e., we just have a Nonlin+ instruction), then the formulas simplify to

Consider a representable architecture $f$ with representation $(\pi,x^{1},\ldots,x^{e})$. Gather all of $\pi$’s backpropagation programs wrt $x^{i}$ into a single (large) program:

whose initial data are the $j$ scalars, $k$ vectors, and $l$ matrices corresponding to $\mathcal{T}_{n}$ in 2.9.1; they are shared among all subprograms $\bm{d}\pi(\xi^{i})$. We call $\bm{d}\bm{\pi}$ the total program of $f$.

9.4 Training Setup

Assume $\varepsilon_{t}$ and $Q^{\Theta}_{t}$ for all $\Theta,t$ are pseudo-Lipschitz and all nonlinearities used in the representing program has pseudo-Lipschitz derivatives.

9.5 Neural Tangent Limit

where the “bar” notation abbreviates application of $Q$ as in Eq. 1.2 and all kets and bras are evaluated in $\bm{d}\bm{\pi}$ (2.9.16) via 2.6.5. To interpret these formulas, we need to tell you two things:

1) Ranges of arguments. Here, $W$ ranges over all matrix parameters and $v$ over all vector parameters of $f$, and all kets are calculated from $\bm{d}\bm{\pi}$ by sampling matrix parameters from $\mathcal{N}(0,1/n)$ and vector parameters from $\mathcal{N}(0,1)$.Again, we can insert hyperparameters like $\sigma_{W}$ and $\sigma_{v}$, but for simplicity we omit them here. The sum in Eq. 2.32 sums over all vectors $h,z,g,y$ in $\pi$ satisfying $h=Wz$ and $g=Wy$ (potentially $h=g$ and $z=y$).

Consider a representable architecture $f$ with representation $(\pi,x^{1},\ldots,x^{e})$ and any training routine in NTP (2.9.11) with memoryless stationary update function $Q$. Adopt 2.9.17. Further assume

$\langle 1\hstretch{0.7}{\talloblong}x^{i}(\xi)\rangle=0$ for every input $\xi$ and output index $i\in[e]$ at initialization.More precisely, the ket $\hstretch{0.7}{\talloblong}x^{i}(\xi)\rangle$ is evaluated in $\bm{d}\bm{\pi}$ (2.9.16) via 2.6.5.

Recall $\bm{f}_{t}$ denotes the function after $t$ steps of updates from random initialization. Then

Note Simple GIA Check may not be satisfied in general architectures, so that we cannot necessarily calculate kets like $\hstretch{0.7}{\talloblong}\bm{d}\bm{h}\rangle$ by assuming matrix parameters and their transposes are independent (even if no transposes occur in $\pi$), e.g., ignoring $\hstretch{0.7}{\talloblong}\bm{d}\bm{h}\dot{\rangle\mkern-3.0mu}\mkern 3.0mu$ in our calculations. Nevertheless, 2.9.19 still holds if one calculates $\hstretch{0.7}{\talloblong}\bm{d}\bm{h}\rangle$ correctly using the rules of 2.6.5.

Assumption 39 ($\bm{x}$ being mean zero) in 2.9.19 is obviously necessary for us to arrive at a Gaussian Process limit at initialization. Assumption 1 is also necessary for two reasons: 1) at initialization, random initialization of scalar parameters would make the initial process a mixture of Gaussian processes (or a GP conditioned on the values of scalar parameters). Even if the scalar parameters are deterministic, the gradient wrt scalar parameters will also be a random process (possibly correlated to the process of the function) at initialization. 2) Function space picture (c.f. 2.8.23 and [45, Remark 3.11]) would no longer hold: One cannot track the evolution of the $\mathring{f}_{t}$ solely by knowing what it is at time $t=0$. Instead, one would need to track the values of the scalar parameters and their gradients as well. So in a sense we will have a “function space + scalar parameters” picture. The complete evolution of $\mathring{f}_{t}$ can then be described by a) the joint process of the function output, scalar parameters, and their gradients at initialization, together with b) an evolution equation involving how they evolve given their previous values in time. This is similar to .

As before, with some more notation we can write down the NT limit for any representable architecture.

where $\overline{\hstretch{0.7}{\talloblong}\bm{d}\bm{h}\rangle{}_{\bm{\chi}_{\leq t}}\langle\bm{z}\hstretch{0.7}{\talloblong}}$ is shorthand for $Q^{W}_{t}\left(\hstretch{0.7}{\talloblong}\bm{d}\bm{h}\rangle{}_{\bm{\chi}_{0}}\langle\bm{z}\hstretch{0.7}{\talloblong},\ldots,\hstretch{0.7}{\talloblong}\bm{d}\bm{h}\rangle{}_{\bm{\chi}_{t}}\langle\bm{z}\hstretch{0.7}{\talloblong}\right)$ and $\overline{\hstretch{0.7}{\talloblong}\bm{d}\bm{v}\rangle\cdot\bm{\chi}_{\leq t}}$ is shorthand for $Q^{v}_{t}\left(\hstretch{0.7}{\talloblong}\bm{d}\bm{v}\rangle\cdot\bm{\chi}_{0},\ldots,\hstretch{0.7}{\talloblong}\bm{d}\bm{v}\rangle\cdot\bm{\chi}_{t}\right)$

With this in mind, the following theorem yields the NT limit of Adam (Eq. $\mathtt{Adam}$) as a corollary.

If the update functions $\bm{Q}$ are memoryless but not necessarily stationary, then 2.9.19 holds with Eq. 2.36 replaced by

For general $\bm{Q}$, not necessarily memoryless, 2.9.19 holds with Eq. 2.36 replaced by

9.6 Maximal Update Limit

Given a representable architecture, the representing program describes the symbolic procedure for computing the output of the network given an assignment of the input and parameters to concrete values. This procedure is well defined for any network width (by construction). Naturally, it remains well defined as explicitly spelled out below even when we “pass to the infinite-width limit.” In short, 2.9.24 specifies what it means for a program to compute something given an assignment of the input and parameters to concrete infinite-width limits (e.g., kets and operators).

For any program $\pi$ and a vector $y$ in $\pi$, $\hstretch{0.7}{\talloblong}y\rangle$ can be thought of as a function, defined via 2.6.5, of $\{\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}\}_{W},\{\hstretch{0.7}{\talloblong}v\rangle\}_{v},\{\mathring{c}\}_{c}$ where $W,v,c$ range over the initial matrices, vectors, and scalars respectively. Consider an assignment $\Theta$ that assigns an operator ${\Theta\{W\}}$ to each initial matrix $W$, ket ${\Theta\{v\}}$ to each initial vector $v$, and a deterministic number $\Theta\{c\}$ to each initial scalar $c$. Then we write

for $\hstretch{0.7}{\talloblong}x\rangle{}_{\mathring{\bm{\chi}}}\langle y\hstretch{0.7}{\talloblong},\hstretch{0.7}{\talloblong}y\rangle$, $\langle y\hstretch{0.7}{\talloblong}z\rangle$ and $\mathring{\theta}$ computed using this assignment (i.e., swapping out $\hstretch{0.7}{\talloblong}W\hstretch{0.7}{\talloblong}$ for $\Theta\{W\}$, $\hstretch{0.7}{\talloblong}v\rangle$ for $\Theta\{v\}$, and $\mathring{c}$ for $\Theta\{c\}$).

where we write $\hstretch{0.7}{\talloblong}\Theta\{W\}\hstretch{0.7}{\talloblong}=\Theta\{W\}$ for each initial matrix $W$, $\hstretch{0.7}{\talloblong}\Theta\{v\}\rangle=\Theta\{v\}$ for each initial vector $v$, and $\Theta\{\mathring{c}\}=\Theta\{c\}$ for each initial scalar $c$.so that $\hstretch{0.7}{\talloblong}\Theta\{\bullet\}\hstretch{0.7}{\talloblong}$ and $\hstretch{0.7}{\talloblong}\Theta\{\bullet\}\rangle$ are just redundant affirmations of the “shape” of $\Theta\{\bullet\}$, rather than saying $\Theta\{\bullet\}$ is some object in some program and $\hstretch{0.7}{\talloblong}\Theta\{\bullet\}\hstretch{0.7}{\talloblong}$ or $\hstretch{0.7}{\talloblong}\Theta\{\bullet\}\rangle$ are their “limits.”

Consider a representable architecture $f$ with representation $(\pi,x^{1},\ldots,x^{e})$ and any training routine in $\mu$P (2.9.12) with memoryless stationary update function $Q$. Adopt 2.9.17.

Then for each $t\geq 0$, $\bm{f}_{t}$ converges almost surely to $\mathring{\bm{f}}_{t}$ computed from the following.

(Forward and Backward Propagation) Let $\Theta_{t}$ be the assignment that assigns $\hstretch{0.7}{\talloblong}W_{t}\hstretch{0.7}{\talloblong}$ to each matrix parameter $W$, $\hstretch{0.7}{\talloblong}v\rangle$ to each vector parameter $v$, and $\mathring{c}_{t}$ to each scalar parameter $c$.

Here $\bm{x}$ is the multi-vector consisting of $x^{1},\ldots,x^{e}$ evaluated on all $\mathcal{M}$ inputs, and all kets and limits are calculated in $\bm{d}\bm{\pi}$ via 2.6.5.

where the tensor operations and summation over $h=Wx$ should be interpreted as in 2.9.18.

(Initialization) \{\hstretch{0.7}{\talloblong}W_{0}\hstretch{0.7}{\talloblong}:\text{matrix parameterW}\} is a set of independent initial operators. Additionally,

Here, we used the notation $W,v,c$ for matrix, vector, and scalar parameters, in contrast to $y$ and $a$ for vector and scalar generated by the program $\pi$. The former are exemplified by weights while the latter by (pre)activations.

In the MLP case with program $\pi$ given in Eq. 2.3, there are no generated scalars, so we can ignore Eq. 2.44. The generated vectors (in $\bm{d}\bm{\pi}$) are $\bm{h}^{1},\ldots,\bm{h}^{L}$ and $\bm{x}^{1},\ldots,\bm{x}^{L}$ as well as their error signals $d\bm{h}^{l},d\bm{x}^{l}$. So Eq. 2.43 reduces to the 2nd and 3rd rows in 2.7.1(1). Finally, the function output is given by averaging $w^{L+1}\odot x^{L}$, so the $\mathring{\bm{f}}_{t}=\Theta_{t}\{\langle 1\hstretch{0.7}{\talloblong}\bm{x}\rangle\}$ in Eq. 2.42 reduces to $\mathring{\bm{f}}_{t}=\langle w_{t}^{L+1}\hstretch{0.7}{\talloblong}\bm{x}^{L}_{t}\rangle$ in 2.7.1(1).

As discussed in 2.7.2, $\mu$P is in a sense the most natural parametrization because its infinite-width limit is just a direct “ket-translation” of the finite-width computations, no matter the architecture. Compare this with NTP, where the Gaussian process and kernel behaviors can only happen under some relative restrictive conditions, like “no scalar parameters” (2.9.19 and 2.9.21).

In the $\mu$-limit for a general architecture, $\mathring{\bm{f}}_{0}$ no longer needs to be 0, unlike 2.7.1. For example, if in an MLP with relu activation, the output is given by the average entry of the last layer activation vector, then obviously $\mathring{\bm{f}}_{0}$ will always be positive.

If $\bm{Q}$ is memoryless but not stationary, then 2.9.25 holds if the bars in LABEL:{eqn:muP_op_update_general}, LABEL:{eqn:muP_ket_update_general} and LABEL:{eqn:muP_scalar_update_general} are interpreted as $Q^{W}_{t},Q^{v}_{t},$ and $Q^{c}_{t}$.

If $\bm{Q}$ is not memoryless, then 2.9.25 holds if LABEL:{eqn:muP_op_update_general}, LABEL:{eqn:muP_ket_update_general} and LABEL:{eqn:muP_scalar_update_general} are replaced with

10 Extensions

Consider the update equation (Eq. 2.4) with decoupled weight decay $\lambda$:We will assume all layers have the same $\lambda$, for simplicity. The generalization to layer-specific $\lambda$ is straightforward.

This decoupled weight decay $\lambda$ is equivalent to a traditional weight decay value $\nicefrac{{\lambda}}{{\eta n^{-c_{l}}}}$. It’s obvious that $\lambda$ should be invariant to the width $n$: if $\lambda\to\infty$ with $n$, then $1-\lambda<0$ eventually; if $\lambda\to 0$ with $n$, then weight decay has no effect in the limit.

The theory in this section covers AdamW , SGD with weight decay, and so on.

2.7.1 holds with Eqs. 2.20, 2.21 and 2.22 replaced with

With decoupled weight decay $\lambda$, 2.4.7 holds with the update equation Eq. 2.9 replaced by

For memoryless nonstationary updates, interpret the bar as $Q_{t}^{l}$. For memoryful nonstationary updates, replace the nonlinear outer product with $\overline{\hstretch{0.7}{\talloblong}d\bm{h}^{l}_{\leq s}\rangle{}_{\bm{\chi}_{\leq s}}\langle\bm{x}^{l-1}_{\leq s}\hstretch{0.7}{\talloblong}}$ (c.f. 2.7.3 for the notation).

Note that when $\lambda=0$, $\bm{h}^{l}_{s}$ is invariant to $s$ as are $\bm{x}^{l}_{s},d\bm{h}^{l}_{s},d\bm{x}^{l}_{s}$, and this exponentially weighted sum just reduces to a simple sum and thus to Eq. 2.9.

In the memoryful case, interpreting the bar in 2.10.3 as Eq. $\mathtt{Adam}$ yields the NT limit of AdamW. The limit theorems (2.9.19 and 2.9.23) for general architecture also hold with analogous modifications.

10.2 Update Clipping and Normalization

We now add update clipping or normalization to Eq. 2.50:Traditionally, gradient clipping for SGD clips the norm of the entire gradient (for the whole network). Naively, one would just do this before applying the update function $Q$. However, in Adam, for example, this would be meaningless because of Adam’s normalization. For general $Q$, the $d_{l}$ terms can be trivially adjusted according to how the global gradient norm scales. So this notion of clipping or normalization (before $Q$) is not very interesting.

if we are doing update normalization, then we set $\nu\leftarrow\|Q_{t}^{l}(n^{d_{l}}g_{0},\ldots,n^{d_{l}}g_{t})\|$;

if we are doing update clipping, then calculate additionally $\nu\leftarrow\min(\nu,\theta^{l})$ where $\theta^{l}$ is a threshold hyperparameter for layer $l$.