Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent

1 Introduction

Machine learning models based on deep neural networks have achieved unprecedented performance across a wide range of tasks Krizhevsky et al. (2012); He et al. (2016); Devlin et al. (2018). Typically, these models are regarded as complex systems for which many types of theoretical analyses are intractable. Moreover, characterizing the gradient-based training dynamics of these models is challenging owing to the typically high-dimensional non-convex loss surfaces governing the optimization. As is common in the physical sciences, investigating the extreme limits of such systems can often shed light on these hard problems. For neural networks, one such limit is that of infinite width, which refers either to the number of hidden units in a fully-connected layer or to the number of channels in a convolutional layer. Under this limit, the output of the network at initialization is a draw from a Gaussian process (GP); moreover, the network output remains governed by a GP after exact Bayesian training using squared loss (Neal, 1994; Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019; Garriga-Alonso et al., 2018). Aside from its theoretical simplicity, the infinite-width limit is also of practical interest as wider networks have been found to generalize better Neyshabur et al. (2015); Novak et al. (2018); Lee et al. (2018); Novak et al. (2019); Neyshabur et al. (2019).

In this work, we explore the learning dynamics of wide neural networks under gradient descent and find that the weight-space description of the dynamics becomes surprisingly simple: as the width becomes large, the neural network can be effectively replaced by its first-order Taylor expansion with respect to its parameters at initialization. For this induced linear model, the dynamics of gradient descent become analytically tractable. While the linearization is only exact in the infinite width limit, we nevertheless find excellent agreement between the predictions of the original network and those of the linearized version even for finite width configurations. The agreement persists across different architectures, optimization methods, and loss functions.

For squared loss, the exact learning dynamics admit a closed-form solution that allows us to characterize the evolution of the predictive distribution in terms of a GP. This result can be thought of as extension of “sample-then-optimize” posterior sampling Matthews et al. (2017) to the training of deep neural networks. Our empirical simulations confirm that the result accurately models the variation in predictions across an ensemble of finite-width models with different random initializations.

1.1 Our Contribution

We begin by building on a recent result by Jacot et al. (2018) that characterizes the exact dynamics of network outputs throughout gradient descent training in the infinite width limit. Their results establish that gradient descent in parameter space corresponds to kernel gradient descent in function space with respect to a new kernel, the Neural Tangent Kernel (NTK). One may ask what this tells us about the nature of the dynamics in parameter space, where training updates are actually made. A key contribution of our work is to show that dynamics in parameter space are equivalent to the training dynamics of a model which is affine in the collection of all network parameters, the weights and biases. This result holds regardless of the choice of loss function. In the case of squared loss, the dynamics admit a closed-form solution as a function of time.

The output of an infinitely wide neural network is Gaussian at initialization (Lee et al., 2018; Matthews et al., 2018), and as mentioned in Jacot et al. (2018), for squared loss it remains Gaussian throughout training¹¹1The setting is full-batch training under gradient flow.. We derive explicit time-dependent expressions for the mean and covariance functions of this GP, and provide a novel interpretation of the result. In particular, it offers a quantitative understanding of the mechanism by which gradient descent differs from Bayesian posterior sampling of the parameters: while both methods generate draws from a GP, gradient descent does not generate samples from the posterior of any probabilistic model. This observation is in contrast to the “sample-then-optimize” framework of Matthews et al. (2017) in which only the top-layer weights are trained and gradient descent does

sample from the Bayesian posterior. These observations establish a framework with which to analyze the long-standing questions as to if, how, and in what contexts gradient descent provides concrete benefits relative to Bayesian inference.

As argued by Chizat & Bach (2018b), these theoretical results may appear too simple to be applicable to realistic neural networks. 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 & Komodakis (2016).

1.2 Additional related work

Daniely et al. (2016) study the relationship between neural networks and kernels at initialization. They bound the difference between the infinite width kernel and the empirical kernel at finite width $n$ , which diminishes as $O (1 / \sqrt{n})$ . Daniely (2017)

uses the same kernel perspective to study stochastic gradient descent(SGD) training of neural networks.

Saxe et al. (2013) 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 their parameters. In contrast, we show that the outputs of sufficiently wide neural networks are linear in their parameters but not usually their inputs.

Du et al. (2018); Allen-Zhu et al. (2018a, b); Zou et al. (2018) 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 neural networks behave linearly in terms of their parameters throughout training.

Mei et al. (2018); Chizat & Bach (2018a); Rotskoff & Vanden-Eijnden (2018); Sirignano & Spiliopoulos (2018) analyze the mean field SGD dynamics of training neural networks 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 ( $1 / n$ ) different from ours ( $1 / \sqrt{n}$ ), which is commonly used in modern networks He et al. (2016); Glorot & Bengio (2010).

Finally, Zhang et al. (2019) observed that some of the layers in trained neural networks are robust to re-initialization, but not to re-randomization. Our framework provides some theoretical support for this empirical finding.

2 Theoretical Results

2.1 Notation and setup

Let $D \subseteq R^{n_{0}} \times R^{k}$ denote the training set and $X = {x : (x, y) \in D}$ and $Y = {y : (x, y) \in D}$ denote the inputs and labels, respectively. Consider a fully-connected feed-forward network with $L$ hidden layers with width $n_{l}$ , for $l = 1, . . ., L$ and a readout layer with $n_{L + 1} = k$ . For each $x \in R^{n_{0}}$ , we use $h^{l} (x), x^{l} (x) \in 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{matrix} h^{l + 1} & = x^{l} W^{l + 1} + b^{l + 1} x^{l + 1} & = ϕ (h^{l + 1}) \end{matrix} and {\begin{matrix} W_{i, j}^{l} & = \frac{σ_{ω}}{\sqrt{n_{l}}} ω_{i j}^{l} b_{j}^{l} & = σ_{b} β_{j}^{l} \end{matrix},

(1)

where $ϕ$ is a point-wise activation function, $W^{l + 1} \in R^{n_{l} \times n_{l + 1}}$ and $b^{l + 1} \in R^{n_{l + 1}}$ are the weights and biases, $ω_{i j}^{l}$ and $b_{j}^{l}$ are the trainable variables, drawn i.i.d. from standard Gaussian $ω_{i j}^{l}, β_{j}^{l} \sim N (0, 1)$ at initialization, and $σ_{ω}^{2}$ and $σ_{b}^{2}$

are weight and bias variances. Note that this parametrization method is non-standard (see SM

E for further analysis), 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. (2018); Park et al. (2018)

. 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. We compare the test performance of these two parameterization methods in Figure

S1 of SM. Our results (linearity in weights, GP predictions) also hold for infinitely wide networks with a standard parameterization (see SM Section E Figure S2).

We set $θ^{l} \equiv [w_{i, j}^{l}, β_{j}^{l}]_{i, j}$ , the collection of parameters mapping to the $l$ -th layer, whose cardinality is $(n_{l - 1} + 1) \times n_{l}$ . Define $θ = vec (\cup_{l = 1}^{L + 1} θ^{l})$ and similarly $θ^{\leq l}$ or $θ^{> l}$ . Denote by $θ_{t}$ the time-dependence of the parameters and by $θ_{0}$ their initial values. We use $f_{t} (x) \equiv h^{L + 1} (x) \in R^{k}$

to denote the output (or logits) of the neural network at time

t

. Let

ℓ (^y, y) : R^{k} \times R^{k} \to 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²²2To simplify the notation for later equations, we use the total loss here instead of the average loss, but for all plots in Section 3, we show the average loss.,

L = \sum (x, y) \in D ℓ (f_{t} (x, θ), y) .

(2)

Let $η$ be the learning rate³³3Note that compared to the conventional parameterization, $η$ is larger by factor of width. The NTK parameterization allows usage of a universal learning rate scale irrespective of network width.. Via continuous time gradient descent, the evolution of the parameters $θ$ and the logits $f$ can be written as

	${˙ θ}_{t}$	$= - η {\nabla_{θ} f_{t} (X)}^{T} \nabla_{f_{t} (X)} L$		(3)
	${˙ f}_{t} (X)$	$= \nabla_{θ} f_{t} (X) {˙ θ}_{t} = - η {^Θ}_{t} (X, X) \nabla_{f_{t} (X)} L$		(4)

where $f_{t} (X) = vec ({[f_{t} (x)]}_{x \in X})$ , the $k | D | \times 1$ vector of concatenated logits for all examples, and $\nabla_{f_{t} (X)} L$ is the gradient of the loss with respect to the model’s output, $f_{t} (X)$ . ${^Θ}_{t} \equiv {^Θ}_{t} (X, X)$ is the tangent kernel at time $t$ , which is a $k | D | \times k | D |$ matrix

{^Θ}_{t}

= \nabla_{θ} f_{t} (X) {\nabla_{θ} f_{t} (X)}^{T} = L + 1 \sum l = 1 \nabla_{θ^{l}} f_{t} (X) {\nabla_{θ^{l}} f_{t} (X)}^{T} .

(5)

One can define the tangent kernel for general arguments, e.g. ${^Θ}_{t} (x, X)$ where $x$ is test input. We refer to $^Θ$ as the empirical tangent kernel.

2.2 Linearized networks

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,

f_{t}^{lin} (x) \equiv f_{0} (x) + \nabla_{θ} f_{0} (x) ω_{t},

(6)

where $ω_{t} \equiv θ_{t} - θ_{0}$ is the change in the parameters from their initial values⁴⁴4Since $f_{0}^{lin} (x) = f_{0} (x)$ , we will often drop the superscript at time 0.. Note that $f_{t}^{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,

	${˙ ω}_{t}$	$= - η {\nabla_{θ} f_{0} (X)}^{T} \nabla_{f_{t}^{% lin} (X)} L$		(7)
	${˙ f}_{t}^{lin} (x)$	$= - η {^Θ}_{0} (x, X) \nabla_{f_{t}^{lin} (X)} L .$		(8)

As $\nabla_{θ} f_{0} (x)$ remains constant throughout training, these dynamics are often quite simple. In the case of an MSE loss, i.e., $ℓ (^y, y) = \frac{1}{2} ∥^y - y ∥_{2}^{2}$ , the ODEs have closed form solutions

	$ω_{t}$	$= - {\nabla_{θ} f_{0} (X)}^{T} {^Θ}_{0}^{- 1} (I - e^{- η {^Θ}_{0} t}) (f_{0} (X) - Y),$		(9)
	$f_{t}^{lin} (X)$	$= (I - e^{- η {^Θ}_{0} t}) Y + e^{- η {^Θ}_{0} t} f_{0} (X) .$		(10)

For an arbitrary point $x$ , $f_{t}^{lin} (x) = μ_{t} (x) + γ_{t} (x)$ , where

	$μ_{t} (x)$	$= {^Θ}_{0} (x, X) {^Θ}_{0}^{- 1} (I - e^{- η {^Θ}_{0} t}) Y$		(11)
	$γ_{t} (x)$	$= f_{0} (x) - {^Θ}_{0} (x, X) {^Θ}_{0}^{- 1} (I - e^{- η {^Θ}_{0} t}) f_{0} (X) .$		(12)

Therefore, we can obtain the time evolution of the linearized neural network without training it. We only need to compute the tangent kernel ${^Θ}_{0}$ and the outputs $f_{0}$ at initialization and use Equations 11, 12, and 9 to compute the dynamics of the outputs and the weights.

2.3 Infinite width limit yields Gaussian processes

As the width of the hidden layers approaches infinity, the Central Limit Theorem (CLT) implies that the outputs at initialization

{f_{0} (x)}_{x \in X}

converge to a multivariate Gaussian in distribution. Informally, this can be proved by induction. Conditioning on activations at layer

l

, each pre-activation (

h_{j}^{l + 1} (x)

) at layer

l + 1

is a sum of

n^{l}

i.i.d. properly normalized random variables (the weights) and a Gaussian distribution (the bias term). One can apply the CLT to conclude that

{h_{j}^{l + 1} (x)}_{j}

are i.i.d. Gaussian; see Poole et al. (2016); Schoenholz et al. (2017); Lee et al. (2018); Xiao et al. (2018); Yang & Schoenholz (2017) for more details, and Matthews et al. (2018); Novak et al. (2019) for a formal treatment.

Therefore, randomly initialized neural networks are in correspondence with a certain class of Gaussian processes (hereinafter referred to as NNGPs), which facilitates a fully Bayesian treatment of neural networks (Lee et al., 2018; Matthews et al., 2018). More precisely, let $f_{t}^{i}$ denote the $i$ -th output dimension and $K$ denote the the sample-to-sample kernel function (of the pre-activation) of the outputs in the infinite width setting,

K^{i, j} (x, x^{'}) = lim min (n_{1 \dots L}) \to \infty E [f_{0}^{i} (x) \cdot f_{0}^{j} (x^{'})],

(13)

then $f_{0} (X) \sim G P (0, K (X, X))$ , where $K^{i, j} (x, x^{'})$ denotes the covariance between the $i$ -th output of $x$ and $j$ -th output of $x^{'}$ , which can be computed recursively (see Lee et al. (2018) Section 2.3 and Supplementary Material (SM) Section A). For an unobserved test input $x$ , the joint output distribution $f ([x, X])$ is also a Gaussian process. Conditioning on the training samples⁵⁵5 This imposes that $h^{L + 1}$ directly corresponds to the network predictions. In the case of softmax readout, variational or sampling methods are required to marginalize over $h^{L + 1}$ . $f (X) = Y$

, the posterior predictive distribution of

f (x)

is also Gaussian

N (μ (x), Σ (x))

	$μ (x)$	$= K (x, X) K^{- 1} Y$
	$Σ (x)$	$= K (x, x) - K (x, X) K^{- 1} K (x, X)^{T},$		(14)

where $K = K (X, X)$ . This is the posterior predictive distribution resulting from exact Bayesian inference in an infinitely wide neural network.

2.3.1 Gaussian processes from gradient descent training

If we freeze the variables $θ^{\leq 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 ${^Θ}_{0}$ will converge to $K$

in probability and Equation

11 will converge to the posterior Equation 2.3 as

t \to \infty

(for further details see SM Section D). This is a realization of the “sample-then-optimize” approach for evaluating the posterior of a Gaussian process proposed in Matthews et al. (2017).

If none of the variables are frozen, in the infinite width setting, ${^Θ}_{0}$ also converges in probability to a deterministic kernel $Θ$ Jacot et al. (2018), which we sometimes refer to as the analytic kernel, and which can also be computed recursively (see SM Section A). Letting the width go to infinity, for any $t$ , the output $f_{t}^{lin} (x)$ is also Gaussian distributed because Equations 11 and 12 describe an affine transform of the Gaussian $[f_{0} (x), f_{0} (X)]$ . The mean( $μ (x)$ ) and the variance( $Σ (x)$ ) are given by

	$μ (x) = Θ (x, X) Θ^{- 1} (I - e^{- η Θ t}) Y$		(15)
	$Σ (x) = K (x, x) - 2 Θ (x, X) Θ^{- 1} (I - e^{- η Θ t}) K (x, X)^{T}$
	$+ Θ (x, X) Θ^{- 1} (I - e^{- η Θ t}) K Θ^{- 1} (I - e^{- η Θ t}) Θ (x, X)^{T} .$		(16)

Unlike the case when only $θ^{L + 1}$ is optimized, Equations 15 and 16 do not admit an interpretation corresponding to the posterior sampling of a probabilistic model⁶⁶6One possible exception is when the NNGP kernel and NTK are the same up to a scalar multiplication.. We contrast the predictive distributions from the NNGP, NTK-GP (i.e. Equations 15 and 16) and ensembles of NNs in Figure 3.

2.4 Infinite width networks are linearized networks

The ODEs (Equation 3, 4) of the original network are unsolvable in general, since ${^Θ}_{t}$ evolves with time. Remarkably, Jacot et al. (2018) showed that

supt∈[0,T]∥^Θt−^Θ0∥F=O(min{n1,…,nL}−1/2)

(17)

under the technical assumption that the integrated functional gradient norm remains stochastically bounded as $n_{1}, \dots, n_{L} \to \infty$ sequentially, i.e.

\int_{0}^{T} {∥ ∥ \nabla_{f_{t} (X)} L ∥ ∥}_{2} d t < \infty,

(18)

where $T$ is the training budget (independent from the widths), i.e. the amount of allowable training time. This assumption was verified in Jacot et al. (2018) in some specific cases. Recent work on the convergence theory of over-parameterized neural networks provides a rigorous proof of Equation 18 in the simultaneous limit (i.e. $min {n_{1}, \dots, n_{L}} \to \infty)$ under discrete-time stochastic gradient descent with MSE loss on various architectures (fully-connected networks, residual CNNs, RNNs) (Du et al., 2018; Allen-Zhu et al., 2018b, a; Zou et al., 2018). Note that the bound in Equation 17 is just a theoretical upper bound and we observe empirically that the exponent $- \frac{1}{2}$ may be improved to $- 1$ ; see Figure 2.

Coupling Equation 18 with Grönwall’s type arguments, in the MSE setting, we can upper bound the discrepancy between the outputs of the original network and those of its linearization,

sup t \in [0, T] ∥ f_{t} (X) - f_{t}^{lin} (X) ∥_{2} = O (sup t \in [0, T] ∥ {^Θ}_{t} - {^Θ}_{0} ∥_{F}),

(19)

which approaches $0$ as the width goes to infinity. Intuitively, the ODE of the original network (Equation 4) can be considered as a $∥ {^Θ}_{t} - {^Θ}_{0} ∥_{F}$ -fluctuation from the linearized ODE (Equation 8). One expects the difference between the solutions of these two ODEs to be upper bounded by some functional of $∥ {^Θ}_{t} - {^Θ}_{0} ∥_{F}$ (refer to SM Section F for the proof). 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 7 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 dynamics derived in Jacot et al. (2018) are equivalent to a random features method, where the features are the gradients of the model with respect to its weights.

2.5 Extensions

The linearization of the wide neural networks and its training dynamics can be generalized in multiple directions.

2.5.1 Momentum

One direction is to go beyond vanilla gradient descent dynamics. We consider momentum updates⁷⁷7Combining the usual two stage update into a single equation.

θ_{k + 1} = θ_{k} + β (θ_{k} - θ_{k - 1}) - η \nabla_{θ} L |_{θ = θ_{k}} .

(20)

The discrete update to the function output becomes

	$f_{k + 1}^{lin} (x) = f_{k}^{lin} (x)$	$- η {^Θ}_{0} (x, X) \nabla_{f_{k}^{lin} (X)} L$
		$+ β (f_{k}^{lin} (x) - f_{k - 1}^{lin} (x))$		(21)

where $f_{t}^{lin} (x)$ is the output of the linearized network after $t$ steps. One can take the continuous time limit as in Qian (1999); Su et al. (2014) and obtain

	${¨ ω}_{t}$	$= ~ β {˙ ω}_{t} - \nabla_{θ} f_{0}^{lin} (X)^{T} \nabla_{f_{t}^{lin} (X)} L$		(22)
	${¨ f}_{t}^{lin} (x)$	$= ~ β {˙ f}_{t}^{lin} (x) - {^Θ}_{0} (x, X) \nabla_{f_{t}^{lin} (X)} L$		(23)

where continuous time relates to steps $t = k \sqrt{η}$ and $~ β = (β - 1) / \sqrt{η}$ . These equations are also amenable to analytic treatment for MSE loss. See Figure 5, 6 and 7 for experimental agreement.

Figure 1: Kernel convergence. Kernels computed from randomly initialized $ReLU$ networks with one and three hidden layers converge to the corresponding analytic kernel as width $n$ and number of Monte Carlo samples $M$ increases. Colors indicate averages over different numbers of Monte Carlo samples.

Figure 2: Relative Frobenius norm change during training. One hidden layer, $ReLU$ networks trained with $η = 1.0$ , $| D | = 128$ . We measure changes of (read-out/non read-out) weights, empirical $^Θ$ and empirical $^K$ after $T = 65, 536$ steps of gradient descent updates for varying width. We see that the change in weights scales as $\frac{1}{\sqrt{n}}$ whereas the change in $^Θ$ and $^K$ is bounded by $O (\frac{1}{\sqrt{n}})$ but close to $O (\frac{1}{n})$ .

2.5.2 Cross-entropy loss

One can extend the loss function to general functions with multiple output dimensions. Unlike for squared error, we do not have a closed form solution to the dynamics equation. However, the equations for the dynamics can be solved using an ODE solver as an initial value problem. For cross-entropy loss with softmax output (see SM Section C),

ℓ (f, y) = - \sum i y_{i} log σ (f^{i}), σ (f^{i}) \equiv \frac{exp (f^{i})}{\sum_{j} exp (f^{j})} .

(24)

the dynamics equation becomes

	${˙ f}_{t}^{lin} (X)$	$= - η {^Θ}_{0} (X, X) (σ (f_{t}^{lin} (X)) - Y)$		(25)
	${˙ f}_{t}^{lin} (x)$	$= - η {^Θ}_{0} (x, X) (σ (f_{t}^{% lin} (X)) - Y) .$		(26)

2.5.3 Beyond fully-connected networks

Although our theoretical analysis has focused on fully-connected architectures, there is good reason to suspect the results to extend to much broader class of models. 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 & Schoenholz (2017), CNNs Xiao et al. (2018), RNNs Chen et al. (2018)

, batch normalization

Yang et al. (2019), and general architectures Yang (2019). We postpone the development of these theoretical extensions in favor of a purely empirical investigation of linearization for a variety of architectures (see Section 3.3).

Figure 3: Dynamics of mean and variance of trained neural network outputs follow analytic dynamics from linearization. Black lines indicate the time evolution of the predictive output distribution from an ensemble of 100 trained neural networks (NNs). The blue region indicates the analytic prediction of the output distribution throughout training (Equations 15, 16). Finally, the red region indicates the prediction that would result from training only the top layer, corresponding to an NNGP (Equations S30, S31). The trained network has 3 hidden layers of width 8192, $tanh$ activation functions, $σ_{w}^{2} = 1.5$ , no bias, and $η = 0.5$
. The output is computed for inputs interpolated between two training points (denoted with black dots)
$x (α) = α x^{(1)} + (1 - α) x^{(2)}$
. The shaded region and dotted lines denote 2 standard deviations (
$\sim 95 %$ quantile) from the mean denoted in solid lines. Training was performed with full-batch gradient descent with dataset size $| D | = 128$ . For dynamics for individual draw of functions, see SM Figure S4

3 Experiments

In this section, we provide empirical support showing that the training dynamics of wide neural networks are well captured by linearized models. We consider fully-connected, convolutional, and wide ResNet architectures trained with full and mini- batch gradient descent using learning rates sufficiently small so that the continuous time approximation holds well. We consider two-class classification on CIFAR-10 (horses and planes) as well as ten-class classification on MNIST and CIFAR-10. When using MSE loss, we treat the binary classification task as regression with one class regressing to $+ 1$ and the other to $- 1$ .

3.1 Convergence of empirical kernel

As in Novak et al. (2019)

, we can use Monte Carlo estimates of the tangent kernel (Equation

5) to probe convergence to the infinite width kernel (analytically computed using Equations S3, S6). For simplicity, we consider random inputs drawn from

N (0, 1)

with

n_{0} = 1024

. In Figure 1, we observe convergence as both width

n

increases and the number of Monte Carlo samples

M

increases. For both NNGP and tangent kernels we observe

∥ {^Θ}^{(n)} - Θ ∥_{F} = O (\frac{1}{\sqrt{n}})

and

∥ {^K}^{(n)} - K ∥_{F} = O (\frac{1}{\sqrt{n}})

as predicted by a CLT Daniely et al. (2016).

Moreover, as the neural network trains the change during training in the NNGP and tangent kernels, and in individual weights, becomes small as width increases, as shown in Figure 2.

Figure 4: Full batch gradient descent on a model behaves similarly to analytic dynamics on its linearization, both for network outputs, and also for individual weights. A binary CIFAR classification task with MSE loss and a $ReLU$ fully connected network with 5 hidden layers of width $n = 2048$ , $η = 0.01$ , $| D | = 256$ , $k = 1$ , $σ_{w}^{2} = 2.0$ , and $σ_{b}^{2} = 0.1$ . All three panes in the first row show dynamics for a randomly selected subset of datapoints or parameters. The first two panes in the second row show that the dynamics of loss and accuracy for training and test points agree well between the original and linearized model. The bottom right pane shows the dynamics of RMSE between the two models on test points. We observe that the empirical kernel $^Θ$ gives more accurate dynamics for finite width networks.

3.2 Predictive output distribution

In the case of an MSE loss, the output distribution remains Gaussian throughout training. In Figure 3, the predictive output distribution for input points interpolated between two training points is shown for an ensemble of neural networks and their corresponding GPs. The interpolation is given by $x (α) = α x^{(1)} + (1 - α) x^{(2)}$ where $x^{(1, 2)}$ are two training inputs with different classes. We observe that the mean and variance dynamics of neural network outputs during gradient descent training follow the analytic dynamics from linearization well (Equations 15, 16). Moreover the NNGP posterior which corresponds to exact Bayesian inference, while similar, is noticeably different from the predictive distribution at the end of gradient descent training. For dynamics for individual function draws, see SM Figure S4.

Figure 5: A convolutional network and its linearization behave similarly when trained using full batch gradient descent with a momentum optimizer Binary CIFAR classification task with MSE loss, $tanh$ convolutional network with 3 hidden layers of channel size $n = 512$ , $3 \times 3$ size filters, average pooling after last convolutional layer, $η = 0.1$ , $β = 0.9$ , $| D | = 128$ , $σ_{w}^{2} = 2.0$ and $σ_{b}^{2} = 0.1$ . The linearized model is trained directly by full batch gradient descent with momentum, rather than by integrating its continuous time analytic dynamics. Panes are the same as in Figure 4.

3.3 Comparison of training dynamics of linearized network to original network

For a particular realization of a finite width network, one can analytically predict the dynamics of the weights and outputs over the course of training using the empirical tangent kernel at initialization. In Figures 5,6,7, we compare these linearized dynamics (Equations 9, 10) with the result of training the actual network. In all cases we see remarkably good agreement. We also observe that for finite networks, dynamics predicted using the empirical kernel $^Θ$ better match the data than those obtained using the infinite-width, analytic, kernel $Θ$ . To understand this we note that $∥ {^Θ}_{T}^{(n)} - {^Θ}_{0}^{(n)} ∥_{F} = O (\frac{1}{n}) < O (\frac{1}{\sqrt{n}}) = ∥ {^Θ}_{0}^{(n)} - Θ ∥_{F}$ , as plotted in Figure 2.

For general loss, e.g. cross-entropy with softmax output, we need to rely on solving the ODE Equations 25 and 26. We use the dopri5

method for ODE integration, which is the default integrator in TensorFlow (

tf.contrib.integrate.odeint). In Figure 4, we see that the learning dynamics for the CIFAR-10 all class classification task with cross-entropy loss are well described by the linearized model. In Figure 6, we tested full MNIST digit classification with cross-entropy loss, and trained with a momentum optimizer. For cross-entropy loss with softmax output, some logits at late times grow indefinitely, in contrast to MSE loss where logits converge to target value. The error between original and linearized model for cross entropy loss becomes much worse at late times if the two models deviate significantly before the logits enter their late-time steady-growth regime (See Figure S5.).

One can directly optimize parameters of $f^{lin}$ instead of solving the ODE induced by the tangent kernel $^Θ$ . Standard neural network optimization techniques such as mini-batching, weight decay, and data augmentation can be directly applied. In Figure 5 and 7, we compared the training dynamics of the linearized and original network while directly training both networks.

As discussed in Section 2.5.3, the linearized dynamics successfully describes the training of networks beyond vanilla fully connected models. To demonstrate the generality of this procedure we show we can predict the learning dynamics of Wide Residual Networks (WRN) Zagoruyko & Komodakis (2016)

. WRNs are a class of model that are popular in computer vision and leverage convolutions, batch normalization, skip connections, and average pooling. In Figure

7, we show a comparison between the linearized dynamics and the true dynamics for a wide residual network trained with MSE loss and SGD with momentum. We slightly modified the block structure described in Figure 7 so that each layer has a constant number of channels (1024 in this case) and otherwise followed the original implementation. As elsewhere, we see strong agreement between the predicted dynamics and the result of training.

Figure 6: A neural network and its linearization behave similarly when both are trained via SGD with momentum on cross entropy loss on MNIST. Experiment is for 10 class MNIST classification using a $ReLU$ fully connected network with 2 hidden layers of width $n = 2048$ , $η = 1.0$ , $β = 0.9$ , $| D | = 50, 000$ , $k = 10$ , $σ_{w}^{2} = 2.0$ , and $σ_{b}^{2} = 0.1$ . Both models are trained using stochastic minibatching with batch size 64. Panes are the same as in Figure 4, except that the top row shows all ten logits for a single randomly selected datapoint.

Figure 7: A wide residual network and its linearization behave similarly when both are trained by SGD with momentum on MSE loss on CIFAR-10. (top) We adopt the network architecture from Zagoruyko & Komodakis (2016). In the residual block, we follow Batch Normalization-ReLU-Conv ordering. We use $N = 1$ , channel size $1024$ , $η = 1.0$ , $β = 0.9$ , $k = 10$ , $σ_{w}^{2} = 1.0$ , and $σ_{b}^{2} = 0.0$ . Both linearized and original model are trained directly on full CIFAR-10 ( $| D | = 50, 000$ ), using stochastic minibatching with batch size 8. (bottom) Output dynamics for a randomly selected subset of train and test points are shown in the first row. The second row shows training and accuracy curves for original and linearized networks.

group name	output size	block type
conv1	32 $\times$ 32	[3 $\times$ 3, channel size]
conv2	32 $\times$ 32	$[\begin{matrix} 3 \times 3, & channel size 3 \times 3, & channel size \end{matrix}]$ $\times$ N
conv3	16 $\times$ 16	$[\begin{matrix} 3 \times 3, & channel size 3 \times 3, & channel size \end{matrix}]$ $\times$ N
conv4	8 $\times$ 8	$[\begin{matrix} 3 \times 3, & channel size 3 \times 3, & channel size \end{matrix}]$ $\times$ N
avg-pool	1 $\times$ 1	[8 $\times$ 8]

3.4 Effects of depth and dataset size

The training dynamics of a neural network match those of its linearization when the width is infinite and the dataset is finite. In previous experiments, we chose sufficiently wide networks to achieve small error between neural networks and their linearization for smaller datasets. Here we investigate how the agreement between the linearized dynamics and the true dynamics behaves as a function of width and dataset size across a wide range of models. We consider Root Mean Squared Error (RMSE) between the predicted outputs and the true outputs of the network over the test set. Generally, the RMSE will increase with time until it plateaus at the end of training at some final value. In Figure 8 we plot this plateau RMSE for a range of models as a function of both width and dataset size. Overall, we observe that as the width grows the error decreases. This decrease goes approximately as $1 / N$ for fully-connected networks, with more ambiguous scaling for convolutional and WRN architectures. Additionally, we see that the error grows approximately linearly in the size of the dataset. Thus, although error grows with dataset this can be counterbalanced by a corresponding increase in the model size.

Figure 8: Error dependence on depth and dataset size. Final value of the RMSE for fully-conected, convolutional, wide residual network as networks becomes wider for varying depth and dataset size. Top row: error in fully connected networks as the depth is varied from 1 to 16 (left) and the dataset size is varied from 32 to 4096 (right). Bottom row: error in convolutional networks (left) as the depth is varied between 1 and 32 and WRN (right) for depths 10 and 16 corresponding to N=1,2 described in Figure 7. Network is critically initialized $σ_{w}^{2} = 2.0$ , $σ_{b}^{2} = 0.1$ , trained with gradient descent on MSE loss.

4 Discussion

We showed theoretically that the learning dynamics in parameter space of deep nonlinear neural networks are exactly described by a linearized model in the infinite width limit. Empirical investigation revealed that this agrees very well with actual training dynamics and predictive distributions across fully-connected, convolutional, and even wide residual network architectures, as well as with different optimizers (gradient descent, momentum, mini-batching) and loss functions (MSE, cross-entropy). Our results suggest that a surprising number of realistic neural networks may be operating in the regime we studied.

In the regime which we study, since the learning dynamics is fully captured by the kernel $^Θ$ and the target signal, studying the properties of $^Θ$ to determine trainability and generalization are interesting future directions. Furthermore, the infinite width limit gives us a simple characterization of both gradient descent and Bayesian inference. Some preliminary observations in Lee et al. (2018) showed that wide neural networks trained with SGD perform similarly to the corresponding GPs as width increase, while Novak et al. (2019) found the opposite in the case of convopostlutional networks without pooling. By studying properties of the NNGP kernel $K$ and the tangent kernel $Θ$ , we may shed light on the inductive bias of gradient descent.

Acknowledgements

We thank Roman Novak and Greg Yang for useful discussions and feedback.

References

Allen-Zhu et al. (2018a) Allen-Zhu, Z., Li, Y., and Song, Z. A convergence theory for deep learning via over-parameterization. arXiv preprint arXiv:1811.03962, 2018a.
Allen-Zhu et al. (2018b) Allen-Zhu, Z., Li, Y., and Song, Z. On the convergence rate of training recurrent neural networks. arXiv preprint arXiv:1810.12065, 2018b.
Chen et al. (2018) Chen, M., Pennington, J., and Schoenholz, S.
Dynamical isometry and a mean field theory of RNNs: Gating enables signal propagation in recurrent neural networks.
In International Conference on Machine Learning, 2018.
Chizat & Bach (2018a) Chizat, L. and Bach, F. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in neural information processing systems, pp. 3040–3050, 2018a.
Chizat & Bach (2018b) Chizat, L. and Bach, F. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv:1812.07956, 2018b.
Cho & Saul (2009) Cho, Y. and Saul, L. K. Kernel methods for deep learning. In Advances in neural information processing systems, 2009.
Daniely (2017) Daniely, A. SGD learns the conjugate kernel class of the network. In Advances in Neural Information Processing Systems 30. 2017.
Daniely et al. (2016) Daniely, A., Frostig, R., and Singer, Y. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. In Advances In Neural Information Processing Systems, 2016.
Devlin et al. (2018) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
Dragomir (2003) Dragomir, S. S. Some Gronwall type inequalities and applications. Nova Science Publishers New York, 2003.
Du et al. (2018) Du, S. S., Lee, J. D., Li, H., Wang, L., and Zhai, X. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv:1811.03804, 2018.
Garriga-Alonso et al. (2018) Garriga-Alonso, A., Aitchison, L., and Rasmussen, C. E. Deep convolutional networks as shallow gaussian processes. arXiv preprint arXiv:1808.05587, 2018.
Glorot & Bengio (2010) Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In International Conference on Artificial Intelligence and Statistics
, pp. 249–256, 2010.
He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition
, pp. 770–778, 2016.
Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 31. 2018.
Karras et al. (2018) Karras, T., Aila, T., Laine, S., and Lehtinen, J. Progressive growing of GANs for improved quality, stability, and variation. In International Conference on Learning Representations, 2018.
Krizhevsky et al. (2012) Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems. 2012.
Lee et al. (2018) Lee, J., Bahri, Y., Novak, R., Schoenholz, S., Pennington, J., and Sohl-dickstein, J. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018.
Matthews et al. (2017) Matthews, A. G. d. G., Hron, J., Turner, R. E., and Ghahramani, Z. Sample-then-optimize posterior sampling for bayesian linear models. In NeurIPS Workshop on Advances in Approximate Bayesian Inference, 2017. URL http://approximateinference.org/2017/accepted/MatthewsEtAl2017.pdf.
Matthews et al. (2018) Matthews, A. G. d. G., Hron, J., Rowland, M., Turner, R. E., and Ghahramani, Z. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
Mei et al. (2018) Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
Neal (1994) Neal, R. M. Priors for infinite networks (tech. rep. no. crg-tr-94-1). University of Toronto, 1994.
Neyshabur et al. (2015) Neyshabur, B., Tomioka, R., and Srebro, N. In search of the real inductive bias: On the role of implicit regularization in deep learning. In International Conference on Learning Representations workshop track, 2015.
Neyshabur et al. (2019) Neyshabur, B., Li, Z., Bhojanapalli, S., LeCun, Y., and Srebro, N. The role of over-parametrization in generalization of neural networks. In International Conference on Learning Representations, 2019.
Novak et al. (2018) Novak, R., Bahri, Y., Abolafia, D. A., Pennington, J., and Sohl-Dickstein, J. Sensitivity and generalization in neural networks: an empirical study. In International Conference on Learning Representations, 2018.
Novak et al. (2019) Novak, R., Xiao, L., Lee, J., Bahri, Y., Yang, G., Hron, J., Abolafia, D. A., Pennington, J., and Sohl-Dickstein, J. Bayesian deep convolutional networks with many channels are gaussian processes. In International Conference on Learning Representations, 2019.
Park et al. (2018) Park, D. S., Smith, S. L., Sohl-dickstein, J., and Le, Q. V.
Optimal SGD hyperparameters for fully connected networks.
In NeurIPS Bayesian Deep Learning Workshop, 2018. URL http://bayesiandeeplearning.org/2018/papers/86.pdf.
Poole et al. (2016) Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., and Ganguli, S. Exponential expressivity in deep neural networks through transient chaos. In Advances In Neural Information Processing Systems, pp. 3360–3368, 2016.
Qian (1999) Qian, N. On the momentum term in gradient descent learning algorithms. Neural networks, 12(1):145–151, 1999.
Rotskoff & Vanden-Eijnden (2018) Rotskoff, G. M. and Vanden-Eijnden, E. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. arXiv preprint arXiv:1805.00915, 2018.
Saxe et al. (2013) Saxe, A. M., McClelland, J. L., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
Schoenholz et al. (2017) Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J. Deep information propagation. International Conference on Learning Representations, 2017.
Sirignano & Spiliopoulos (2018) Sirignano, J. and Spiliopoulos, K. Mean field analysis of neural networks. arXiv preprint arXiv:1805.01053, 2018.
Su et al. (2014) Su, W., Boyd, S., and Candes, E. A differential equation for modeling nesterov’s accelerated gradient method: Theory and insights. In Advances in Neural Information Processing Systems, pp. 2510–2518, 2014.
van Laarhoven (2017) van Laarhoven, T. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.05350, 2017.
Williams (1997) Williams, C. K. Computing with infinite networks. In Advances in neural information processing systems, pp. 295–301, 1997.
Xiao et al. (2018) Xiao, L., Bahri, Y., Sohl-Dickstein, J., Schoenholz, S., and Pennington, J. Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, 2018.
Yang (2019) Yang, G. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760, 2019.
Yang & Schoenholz (2017) Yang, G. and Schoenholz, S. Mean field residual networks: On the edge of chaos. In Advances in Neural Information Processing Systems. 2017.
Yang et al. (2019) Yang, G., Pennington, J., Rao, V., Sohl-Dickstein, J., and Schoenholz, S. S. A mean field theory of batch normalization. In International Conference on Learning Representations, 2019.
Zagoruyko & Komodakis (2016) Zagoruyko, S. and Komodakis, N. Wide residual networks. In British Machine Vision Conference, 2016.
Zhang et al. (2019) Zhang, C., Bengio, S., and Singer, Y. Are all layers created equal? arXiv preprint arXiv:1902.01996, 2019.
Zou et al. (2018) Zou, D., Cao, Y., Zhou, D., and Gu, Q. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint arXiv:1811.08888, 2018.

Appendix A Computing NTK and NNGP Kernel

For completeness, we reproduce, informally, the recursive formula of the NNGP kernel and the tangent kernel from Lee et al. (2018) and Jacot et al. (2018), respectively. Let the activation function $ϕ : R \to R$ be absolutely continuous. Let $T$ and $˙ T$ be functions from $2 \times 2$ positive semi-definite matrices $Σ$ to $R$ given by

(S1)

In the infinite width limit, the NNGP and tangent kernel can be computed recursively. Let $x, x^{'}$ be two inputs in $R^{n_{0}}$ . Then $h^{l} (x)$ and $h^{l} (x^{'})$ converge in distribution to a joint Gaussian as $min {n_{1}, \dots, n_{l - 1}}$ . The mean is zero and the variance $K^{l} (x, x^{'})$ is

K^{l} (x, x^{'}) = {~ K}^{l} (x, x^{'}) \otimes {I d}_{n_{l}}

(S2)

(S3)

with base case

K^{1} (x, x^{'}) = σ_{ω}^{2} \cdot \frac{1}{n_{0}} x^{T} x^{'} + σ_{b}^{2} .

(S4)

Using this one can also derive the tangent kernel for gradient descent training. We will use induction to show that

Θ^{l} (x, x^{'}) = {~ Θ}^{l} (x, x^{'}) \otimes {I d}_{n_{l}}

(S5)

where

{~ Θ}^{l} (x, x^{'}) = {~ K}^{l} (x, x^{'}) + σ_{ω}^{2} {~ Θ}^{l - 1} (x, x^{'}) ˙ T ([\begin{matrix} {~ K}^{l - 1} (x, x) & {~ K}^{l - 1} (x, x^{'}) {~ K}^{l - 1} (x, x^{'}) & {~ K}^{l - 1} (x^{'}, x^{'}) \end{matrix}])

(S6)

with ${~ Θ}^{1} = {~ K}^{1}$ . Let

J^{l} (x) = \nabla_{θ^{\leq l}} h_{0}^{l} (x) = [\nabla_{θ^{l}} h_{0}^{l} (x), \nabla_{θ^{< l}} h_{0}^{l} (x)] .

(S7)

Then

J^{l} (x) J^{l} (x^{'})^{T}

= \nabla_{θ^{l}} h_{0}^{l} (x) \nabla_{θ^{l}} h_{0}^{l} (x^{'})^{T} + \nabla_{θ^{< l}} h_{0}^{l} (x) \nabla_{θ^{< l}} h_{0}^{l} (x^{'})^{T}

(S8)

Letting $n_{1}, \dots, n_{l - 1} \to \infty$ sequentially, the first term converges to the NNGP kernel $K^{l} (x, x^{'})$

. By applying the chain rule and the induction step (letting

n_{1}, \dots, n_{l - 2} \to \infty

sequentially), the second term is

$\nabla_{θ^{< l}} h_{0}^{l} (x) \nabla_{θ^{< l}} h_{0}^{l} (x^{'})^{T}$	$= \frac{\partial h_{0}^{l} (x)}{\partial h_{0}^{l - 1} (x)} \nabla_{θ^{\leq l - 1}} h_{0}^{l - 1} (x) \nabla_{θ^{\leq l - 1}} h_{0}^{l - 1} (x^{'})^{T} {\frac{\partial h_{0}^{l} (x^{'})}{\partial h_{0}^{l - 1} (x^{'})}}^{T}$	(S9)
	$\to \frac{\partial h_{0}^{l} (x)}{\partial h_{0}^{l - 1} (x)} {~ Θ}^{l - 1} (x, x^{'}) \otimes {I d}_{n_{l - 1}} {\frac{\partial h_{0}^{l} (x^{'})}{\partial h_{0}^{l - 1} (x^{'})}}^{T} (n_{1}, \dots, n_{l - 2} \to \infty)$	(S10)
	$\to σ_{ω}^{2} (E ϕ^{'} (h_{0, i}^{l - 1} (x)) ϕ^{'} (h_{0, i}^{l - 1} (x^{'})) {~ Θ}^{l - 1} (x, x^{'})) \otimes {I d}_{n_{l}} (n_{l - 1} \to \infty)$	(S11)
	$= (σ_{ω}^{2} {~ Θ}^{l - 1} (x, x^{'}) ˙ T ([\begin{matrix} {~ K}^{l - 1} (x, x) & {~ K}^{l - 1} (x, x^{'}) {~ K}^{l - 1} (x, x^{'}) & {~ K}^{l - 1} (x^{'}, x^{'}) \end{matrix}])) \otimes {I d}_{n_{l}}$	(S12)

Appendix B Tangent kernel for $ReLU$ and $erf$

For $ReLU$ and $erf$ activation functions, the tangent kernel can be computed analytically. We begin with the case $ϕ = {ReLU}$ ; using the formula from Cho & Saul (2009), we can compute $T$ and $˙ T$ in closed form. Let $Σ$ be a $2 \times 2$ PSD matrix. We will use

k_{n} (x, y) = \int ϕ^{n} (x \cdot w) ϕ^{n} (y \cdot w) e^{- ∥ w ∥^{2} / 2} d w \cdot (2 π)^{- d / 2} = \frac{1}{2 π} ∥ x ∥^{n} ∥ y ∥^{n} J_{n} (θ)

(S13)

where

	$ϕ (x)$	$= max (x, 0), θ (x, y) = arccos (\frac{x \cdot y}{∥ x ∥ ∥ y ∥}),$
	$J_{0} (θ)$	$= π - θ, J_{1} (θ) = sin θ + (π - θ) cos θ = \sqrt{1 - (\frac{x \cdot y}{∥ x ∥ ∥ y ∥})^{2}} + (π - θ) (\frac{x \cdot y}{} ∥ x ∥ ∥ y ∥) .$		(S14)

Let $d = 2$ and $u = (x \cdot w, y \cdot w)^{T}$ . Then $u$ is a mean zero Gaussian with $Σ = [[x \cdot x, x \cdot y]; [x \cdot y, y \cdot y]]$ . Then

	$T (Σ)$	$= k_{1} (x, y) = \frac{1}{2 π} ∥ x ∥ ∥ y ∥ J_{1} (θ)$		(S15)
	$˙ T (Σ)$	$= k_{0} (x, y) = \frac{1}{2 π} J_{0} (θ)$		(S16)

For $ϕ = erf$ , let $Σ$ be the same as above. Following Williams (1997), we get

	$T (Σ)$	$= \frac{2}{π} {sin}^{- 1} (\frac{2 x \cdot y}{\sqrt{(1 + 2 x \cdot x) (1 + 2 y \cdot y)}})$		(S17)
	$˙ T (Σ)$	$= \frac{4}{π} d e t (I + 2 Σ)^{- 1 / 2}$		(S18)

Appendix C Multi-dimensional output and cross-entropy loss

For completeness, we include the derivation for cross entropy loss for softmax output,

ℓ (f, y) = - \sum i y^{i} log σ (f^{i}), σ (f^{i}) \equiv \frac{exp (f^{i})}{\sum_{j} exp (f^{j})} .

(S19)

Recall that $\frac{\partial ℓ}{\partial {^y}^{i}} = σ ({^y}^{i}) - y^{i}$ . For general input point $x$ and for an arbitrary parameterized function $f^{i} (x)$ parameterized by $θ$ , gradient flow dynamics is given by

	${˙ f}_{t}^{i} (x) = \nabla_{θ} f_{t}^{i} (x) \frac{d θ}{d t}$	$= - η \nabla_{θ} f_{t}^{i} (x) \sum j \sum (z, y) \in D [\nabla_{θ} f_{t}^{j} (z)^{T} \frac{\partial ℓ (f_{t}, y)}{\partial {^y}^{j}}]$		(S20)
		$= - η \sum (z, y) \in D \sum j \nabla_{θ} f_{t}^{i} (x) \nabla_{θ} f_{t}^{j} (z)^{T} (σ (f_{t}^{j} (z)) - y^{j})$		(S21)

Let $Θ^{i j} (x, X) = \nabla_{θ} f^{i} (x) \nabla_{θ} f^{j} (X)^{T}$ . The above is

	$_{t} (X)$	$= - η Θ_{t} (X, X) (σ (f_{t} (X)) - Y)$		(S22)
	$_{t} (x)$	$= - η Θ_{t} (x, X) (σ (f_{t} (X)) - Y) .$		(S23)

The linearization is

	${_{t}}^{lin} (X)$	$= - η Θ_{0} (X, X) (σ (f_{t}^{% lin} (X)) - Y)$		(S24)
	${_{t}}^{lin} (x)$	$= - η Θ_{0} (x, X) (σ (f_{t}^{lin} (X)) - Y) .$		(S25)

Appendix D Gradient flow dynamics for training only the readout-layer

The connection between Gaussian processes and Bayesian wide neural networks can be extended to the setting when only the readout layer parameters are being optimized. More precisely, we show that when training only the readout layer, the outputs of the network form a Gaussian process (over an ensemble of draws from the parameter prior) throughout training, where that output is an interpolation between the GP prior and GP posterior.

Note that for any $x, x^{'} \in R^{n_{0}}$ , in the infinite width limit $¯ x (x) \cdot ¯ x (x^{'}) =^K (x, x^{'}) \to K (x, x^{'})$ in probability, where for notational simplicity we assign $¯ x (x) = [\frac{σ_{w} x^{L} (x)}{\sqrt{n_{L}}}, σ_{b}]$ . The regression problem is specified with mean-squared loss

L = \frac{1}{2} ∥ f (X) - Y ∥_{2}^{2} = \frac{1}{2} ∥ ¯ x (X) θ^{L + 1} - Y ∥_{2}^{2},

(S26)

and applying gradient flow to optimize the readout layer (and freezing all other parameters),

{˙ θ}^{L + 1} = - η {¯ x (X)}^{T} (¯ x (X) θ^{L + 1} - Y),

(S27)

where $η$ is the learning rate. The solution to this ODE gives the evolution of the output of an arbitrary $x^{*}$ . So long as the empirical kernel $¯ x (X) ¯ x (X)^{T}$ is invertible, it is

(S28)

For any $x, x^{'} \in R^{n_{0}}$ , letting $n_{l} \to \infty$ for $l = 1, \dots, L$ , one has the convergence in distribution in probability and distribution respectively

(S29)

Moreover $¯ x (X) θ_{0}^{L + 1}$ and the term containing $f_{0} (X)$ are the only stochastic term over the ensemble of network initializations, therefore for any $t$ the output $f (x^{*})$ throughout training converges to a Gaussian distribution in the infinite width limit, with

	$E [f_{t} (x^{*})]$	$= K (x^{*}, X) K^{- 1} (I - e^{- η K t}) Y,$		(S30)
	$V a r [f_{t} (x^{*})]$	$= K (x^{}, x^{}) - K (x^{}, X) K^{- 1} (I - e^{- η K t}) K (x^{}, X)^{T} .$		(S31)

Thus the output of the neural network is also a GP and the asymptotic solution (i.e. $t \to \infty$ ) is identical to the posterior of the NNGP (Equation 2.3). Therefore, in the infinite width case, the optimized neural network is performing posterior sampling if only the readout layer is being trained. This result is a realization of sample-then-optimize equivalence identified in Matthews et al. (2017).

Appendix E Results for NTK parameterization transfer to standard parameterization

Here we present a sketch for why the linearization results, derived for NTK parameterized networks, also apply to networks with a standard parameterization

{\begin{matrix} h^{l + 1} & = x^{l} W^{l + 1} + b^{l + 1} x^{l + 1} & = ϕ (h^{l + 1}) \end{matrix} and ⎧ ⎨ ⎩ \begin{matrix} W_{i, j}^{l} & = ω_{i j}^{l} \sim N (0, \frac{σ_{ω}^{2}}{n_{l}}) b_{j}^{l} & = β_{j}^{l} \sim N (0, σ_{b}^{2}) \end{matrix} .

(S32)

The NTK parameterization in Equation 1 is not commonly used for training neural networks. While the function that the network represents is the same for both NTK and standard parameterization, training dynamics under gradient descent are generally different for the two parameterizations. However, for a particular choice of layer-dependent learning rate training dynamics also become identical. Let $η_{NTK, w}^{l}$ and $η_{NTK, b}^{l}$ be layer-dependent learning rate for $W^{l}$ and $b^{l}$ in the NTK parameterization, and $η_{std} = \frac{1}{n_{max}} η_{0}$ be the learning rate for all parameters in the standard parameterization, where $n_{max} = {max}_{l} n_{l}$ . Recall that gradient descent training in standard neural networks requires a learning rate that scales with width like $\frac{1}{n_{max}}$ , so $η_{0}$ defines a width-invariant learning rate (Park et al., 2018). If we choose

η_{NTK, w}^{l} = \frac{n_{l}}{n_{max} σ_{ω}^{2}} η_{0}, and η_{NTK, b}^{l} = \frac{1}{}

Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent

Authors

Disentangling trainability and generalization in deep learning

On the Optimization Dynamics of Wide Hypernetworks

The Principles of Deep Learning Theory

Asymptotics of Wide Networks from Feynman Diagrams

The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective

Information in Infinite Ensembles of Infinitely-Wide Neural Networks

Implicit Acceleration and Feature Learning in Infinitely Wide Neural Networks with Bottlenecks

1 Introduction

1.1 Our Contribution

1.2 Additional related work

2 Theoretical Results

2.1 Notation and setup

2.2 Linearized networks

2.3 Infinite width limit yields Gaussian processes

2.3.1 Gaussian processes from gradient descent training

2.4 Infinite width networks are linearized networks

2.5 Extensions

2.5.1 Momentum

2.5.2 Cross-entropy loss

2.5.3 Beyond fully-connected networks

3 Experiments

3.1 Convergence of empirical kernel

3.2 Predictive output distribution

3.3 Comparison of training dynamics of linearized network to original network

3.4 Effects of depth and dataset size

4 Discussion

Acknowledgements

References

Appendix A Computing NTK and NNGP Kernel

Appendix B Tangent kernel for $ReLU$ and $erf$

Appendix C Multi-dimensional output and cross-entropy loss

Appendix D Gradient flow dynamics for training only the readout-layer

Appendix E Results for NTK parameterization transfer to standard parameterization

Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent

Authors

Disentangling trainability and generalization in deep learning

On the Optimization Dynamics of Wide Hypernetworks

The Principles of Deep Learning Theory

Asymptotics of Wide Networks from Feynman Diagrams

The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective

Information in Infinite Ensembles of Infinitely-Wide Neural Networks

Implicit Acceleration and Feature Learning in Infinitely Wide Neural Networks with Bottlenecks

This week in AI

1 Introduction

1.1 Our Contribution

1.2 Additional related work

2 Theoretical Results

2.1 Notation and setup

2.2 Linearized networks

2.3 Infinite width limit yields Gaussian processes

2.3.1 Gaussian processes from gradient descent training

2.4 Infinite width networks are linearized networks

2.5 Extensions

2.5.1 Momentum

2.5.2 Cross-entropy loss

2.5.3 Beyond fully-connected networks

3 Experiments

3.1 Convergence of empirical kernel

3.2 Predictive output distribution

3.3 Comparison of training dynamics of linearized network to original network

3.4 Effects of depth and dataset size

4 Discussion

Acknowledgements

References

Appendix A Computing NTK and NNGP Kernel

Appendix B Tangent kernel for ReLU and erf

Appendix C Multi-dimensional output and cross-entropy loss

Appendix D Gradient flow dynamics for training only the readout-layer

Appendix E Results for NTK parameterization transfer to standard parameterization

Appendix B Tangent kernel for $ReLU$ and $erf$