Tensor Programs IIb: Architectural Universality of Neural Tangent Kernel Training Dynamics

Introduction

(Jacot et al., 2018)’s pioneering work showed that a multi-layer perceptron (MLP) trained by gradient descent (GD) evolves like a linear model. This spurred a flurry of research papers using this insight to tackle the core questions in deep learning theory, from optimization to generalization in both finite and infinite width regimes. (Jacot et al., 2018)’s argument consists of two observations:

For the output of a network $f(\xi;w)$ with parameters $w$ given example $\xi$, (Jacot et al., 2018) identified the kernel $\mathcal{K}(\xi,\bar{\xi})=\langle\nabla f(\xi;w),\nabla f(\bar{\xi},w)\rangle$, known as the Neural Tangent Kernel (NTK). They showed that if $f$ is parametrized and initialized appropriately, then $\mathcal{K}$ converges to a deterministic kernel $\mathring{\mathcal{K}}$ as the width of the network tends to infinity.

As the infinitely wide network is trained by gradient descent, the NTK remains frozen in its initial state, and the network evolves as by kernel gradient descent with kernel $\mathring{\mathcal{K}}$

In (Yang, 2020a), the ntkInit property was proven to hold for standard architectures, meaning any composition of MLPs, recurrent neural networks (RNN), LSTMs (Hochreiter & Schmidhuber, 1997), gated recurrent unit (GRU) (Cho et al., 2014), convolutions (Fukushima, 1980, 1975; Lecun et al., 1998, 2000; Rumelhart et al., 1986), residual connections (He et al., 2016; Huang et al., 2017), batch normalization (Ioffe & Szegedy, 2015), graph neural networks (Bruna et al., 2014; Defferrard et al., 2016; Duvenaud et al., 2015; Henaff et al., 2015; Kipf & Welling, 2017) and attention (Bahdanau et al., 2015; Vaswani et al., 2017), along with arbitrary weight sharing between components. More generally, it holds for any architecture expressible in a so-called Tensor Program (Yang, 2019b, a, 2020a, 2020b), of which the standard architectures are a subset. However, their reasoning is limited to initialization only.

A statement is architecturally universal if it holds for any reasonable neural architecture. This is an informal property, but here we will formalize it by taking reasonable to be “expressable in Tensor Programs.” By the expressiveness of such programs (Yang, 2019a, 2020a), architectural universality is a fairly robust notion that covers present (and, we expect, future) architectures comprehensively. In this terminology, (Yang, 2020a) showed that ntkInit is architecturally universal.

We show the architectural universality of the entire NTK theory by proving ntkTrain for the same architectures discussed above, including all standard architectures. In the process, we introduce a new graphical form of Tensor Programs that is both required in our proofs and useful for the pedagogy of Tensor Programs.

This paper follows (Yang, 2019b, a, 2020a, 2020b; Yang & Hu, 2020) in the series. While we number this paper “IIb” right after (Yang, 2020a), we actually need the complete theoretical foundation developed in III (Yang, 2020b). See Footnote 22 for more details.

Background

Let $f(\xi;w)\in\mathbb{R}$ denote the (scalar) output of a neural network parameterized by $w$, given example $\xi$. To understand how the output changes with a slight change in the network parameters $w_{0}-\delta w$, we may naively expand the network function using the first order Taylor expansion around a base point $w_{0}$:

Under the SGD algorithm, the weight update $\delta w$ is given by the gradient $\delta w=-\eta\chi(\hat{\xi})\nabla_{w}f(\hat{\xi};w_{0})$ where $\chi(\hat{\xi})$ is the loss derivative, $\hat{\xi}$ is a sample from the training set, and $\eta$ is the learning rate. Plugging into Eq. 1, we get:

where $\mathcal{K}(\xi,\hat{\xi})=\langle\nabla_{w}f(\xi;w_{0}),\nabla_{w}f(\hat{\xi};w_{0})\rangle$ is the NTK. The NTK theory of infinitely wide neural networks as first proposed by (Jacot et al., 2018) boils down to the the following observations: When the width of $f$ tend to infinity, the NTK $\mathcal{K}$ converges to a fixed kernel $\mathring{\mathcal{K}}$ at random initialization, independent of the specific instantiation of the weights, and remains frozen during the optimization process. Eq. 2 then gives an accurate description of the output evolution with if we substitue $\mathcal{K}$ with $\mathring{\mathcal{K}}$. The seemingly complex optimization trajectory of SGD therefore reduce to the convex trajectory of kernel gradient descent with a time-independent kernel $\mathring{\mathcal{K}}$. Consider the output of the network $f\in\mathbb{R}^{D}$ on the full training dataset. As shown in (Jacot et al., 2018), when the $L2$ loss is used the evolution of the output $f_{t}$ at time $t$ under continuous time GD (i.e. gradient flow) takes a simple form:

where $\mathring{\mathcal{K}}\in\mathbb{R}^{D\times D}$ is the full NTK matrix evaluated on the training data, $f^{\star}$ is the label function, and $f_{0}$ is the output at initialization. Hence, provided $\mathring{\mathcal{K}}$ is full rank, as $t\to\infty$ we have that $f_{t}\to f^{\star}$, and the network can fit the training data perfectly.

A common theme in showing ntkTrain for MLP is to derive high-probability bounds on the deviation of the NTK $\mathcal{K}$ from its initial value after training (e.g. Allen-Zhu et al. (2018); Du et al. (2018); Zou et al. (2018)).111 In the original NTK paper (Jacot et al., 2018), the limit is taken as each layer width goes to infinity sequentially, which already doesn’t make sense for weight-tied architectures like RNNs. Obtaining these bounds usually requires developing ad hoc methods on a per-architecture basis, hindering the scalability of the method to other settings. In the present work we take a more holistic approach, leveraging the recently developed Tensor Programs framework (Yang, 2019b, a, 2020a, 2020b). It consists of two layers of arguments: 1) The bottom layer analyzes how the distribution of (pre-)activations change throughout the course of training; this crucially leverages the mathematical machinery of the Tensor Programs Master Theorem.222In particular, we need to use the Master Theorem in (Yang, 2020b), so (Yang, 2020a) could not have obtained ntkTrain at the same time as ntkInit. 2) The top layer packages these insights systematically via the notion of paths so as to apply to any architecture expressible by a Tensor Program. We will illustrate 1) through examples in Section 3 and 2) through figures in Section 5.1.

In this paper, we will consider the architecture (including depth), data, and training time to be fixed as width $n\to\infty$.333They will affect the rate of convergence to the infinite-width limit, but since we are only concerned with whether convergence occurs, they do not appear in our theorem statements here. We describe common notations used in the remainder of the paper. For simplicity, we will consider SGD with batch size 1 and learning rate $\eta$ (often set to 1 WLOG).444This generalizes readily to any batch size and learning rate. We use $\xi_{t}$ to denote the input and $\mathcal{L}_{t}$ to denote the loss function (absorbing the label) at step $t$. More generally, subscript $t$ on any symbol means time $t$. However, for brevity, we abuse notation and shorthand $f_{t}$ for $f_{t}(\xi_{t})$, and, for any (pre-)activation $x$, $x_{t}$ for $x_{t}(\xi_{t})$.555We will not refer to the function $x_{t}:\xi\to x_{t}(\xi)$ (likewise for $f_{t},\chi_{t}$), so this abuse of notation should cause no confusion. We will also write $\chi_{t}$ for the loss derivative $\mathcal{L}_{t}^{\prime}(f_{t})$. For any vector $x(\xi)$ we define \delta x_{t+1}(\xi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\big{(}x_{t+1}(\xi)-x_{t}(\xi)\big{)} and dx(\xi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\frac{\partial f(\xi)}{\partial x(\xi)}. We will track the evolution of $f$ on an arbitrary input $\tilde{\xi}$.666It might help to think of $\tilde{\xi}$ as some test sample, but it can also fall in the training set. Similar to above, we shorthand $\tilde{x}_{t},\tilde{f}_{t}$ for $x_{t}(\tilde{\xi}),f_{t}(\tilde{x})$.

Motivating Examples

The purpose of this section is to illustrate our key ideas via simple, intuitive examples without diving into the specifics of Tensor Programs. In the process, we will gain insight into how randomness from initialization propagates over the course of training. As these examples intend to provide the reader with the proper intuition, we use informal arguments alone and relegate all formal statements to the appendix. For brevity, we will gloss over minor details or routine calculations, but interested readers can see Appendix A for these omissions.

It turns out that the random initialization and the overparametrization of weights cause each (pre-)activation vector $x_{t}(\xi)\in\mathbb{R}^{n}$, its gradient $dx_{t}(\xi)\in\mathbb{R}^{n}$, and its (scaled) change $\delta x_{t}(\xi)\in\mathbb{R}^{n}$ every time step $t$ to have roughly iid coordinates, not just initially but throughout training.777This is a consequence of the Tensor Program Master Theorem. Then, as we shall demonstrate through the examples below, to track the evolution of the neural network function, it suffices to track the evolution of the coordinate distributions of $x(\xi)$, $dx(\xi)$, $\delta x(\xi)$. We write $Z^{x(\xi)}$, $Z^{dx(\xi)}$, $Z^{\delta x(\xi)}\in\mathbb{R}$ for the random variables corresponding to such coordinate distributions.888As we will explain below, different $Z$s may correlate, reflecting correlations between corresponding vectors.

Our goal is to derive, from these insights, {restatable}claimclaimNTK In the large width limit, $\tilde{f}_{t}=f_{t}(\tilde{\xi})$ changes by

𝑡1subscript~𝑓𝑡subscript̊𝜒𝑡̊𝒦~𝜉subscript𝜉𝑡\displaystyle\lim_{n\to\infty}\tilde{f}_{t+1}-\tilde{f}_{t}=-\mathring{\chi}_{t}\mathring{\mathcal{K}}(\tilde{\xi},\xi_{t}) (3) at step $t$, where $\mathring{\mathcal{K}}$ is the limiting NTK of the architecture and $\mathring{\chi}_{t}=\mathcal{L}_{t}^{\prime}(\lim_{n}f_{t})$ is the loss derivative.

We start with an example derivation for 1-hidden-layer MLP, before moving on to 2-hidden-layers, where the mathematics quickly become much more involved.

1 1 Hidden Layer

Consider a 1-hidden-layer network with nonlinearity $\phi$:

where $\xi\in\mathbb{R}^{d},~{}U=\frac{u}{\sqrt{d}}\in\mathbb{R}^{n\times d},~{}V=\frac{v}{\sqrt{n}}\in\mathbb{R}^{n\times 1}$, for trainable parameter tensor $u$, initialized iid from $\mathcal{N}(0,1)$. In the interest of clarity we assume the output layer is not trained, and $d=1$.

For a vector $v\in\mathbb{R}^{n}$, let $v=\Theta(n^{a})$ mean that “$v$ has coordinates of order $n^{a}$ when $n$ is large”999 More rigorously, we mean that $\|v\|^{2}/n=\Theta(n^{2a})$. Note this is different from the common interpretation that $\|v\|=\Theta(n^{a})$. ; likewise for $o(n^{a})$, etc. Recall the notations $x_{t}=x_{t}(\xi_{t}),\tilde{x}_{t}=x_{t}(\tilde{\xi}),\delta x_{t}=\sqrt{n}(x_{t}-x_{t-1})$ and likewise for $h_{t},\tilde{h}_{t},\delta h_{t}$. The key insights are as follows:

It turns out $\tilde{x}_{t+1}-\tilde{x}_{t}=\Theta(1/\sqrt{n})=o(1)$ so $\delta\tilde{x}_{t+1}=\sqrt{n}(\tilde{x}_{t+1}-\tilde{x}_{t}))$ has $\Theta(1)$ coordinates. Likewise for $\delta\tilde{h}_{t+1}$. Consequently, for any $t$ and input $\xi$, by telescoping,

subscriptℎ0𝜉𝑜1\displaystyle=h_{0}(\xi)+o(1). (4) Using $\nabla_{u}f=\frac{1}{\sqrt{n}}dh_{t}\xi_{t}^{\top}$ and $dh_{t}=\phi^{\prime}(h_{t})\odot v$, we have:

𝑡1direct-productsubscript𝜒𝑡superscriptsubscript𝜉𝑡top~𝜉superscriptitalic-ϕ′subscriptℎ𝑡𝑣\displaystyle\delta\tilde{h}_{t+1}=-\chi_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(h_{t})\odot v. (5) Also, \delta\tilde{x}_{t+1}=\sqrt{n}\big{(}\phi(\tilde{h}_{t}+\frac{\delta\tilde{h}_{t+1}}{\sqrt{n}})-\phi(\tilde{h}_{t})\big{)}. Since $\delta\tilde{h}_{t+1}=\Theta(1)$, by intuitive Taylor expansion, we have

𝑡1\displaystyle\delta\tilde{x}_{t+1} $\displaystyle\approx\phi^{\prime}(\tilde{h}_{t})\odot\delta\tilde{h}_{t+1}.$ (6) The change in the output on example $\tilde{\xi}$ from step $t$ to step $t+1$ is given by:

By definition $v$ has iid coordinates. It turns out $h_{t}(\xi),\delta h_{t}(\xi)$ (likewise for $x$) all have approx. iid coordinates of size $\Theta(1)$ as well.101010Technically, they have iid coordinates only after conditioning on the initial function (GP) $f_{0}$. Likewise, when we take expectation in this example (e.g. Eqs. 9 and LABEL:{eqn:dot}), it’s a conditional expectation of this kind. See Section D.1.1 for a rigorous treatment. However, to convey the main intuition, we gloss over this technicality here. Let $Z^{\delta\tilde{x}_{t}},Z^{v}$ denote the random variables encoding the corresponding coordinate distributions; likewise for the other vectors. Note that $Z^{\delta\tilde{x}_{t}},Z^{v}$ will in general be correlated, reflecting the coordinatewise correlation between $v$ and $\delta\tilde{x}_{t}$.

𝑡1subscript~𝑓𝑡𝔼superscript𝑍𝑣superscript𝑍𝛿subscript~𝑥𝑡1\displaystyle\phantom{{}={}}\lim_{n\to\infty}\tilde{f}_{t+1}-\tilde{f}_{t}=\operatorname*{\mathbb{E}}Z^{v}Z^{\delta\tilde{x}_{t+1}} (8) $\displaystyle=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}\phi^{\prime}(Z^{\tilde{h}_{t}})\phi^{\prime}(Z^{h_{t}})(Z^{v})^{2}.$ (9) where $\mathring{\chi}_{t}=\mathcal{L}_{t}^{\prime}(\lim_{n}f_{t})$ as in Section 3.

By Eq. 4, in the $n\to\infty$ limit, $Z^{\tilde{h}_{t}}=Z^{h_{0}(\tilde{\xi})}$ and $Z^{h_{t}}=Z^{h_{0}(\xi_{t})}$. They are independent from $Z^{v}$ and jointly Gaussian with variances $\|\tilde{\xi}\|^{2},\|\xi_{t}\|^{2}$ and covariance $\tilde{\xi}^{\top}\xi_{t}$. So (using the initialization of $v$ to simplify $\operatorname*{\mathbb{E}}(Z^{v})^{2}=1$),

𝑡1subscript~𝑓𝑡\displaystyle\lim_{n\to\infty}\tilde{f}_{t+1}-\tilde{f}_{t} $\displaystyle=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}\phi^{\prime}(Z^{h_{t}})\phi^{\prime}(Z^{\tilde{h}})$ (10) This can easily be seen to be Eq. 3 (recall we assumed for simplicity the output layer $V$ is not trained).

Our strategy so far has been computing the form of $Z^{\delta\tilde{x}_{t}}$ and plugging it into Eq. 8 to compute the dynamics of the output in the limit. Note that our approach differs from previous work which mainly focus on proving a bound on the change of the NTK post training. As the architecture gets more complex, bounding the NTK movement becomes quite complex, but our approach easily scales due to the automation provided by the Tensor Programs framework (see Section 4).

In the previous example, the coordinate distribution $Z^{\tilde{x}_{t}}$ took a fairly simple form, which allowed us to intuitively compute the expectation $\operatorname*{\mathbb{E}}Z^{\tilde{x}_{t}}Z^{v}$. Before introducing a method for computing coordinate distributions in a general architecture, we move on to a slightly more involved architecture, with the intention of highlighting the intuition behind the general case.

2 2 Hidden Layers

In this example we consider a model of the form:

where $U=\frac{u}{\sqrt{d}}\in\mathbb{R}^{n\times d},~{}W=\frac{w}{\sqrt{n}}\in\mathbb{R}^{n\times n},~{}V=\frac{v}{\sqrt{n}}\in\mathbb{R}^{n\times 1}$, for trainable parameters $u,w$, initialized iid from a normal distribution. As before we assume the last layer is not trained, and $d=1$.

Again, we want to establish Section 3 for this model. As in the 1-hidden-layer example, the dynamics of the output in the limit is still given by Eq. 8. This time around, the second hidden layer adds nontrivial complexity when evaluating the expectation $\operatorname*{\mathbb{E}}Z^{v}Z^{\delta\tilde{x}_{t+1}}$. As we shall see, this complexity arises from the dependency of $\tilde{x}$ on the $n\times n$ matrices $w$ and $w^{\top}$, which will make it wrong to naively apply LLN arguments. Resolving this complexity will pave the way to a general strategy which we will then be able to apply in any arbitrary architecture. We now apply the same insights as in the 1-hidden-layer MLP. Namely:

Eq. 4 continues to hold with $h$ replaced by any of $\{x,h,z,g\}$.

After some brief calculations, with $dh_{t}$ denoting the scaled gradient $\sqrt{n}\nabla_{h_{t}}f$,

𝑡1\displaystyle\delta\tilde{g}_{t+1} \displaystyle\approx-\chi_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(g_{t})\odot\big{(}W^{\top}dh_{t}\big{)} (11) $\displaystyle\delta\tilde{h}_{t+1}$ $\displaystyle\approx W\delta\tilde{z}_{t+1}-\chi_{t}\frac{z_{t}^{\top}\tilde{z}_{t}}{n}\phi^{\prime}(h_{t})\odot v.$ (12) • As in the 1-hidden-layer case, for all $x\in\{g,z,h,x\}$, $x_{t}(\xi),\delta x_{t}(\xi)$ have iid coordinates of size $\Theta(1)$, as does $v$ by definition.111111Again, this is technically true only after conditioning on $f_{0}$; see Footnote 10. Let $Z^{\delta\tilde{x}_{t}},Z^{v}$ denote the (generally correlated) random variables encoding the corresponding coordinate distributions.

As in Eq. 7, by naive Taylor expansion we have:

𝑡1\displaystyle\delta\tilde{z}_{t+1} $\displaystyle\approx\phi^{\prime}(\tilde{g}_{t})\odot\delta\tilde{g}_{t+1}.$ (13) • Eqs. 7 and 6 in the 1-hidden-layer case continue to hold here. Then by Eq. 12 and Law of Large Numbers,

𝑡1subscript~𝑓𝑡𝔼superscript𝑍𝑣superscript𝑍𝛿subscript~𝑥𝑡1\displaystyle\phantom{{}={}}\lim_{n\to\infty}\tilde{f}_{t+1}-\tilde{f}_{t}=\operatorname*{\mathbb{E}}Z^{v}Z^{\delta\tilde{x}_{t+1}} (14) $\displaystyle=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}[Z^{z_{t}}Z^{\tilde{z}}]\operatorname*{\mathbb{E}}[\phi^{\prime}(Z^{h_{t}})\phi^{\prime}(Z^{\tilde{h}})]$ (15) $\displaystyle\phantom{{}={}}-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}[\phi^{\prime}(Z^{\tilde{h}_{t}})Z^{W\delta\tilde{z}_{t+1}}Z^{v}].$ (16) where $\mathring{\chi}_{t}=\mathcal{L}_{t}^{\prime}(\lim_{n}f_{t})$ as in Section 3.

In this expression, the first term (Eq. 15) can easily be seen to correspond to the contribution from $w$ to the NTK. It remains to show that the second (Eq. 16) corresponds to the contribution from $u$.

To do this, we must reason about the coordinate distribution of $W\delta\tilde{z}_{t+1}$ (encoded by random variable $Z^{W\delta\tilde{z}_{t+1}}$) and compute the expectation in Eq. 16. To understand why this represents a greater challenge than it might first appear, note that from $\delta\tilde{z}_{t+1}\approx\phi^{\prime}(\tilde{g}_{t})\odot\delta\tilde{g}_{t+1}$ (Eq. 13), the term ${W\delta\tilde{z}_{t+1}}$ hides within itself a dependency on $W^{\top}dh_{t}$ through $\delta\tilde{g}$ (Eq. 11). While at $t=0$, we may assume $W^{\top}$ be independent of $W$ and obtain the correct results (Gradient Independent Assumption (Yang & Schoenholz, 2017; Yang, 2020a)), this is no longer the case for $t>0$: $Z^{W\delta\tilde{z}_{t+1}}$ will be nontrivially correlated with $\phi^{\prime}(Z^{\tilde{h}_{t}})$ and $Z^{v}$ (which would be false if $W^{\top}$ can be assumed independent of $W$). We will give some intuition why later in Eq. 20. Now, what is this dependence exactly? {restatable}claimclaimZdot Based on the above discussion and some easy calculations, $\delta\tilde{z}_{t+1}$ can be written as $\Phi(W^{\top}dh_{t})$ for some $\Phi:\mathbb{R}\to\mathbb{R}$ applied coordinatewise (which will depend on other vectors not of the form $W^{\top}\bullet$). Then it turns out††footnotemark:

𝑡1𝐺superscript𝑍𝑑subscriptℎ𝑡𝔼superscript𝑍𝛿subscript~𝑧𝑡1superscript𝑍superscript𝑊top𝑑subscriptℎ𝑡\displaystyle Z^{W\delta\tilde{z}_{t+1}}=G+Z^{dh_{t}}\operatorname*{\mathbb{E}}\frac{\partial Z^{\delta\tilde{z}_{t+1}}}{\partial Z^{W^{\top}dh_{t}}}, (17) where $G$ is some Gaussian variable independent from $Z^{v}$, and \frac{\partial Z^{\delta\tilde{z}_{t+1}}}{\partial Z^{W^{\top}dh_{t}}}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\Phi^{\prime}(Z^{W^{\top}dh_{t}}).

Thus, from Eqs. 13 and 11, it follows that:

𝑡1\displaystyle Z^{\delta\tilde{z}_{t+1}} $\displaystyle=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(Z^{\tilde{g}_{t}})\phi^{\prime}(Z^{g_{t}})Z^{W^{\top}dh_{t}}$ (18) $\displaystyle\operatorname*{\mathbb{E}}\frac{\partial Z^{\delta\tilde{z}_{t+1}}}{\partial Z^{W^{\top}dh_{t}}}$ $\displaystyle=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}[\phi^{\prime}(Z^{\tilde{g}_{t}})\phi^{\prime}(Z^{g_{t}})].$ (19) Plugging into Eqs. 14 and LABEL:{eqn:dot}, followed by some straightforward calculation, then yields Section 3.

Eq. 17 may appear cryptic at first, so let’s give some intuition using an example. Suppose $\Phi$ in Section 3.2 is actually identity. For brevity, we set $\mathbf{x}=dh_{t},\mathbf{y}=\delta\tilde{z}_{t+1}\in\mathbb{R}^{n}$. Then, following straightforward calculation, $W\mathbf{y}=WW^{\top}\mathbf{x}$ has coordinates

subscript𝛾𝛼subscript𝐱𝛾superscriptsubscript𝛽1𝑛subscript𝑊𝛼𝛽subscript𝑊𝛾𝛽superscriptsubscript𝛽1𝑛superscriptsubscript𝑊𝛼𝛽2subscript𝐱𝛼\displaystyle(W\mathbf{y})_{\alpha}=\sum_{\gamma\neq\alpha}\mathbf{x}_{\gamma}\sum_{\beta=1}^{n}W_{\alpha\beta}W_{\gamma\beta}+\sum_{\beta=1}^{n}(W_{\alpha\beta})^{2}\mathbf{x}_{\alpha} (20) Now, the second sum converges via LLN to $\mathbf{x}_{\alpha}$ as $n\to\infty$. On the other hand, the first sum will converge via CLT to $\mathcal{N}(0,\lim\|\mathbf{x}\|^{2}/n)$. Thus, in terms of $Z$s, we have

𝐺superscript𝑍𝐱𝐺superscript𝑍𝐱𝔼superscriptΦ′\displaystyle Z^{W\mathbf{y}}=G+Z^{\mathbf{x}}=G+Z^{\mathbf{x}}\operatorname*{\mathbb{E}}\Phi^{\prime} (21) for some Gaussian $G$; this corresponds directly to Eq. 17.121212 This example was worked out in (Yang, 2020a, b) as well, though in different contexts. Readers needing more explanation may see those works. For general $\Phi$, a similar intuition applies after Taylor expansion of $\Phi$.

This 2-hidden-layer example proceeded much the same as the 1-hidden-layer case, with the main exception of analyzing the interaction of the $n\times n$ Gaussian matrix $W$ and $W^{\top}$ (Eq. 16) that occurs after taking at least 1 step of SGD. This was absent in the 1-hidden-layer case because each weight matrix has at most one side tending to infinity. Such analysis is crucial to obtaining the right results, as assuming $W^{\top}$ be independent from $W$ would imply $f$ does not move from initialization.131313One can see this easily by modifying our calculations above.

It turns out these two examples have essentially covered all of the core ideas needed to extend the analysis into arbitrary architectures. To formalize and scale up our calculations, we now turn to the Tensor Programs framework.

Tensor Programs

So far, our results have been obtained by unrolling the SGD updates on toy models with specific architectures, and using informal arguments. Obviously, these computations quickly become unmanageable when the architecture becomes more complex. The sheer amount of architectural innovations that have sprung up in recent years requires us to adopt a much more general formulation of our results. To that end, we adopt the Tensor Programs (TP) framework developed in (Yang, 2019a, 2020a, 2020b). In a nutshell, it provides a language for describing typical computations done in the context of neural networks, such as forward and backward propagation. It is simultaneously simple and expressive, covering all standard architectures (Yang, 2019a, 2020a). Here we review two basic forms of Tensor Programs, $\textsc{Netsor}\top$ and $\textsc{Netsor}\top^{+}$.

A $\textsc{Netsor}\top$ program is just a sequence of $\mathbb{R}^{n}$ vectors inductively generated via one of the following instructions from an initial set $\mathcal{V}$ of random $\mathbb{R}^{n}$ vectors and a set $\mathcal{W}$ of random $n\times n$ matrices

For $x^{1},\ldots,x^{k}\in\mathbb{R}^{n}$ in the program and any $\psi:\mathbb{R}^{k}\to\mathbb{R}$, we can generate $\psi(x^{1},\ldots,x^{k})\in\mathbb{R}^{n}$

Given $W\in\mathbb{R}^{n\times n}$ and $x\in\mathbb{R}^{n}$, we can generate $Wx\in\mathbb{R}^{n}$ or $W^{\top}x\in\mathbb{R}^{n}$

We propose to represent a $\textsc{Netsor}\top$ program as a computational graph, where each node in the graph represents vectors (initial or generated), each (dashed) edge represents a MatMul, and each gate represents a Nonlin. For example, Fig. 1 shows the computation graphs expressing (the forward passes of) an MLP and an RNN. We can also express the backpropagation as well (see Fig. 6). Graphically, the initial vectors are the empty nodes with only one edge coming out, toward the direction of computation. The matrices correspond to (the labels of) the dashed edges. We can also define the output vectors to correspond to the nodes that have only one edge coming out, against the direction of computation.

Each $\textsc{Netsor}\top$ program can be thought of as computing a function $(\mathbb{R}^{n})^{\mathcal{V}}\times(\mathbb{R}^{n\times n})^{\mathcal{W}}\to\mathbb{R}^{\mathcal{Y}}$ taking an instantiation of the initial vectors $\mathcal{V}$ and matrices $\mathcal{W}$ and computing the values of all output vectors $\mathcal{Y}$. We can say a program represents a network if it computes the body of it (without the input and output layers), as exemplified by Fig. 1. This is formalized below.

Consider a neural network $f:(\mathbb{R}^{d})^{k}\to\mathbb{R}$ with input embedding matrices $U^{1},\ldots,U^{k}\in\mathbb{R}^{n\times d}$ (not necessarily distinct) and readout matrix $V\in\mathbb{R}^{n}$, so that $f(\xi^{1},\ldots,\xi^{k})=V^{\top}\Phi(U^{1}\xi^{1},\ldots,U^{k}\xi^{k};\Theta)$ for some function $\Phi(x^{1},\ldots,x^{k};\Theta)$ with parameters $\Theta$. We say a $\textsc{Netsor}\top$ program represents $f$ if it computes $\Phi$ (under some correspondence of $\mathcal{V}\cup\mathcal{W}$ to $\{x^{1},\ldots,x^{k}\}\cup\Theta$).141414 We only consider $f$ with scalar output for simplicity but generalization to multi-dimensional output is straightforward.

For example, the programs in Fig. 1 resp. represent a 3-hidden-layer MLP and an RNN running for 3 steps. Note that the initial vectors correspond to a combination of input embeddings (e.g. $W^{1}\xi$) and vector parameters (e.g. biases) and the matrices correspond to matrix parameters (e.g. weights).

Typically, the vectors (resp. matrices) in a program will be sampled iid like $\mathcal{N}(0,1)$ (resp. $\mathcal{N}(0,1/n)$), corresponding to the “standard” initialization of neural networks.151515 In the original definition of (Yang, 2019a, 2020a, 2020b), the vectors can have correlations between them, but we can always rewrite these vectors as linear image of another set of uncorrelated vectors. In such cases, when $n\to\infty$, a program behaves as follows, in a gist:

Any vector $x\in\mathbb{R}^{n}$ in the program has roughly iid coordinates. We write $Z^{x}$ for the random variable encoding this coordinate distribution. This $Z^{x}$ may be correlated with $Z^{y}$ for other vector $y$ in the program, such that, for example, $\lim_{n\to\infty}x^{\top}y/n=\operatorname*{\mathbb{E}}Z^{x}Z^{y}$.

$Z^{\psi(x^{1},\ldots,x^{k})}=\psi(Z^{x^{1}},\ldots,Z^{x^{k}})$.

Consider a matrix $W\in\mathbb{R}^{n\times n}$ in the program and any set of $\mathbb{R}^{n}$ vectors $\mathcal{\mathbf{X}}$ not dependent on vectors of the form $W^{\top}\bullet$. Then the set of random variables $\{Z^{Wx}:x\in\mathcal{\mathbf{X}}\}$ are jointly Gaussian with mean zero and covariance $\mathrm{Cov}(Z^{Wx},Z^{W\tilde{x}})=\operatorname*{\mathbb{E}}Z^{x}Z^{\tilde{x}}$ for any $x,\tilde{x}\in\mathcal{\mathbf{X}}$. If $\bar{W}\neq W$ is another matrix in the program and $\mathcal{\mathbf{Y}}$ is a set of such vectors w.r.t. $\bar{W}$, then the set $\{Z^{Wx}:x\in\mathcal{\mathbf{X}}\}$ is independent from $\{Z^{\tilde{W}y}:y\in\mathcal{\mathbf{Y}}\}$.

For general $x$, $Z^{Wx}$ decomposes into a sum of a Gaussian part, identical to $Z^{Wx}$ in the above case, and a correction term. This decomposition is a generalization of Eq. 17.

\textsc{Netsor}\top^{+} Programs (Yang, 2019a, 2020a) showed that $\textsc{Netsor}\top$ suffices to express the forward and backward passes of most architectures such as ResNet (with Batchnorm), but Transformer and other standard architectures require adding to $\textsc{Netsor}\top$ a new “averaging” instruction161616For readers familiar with Tensor Programs, this is equivalent in expressivity to Moment, which is a composition of a Nonlin and this averaging instruction. that returns the “empirical average” $\frac{1}{n}\sum_{\alpha=1}^{n}x_{\alpha}$ of a vector $x\in\mathbb{R}^{n}$. In the $n\to\infty$ limit, this scalar converges to $\operatorname*{\mathbb{E}}Z^{x}$ as would be expected from the intuitions above. This extension of $\textsc{Netsor}\top$ (called $\textsc{Netsor}\top^{+}$) also allows us to express the network output and loss derivative (e.g. in contrast to Fig. 1), which will be a technical requirement for unrolling the entirety of SGD training inside a single program, a key step in the proof of our main result. See discussion of proof formalization in Section 5. We can say a $\textsc{Netsor}\top^{+}$ program represents a network $f$ if it computes the body of $f$.

Universality of Kernel Dynamics

(Yang, 2019a, 2020a) showed that any neural network of standard architecture is represented by a $\textsc{Netsor}\top^{+}$ program. Moreover,

For a neural network as in 5.2 below, its Neural Tangent Kernel at initialization has a well-defined infinite-width limit $\mathring{\mathcal{K}}$.

Suppose a neural network $f$ is represented by a $\textsc{Netsor}\top^{+}$ program (in the sense of 4.2) whose Nonlin all have polynomially bounded derivatives.171717More generally, we can allow any pseudo-Lipschitz function here, but for simplicity we go with the statement in the main text. Adopt the NTK parametrization: for every matrix parameter $W\in\mathbb{R}^{n\times n}$ of $f$, we factor $W=\frac{1}{\sqrt{n}}w$ where $w$ is the trainable parameter; likewise, for each input layer matrix $U^{i}\in\mathbb{R}^{n\times d}$, we factor $U^{i}=\frac{1}{\sqrt{d}}u^{i}$, and likewise the output matrix $V=\frac{1}{\sqrt{n}}v$, such that $u^{i},v$ are trainable. Finally, we randomly initialize all trainable parameters iid as $\mathcal{N}(0,1)$.

Our main result is to show that the SGD training of such a neural network described in 5.2 reduces to kernel gradient descent with kernel $\mathring{\mathcal{K}}$ in the infinite-width limit.

Consider training a network $f$ described in 5.2 via SGD with batch-size 1 and (WLOG) learning rate 1. Let $\xi_{t}$ be the input and $\mathcal{L}_{t}:\mathbb{R}\to\mathbb{R}$ be the loss function (absorbing the label) at time $t$. Suppose $\mathcal{L}_{t}$ is continuous for all $t$. Then, for any $\xi$ and $t$, $f_{t}(\xi)$ converges almost surely to a random variable $\mathring{f}_{t}(\xi)$ as width $\to\infty$, such that

𝑡1𝜉subscript̊𝑓𝑡𝜉̊𝒦𝜉subscript𝜉𝑡superscriptsubscriptℒ𝑡′subscript̊𝑓𝑡subscript𝜉𝑡\displaystyle\mathring{f}_{t+1}(\xi)-\mathring{f}_{t}(\xi)=-\mathring{\mathcal{K}}(\xi,\xi_{t})\mathcal{L}_{t}^{\prime}(\mathring{f}_{t}(\xi_{t})) (22) where $\mathring{\mathcal{K}}$ is the infinite-width NTK (at initialization) of the neural network.

The full proof of 5.3 is given Appendix D.

We briefly mention several ways our result can be easily extended. 0) Different batch sizes, learning rate schedules, and nonscalar outputs. 1) Variants of NTK parametrization. We can deal with any parametrization that scales the same way as NTK parametrization, e.g. weights are sampled like $N(0,\sigma^{2})$ for any $\sigma$, with the multipliers $\gamma/fanin$ for any $\gamma$. 2) Variable width. In real networks, the width of different layers can often be different (e.g. in ResNet). Our result can be extended to the case where the widths tend to infinity at a fixed ratio, using the variable-width version of Tensor Programs (Yang, 2020b). 3) Unsupervised and other learning settings can be covered because their training and testing computation can be written into Tensor Programs. 4) Weight decay, momentum, and other optimizer tricks can be covered as well as they can be straightforwardly written into Tensor Programs, but in general the kernel will change from step to step in contrast to 5.3.

1 Proof Sketch of Special Case

To convey the main idea, we give a proof sketch of a simplified problem: we assume 1) the input, output layers and biases are not trained (the network has only $\mathbb{R}^{n\times n}$ matrices as trainable parameters); 2) the forward pass does not contain both a weight matrix $W$ and its transpose $W^{\top}$ (but a single matrix $W$ can still be used multiple times without being transposed); 3) input space is $\mathbb{R}^{d}$ (with $k=1$), and $f=V^{\top}x\in\mathbb{R}$; 4) the output vector $x$ is a G-var; 5) the network is represented by a $\textsc{Netsor}\top$ (instead of $\textsc{Netsor}\top^{+}$) program; 181818This is not common in deep learning, but appears in some weight-tied autoencoders (Li & Nguyen, 2019). In the appendix, we prove the general case with these simplifications lifted.

It turns out, every $\textsc{Netsor}\top$ can be simplified into a standard form of sorts, which greatly facilitates our proof.

In a $\textsc{Netsor}\top$ program, a G-var191919“G” because G-vars often are roughly Gaussian vectors is an initial vector or a vector created by MatMul, while an X-var is a vector created by Nonlin.202020Var is short for variable, as the vectors are considered variables in the program. In previous works, H-var refers to any vector in the program; we will not use this terminology. We define a reduced $\textsc{Netsor}\top$ program as a program in which only G-vars are allowed as inputs to a Nonlin, while only an X-var is allowed as input to a MatMul.

Observe that any $\textsc{Netsor}\top$ program may be trivially expressed as a reduced $\textsc{Netsor}\top$ program by: 1) collapsing chains of non-linearities which appear consecutively, and 2) insert a Nonlin operation with $\psi(x)=x$ in between consecutive G-vars. Hence, we may safely assume that $f$ is representable by a reduced $\textsc{Netsor}\top$ program.

The examples of Sections 3.1 and 3.2 exposed several insights, such as the iid-coordinates intuition, important for proving 5.3. Now we discuss the one remaining key idea for scaling up to general architectures:

[Paths]defnPaths In a $\textsc{Netsor}\top$ program, a path $p$ starts with an X-var and ends with a G-var, alternating between X- and G-vars along the path. We write $p^{0}$ for the starting X-var, $p^{1}$ for the following G-var, and so on, as well as $p^{-1}$ for the ending G-var (see Fig. 2 for a graphical illustration). For odd $i$, let $W^{p^{i}}$ denote the defining matrix of G-var $p^{i}$. For two equal length paths $p,q$, we write $p\cong q$ (path $p$ is isomorphic to path $q$) if for all odd $i$, $W^{p^{i}}$ is the same matrix as $W^{q^{i}}$.212121Here we are talking about equality of symbols rather than equality of values of those symbols. In other words, we say path $p$ is isomorphic to path $q$ if their sequences of MatMul matrices are identical, (but the Nonlin don’t have to be, see Fig. 3 for a graphical illustration). Let $|p|$ denote the number of vectors in $p$ (this is always an even number).

The collection of paths $p$ starting with an X-var $p^{0}=x$ and ending with a G-var $h$ describes all possible pathways of backpropagating an error signal $dh$ at $h$ to an error signal $dx$ at $x$. Simultaneously, it also describes all possible pathways of forward propagating a change in $x$ to a change in $h$.

Because the gradient $\nabla_{W}f$ of a weight $W\in\mathbb{R}^{n\times n}$ is the sum of outer products $\sum_{h,x:h=Wx}dh\otimes x$, summing over all G-vars $h$ and X-vars $x$ in the program with $h=Wx$ (where $dh$ denotes $\nabla_{h}f$), we also have

𝑓superscript𝑝1superscript𝑝1superscript𝑝2superscript𝑝2superscript𝑝3…superscript𝑝2ℎ\displaystyle(J^{p})^{\top}=\frac{\partial f}{\partial p^{-1}}\times\frac{\partial p^{-1}}{\partial p^{-2}}\times\frac{\partial p^{-2}}{\partial p^{-3}}\times...\times\frac{\partial p^{2}}{\partial h} (24) i.e, $J^{p}$ denotes the error signal at $h$ from backpropagation through path $p$, and $p$ ranges over all paths starting with $p^{0}=x,p^{1}=h$ and ending with the output node of the underlying program. Recall $W$ factors as $\frac{1}{\sqrt{n}}w$ where $w$ is the trainable parameter, not $W$. By the discussion above, updating $w$ with $\nabla_{w}f=\frac{1}{\sqrt{n}}\nabla_{W}f$ causes $f$ to change by

When every parameter $w$ is randomly initialized iid as $\mathcal{N}(0,1)$, it turns out that $\langle J^{p},J_{\bar{p}}\rangle$ will go to 0 as $n\to\infty$ unless $p\cong\bar{p}$ (Fig. 3). If one think of $J^{p}$ as a product of random Gaussian matrices (interleaved with other matrices), then this is akin to the fact that, for a mixed moment M\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\operatorname*{\mathbb{E}}\prod_{i=1}^{r}Z_{\gamma(i)},\gamma:[r]\to[k] of standard iid Gaussians $Z_{1},\ldots,Z_{k}$, $M$ is nonzero iff every $Z_{i}$ appears an even number of times in the product. This means we can replace the 4 nested sums in Eq. 25 with the single sum $\sum_{p\cong\bar{p}}$ and rewrite $x=p^{0},\bar{x}=\bar{p}^{0}$.

We have suppressed dependence on input in Eq. 25. Being more explicit about it and performing updates on all weights, we have

𝑡1𝜉subscript𝑓𝑡𝜉superscriptsubscriptℒ𝑡′subscript𝑓𝑡subscript𝜉𝑡subscript𝑝¯𝑝superscript𝐽𝑝𝜉superscript𝐽¯𝑝subscript𝜉𝑡superscript𝑝0𝜉superscript¯𝑝0subscript𝜉𝑡\begin{gathered}f_{t+1}(\xi)-f_{t}(\xi)\\ \approx-\mathcal{L}_{t}^{\prime}(f_{t}(\xi_{t}))\sum_{p\cong\bar{p}}\langle J^{p}(\xi),J^{\bar{p}}(\xi_{t})\rangle\langle p^{0}(\xi),\bar{p}^{0}(\xi_{t})\rangle.\end{gathered} (26) after taking a gradient step on input $\xi_{t}$ at initialization $t=0$. (Here $p^{0}(\xi)$ denotes the vector $p^{0}$ as a function of $\xi$ at initialization). However, Eq. 25 holds for general $t$ as well: The key insight is similar to Eq. 4 in the 1-hidden-layer example, that vectors $p^{0}(\xi)$, $J^{p}(\xi)$, etc change vanishingly from their initial values as $n\to\infty$, after any number of SGD steps. Because our arguments above only depend on inner products between vectors, this means that the error in Eq. 26 for $t>0$ vanishes as $n\to\infty$. Finally, at least heuristically, it is straightforward to show the RHS of Eq. 26 is the NTK via a series of calculations exemplified by those in Sections 3.1 and 3.2, to get Eq. 22.

While the core ideas discussed above are intuitive, making them rigorous at face value would be quite challenging. Instead we use the machinery offered by the Tensor Programs framework. The mechanics of the proof then goes as follows: 1) First we unroll SGD of $f$ into a $\textsc{Netsor}\top^{+}$ program.222222 We note that this formalization crucially relies on $\textsc{Netsor}\top^{+}$ and its Master Theorem from (Yang, 2020b) because the SGD unrolling cannot be done in $\textsc{Netsor}\top$. The reason is that we need to express the output and loss derivatives of the network, which are scalars (or at least finite dimensional), and that cannot be done in a $\textsc{Netsor}\top$ program. Furthermore, the Master Theorem from (Yang, 2020a) only pertains to a specific type of programs that look like the first backpropagation after initialization. Thus, it cannot deal with the complete unrolling of SGD as we do here, which requires the more advanced Master Theorem from (Yang, 2020b). This is similar to the equations in Sections 3.1 and 3.2; the key here is to express $\delta x_{t+1}=\sqrt{n}(x_{t+1}-x_{t})$ as a vector in the program, for any (pre-)activation $x$. 2) We apply the $\textsc{Netsor}\top^{+}$ Master Theorem (Yang, 2020b) to this program. This yields the coordinate distribution of each vector. The core insights here are demonstrated by the calculation with the $Z$ random variables in Sections 3.1 and 3.2. 3) Finally, we need to show that (the rigorous version of) Eq. 26 indeed recovers the NTK and agrees with Eq. 22. This is done via an inductive (symbolic) computation, and the path concept in this section plays a key role here.

Related Works

The connection between kernel methods and neural networks has had a long history before its recent resurgence. The Gaussian Process (NNGP) view of wide neural networks, which characterizes the behaviour of training only the last layer of a wide neural network, has been studied in (Daniely et al., 2016; Hazan & Jaakkola, 2015; Roux & Bengio, 2007; Lee et al., 2018; Matthews et al., 2018; Hinton & Neal, 1995; Novak et al., 2019). Since the original NTK paper (Jacot et al., 2018), many works have informally derived the infinite-width NTK for various architectures such as CNNs (Arora et al., 2019), RNN (Alemohammad et al., 2020), attention (Hron et al., 2020), ensembles (Littwin et al., 2020b) and graph neural networks (Du et al., 2019), but none of them formally proved ntkInit or ntkTrain for those architectures. Finite width corrections to the NTK were derived for fully connected networks in (Hanin & Nica, 2019; Littwin et al., 2020a). The validity of the NTK theory was empirically studied in (Lee et al., 2019) for a variety of architectures.

The Tensor Program framework (Yang, 2019a, 2020a, 2020b) was introduced in an attempt to unify and generalize the NNGP/NTK theory to a broad range of architectures, eliminating the need to re-develop the theory for each new architecture. For example, (Yang, 2019a) proved the architectural universality of NNGP correspondence, while (Yang, 2020a) proved that of ntkInit. On the other hand, (Yang, 2020b) developed the most general machinery for Tensor Programs and as a corollary constructed a comprehensive theory of nonlinear random matrix theory, that, for example, can calculate the singular value distribution of a wide neural network of any architecture. Our proofs depend on the machinery of (Yang, 2020b) crucially, as discussed in Section 5.

Conclusion

New theories of deep learning almost always start with MLPs, rightly so as they are the simplest case and can often reveal the key insights more clearly. Of course, as deep learning itself is an applied field, one should always ask whether insights on MLPs extend to more general architectures, i.e. whether there is an architecturally universal extension of a proposed theory of MLPs. This is not always easy to answer.

In this paper, we showed that the NTK theory is architecturally universal, but more importantly, we showed that the Tensor Programs technique is a very powerful tool for answering the above question as a matter of routine. Looking forward, we hope to apply it to generate more novel and general insights.

References

The appendix is organized as follows: In Appendix A we expands upon the examples given in Section 3, while adding some additional details. In Appendix B we introduce the formal version of the $\textsc{Netsor}\top$,$\textsc{Netsor}\top^{+}$ programs. In Appendix C we introduce the graphical notation of $\textsc{Netsor}\top^{+}$ and demonstrate other examples of architectures or computations expressible in Tensor Programs. In Appendix D we prove our main result.

For the readers convenience we restate the notations described in Section 2, along with some additional ones which will be used throughout the appendix. We will consider SGD with batch size 1 and learning rate of 1 (WLOG).We use $\xi_{t}$ to denote the input and $\mathcal{L}_{t}$ to denote the loss function (absorbing the label) at step $t$. More generally, subscript $t$ on any symbol means time $t$. However, for brevity, we abuse notation and shorthand $f_{t}$ for $f_{t}(\xi_{t})$, and, for any (pre-)activation $x$, $x_{t}$ for $x_{t}(\xi_{t})$. We will also write $\chi_{t}$ for the loss derivative $\mathcal{L}_{t}^{\prime}(f_{t})$. For any vector $x(\xi)$ we define \delta x_{t+1}(\xi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\big{(}x_{t+1}(\xi)-x_{t}(\xi)\big{)} and dx(\xi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\frac{\partial f(\xi)}{\partial x(\xi)}. We will track the evolution of $f$ on an arbitrary input $\tilde{\xi}$.232323It might help to think of $\tilde{\xi}$ as some test sample, but it can also fall in the training set. Similar to above, we shorthand $\tilde{x}_{t},\tilde{f}_{t}$ for $x_{t}(\tilde{\xi}),f_{t}(\tilde{x})$. In general, omitting the time index $t$ for any time dependent quantity implies its value at initialization. (i.e $x(\xi)=x_{0}(\xi),f(\xi)=f_{0}(\xi)$). Finally, we use $\triangleq$ to imply equality of symbols (i.e $W^{1}\triangleq W^{2}$ iff $W^{1},W^{2}$ represent the same variable, as opposed to equality in value).

Appendix A Additional Examples

In this section we flesh out the examples given in Section 3 of the main text with the purpose of adding additional clarity, while maintaining the intuitive arguments as presented in each example to perform these calculations. The rigorous justification for these calculations will be given in the following section with the formal introduction of the Tensor Program framework.

Recall that our objective is to derive Section 3 by tracking the coordinate distribution of each (pre-)activations vector x(\xi),dx(\xi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\frac{\partial f(\xi)}{\partial x(\xi)},\delta x(\xi)\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\big{(}x_{t+1}(\xi)-x_{t}(\xi)\big{)}. \claimNTK*

In our calculations we will rely on the following rules relating to the coordinates of any $\mathbb{R}^{n}$ (pre-)activation vector $x(\xi)$ in the large width regime, which we will later formalize:

$x_{t+1}(\xi)-x_{t}(\xi)$ has $\Theta(\frac{1}{\sqrt{n}})$ coordinates.

$\delta x_{t+1}(\xi)$ has $\Theta(1)$ coordinates.

$x_{t}(\xi)=x(\xi)+o(1)$. Consequently $Z^{x_{t}(\xi)}=Z^{x(\xi)}$.

If x(\xi)=\phi\big{(}y(\xi)\big{)} for some vector $y(\xi)\in\mathbb{R}^{n}$, then by Taylor approximation \delta x_{t+1}(\xi)=\sqrt{n}\big{(}\phi(y_{t}(\xi)+\frac{\delta y_{t+1}(\xi)}{\sqrt{n}})-\phi(y_{t}(\xi))\big{)}\approx\phi^{\prime}\big{(}y_{t}(\xi)\big{)}\odot\delta y_{t+1}(\xi). Consequently $Z^{\delta x_{t+1}(\xi)}=\phi^{\prime}(Z^{y_{t}(\xi)})Z^{\delta y_{t+1}(\xi)}$.

We write $Z^{x}$ to denote the limit coordinate distribution of $x\in\mathbb{R}^{n}$ conditioned on the output function $f$ at initialization. Consequently we write $\operatorname*{\mathbb{E}}X^{x}Z^{y}$ to express a conditional expectation given the output function $f$. See Appendix B for the formal statement.

where $\xi\in\mathbb{R}^{d},~{}U=\frac{u}{\sqrt{d}}\in\mathbb{R}^{n\times d},~{}V=\frac{v}{\sqrt{n}}\in\mathbb{R}^{n\times 1}$, for trainable parameter tensor $u$, initialized iid from $\mathcal{N}(0,1)$, and $d=1$.

The infinite width NTK of this architecture is given by:

𝑓𝜉ℎ𝜉direct-productsuperscriptitalic-ϕ′ℎ𝜉𝑣\displaystyle\mathrel{\raisebox{-1.29167pt}{$\mathbin{\overset{\text{\tiny{def}}}{=}}$}}\sqrt{n}\frac{\partial f(\xi)}{\partial h(\xi)}=\phi^{\prime}(h(\xi))\odot v. (28) Therefore, by LLN it follows:242424While in A.1, we said $\operatorname*{\mathbb{E}}$ denotes expectation conditioned on $\lim f_{0}$, the NTK here does not actually depend on $\lim f_{0}$.

To show that Section 3 holds with the kernel in Eq. 29, we track the coordinate distribution $Z^{\delta\tilde{x}_{t+1}}$ at each step of SGD. At step $t$, the update to the weights $u_{t+1}-u_{t}$ is given by the gradient of the loss with respect to $u_{t}$:

𝑡1subscript𝑢𝑡subscript𝜒𝑡𝑑subscriptℎ𝑡superscriptsubscript𝜉𝑡top𝑛𝑑subscriptℎ𝑡direct-productsuperscriptitalic-ϕ′subscriptℎ𝑡𝑣\displaystyle u_{t+1}-u_{t}=-\chi_{t}\frac{dh_{t}\xi_{t}^{\top}}{\sqrt{n}},~{}~{}~{}dh_{t}=\phi^{\prime}(h_{t})\odot v (30) Recall that $\delta\tilde{h}_{t+1}=\sqrt{n}(\tilde{h}_{t+1}-\tilde{h}_{t})=\sqrt{n}(u_{t+1}\tilde{\xi}-u_{t}\tilde{\xi})$ and $\delta\tilde{x}_{t+1}=\sqrt{n}(\tilde{x}_{t+1}-\tilde{x}_{t})$. It therefore follows:

𝑡1direct-productsubscript𝜒𝑡superscriptsubscript𝜉𝑡top~𝜉superscriptitalic-ϕ′subscriptℎ𝑡𝑣𝛿subscript~𝑥𝑡1𝑛italic-ϕsubscript~ℎ𝑡𝛿subscript~ℎ𝑡1𝑛italic-ϕsubscript~ℎ𝑡\displaystyle\delta\tilde{h}_{t+1}=-\chi_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(h_{t})\odot v,~{}~{}~{}\delta\tilde{x}_{t+1}=\sqrt{n}\big{(}\phi(\tilde{h}_{t}+\frac{\delta\tilde{h}_{t+1}}{\sqrt{n}})-\phi(\tilde{h}_{t})\big{)}. (31) Since $h_{t}=\Theta(1)$ and $\delta\tilde{h}_{t+1}=\Theta(1)$, for large $n$ we may Taylor expand $\phi$ to first order around $\tilde{h}_{t}$:

𝑡1\displaystyle\delta\tilde{x}_{t+1} \displaystyle\approx\sqrt{n}\big{(}\phi(\tilde{h}_{t})+\frac{1}{\sqrt{n}}\phi^{\prime}(\tilde{h}_{t})\odot\delta\tilde{h}_{t+1}-\phi(\tilde{h}_{t})\big{)} (32) $\displaystyle=\phi^{\prime}(\tilde{h}_{t})\odot\delta\tilde{h}_{t+1}$ (33) $\displaystyle=-\chi_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(\tilde{h}_{t})\odot\phi^{\prime}(h_{t})\odot v.$ (34) Again since $\delta h_{t}(\xi)=\Theta(1)$, it follows that $h_{t}(\xi)=h(\xi)+\sum_{s=1}^{t}\frac{\delta h_{s}(\xi)}{\sqrt{n}}=h(\xi)+o(1)$. Hence, in the infinite width limit the coordinate distribution of $h_{t}(\xi)$ is identical to the coordinate distribution of $h(\xi)$ (i.e $Z^{h_{t}(\xi)}=Z^{h(\xi)}$). Using Eq. 34, the coordinate distribution of $\delta\tilde{x}_{t+1}$ is given by:

𝑡1subscript̊𝜒𝑡superscriptsubscript𝜉𝑡top~𝜉superscriptitalic-ϕ′superscript𝑍subscript~ℎ𝑡superscriptitalic-ϕ′superscript𝑍subscriptℎ𝑡superscript𝑍𝑣\displaystyle Z^{\delta\tilde{x}_{t+1}}=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(Z^{\tilde{h}_{t}})\phi^{\prime}(Z^{h_{t}})Z^{v}. (35) In the large width limit, the change in the output is simply given by $\tilde{f}_{t+1}-\tilde{f}_{t}=\frac{v^{\top}\delta\tilde{x}_{t+1}}{n}=\operatorname*{\mathbb{E}}Z^{\delta\tilde{x}_{t+1}}Z^{v}$. Using Eq. 35 and the independence of $Z^{v}$ from the other random variables,252525 Again, as in Footnote 10, the expectations in Sections A.1 and A.2 should more rigorously be interpreted as expectation conditional on the function values of $\lim f_{0}$.

A.2 2 hidden layers

where $U=\frac{u}{\sqrt{d}}\in\mathbb{R}^{n\times d},~{}W=\frac{w}{\sqrt{n}}\in\mathbb{R}^{n\times n},~{}V=\frac{v}{\sqrt{n}}\in\mathbb{R}^{n\times 1}$, for trainable parameters $u,w$, initialized iid from a normal distribution. As before we assume the last layer is not trained, and $d=1$.

subscript∇𝑢𝑓𝜉subscript∇𝑢𝑓~𝜉subscript∇𝑤𝑓𝜉subscript∇𝑤𝑓~𝜉\displaystyle=\langle\nabla_{u}f(\xi),\nabla_{u}f(\tilde{\xi})\rangle+\langle\nabla_{w}f(\xi),\nabla_{w}f(\tilde{\xi})\rangle (39) $\displaystyle=\xi^{\top}\tilde{\xi}\frac{dg(\xi)^{\top}dg(\tilde{\xi})}{n}+\frac{z(\xi)^{\top}z(\tilde{\xi})}{n}\frac{dh(\xi)^{\top}dh(\tilde{\xi})}{n}$ (40) $\displaystyle dh(\xi)$ \displaystyle\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\frac{\partial f(\xi)}{\partial h(\xi)}=\phi^{\prime}\big{(}h(\xi)\big{)}\odot v (41) $\displaystyle dg(\xi)$ \displaystyle\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\frac{\partial f(\xi)}{\partial g(\xi)}=\phi^{\prime}\big{(}g(\xi)\big{)}\odot\big{(}W^{\top}dh(\xi)\big{)}. (42) Naively using LLN on Eq. 39 (and $Z^{v}$ being independent from everything else) should result in:

superscript𝜉top~𝜉𝔼superscript𝑍𝑑𝑔𝜉superscript𝑍𝑑𝑔~𝜉𝔼superscript𝑍𝑧𝜉superscript𝑍𝑧~𝜉𝔼superscriptitalic-ϕ′superscript𝑍ℎ𝜉superscriptitalic-ϕ′superscript𝑍ℎ~𝜉\displaystyle\mathcal{K}(\xi,\tilde{\xi})=\xi^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}\big{[}Z^{dg(\xi)}Z^{dg(\tilde{\xi})}\big{]}+\operatorname*{\mathbb{E}}\big{[}Z^{z(\xi)}Z^{z(\tilde{\xi})}\big{]}\operatorname*{\mathbb{E}}\big{[}\phi^{\prime}(Z^{h(\xi)})\phi^{\prime}(Z^{h(\tilde{\xi})})\big{]}. (43) Evaluating the term \operatorname*{\mathbb{E}}\big{[}Z^{dg(\xi)}Z^{dg(\tilde{\xi})}\big{]} however presents a challenge since $dh(\xi)$ depends on both $W$ and $W^{\top}$. As it turns out, at initialization we may naively assume that $W^{\top},W$ are independent (formally known in the literature as gradient independence assumption, or GIA) 262626For a rigorous justification of the GIA assumption see (Yang, 2020a), we arrive using simple LLN arguments to:

Plugging Eq. 44 into Eq. 43 we arrive at the correct expression for the infinite width NTK.

To show that Section 3 holds at any step $t$ (where we may not assume that GIA holds), we track the distributions of the vectors $g(\xi),z(\xi),h(\xi),x(\xi)$ throughout training.

At any step $t$ the weights are updated according to:

𝑡1subscript𝑢𝑡subscript𝜒𝑡𝑑subscript𝑔𝑡superscriptsubscript𝜉𝑡top𝑛subscript𝑤𝑡1subscript𝑤𝑡subscript𝜒𝑡𝑑subscriptℎ𝑡superscriptsubscript𝑧𝑡top𝑛\displaystyle u_{t+1}-u_{t}=-\chi_{t}\frac{dg_{t}\xi_{t}^{\top}}{\sqrt{n}},~{}~{}~{}w_{t+1}-w_{t}=-\chi_{t}\frac{dh_{t}z_{t}^{\top}}{n}. (45) The update \delta\tilde{g}_{t+1}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}(\tilde{g}_{t+1}-\tilde{g}_{t}),\delta\tilde{z}_{t+1}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}(\tilde{z}_{t+1}-\tilde{z}_{t}) are given by:

As before, with large $n$ we have that $\bullet_{t+1}(\xi)-\bullet_{t}(\xi)\sim\Theta(\frac{1}{\sqrt{n}})$ and $\delta\bullet_{t+1}(\xi)\sim\Theta(1)$ coordinates for $\bullet$ replaced by $\{g,z,h,x\}$. And so after Taylor expanding $\phi(\tilde{g}_{t}+\frac{\delta\tilde{g}_{t+1}}{\sqrt{n}})$ around $\tilde{g}_{t}$:

𝑡1direct-productsuperscriptitalic-ϕ′subscript~𝑔𝑡𝛿subscript~𝑔𝑡1\displaystyle\delta\tilde{z}_{t+1}\approx\phi^{\prime}(\tilde{g}_{t})\odot\delta\tilde{g}_{t+1}. (47) In a similar fashion, using Eqs. 41 and 46 the updates $\delta\tilde{z}_{t+1},\delta\tilde{x}_{t+1},\delta\tilde{x}_{t+1}$ take the form:

𝑡1\displaystyle\delta\tilde{z}_{t+1} \displaystyle\approx\phi^{\prime}(g_{t})\odot\delta\tilde{g}_{t+1}=-\chi_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(g_{t})\odot\phi^{\prime}(\tilde{g}_{t})\odot\big{(}W^{\top}dh_{t}\big{)} (48) $\displaystyle\delta\tilde{h}_{t+1}$ $\displaystyle\approx W\delta\tilde{z}_{t+1}-\chi_{t}\frac{z_{t}^{\top}\tilde{z}_{t}}{n}\phi^{\prime}(h_{t})\odot v-\frac{1}{\sqrt{n}}\sum_{s=0}^{t}\chi_{s}\frac{z_{s}^{\top}\delta\tilde{z}_{t+1}}{n}\phi^{\prime}(h_{s})\odot v$ (49) $\displaystyle\delta\tilde{x}_{t+1}$ $\displaystyle\approx\phi^{\prime}(\tilde{h}_{t})\odot\delta\tilde{h}_{t+1}.$ (50) where we used Eq. 45 and

𝑡1\displaystyle\delta\tilde{h}_{t+1} $\displaystyle=W_{t+1}\delta\tilde{z}_{t+1}+\sqrt{n}(W_{t+1}-W_{t})\tilde{z}_{t}$ (51) $\displaystyle=W\delta\tilde{z}_{t+1}+\sum_{s=0}^{t}(W_{s+1}-W_{s})\delta\tilde{z}_{t+1}+\sqrt{n}(W_{t+1}-W_{t})\tilde{z}_{t}$ (52) to get Eq. 49. Based on Eqs. 41, 48, 49 and 50, the corresponding coordinate distributions take the form:

𝑡1\displaystyle Z^{\delta\tilde{z}_{t+1}} $\displaystyle=-\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\phi^{\prime}(Z^{g_{t}})\phi^{\prime}(Z^{\tilde{g}_{t}})Z^{W^{\top}dh_{t}}$ (54) $\displaystyle Z^{\delta\tilde{h}_{t+1}}$ $\displaystyle=Z^{W\delta\tilde{z}_{t+1}}-\mathring{\chi}_{t}\operatorname*{\mathbb{E}}[Z^{z_{t}}Z^{\tilde{z}}]\phi^{\prime}(Z^{h_{t}})Z^{v}$ (55) $\displaystyle Z^{\delta\tilde{x}_{t+1}}$ $\displaystyle=\phi^{\prime}(Z^{\tilde{h}_{t}})Z^{\delta\tilde{h}_{t+1}}.$ (56) As before, the functional update is given by $\lim_{n\to\infty}\tilde{f}_{t+1}-\tilde{f}_{t}=\lim_{n\to\infty}\frac{v^{\top}\delta\tilde{x}_{t+1}}{n}=\operatorname*{\mathbb{E}}Z^{v}Z^{\delta\tilde{x}_{t+1}}$. Plugging Eqs. 56 and 55:

𝑡1subscript̊𝜒𝑡𝔼superscript𝑍subscript𝑧𝑡superscript𝑍~𝑧𝔼superscriptitalic-ϕ′superscript𝑍subscriptℎ𝑡superscriptitalic-ϕ′superscript𝑍subscript~ℎ𝑡𝔼superscriptitalic-ϕ′superscript𝑍subscriptℎ𝑡superscript𝑍𝑊𝛿subscript~𝑧𝑡1superscript𝑍𝑣\displaystyle\operatorname*{\mathbb{E}}Z^{v}Z^{\delta\tilde{x}_{t+1}}=-\mathring{\chi}_{t}\operatorname*{\mathbb{E}}\big{[}Z^{z_{t}}Z^{\tilde{z}}\big{]}\operatorname*{\mathbb{E}}\big{[}\phi^{\prime}(Z^{h_{t}})\phi^{\prime}(Z^{\tilde{h}_{t}})\big{]}-\operatorname*{\mathbb{E}}\big{[}\phi^{\prime}(Z^{h_{t}})Z^{W\delta\tilde{z}_{t+1}}Z^{v}\big{]}. (57) To compute the second term of the RHS of Eq. 57, we use Section 3.2, reproduced below. \claimZdot*

Applying Section 3.2 to get the expression for $Z^{W\delta\tilde{z}_{t+1}}$:

𝑡1\displaystyle Z^{W\delta\tilde{z}_{t+1}} $\displaystyle=G+Z^{dh_{t}}\operatorname*{\mathbb{E}}\frac{\partial Z^{\delta\tilde{z}_{t+1}}}{\partial Z^{W^{\top}dh_{t}}}$ (58) \displaystyle=G-\phi^{\prime}(Z^{h_{t}})Z^{v}\mathring{\chi}_{t}\xi_{t}^{\top}\tilde{\xi}\operatorname*{\mathbb{E}}\big{[}\phi^{\prime}(Z^{g_{t}})\phi^{\prime}(Z^{\tilde{g_{t}}})\big{]} (59) As before, for $h\in\{g,z,h,x\}$ it holds that $h_{t}(\xi)=h(\xi)+\sum_{s=1}^{t}\frac{\delta h_{s}(\xi)}{\sqrt{n}}=h(\xi)+o(1)$, and $Z^{h_{t}(\xi)}=Z^{h(\xi)}$. Plugging Eq. 58 into Eq. 57 yields Section 3.

Appendix B Tensor Programs: the Formal Version

We briefly review the formal definition of Tensor Programs below, but readers needing more explanation and intuition should see (Yang, 2020b). We will directly describe $\textsc{Netsor}\top^{+}$ programs, which generalizes $\textsc{Netsor}\top$.

A $\textsc{Netsor}\top^{+}$ program is a sequence of $\mathbb{R}^{n}$-vectors and $\mathbb{R}$-scalars inductively generated via one of the following ways from an initial set $\mathcal{C}$ of random scalars, $\mathcal{V}$ of random $\mathbb{R}^{n}$ vectors, and a set $\mathcal{W}$ of random $\mathbb{R}^{n\times n}$ matrices (which will be sampled with iid Gaussian entries in B.2)

Given $\phi:\mathbb{R}^{k}\times\mathbb{R}^{l}\to\mathbb{R}$, previous scalars $\theta_{1},\ldots,\theta_{l}\in\mathbb{R}$ and vectors $x^{1},\ldots,x^{k}\in\mathbb{R}^{n}$, we can generate a new vector

where $\psi(-;\theta_{1},\ldots,\theta_{l})$ applies coordinatewise to each “$\alpha$-slice” $(x_{\alpha}^{1},\ldots,x_{\alpha}^{k})$.

Given same setup as above, we can also generate a new scalar

A $\textsc{Netsor}\top$ program is just a $\textsc{Netsor}\top^{+}$ program without scalars, without the usage of Moment, and without parameters $\theta_{1},\ldots,\theta_{l}$ in Nonlin+.

We will typically randomly sample the initial matrices, vectors, and scalars of the program as follows.

1) For each initial $W\in\mathcal{W}$, we sample iid $W_{\alpha\beta}\sim\mathcal{N}(0,\sigma_{W}^{2}/n)$ for some variance $\sigma_{W}^{2}$ associated to $W$, independent of other $W^{\prime}\in\mathcal{W}$; 2) for some multivariate Gaussian $Z^{\mathcal{V}}=\left\{Z^{h}:h\in\mathcal{V}\right\}\in\mathbb{R}^{\mathcal{V}}$, we sample the initial set of vectors $\mathcal{V}$ like $\left\{h_{\alpha}:h\in\mathcal{V}\right\}\sim Z^{\mathcal{V}}$ iid for each $\alpha\in[n]$. 3) For each initial scalar $\theta\in\mathcal{C}$, we require $\theta\xrightarrow{\mathrm{\mathrm{a.s.}}}\mathring{\theta}$ for some deterministic $\mathring{\theta}\in\mathbb{R}$.

The following constructs a random variable $Z^{h}$ for every vector $h$ and a deterministic scalar $\mathring{\theta}$ for every scalar $\theta$ in the program. The interpretation is that $h$ will have iid coordinates distributed like $Z^{h}$, and $\theta$ will converge to $\mathring{\theta}$ as $n\to\infty$.

Given a $\textsc{Netsor}\top^{+}$ program, we recursively define $Z^{h}$ for each vector $h$ and $\mathring{\theta}$ for each scalar $\theta$ as follows.