PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, Jeffrey Pennington

arXiv:1902.06720v4stat.MLcs.LG

TL;DR

Precisely characterizing neural-network training dynamics is difficult because high-dimensional non-convex loss surfaces complicate analysis. This paper shows that wide networks can be modeled through a first-order linearization at initialization, which agrees well with empirical training behavior across architectures, optimizers, and loss functions.

  • Problem

    Characterizing gradient-based training dynamics is challenging because neural networks are governed by high-dimensional, non-convex loss surfaces.

    Source: §1 · p. 1

  • Method

    The paper analyzes wide-network training by replacing the network with its first-order Taylor expansion at initialization and derives a Gaussian-process description of predictions for squared loss.

    Source: §1 · pp. 1–2

  • Results

    The linearized model agrees well with actual training dynamics and predictive distributions across fully-connected, convolutional, and wide residual networks, optimizers, and loss functions.

    Source: §4 · p. 9

  • Takeaways & Limitations

    The results suggest that a surprising number of realistic neural networks may operate in the regime studied.

    Source: §4 · p. 9

    Some layers of modern neural networks may operate far from the linearized regime.

    Source: §4 · p. 9

Abstract

from arXiv · show

A longstanding goal in deep learning research has been to precisely characterize training and generalization. However, the often complex loss landscapes of neural networks have made a theory of learning dynamics elusive. In this work, we show that for wide neural networks the learning dynamics simplify considerably and that, in the infinite width limit, they are governed by a linear model obtained from the first-order Taylor expansion of the network around its initial parameters. Furthermore, mirroring the correspondence between wide Bayesian neural networks and Gaussian processes, gradient-based training of wide neural networks with a squared loss produces test set predictions drawn from a Gaussian process with a particular compositional kernel. While these theoretical results are only exact in the infinite width limit, we nevertheless find excellent empirical agreement between the predictions of the original network and those of the linearized version even for finite practically-sized networks. This agreement is robust across different architectures, optimization methods, and loss functions.

1 Introduction

The paper argues that, as neural networks widen, gradient-descent training is governed by the network’s first-order Taylor expansion around initialization. For squared loss, predictions follow Gaussian-process dynamics, and experiments report strong agreement between the linearized and original networks at finite widths.

1 Introduction

  • Finite-width linearized models accurately characterize learning dynamics and prediction distributions across architectures, optimization methods, and losses, including wide residual networks.The reported agreement includes ensembles of finite-width models with different random initializations.Source: §1 · p. 2, §1.1 · p. 3
  • With squared loss, training dynamics have a closed-form solution and predictions evolve as a Gaussian process at infinite width, with explicit time-dependent expressions; gradient descent does not sample any model’s parameter posterior.The paper reports empirical confirmation that the GP result models prediction variation across finite-width models with different random initializations.Source: §1 · p. 2
  • For any loss, sufficiently wide network training is equivalent to training a model affine in all weights and biases; below a threshold learning rate, trajectories remain within an O-neighborhood of the linearized trajectory.The linearization result holds in both standard and NTK parameterizations.Source: §1 · pp. 1–2
  • The authors compute analytic neural tangent kernels for fully connected networks with ReLU or erf nonlinearities.Source: §1 · p. 2

1.1 Related work

  • Building on Jacot et al. [13]’s correspondence between infinite-width output dynamics and kernel gradient descent, the paper shows network outputs are linear in parameter updates, but not usually in inputs.This distinguishes the result from deep linear networks, which are linear in their inputs.Source: §1.1 · pp. 2–3

2 Theoretical results

In the infinite-width limit, neural-network training dynamics are captured by a first-order linearization, and gradient-trained predictions follow a Gaussian process with a kernel determined by the training dynamics. Under a learning-rate condition, the original network converges to those linearized predictions.

2.1 Notation and setup for architecture and training dynamics

  • The NTK parameterization normalizes forward and backward dynamics, and its predictions and training dynamics match standard parameterization up to width-dependent learning-rate scaling.The paper states that the linearity and Gaussian-process results also hold for standard parameterization.Source: §2.1 · p. 3

2.2 Linearized networks have closed form training dynamics for parameters and outputs

  • Replacing network outputs with their first-order Taylor expansion yields linearized dynamics; for MSE, closed-form solutions require only the initial tangent kernel and outputs.The initial parameter gradients remain fixed throughout training, so the dynamics can be computed without running gradient descent.Source: §2.2 · p. 4

2.3 Infinite width limit yields a Gaussian process

  • At infinite width, randomly initialized network outputs form a Gaussian process, whose conditioning on training data gives the posterior predictive distribution of an infinitely wide network.Its covariance kernel is computed recursively from the network architecture.Source: §2.3 · pp. 4–5

2.3.1 Gaussian processes from gradient descent training

  • For every training time, the infinite-width linearized network has Gaussian test predictions, but full-network gradient training generally does not correspond to Bayesian posterior sampling.When only the readout layer is optimized, the limiting predictions instead converge to the NNGP posterior as training time tends to infinity.Source: §2.3.1 · p. 5

2.4 Infinite width networks are linearized networks

  • With η < 2(λmin(Θ) + λmax(Θ))^-1, infinite-width network outputs converge to the Gaussian predictions of the linearized dynamics.The theorem assumes equal hidden-layer widths, λmin(Θ) > 0, and inputs with norm at most 1.Source: §2.4 · p. 6
  • Although individual weights and hidden activations move by vanishing amounts at infinite width, their collective changes produce a finite output change.The weight-gradient features also give a random-features interpretation of the individual-instance dynamics.Source: §2.4 · pp. 6–7

2.5 Extensions to other optimizers, architectures, and losses

  • The analysis also extends to multi-output networks, cross-entropy loss, and gradient descent with momentum.Broader architecture extensions are discussed as likely but deferred to empirical investigation.Source: §2.5 · p. 7

3 Experiments

Across the tested architectures and training settings, neural-network training dynamics closely follow those of linearized models. With MSE loss, the predictive output distribution remains Gaussian during training, though its final distribution differs from the NNGP posterior.

Comparison of training dynamics of linearized network to original network

  • Finite-width networks’ outputs, weights, and losses closely follow linearized dynamics, with empirical tangent kernels predicting them better than infinite-width kernels.The comparison includes full-batch gradient descent on a ReLU fully connected network.Source: Figure 3 · p. 8, §3 · p. 8

Predictive output distribution

  • The mean and variance of network outputs follow linearization’s analytic dynamics, while the final predictive distribution differs noticeably from the NNGP posterior.For MSE loss, the output distribution remains Gaussian throughout training.Source: §3 · pp. 7–8

Direct optimization of the linearized model

  • With cross-entropy and softmax, logits can grow indefinitely, and linearized-model error worsens late if the models diverge before that growth regime.This behavior contrasts with MSE, where logits converge to target values.Source: §3 · p. 8

Generalization beyond fully connected models

  • A wide residual network with 1,024 channels closely matches its linearization when trained directly on full CIFAR-10 with SGD and momentum.The comparison uses MSE loss and batch size 8.Source: §3 · p. 8, Figure 4 · p. 9

Effects of dataset size

  • As dataset size grows, linearization error increases, but increasing model width counterbalances that error.The experiments observe decreasing error as width grows and increasing error with dataset size.Source: §3 · pp. 8–9

4 Discussion

In the infinite-width limit, deep networks’ learning dynamics are exactly described by a linearized model, with empirical agreement across architectures, optimizers, and loss functions. The discussion identifies kernel properties as a path to understanding trainability and generalization, while noting unresolved finite-width performance gaps.

4 Discussion

  • In the infinite-width limit, a linearized model exactly describes deep nonlinear networks’ parameter-space learning dynamics, and empirical results agree well across architectures, optimizers, and loss functions.The authors suggest that many realistic networks may operate in this regime, consistent with evidence that networks are robust to re-initialization but not layer re-randomization (Zhang et al. [43]).Source: §4 · p. 9
  • Studying the kernel ˆΘ and target signal may clarify trainability and generalization, while examining NNGP kernel K and tangent kernel Θ may illuminate gradient descent’s inductive bias.The authors identify these as future research directions.Source: §4 · p. 9
  • Lee et al. [5] found wide SGD-trained networks approach corresponding GPs as width increases, while Novak et al. [7] found finite- versus infinite-width performance depends strongly on architecture.Novak et al. [7] reported infinite-width networks match or outperform finite-width ones for many fully connected or locally connected architectures, but the reverse for convolutional networks without pooling; the factors behind these gaps remain open. The authors see infinite-width behavior as a basis for systematically understanding finite-width networks and networks trained with large learning rates.Source: §4 · p. 9

Supplementary Material … C Neural Tangent kernel for ReLU and erf

The supplementary material tests agreement between neural networks and their linearizations under additional training conditions, extends the dynamics to momentum and general losses, and gives analytical tangent kernels for ReLU and erf activations.

A Additional figures

  • Across 100 initializations, the plotted lines show the different functions learned by the neural networks.The configuration matches Figure 2.Source: Figure S1 · p. 13
  • The original networks and their linearizations behave similarly with full-batch momentum on CIFAR and with minibatch SGD momentum on MNIST.The CIFAR experiment uses MSE with a tanh convolutional network; the MNIST experiment uses cross-entropy with a ReLU fully connected network.Source: Figure S2 · p. 13, Figure S3 · p. 14
  • With cross-entropy, logits can diverge at late times; wider networks delay the onset, while early deviation in narrower networks can lead to large late deviation.The comparison is for a fully connected tanh network with four layers on binary CIFAR.Source: Figure S4 · p. 14
  • Final RMSE is evaluated across varying depth and dataset size for fully connected networks, and across depth for convolutional networks and wide residual networks.The figure describes critically initialized networks trained with gradient descent on MSE.Source: Figure S5 · p. 15, A Additional figures · p. 15
  • The supplementary experiments measure training-time relative Frobenius norm changes in read-out and non-read-out weights, the empirical tangent kernel, and the empirical kernel.The measurements include one-hidden-layer ReLU networks at varying widths after 216 gradient-descent steps and three-layer tanh networks.Source: Figure S6 · p. 15

B.1 Momentum

  • Momentum updates yield discrete and continuous-time dynamics for the linearized network, with the continuous-time limit following Qian and Su et al. [45].The equations are amenable to analytic treatment for MSE loss.Source: §B.1 · pp. 15–16

B.2 Multi-dimensional output and cross-entropy loss

  • For multiple outputs and general losses such as softmax cross-entropy, the dynamics lack a closed-form solution and can instead be solved as an ODE initial-value problem.The described implementation uses TensorFlow's default dopri5 integrator.Source: §B.2 · p. 16

C Neural Tangent kernel for ReLU and erf

D Gradient flow dynamics for training only the readout-layer · E Computing NTK and NNGP Kernel

With mean-squared loss and all but the readout layer frozen, infinite-width training yields Gaussian-process outputs that interpolate between the prior and posterior. The accompanying section describes recursive computation of the NNGP and tangent kernels in the infinite-width limit.

D Gradient flow dynamics for training only the readout-layer

  • In the infinite-width limit, the network output throughout readout-layer training is Gaussian over initializations and interpolates between the GP prior and posterior.The time-dependent output is Gaussian because the initial-output term is the only stochastic term over network initializations.Source: D Gradient flow dynamics for training… · p. 17
  • Training only the readout layer uses mean-squared loss and gradient flow, with all other parameters frozen.The stated solution applies so long as the empirical kernel on the training inputs is invertible.Source: D Gradient flow dynamics for training… · p. 17
  • As training time tends to infinity, the output matches the NNGP posterior, making readout-only optimization posterior sampling and realizing sample-then-optimize equivalence.Source: D Gradient flow dynamics for training… · p. 17

E Computing NTK and NNGP Kernel

F Results in function space for NTK parameterization transfer to standard parameterization

The section argues that sufficiently wide standard-parameterized networks inherit NTK function-space linearization when layer-dependent learning rates make their training dynamics identical, and points to a formal infinite-width proof.

F Results in function space for NTK parameterization transfer to standard parameterization

  • Across datasets, activations, and loss functions, NTK- and standard-parameterized models achieve similar final performance at similar normalized learning rates.Figure S7 compares models trained with (stochastic) gradient descent; the reported normalized rate is ηstd = ηNTK/n.Source: Figure S7 · p. 19, F Results in function space for NTK… · p. 19
  • Although both parameterizations represent the same function, their gradient-descent dynamics generally differ unless a particular layer-dependent learning-rate choice is used.The section sketches this transfer and refers to §G for a formal proof of standard-network convergence to its linearization at infinite width.Source: F Results in function space for NTK… · pp. 18–19
  • Standard-parameterization gradient descent uses a learning rate scaling as 1/nmax, with η0 defining a width-invariant learning rate.Here nmax is the maximum layer width.Source: F Results in function space for NTK… · p. 19
  • With matched layer-dependent learning rates, standard and NTK parameterizations have identical training dynamics, so sufficiently wide standard networks remain linear in their parameters throughout training.The transfer requires only minor modifications to the per-layer Jacobian contributions.Source: F Results in function space for NTK… · p. 19

G Convergence of neural network to its linearization, and stability of NTK under gradient descent

Under stated assumptions, the section establishes global convergence and NTK stability for sufficiently wide networks trained by full-batch gradient descent or gradient flow. It also reports that experimental output dynamics closely match the linearized model, while individual-weight dynamics are similar.

G Convergence of neural network to its linearization, and stability of NTK under gradient descent

  • The analysis assumes equal hidden-layer widths and a full-rank analytic NTK, with a compact training set containing distinct inputs.The full-rank condition bounds the NTK eigenvalues away from zero and infinity.Source: G Convergence of neural network to its… · pp. 19–20
  • In the reported ReLU-network experiment, output dynamics of the original and linearized models are nearly identical, while individual-weight dynamics are similar.The experiment uses five hidden layers of width 2048 and MSE loss; training and test loss and accuracy also agree well.Source: Figure S8 · p. 20
  • The proofs for both training procedures rely on the Jacobian being locally Lipschitz.The section states this as Lemma 1.Source: G Convergence of neural network to its… · p. 21
  • The section establishes global convergence of a neural network and NTK stability under full-batch gradient descent, with probability at least 1 − δ0 over random initialization when η = η0 < ηcritical under Assumptions [1-4].Source: G Convergence of neural network to its… · pp. 19, 21
  • The results extend to other architectures or functions when the empirical NTK converges in probability to a positive-definite limit and the Jacobian is locally Lipschitz.These are the two conditions stated for the extension.Source: G Convergence of neural network to its… · p. 21

G.1 Proof of Theorem G.1

The proof proceeds by induction on Equation S49, using high-probability estimates and the gradient-descent update. It then applies the mean value theorem and initialization-kernel convergence to establish the remaining estimates under a learning-rate condition.

G.1 Proof of Theorem G.1

  • Equation S49 is proved by induction, with the base case immediate and Equations S47 and S53 holding with high probability for sufficiently large width.The induction step uses the second estimate of Equation S47.Source: §G.1 · p. 22
  • The induction yields the first estimate of Equation S49 at the next step and bounds each iterate’s distance from initialization by 3KR0.The bound applies for j = 0, …, t + 1.Source: §G.1 · p. 22
  • The second estimate is derived using the mean value theorem at an interpolated parameter between successive iterates and the gradient-descent update.The interpolation point lies between θt and θt+1.Source: §G.1 · p. 22
  • Initialization-kernel convergence in probability [37] and Lemma 1 support the probability bound, under the condition η0 < 2/(λmin + λmax).The passage gives the condition as η0 < 2/(λmin + λmax).Source: §G.1 · p. 22

G.2 Proof of Theorem G.2

The proof establishes high-probability conditions using initialization and the arguments of Section G.1, then argues by contradiction that t1 is infinite.

G.2 Proof of Theorem G.2

  • The first step establishes constants R0 and n0 for which a random-initialization event holds with probability at least 1 − δ0/10 for every n ≥ n0.Source: §G.2 · p. 23
  • Using Section G.1 arguments and λmin in Lemma 1, the proof establishes an n1 such that a stated event holds with probability at least 1 − δ0/10 for n ≥ n1.Source: §G.2 · p. 23
  • The proof claims t1 = ∞ by contradiction: assuming finite t1 keeps θt within B(θ0, Cn−1) through t1, contradicting its definition.The proof also notes that Equation S78 matches the first equation of Equation S51.Source: §G.2 · p. 23

G.3 Proof of Lemma 1 … J Details on Wide Residual Network

These supplementary sections establish bounds used in the proof, state NTK-parameterization and MSE-proof extensions, report empirical kernel convergence, and specify the Wide Residual Network block ordering.

G.3 Proof of Lemma 1

  • The proof of Lemma 1 uses operator-norm bounds for random Gaussian matrices [48] to bound initialized and nearby layer-weight norms with high probability.Source: §G.3 · pp. 23–24
  • Layer activations and backpropagated deltas vary at most linearly with parameter distance, with a constant depending on the stated model quantities.The estimates are obtained by induction, starting at l=1 for activations and l=L for deltas.Source: §G.3 · p. 24, §G.4 · p. 24

G.4 Remarks on NTK parameterization

  • The appendix states NTK-parameterization analogues of Theorem G.1 and Lemma 1, including a local-Lipschitzness-of-the-Jacobian lemma.The theorem assumes the stated assumptions and a gradient-descent learning rate η0 below ηcritical.Source: §G.4 · pp. 24–25

H Bounding the discrepancy between the original and the linearized network: MSE loss

I Convergence of empirical kernel

  • O(1/sqrt(n)) convergence holds for both empirical tangent-kernel and NNGP estimates as width increases; Figure S9 also shows convergence improving with width and Monte Carlo samples.The reported rates are ∥ˆΘ(n) − Θ∥F = O(1/sqrt(n)) and ∥ˆK(n) − K∥F = O(1/sqrt(n)).Source: Figure S9 · p. 27, I Convergence of empirical kernel · p. 27

J Details on Wide Residual Network

  • The Wide Residual Network uses Batch Normalization–ReLU–Conv ordering within each residual block [14].Source: Table S1 · p. 27

Related papers

Loading 1902.06720v4…