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

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:

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.

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.

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$:

2 2 Hidden Layers

In this example we consider a model of the form:

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$,

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

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$.

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

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^{+}$.

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.

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:

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

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.

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}}$.

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.

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

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$.

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

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.

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*

$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)}$.

The infinite width NTK of this architecture is given by:

A.2 2 hidden layers

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:

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$.

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.

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.

If $h\in\mathcal{V}$, then $Z^{h}$ is defined as in B.2. We also set \hat{Z}^{h}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}Z^{h} and \dot{Z}^{h}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}0.

Given same setup as above and scalar $\theta=\frac{1}{n}\sum_{\alpha=1}^{n}\psi(x_{\alpha}^{1},\ldots,x_{\alpha}^{k};\theta_{1},\ldots,\theta_{l})$, then

Here $\mathring{\theta}_{1},\ldots,\mathring{\theta}_{l}$ are deterministic, so the expectation is taken over $Z^{x^{1}},\ldots,Z^{x^{k}}$.

Z^{Wx}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\hat{Z}^{Wx}+\dot{Z}^{Wx} for every matrix $W$ (with $\mathcal{N}(0,\sigma_{W}^{2}/n)$ entries) and vector $x$, where

$\hat{Z}^{Wx}$ is a Gaussian variable with zero mean. Let $\mathcal{V}_{W}$ denote the set of all vectors in the program of the form $Wy$ for some $y$. Then $\{\hat{Z}^{Wy}:Wy\in\mathcal{V}_{W}\}$ is defined to be jointly Gaussian with zero mean and covariance

Furthermore, $\{\hat{Z}^{Wy}:Wy\in\mathcal{V}_{W}\}$ is mutually independent from $\{\hat{Z}^{v}:v\in\mathcal{V}\cup\bigcup_{\bar{W}\neq W}\mathcal{V}_{\bar{W}}\}$, where $\bar{W}$ ranges over $\mathcal{W}\cup\{A^{\top}:A\in\mathcal{W}\}$.

The following theorem ties the symbolic nature of the $Z$s to the analytic nature of a Tensor Program.

\textsc{Netsor}\top^{+} Master Theorem, c.f. Theorem E.15 of (Yang, 2020b)). Fix a Tensor Program initialized accordingly to B.2. Adopt B.8. Then

for any vectors $h^{1},\ldots,h^{k}$ in the program, where $Z^{h^{i}}$ are as defined in LABEL:{defn:netsortplusKeyIntuit}.

Any scalar $\theta$ in the program tends to $\mathring{\theta}$ almost surely, where $\mathring{\theta}$ is as defined in LABEL:{defn:netsortplusKeyIntuit}.

The partial derivative in ZDot should be interpreted as follows. By a simple inductive argument, $Z^{x}$ for every vector $x$ in the program is defined uniquely as a deterministic function $\varphi(\hat{Z}^{x^{1}},\ldots,\hat{Z}^{x^{k}})$ of some $x^{1},\ldots,x^{k}$ in $\mathcal{V}$ or introduced by MatMul (notationally, we are suppressing the possible dependence on limit scalars $\mathring{\theta}_{1},\ldots,\mathring{\theta}_{l}$). For instance, if in a program we have $A\in\mathcal{W},v\in\mathcal{V}$, $y=Av,x=A^{\top}y$, then $Z^{x}=\hat{Z}^{x}+\hat{Z}^{v}$, so $\varphi$ is given by $\varphi(a,b)=a+b$. Then

are, roughly speaking, functions whose weak derivatives are polynomially bounded.

Here are some basic properties of pseudo-Lipschitz functions:

A pseudo-Lipschitz function is polynomially bounded.

A composition of pseudo-Lipschitz functions of degrees $d_{1}$ and $d_{2}$ is pseudo-Lipschitz of degree $d_{1}+d_{2}$.

A pseudo-Lipschitz function is Lipschitz on any compact set.

We adopt the following assumption for the Master Theorem B.4.

Statement 1 in B.8 essentially says that if we have scalars $\theta_{1},\ldots,\theta_{l}$ in the program, then we can produce a new scalar by applying a continuous function (a weaker restriction than a pseudo-Lipschitz function) to them. Indeed, if $\theta_{1},\ldots,\theta_{l}$ converge almost surely, then this new scalar does too. In our setting, statement 1 is used to allow any loss function whose derivative is continuous.

Other versions of the Master Theorem can be found in (Yang, 2020b), for example, versions where the we do not assume any smoothness condition at all on the nonlinearities beyond that they be polynomially bounded, in exchange for assuming what’s called a rank stability condition. This rank stability should be generically true, but checking it rigorously is subtle, so we are content with the pseudo-Lipschitz condition in this paper.

Appendix C More Diagrams

We can augment the graphical form of $\textsc{Netsor}\top$ to accomodate the Moment instruction in $\textsc{Netsor}\top^{+}$. See Fig. 4 for an example for layernorm and attention. In short, we denote scalar variables with a square, in contrast to the circle for vector variables, and we use a “bar-gate” to denote the Moment, where the function in the gate corresponds to $\psi$ in Moment.

In addition, for more examples of the expressivity of $\textsc{Netsor}\top$, Figs. 5 and 6 demonstrate convolution and MLP backpropagation in $\textsc{Netsor}\top$.

Appendix D Proof of Main Result

We dedicate the following section to prove 5.3. We will begin by proving a simplified version under the same assumptions as Section 5.1, as reproduced below:

The forward pass does not use both a matrix and its transpose (in different MatMuls).

We assume the last layer embedding is a G-var.

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.

\textsc{Netsor}\top^{+} Program SGD is comprised of a sequence of forward and backward passes computed on some architecture. WLOG, let $\pi_{0}$ denote the reduced program implementing the body of network $f$, and let $x(\xi)$ denote the final embedding such that $f(\xi)=V^{\top}x(\xi)$, we will now show how the SGD procedure on $\pi_{0}$ can be implemented by a $\textsc{Netsor}\top^{+}$ program.

Given $\mathtt{v}$ and (the columns of) $X$ as vectors and $\mathtt{f}$ as scalars in a program, $\hat{v}$ may be defined in the same program via Nonlin, where $\frac{\Sigma^{+}\mathtt{f}}{\sqrt{n}}$ and $\Sigma^{+}\mu$ (both finite-dimensional) provide coefficients for the linear combination over (columns of) $X$. Formally, to express the evolution of $f$ conditioned on $f_{0}=\mathtt{f}$ at initialization, the program will calculate the first forward pass up to $X$, calculate the loss derivatives $\chi$ assuming $f_{0}=\mathtt{f}$, and then proceed with the backward pass and later forward/backward passes with $v$ replaced by $\hat{v}$.

The loss derivative $\chi(\xi)=\frac{\partial\mathcal{L}(f(\xi))}{\partial(f(\xi))}$ after the first forward pass given $f(\xi)$ can be implemented with Moment instructions using $\psi(;f(\xi))=\mathcal{L}^{\prime}(f(\xi))$.

D.1.2 Implementing SGD

To write the full unrolled SGD as a $\textsc{Netsor}\top^{+}$ program, we will need to implement the error signal dg_{t}\mathrel{\raisebox{-1.29167pt}{\mathbin{\overset{\text{\tiny{def}}}{=}}}}\sqrt{n}\frac{\partial f_{t}}{\partial g_{t}} for each G-var $g$ at time $t$. To accomplish this, we recall the notion of paths in program $\pi_{0}$: \Paths* Note that a path $p$ represents a series of nodes independent of an input, and can be instantiated as $p(\xi)$ by an input $\xi$, resulting in a series of instantiated G-vars and X-vars $p^{i}(\xi)$.

For any G-var $g=Wh$, we can write the error term $dg$ as the summation of errors signals over paths $p$:

Note that Eq. 78 can be written in $\textsc{Netsor}\top$ language using MatMul instructions using the transposed weights, and Nonlin instructions using $\psi^{\prime}$, which is pseudo-Lipschitz by D.1.

Eq. 82 may be implemented as a Nonlin+ instruction:

the scalar type outputs $f_{t}(\xi)$ at $t>0$ for any input $\xi$ can be implemented using the Moment instruction. The loss derivative $\chi_{t},t>0$ given $f_{t}$ can be implemented with Moment instructions using \psi(-;f(\xi))=\mathcal{L}^{\prime}\big{(}f(\xi)\big{)} where $f(\xi)$ is treated as a scalar type as in the first forward pass.

$dg(\xi)_{t+1}$ is implemented using Moment and Nonlin+ instructions.

According to the $\textsc{Netsor}\top^{+}$ rules as specified in B.3, we have the following identities:

D.2 Deriving The NTK

Instantiate paths $p$ and $q$ on two inputs $\xi,\xi^{\prime}$ by $p=p(\xi),q=q(\xi^{\prime})$ (abusing notation slightly). We define an inner product between them as follows:

For each weight $W\in\mathcal{W}$, the gradient of the output with respect to $w$ is given by:

Here, $g,h$ represent nodes in program $\pi_{0}$ that can be instantiated by an input $\xi$. The NTK of $f$ can be expressed as:

Using Eqs. 74 and 78, for any G-var $g=Wh$, we can write the error term $dg$ as the summation of errors signals over paths $p$:

D.3 Getting Section 3

We write $Z^{x}\equiv Z^{y}\mod\hat{Z}^{W\bullet}$ to denote that $Z^{x}-Z^{y}$ is a linear combination of $\hat{Z}^{Wu}$ for various vectors $u$.

At any time $t$ and G-var $g=Wh$, the following holds:

D.3.1 Base Case

D.3.2 Inductive case

Suppose $g=Wh$, where $h=\psi(h^{1},...,h^{k})$, we then have using Eq. 91:

Hence, we have proven the induction hypothesis.

D.3.3 Proving Section 3 using the induction hypothesis

WLOG assume $f(\xi)=V^{\top}x(\xi)$ for some G-var $x(\xi)$. Using the induction hypothesis and the Master Theorem (B.4), we have that:

D.4 Relaxing A1., A2., A3. and A4.

We now briefly discuss the case where A1., A2., A3. and A4. are relaxed, as well as the case where $f$ is represented by a $\textsc{Netsor}\top^{+}$ program. As the proof of the general case follows roughly the same logic as in D.1, we only discuss the meaningful differences in each case.

Recall the input and output layers are parameterized by $\{u^{i}\},v$ which now depend on $t$. The output evolution is now given by:

When both $W,W^{\top}$ are allowed in the forward pass, the update equations for each $w_{t}$ take the form:

If $g=W^{\top}h$ then Eq. 107 holds with $\mod\hat{Z}^{W^{\top}\bullet}$ replacing $\mod\hat{Z}^{W\bullet}$.

Some additional complications need to be resolved. Specifically, with setup D.1 we have used in two places the fact that no transpose is used in the forward pass to prove the induction hypothesis (see Eqs. 112 and 120). To prove the induction, and assuming $g=Wh$, we now have instead of Eq. 112 (using Eq. 111):

D.4.3 Multiple outputs and arbitrary batchsize

We have used a scalar output and a batchsize of 1 throughout this paper. However, extending to multiple (finite) outputs and an arbitrary batchsize requires no additional arguments besides some additional notations. For example, the definition of path should now be altered to express dependency on multiple samples (if batchnorm is used for example). The proof however follows roughly the same logic in D.1.

D.4.4 X-var embedding

We assumed in our proof that $x$, which represents the final embedding of $f$ is a G-var. However, extending the proof to the case where $x$ is an X-var is straightforward. Let $f(\xi)=V^{\top}x(\xi)$ where $x=\psi(h^{1},...,h^{k})$ and $h^{1},...,h^{k}$ are G-vars. Using the induction hypothesis, along with Eq. 93 yields:

\textsc{Netsor}\top^{+} If the network is more generally represented by a $\textsc{Netsor}\top^{+}$ program instead of just a $\textsc{Netsor}\top$ program, then our proof can be very simply modified to accommodate as follows: The new operation allowed in such a network is the production of a scalar through Moment, say $a=\frac{1}{n}\sum_{\alpha=1}^{n}\psi(x^{1}_{\alpha},\ldots,x^{k}_{\alpha};\theta^{1},\ldots,\theta^{l})$. By a similar inductive argument as before, we will see that 1) $x^{i}_{t}=x^{i}_{0}+o(1)$ for all $i\in[k]$ and $\theta^{j}_{t}=\theta^{j}_{0}+o(1)$ for all $j\in[l]$, so that $a_{t}=a_{0}+o(1)$; 2) in the backward pass, any backpropagation through $a$ will zero out: For example, if $a$ is only used later in a Nonlin $z=\mathtt{\psi}(y;a)$, then $\frac{1}{\sqrt{n}}\nabla_{a}f=\langle dz,\partial_{a}\mathtt{\psi}(y;a)\rangle/n$ will converge to 0 because of GIA (as $dz$ is linear in the final layer), and the error signal at $x^{i}$ times $\sqrt{n}$ is the constant vector with entries $\frac{1}{\sqrt{n}}\nabla_{a}f$, which is $o(1)$.

Therefore, we can treat any scalar produced through Moment as a constant fixed at initialization, and the notion of path from before carries over here without change (by assuming all nonlinearities with scalar parameters to be parameterless nonlinearities where the parameters are fixed). Then the same reasoning follows.