Wide neural networks of any depth evolve as linear models under gradient descent^*

We build on recent work by Jacot et al (2018) that characterize the exact dynamics of network outputs throughout gradient descent training in the infinite width limit. Their results establish that full batch gradient descent in parameter space corresponds to kernel gradient descent in function space with respect to a new kernel, the NTK. We examine what this implies about dynamics in parameter space, where training updates are actually made.

Daniely et al (2016) study the relationship between NNs and kernels at initialization. They bound the difference between the infinite width kernel and the empirical kernel at finite width n, which diminishes as $\mathcal{O}\left(1/\sqrt{n}\right)$. Daniely (2017) uses the same kernel perspective to study stochastic gradient descent (SGD) training of NNs.

Saxe et al (2014) study the training dynamics of deep linear networks, in which the nonlinearities are treated as identity functions. Deep linear networks are linear in their inputs, but not in their parameters. In contrast, we show that the outputs of sufficiently wide NNs are linear in the updates to their parameters during gradient descent, but not usually their inputs.

Du et al (2019), Allen-Zhu et al (2019, 2018), Zou et al (2019) study the convergence of gradient descent to global minima. They proved that for i.i.d. Gaussian initialization, the parameters of sufficiently wide networks move little from their initial values during SGD. This small motion of the parameters is crucial to the effect we present, where wide NNs behave linearly in terms of their parameters throughout training.

Mei et al (2018), Chizat and Bach (2018), Rotskoff and Vanden-Eijnden (2018), Sirignano and Spiliopoulos (2018) analyze the mean field SGD dynamics of training NNs in the large-width limit. Their mean field analysis describes distributional dynamics of network parameters via a PDE. However, their analysis is restricted to one hidden layer networks with a scaling limit $\left(1/n\right)$ different from ours $\left(1/\sqrt{n}\right)$, which is commonly used in modern networks (He et al 2016, Glorot and Bengio 2010).

Chizat et al (2018) ⁵ argued that infinite width networks are in 'lazy training' regime and maybe too simple to be applicable to realistic NNs. Nonetheless, we empirically investigate the applicability of the theory in the finite-width setting and find that it gives an accurate characterization of both the learning dynamics and posterior function distributions across a variety of conditions, including some practical network architectures such as the wide residual network (Zagoruyko and Komodakis 2016).

Let $\mathcal{D}\subseteq {\mathbb{R}}^{{n}_{0}}{\times}{\mathbb{R}}^{k}$ denote the training set and $\mathcal{X}=\left\{x:\left(x,y\right)\in \mathcal{D}\right\}$ and $\mathcal{Y}=\left\{y:\left(x,y\right)\in \mathcal{D}\right\}$ denote the inputs and labels, respectively. Consider a fully-connected feed-forward network with L hidden layers with widths n_l , for l = 1,..., L and a readout layer with n_L+1 = k. For each $x\in {\mathbb{R}}^{{n}_{0}}$, we use ${h}^{l}\left(x\right),{x}^{l}\left(x\right)\in {\mathbb{R}}^{{n}_{l}}$ to represent the pre- and post-activation functions at layer l with input x. The recurrence relation for a feed-forward network is defined as

$$\begin{equation}\begin{cases}^{l+1}\quad \hfill & ={x}^{l}{W}^{l+1}+{b}^{l+1}\hfill \\ {x}^{l+1}\quad \hfill & =\phi \left({h}^{l+1}\right)\hfill \end{cases}\quad \mathrm{a}\mathrm{n}\mathrm{d}\quad \begin{cases}_{i,j}^{l}\quad \hfill & =\frac{{\sigma }_{\omega }}{\sqrt{{n}_{l-1}}}{\omega }_{ij}^{l}\hfill \\ {b}_{j}^{l}\quad \hfill & ={\sigma }_{b}{\beta }_{j}^{l}\hfill \end{cases},\end{equation} \tag{ 1 }$$

where ϕ is a point-wise activation function, ${W}^{l+1}\in {\mathbb{R}}^{{n}_{l}{\times}{n}_{l+1}}$ and ${b}^{l+1}\in {\mathbb{R}}^{{n}_{l+1}}$ are the weights and biases, ${\omega }_{ij}^{l}$ and ${b}_{j}^{l}$ are the trainable variables, drawn i.i.d. from a standard Gaussian ${\omega }_{ij}^{l},{\beta }_{j}^{l}\sim \mathcal{N}\left(0,1\right)$ at initialization, and ${\sigma }_{\omega }^{2}$ and ${\sigma }_{b}^{2}$ are weight and bias variances. Note that this parametrization is non-standard, and we will refer to it as the NTK parameterization. It has already been adopted in several recent works (van Laarhoven 2017, Karras et al 2018, Jacot et al 2018, Du et al 2019, Park et al 2019). Unlike the standard parameterization that only normalizes the forward dynamics of the network, the NTK-parameterization also normalizes its backward dynamics. We note that the predictions and training dynamics of NTK-parameterized networks are identical to those of standard networks, up to a width-dependent scaling factor in the learning rate for each parameter tensor. As we derive, and support experimentally, in the supplementary material (SM) (https://stacks.iop/JSTAT/2020/124002/mmedia) sections F and G, our results (linearity in weights, GP predictions) also hold for networks with a standard parameterization.

We define ${\theta }^{l}\equiv \mathrm{v}\mathrm{e}\mathrm{c}\left(\left\{{W}^{l},{b}^{l}\right\}\right)$, the $\left(\left({n}_{l-1}+1\right){n}_{l}\right){\times}1$ vector of all parameters for layer l. $\theta =\mathrm{v}\mathrm{e}\mathrm{c}\left({\cup }_{l=1}^{L+1}{\theta }^{l}\right)$ then indicates the vector of all network parameters, with similar definitions for θ^⩽l and θ^>l . Denote by θ_t the time-dependence of the parameters and by θ₀ their initial values. We use ${f}_{t}\left(x\right)\equiv {h}^{L+1}\left(x\right)\in {\mathbb{R}}^{k}$ to denote the output (or logits) of the neural network at time t. Let $\ell \left(\hat{y},y\right):{\mathbb{R}}^{k}{\times}{\mathbb{R}}^{k}\to \mathbb{R}$ denote the loss function where the first argument is the prediction and the second argument the true label. In supervised learning, one is interested in learning a θ that minimizes the empirical loss ⁶ , $\mathcal{L}={\sum }_{\left(x,y\right)\in \mathcal{D}}\ell \left({f}_{t}\left(x,\theta \right),y\right)$.

Let η be the learning rate ⁷ . Via continuous time gradient descent, the evolution of the parameters θ and the logits f can be written as

$$\begin{equation}{\dot {\theta }}_{t}=-\eta {{\nabla }_{\theta }{f}_{t}\left(\mathcal{X}\right)}^{T}{\nabla }_{{f}_{t}\left(\mathcal{X}\right)}\mathcal{L}\end{equation} \tag{ 2 }$$

$$\begin{equation}{\dot {f}}_{t}\left(\mathcal{X}\right)={\nabla }_{\theta }{f}_{t}\left(\mathcal{X}\right){\dot {\theta }}_{t}=-\eta {\hat{{\Theta}}}_{t}\left(\mathcal{X},\mathcal{X}\right){\nabla }_{{f}_{t}\left(\mathcal{X}\right)}\mathcal{L}\end{equation} \tag{ 3 }$$

where ${f}_{t}\left(\mathcal{X}\right)=\mathrm{v}\mathrm{e}\mathrm{c}\left({\left[{f}_{t}\left(x\right)\right]}_{x\in \mathcal{X}}\right)$, the $k\vert \mathcal{D}\vert {\times}1$ vector of concatenated logits for all examples, and ${\nabla }_{{f}_{t}\left(\mathcal{X}\right)}\mathcal{L}$ is the gradient of the loss with respect to the model's output, ${f}_{t}\left(\mathcal{X}\right)$. ${\hat{{\Theta}}}_{t}\equiv {\hat{{\Theta}}}_{t}\left(\mathcal{X},\mathcal{X}\right)$ is the tangent kernel at time t, which is a $k\vert \mathcal{D}\vert {\times}k\vert \mathcal{D}\vert$ matrix

$$\begin{equation}{\hat{{\Theta}}}_{t}={\nabla }_{\theta }{f}_{t}\left(\mathcal{X}\right){{\nabla }_{\theta }{f}_{t}\left(\mathcal{X}\right)}^{T}=\sum _{l=1}^{L+1}{\nabla }_{{\theta }^{l}}{f}_{t}\left(\mathcal{X}\right){{\nabla }_{{\theta }^{l}}{f}_{t}\left(\mathcal{X}\right)}^{T}.\end{equation} \tag{ 4 }$$

One can define the tangent kernel for general arguments, e.g. ${\hat{{\Theta}}}_{t}\left(x,\mathcal{X}\right)$ where x is test input. At finite-width, $\hat{{\Theta}}$ will depend on the specific random draw of the parameters and in this context we refer to it as the empirical tangent kernel.

The dynamics of discrete gradient descent can be obtained by replacing ${\dot {\theta }}_{t}$ and ${\dot {f}}_{t}\left(\mathcal{X}\right)$ with (θ_i+1 − θ_i ) and $\left({f}_{i+1}\left(\mathcal{X}\right)-{f}_{i}\left(\mathcal{X}\right)\right)$ above, and replacing ${\text{e}}^{-\eta {\hat{{\Theta}}}_{0}t}$ with $\left(1-{\left(1-\eta {\hat{{\Theta}}}_{0}\right)}^{i}\right)$ below.

In this section, we consider the training dynamics of the linearized network. Specifically, we replace the outputs of the neural network by their first order Taylor expansion,

$$\begin{equation}{f}_{t}^{\text{lin}}\left(x\right)\equiv {f}_{0}\left(x\right)+{\left.{\nabla }_{\theta }{f}_{0}\left(x\right)\right\vert }_{\theta ={\theta }_{0}}{\omega }_{t},\end{equation} \tag{ 5 }$$

where ω_t ≡ θ_t − θ₀ is the change in the parameters from their initial values. Note that ${f}_{t}^{\text{lin}}$ is the sum of two terms: the first term is the initial output of the network, which remains unchanged during training, and the second term captures the change to the initial value during training. The dynamics of gradient flow using this linearized function are governed by,

$$\begin{equation}{\dot {\omega }}_{t}=-\eta {{\nabla }_{\theta }{f}_{0}\left(\mathcal{X}\right)}^{T}{\nabla }_{{f}_{t}^{\text{lin}}\left(\mathcal{X}\right)}\mathcal{L}\end{equation} \tag{ 6 }$$

$$\begin{equation}{\dot {f}}_{t}^{\text{lin}}\left(x\right)=-\eta {\hat{{\Theta}}}_{0}\left(x,\mathcal{X}\right){\nabla }_{{f}_{t}^{\text{lin}}\left(\mathcal{X}\right)}\mathcal{L}.\end{equation} \tag{ 7 }$$

As ∇_θ f₀(x) remains constant throughout training, these dynamics are often quite simple. In the case of an MSE loss, i.e. $\ell \left(\hat{y},y\right)=\frac{1}{2}{\Vert}\hat{y}-y{{\Vert}}_{2}^{2}$, the ODEs have closed form solutions

$$\begin{equation}{\omega }_{t}=-{{\nabla }_{\theta }{f}_{0}\left(\mathcal{X}\right)}^{T}{\hat{{\Theta}}}_{0}^{-1}\left(I-{\mathrm{e}}^{-\eta {\hat{{\Theta}}}_{0}t}\right)\left({f}_{0}\left(\mathcal{X}\right)-\mathcal{Y}\right),\end{equation} \tag{ 8 }$$

$$\begin{equation}{f}_{t}^{\text{lin}}\left(\mathcal{X}\right)=\left(I-{\mathrm{e}}^{-\eta {\hat{{\Theta}}}_{0}t}\right)\mathcal{Y}+{\text{e}}^{-\eta {\hat{{\Theta}}}_{0}t}\enspace {f}_{0}\left(\mathcal{X}\right).\end{equation} \tag{ 9 }$$

For an arbitrary point x, ${f}_{t}^{\text{lin}}\left(x\right)={\mu }_{t}\left(x\right)+{\gamma }_{t}\left(x\right)$, where

$$\begin{equation}{\mu }_{t}\left(x\right)={\hat{{\Theta}}}_{0}\left(x,\mathcal{X}\right){\hat{{\Theta}}}_{0}^{-1}\left(I-{\text{e}}^{-\eta {\hat{{\Theta}}}_{0}t}\right)\mathcal{Y}\end{equation} \tag{ 10 }$$

$$\begin{equation}{\gamma }_{t}\left(x\right)={f}_{0}\left(x\right)-{\hat{{\Theta}}}_{0}\left(x,\mathcal{X}\right){\hat{{\Theta}}}_{0}^{-1}\left(I-{\mathrm{e}}^{-\eta {\hat{{\Theta}}}_{0}t}\right){f}_{0}\left(\mathcal{X}\right).\end{equation} \tag{ 11 }$$

Therefore, we can obtain the time evolution of the linearized neural network without running gradient descent. We only need to compute the tangent kernel ${\hat{{\Theta}}}_{0}$ and the outputs f₀ at initialization and use equations (8), (10) and (11) to compute the dynamics of the weights and the outputs.

As the width of the hidden layers approaches infinity, the central limit theorem (CLT) implies that the outputs at initialization ${\left\{{f}_{0}\left(x\right)\right\}}_{x\in \mathcal{X}}$ converge to a multivariate Gaussian in distribution. Informally, this occurs because the pre-activations at each layer are a sum of Gaussian random variables (the weights and bias), and thus become a Gaussian random variable themselves. See Poole et al (2016), Schoenholz et al (2017), Lee et al (2018), Xiao et al (2018), Yang and Schoenholz (2017) for more details, and Matthews et al (2018b), Novak et al (2019a) for a formal treatment.

Therefore, randomly initialized NNs are in correspondence with a certain class of GPs (hereinafter referred to as NNGPs), which facilitates a fully Bayesian treatment of NNs (Lee et al 2018, Matthews et al 2018a). More precisely, let ${f}_{t}^{i}$ denote the ith output dimension and $\mathcal{K}$ denote the sample-to-sample kernel function (of the pre-activation) of the outputs in the infinite width setting,

$$\begin{equation}{\mathcal{K}}^{i,j}\left(x,{x}^{\prime }\right)=\underset{\mathrm{min}\left({n}_{1},\dots ,{n}_{L}\right)\to \infty }{\mathrm{lim}}\enspace \mathbb{E}\left[{f}_{0}^{i}\left(x\right)\cdot {f}_{0}^{j}\left({x}^{\prime }\right)\right],\end{equation} \tag{ 12 }$$

then ${f}_{0}\left(\mathcal{X}\right)\sim \mathcal{N}\left(0,\mathcal{K}\left(\mathcal{X},\mathcal{X}\right)\right)$, where ${\mathcal{K}}^{i,j}\left(x,{x}^{\prime }\right)$ denotes the covariance between the ith output of x and jth output of x', which can be computed recursively (see Lee et al 2018, section 2.3 and SM section E). For a test input $x\in {\mathcal{X}}_{\mathrm{T}}$, the joint output distribution $f\left(\left[x,\mathcal{X}\right]\right)$ is also multivariate Gaussian. Conditioning on the training samples ⁸ , $f\left(\mathcal{X}\right)=\mathcal{Y}$, the distribution of $\left.f\left(x\right)\right\vert \mathcal{X},\mathcal{Y}$ is also a Gaussian $\mathcal{N}\left(\mu \left(x\right),{\Sigma}\left(x\right)\right)$,

$$\begin{equation}\mu \left(x\right)=\mathcal{K}\left(x,\mathcal{X}\right){\mathcal{K}}^{-1}\mathcal{Y},\qquad {\Sigma}\left(x\right)=\mathcal{K}\left(x,x\right)-\mathcal{K}\left(x,\mathcal{X}\right){\mathcal{K}}^{-1}\mathcal{K}{\left(x,\mathcal{X}\right)}^{T},\end{equation} \tag{ 13 }$$

and where $\mathcal{K}=\mathcal{K}\left(\mathcal{X},\mathcal{X}\right)$. This is the posterior predictive distribution resulting from exact Bayesian inference in an infinitely wide neural network.

2.3.1. GPs from gradient descent training

If we freeze the variables θ^⩽L after initialization and only optimize θ^L+1, the original network and its linearization are identical. Letting the width approach infinity, this particular tangent kernel ${\hat{{\Theta}}}_{0}$ will converge to $\mathcal{K}$ in probability and equation (10) will converge to the posterior equation (13) as t → ∞ (for further details see SM section D). This is a realization of the 'sample-then-optimize' approach for evaluating the posterior of a GP proposed in Matthews et al (2017).

If none of the variables are frozen, in the infinite width setting, ${\hat{{\Theta}}}_{0}$ also converges in probability to a deterministic kernel Θ (Jacot et al 2018, Yang 2019), which we sometimes refer to as the analytic kernel, and which can also be computed recursively (see SM section E). For ReLU and erf nonlinearity, Θ can be exactly computed (SM section C) which we use in section 3. Letting the width go to infinity, for any t, the output ${f}_{t}^{\text{lin}}\left(x\right)$ of the linearized network is also Gaussian distributed because equations (10) and (11) describe an affine transform of the Gaussian $\left[{f}_{0}\left(x\right),{f}_{0}\left(\mathcal{X}\right)\right]$. Therefore

Corollary 1. For every test points in $x\in {\mathcal{X}}_{\mathrm{T}}$, and t ⩾ 0, ${f}_{t}^{\text{lin}}\left(x\right)$ converges in distribution as width goes to infinity to a Gaussian with mean and covariance given by ⁹

$$\begin{equation}\mu \left({\mathcal{X}}_{\mathrm{T}}\right)={\Theta}\left({\mathcal{X}}_{\mathrm{T}},\mathcal{X}\right){{\Theta}}^{-1}\left(I-{\mathrm{e}}^{-\eta {\Theta}t}\right)\mathcal{Y},\end{equation} \tag{ 14 }$$

$$\begin{align}\hfill {\Sigma}\left({\mathcal{X}}_{\mathrm{T}},{\mathcal{X}}_{\mathrm{T}}\right)& =\mathcal{K}\left({\mathcal{X}}_{\mathrm{T}},{\mathcal{X}}_{\mathrm{T}}\right)+{\Theta}\left({\mathcal{X}}_{\mathrm{T}},\mathcal{X}\right){{\Theta}}^{-1}\left(I-{\mathrm{e}}^{-\eta {\Theta}t}\right)\mathcal{K}\left(I-{\mathrm{e}}^{-\eta {\Theta}t}\right){{\Theta}}^{-1}{\Theta}\left(\mathcal{X},{\mathcal{X}}_{\mathrm{T}}\right)\hfill \\ \hfill & \quad -\left({\Theta}\left({\mathcal{X}}_{\mathrm{T}},\mathcal{X}\right){{\Theta}}^{-1}\left(I-{\mathrm{e}}^{-\eta {\Theta}t}\right)\mathcal{K}\left(\mathcal{X},{\mathcal{X}}_{\mathrm{T}}\right)+\mathrm{h}.\mathrm{c}.\right).\hfill \end{align} \tag{ 15 }$$

Therefore, over random initialization, ${\mathrm{lim}}_{t\to \infty }{\mathrm{lim}}_{n\to \infty }\enspace {f}_{t}^{\text{lin}}\left(x\right)$ has distribution

$$\begin{align}\hfill & \mathcal{N}\left({\Theta}\left({\mathcal{X}}_{\mathrm{T}},\mathcal{X}\right){{\Theta}}^{-1}\mathcal{Y},\mathcal{K}\left({\mathcal{X}}_{\mathrm{T}},{\mathcal{X}}_{\mathrm{T}}\right)+{\Theta}\left({\mathcal{X}}_{\mathrm{T}},\mathcal{X}\right){{\Theta}}^{-1}\mathcal{K}{{\Theta}}^{-1}{\Theta}\left(\mathcal{X},{\mathcal{X}}_{\mathrm{T}}\right)\right.\hfill \\ \hfill & \qquad \left.-\left({\Theta}\left({\mathcal{X}}_{\mathrm{T}},\mathcal{X}\right){{\Theta}}^{-1}\mathcal{K}\left(\mathcal{X},{\mathcal{X}}_{\mathrm{T}}\right)+\mathrm{h}.\mathrm{c}.\right)\right).\hfill \end{align} \tag{ 16 }$$

Unlike the case when only θ^L+1 is optimized, equations (14) and (15) do not admit an interpretation corresponding to the posterior sampling of a probabilistic model ¹⁰ . We contrast the predictive distributions from the NNGP, NTK-GP (i.e. equations (14) and (15)) and ensembles of NNs in figure 2.

Infinitely-wide NNs open up ways to study deep NNs both under fully Bayesian training through the GP correspondence, and under GD training through the linearization perspective. The resulting distributions over functions are inconsistent (the distribution resulting from GD training does not generally correspond to a Bayesian posterior). We believe understanding the biases over learned functions induced by different training schemes and architectures is a fascinating avenue for future work.

Equations (2) and (3) of the original network are intractable in general, since ${\hat{{\Theta}}}_{t}$ evolves with time. However, for the mean squared loss, we are able to prove formally that, as long as the learning rate $\eta {< }{\eta }_{\text{critical}}{:=}2{\left({\lambda }_{\text{min}}\left({\Theta}\right)+{\lambda }_{\text{max}}\left({\Theta}\right)\right)}^{-1}$, where λ_min/max(Θ) is the min/max eigenvalue of Θ, the gradient descent dynamics of the original neural network falls into its linearized dynamics regime.

Theorem 2.1 (Informal). Let n₁ = ⋯ = n_L = n and assume λ_min(Θ) > 0. Applying gradient descent with learning rate η < η_critical (or gradient flow), for every $x\in {\mathbb{R}}^{{n}_{0}}$ with ||x||₂ ⩽ 1, with probability arbitrarily close to 1 over random initialization,

$$\begin{equation}\underset{t{\geqslant}0}{\mathrm{sup}}{\Vert}{f}_{t}\left(x\right)-{f}_{t}^{\text{lin}}\left(x\right){{\Vert}}_{2},\qquad \underset{t{\geqslant}0}{\mathrm{sup}}\frac{{\Vert}{\theta }_{t}-{\theta }_{0}{{\Vert}}_{2}}{\sqrt{n}},\qquad \underset{t{\geqslant}0}{\mathrm{sup}}{\Vert}{\hat{{\Theta}}}_{t}-{\hat{{\Theta}}}_{0}{{\Vert}}_{F}=\mathcal{O}\left({n}^{-\frac{1}{2}}\right),\quad \mathrm{a}\mathrm{s}\enspace n\to \infty .\end{equation} \tag{ 17 }$$

Therefore, as n → ∞, the distributions of f_t (x) and ${f}_{t}^{\text{lin}}\left(x\right)$ become the same. Coupling with corollary 1, we have

Theorem 2.2. If η < η_critical, then for every $x\in {\mathbb{R}}^{{n}_{0}}$ with ||x||₂ ⩽ 1, as n → ∞, f_t (x) converges in distribution to the Gaussian with mean and variance given by equations (14) and (15).

We refer the readers to figure 2 for empirical verification of this theorem. The proof of theorem 2.1 consists of two steps. The first step is to prove the global convergence of overparameterized NNs (Du et al 2019, Allen-Zhu et al 2019, 2018, Zou et al 2019) and stability of the NTK under gradient descent (and gradient flow); see SM section G. This stability was first observed and proved in Jacot et al (2018) in the gradient flow and sequential limit (i.e. letting n₁ → ∞, ..., n_L → ∞ sequentially) setting under certain assumptions about global convergence of gradient flow. In section G, we show how to use the NTK to provide a self-contained (and cleaner) proof of such global convergence and the stability of NTK simultaneously. The second step is to couple the stability of NTK with Grönwall's type arguments (Dragomir 2003) to upper bound the discrepancy between f_t and ${f}_{t}^{\text{lin}}$, i.e. the first norm in equation (17). Intuitively, the ODE of the original network (equation (3)) can be considered as a ${\Vert}{\hat{{\Theta}}}_{t}-{\hat{{\Theta}}}_{0}{{\Vert}}_{F}$-fluctuation from the linearized ODE (equation (7)). One expects the difference between the solutions of these two ODEs to be upper bounded by some functional of ${\Vert}{\hat{{\Theta}}}_{t}-{\hat{{\Theta}}}_{0}{{\Vert}}_{F}$; see section H. Therefore, for a large width network, the training dynamics can be well approximated by linearized dynamics.

Note that the updates for individual weights in equation (6) vanish in the infinite width limit, which for instance can be seen from the explicit width dependence of the gradients in the NTK parameterization. Individual weights move by a vanishingly small amount for wide networks in this regime of dynamics, as do hidden layer activations, but they collectively conspire to provide a finite change in the final output of the network, as is necessary for training. An additional insight gained from linearization of the network is that the individual instance dynamics derived in Jacot et al (2018) can be viewed as a random features method ¹¹ where the features are the gradients of the model with respect to its weights.

Our theoretical analysis thus far has focused on fully-connected single-output architectures trained by full batch gradient descent. In SM section B we derive corresponding results for: networks with multi-dimensional outputs, training against a cross entropy loss, and gradient descent with momentum.

In addition to these generalizations, there is good reason to suspect the results to extend to much broader class of models and optimization procedures. In particular, a wealth of recent literature suggests that the mean field theory governing the wide network limit of fully-connected models (Poole et al 2016, Schoenholz et al 2017) extends naturally to residual networks (Yang and Schoenholz 2017), CNNs (Xiao et al 2018), RNNs (Chen et al 2018), batch normalization (Yang et al 2019), and to broad architectures (Yang 2019). We postpone the development of these additional theoretical extensions in favor of an empirical investigation of linearization for a variety of architectures.