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

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.06720v4Feb 18, 2019stat.MLcs.LG

TL;DR

The paper addresses the difficulty of characterizing deep-network training by analyzing wide neural networks in the infinite-width limit. It replaces their dynamics with a first-order linearized model and derives Gaussian-process predictive dynamics for squared loss, finding strong agreement with finite networks across settings. The paper also identifies scope boundaries, including cases where some modern-network layers may lie far from the linearized regime and where gradient-descent distributions are not Bayesian posteriors.

Problem
Deep neural networks have complex, high-dimensional non-convex loss landscapes, making their training dynamics difficult to characterize theoretically.
Method
The paper analyzes wide networks by replacing them with their first-order Taylor expansion around initialization and characterizing squared-loss dynamics through Gaussian processes.
Results
The linearized theory agrees well with original-network learning dynamics and predictive distributions across architectures, optimization methods, and loss functions, including finite-width settings.
Takeaways & Limitations
Wide-network training can be studied through analytically tractable linearized dynamics, while its evolving predictions under squared loss admit a Gaussian-process description.
Takeaways & Limitations
Some modern neural-network layers may operate far from the linearized regime, and gradient-descent predictive distributions do not generally correspond to Bayesian posteriors.

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 makes wide-network training analytically tractable by showing that, at large width, nonlinear networks follow a linearized parameter-space model. It further connects gradient-descent predictions under squared loss to Gaussian-process dynamics and finds strong finite-width agreement across settings.

Core contribution: At large width, a neural network can be replaced by its first-order Taylor expansion around initialization, making gradient-descent dynamics analytically tractable.The linearization describes the network’s dependence on parameter updates during training.
Empirical support: The theory accurately characterizes learning dynamics and predictive distributions for finite-width networks across architectures, optimizers, and loss functions.Reported settings include fully connected, convolutional, and wide residual networks, with gradient descent, momentum, mini-batching, MSE, and cross-entropy.
Core contribution: For squared loss, exact training dynamics admit a closed-form solution that characterizes the evolving predictive distribution as a Gaussian process.The result extends the analysis of output dynamics throughout gradient-descent training.
Theoretical guarantees: The parameter-space result holds for arbitrary loss functions, while sufficiently small learning rates keep sufficiently wide-network trajectories near the linearized trajectory.The paper identifies a critical learning-rate threshold for this regime under its formal conditions.
Interpretation: The work distinguishes gradient-descent dynamics from Bayesian posterior sampling, despite both producing Gaussian-process draws in the relevant settings.The paper notes that gradient descent does not generally correspond to sampling from a Bayesian posterior.

2 Theoretical results

The paper shows that sufficiently wide neural networks follow dynamics equivalent to a first-order linearization around initialization, enabling tractable parameter and output evolution. In the infinite-width limit, trained outputs become Gaussian processes, while finite-width networks closely follow these predictions under specified conditions.

Linearized networks: The linearized network replaces the original output with its first-order Taylor expansion around initialization, separating the fixed initial output from the training-induced change.Because the initialization gradient remains constant, the linearized dynamics can be computed from the initial outputs and tangent kernel.
Output distribution dynamics: As width tends to infinity, the linearized network's test outputs converge in distribution to a Gaussian with time-dependent mean and covariance.This follows because the initialization outputs are Gaussian in the infinite-width limit and the linearized outputs are affine transformations of them.
Gaussian-process interpretation: Only optimizing the readout layer makes the network and its linearization identical and recovers the Gaussian-process posterior in the infinite-width limit, unlike joint optimization.With all parameters trainable, the resulting NTK-GP distribution does not generally have a posterior-sampling interpretation.
Sufficient conditions for linearization: For sufficiently large width and learning rate η < ηcritical, gradient descent trajectories of the original network remain near those of the linearized network.The stated threshold is ηcritical = 2(λmin(Θ) + λmax(Θ))^-1, under the assumption λmin(Θ) > 0.
Extensions: The analysis extends beyond the core setting to multidimensional outputs, cross-entropy loss, and gradient descent with momentum.The principal theoretical analysis focuses on fully connected single-output networks trained by full-batch gradient descent.

3 Experiments

Experiments show that linearized models closely reproduce wide-network training dynamics across architectures, optimizers, batching schemes, and losses. Approximation error decreases with width but increases with dataset size, while squared-loss predictive distributions follow the predicted Gaussian-process dynamics.

Cross-architecture agreement: Across fully connected, convolutional, and wide residual architectures, linearized models closely match original network training dynamics.The comparisons include full- and mini-batch training and multiple datasets.
Predictive distributions: For squared loss, neural-network output mean and variance dynamics follow the analytic linearization dynamics, while the final gradient-descent distribution differs from the NNGP posterior.The predictive output distribution remains Gaussian throughout training, but gradient-descent training and exact Bayesian inference produce noticeably different distributions.
Optimization methods: Directly optimizing the linearized model supports standard techniques including mini-batching, weight decay, data augmentation, and momentum optimization.Experiments directly train the linearized model rather than only integrating its continuous-time tangent-kernel dynamics.
Loss-function effects: Cross-entropy training can produce late-time logit divergence, especially when the original and linearized models already differ substantially early in training.Greater width delays the onset of logit deviation.
Width and dataset size: As width grows, linearization error decreases; as dataset size grows, error increases, which can be counterbalanced by increasing model size.The empirical tangent kernel gives more accurate dynamics for finite-width networks.

4 Discussion

The discussion concludes that linearized dynamics describe a broad range of wide and finite networks, while emphasizing that some modern-network layers may lie outside this regime. Kernel properties remain a route toward understanding trainability, generalization, and gradient-descent inductive bias.

4 Discussion: Across fully connected, convolutional, and wide residual networks, linearized dynamics agree well with training and predictive distributions under varied optimizers and losses.The reported settings include gradient descent, momentum, mini-batching, MSE, and cross-entropy.
4 Discussion: The theory suggests that properties of the tangent kernel and NNGP kernel may clarify trainability, generalization, and gradient-descent inductive bias.In the studied regime, learning dynamics are captured by the tangent kernel and target signal.
4 Discussion: Some layers of modern neural networks may operate far from the linearized regime, and performance gaps between finite- and infinite-width models remain architecture-dependent.The text identifies convolutional networks without pooling as one setting where infinite-width performance can be worse than finite-width performance.

A Additional figures

Additional figures test linearization across initializations, convolutional and fully connected architectures, momentum, cross-entropy, and varying depth, width, and dataset size. They also examine momentum dynamics and analytic or numerical treatment for general losses.

Figure S1: Figure S1 plots functions learned from 100 different initializations using the Figure 2 configuration.The figure samples output functions to examine variation across random initializations.
Figure S2: Figure S2 compares a tanh convolutional network with its linearization under full-batch gradient descent with momentum on binary CIFAR classification.The networks have three hidden layers with channel size n = 512, and the linearized model is trained directly.
Figure S3: Figure S3 compares a ReLU fully connected network and its linearization under momentum SGD with cross-entropy loss on 10-class MNIST.The experiment uses two hidden layers of width n = 2048 and the full dataset of |D| = 50,000.
Figure S4: Under cross-entropy, logit deviation can grow at late times; wider networks delay this deviation, whereas early mismatch in narrower networks predicts larger later discrepancy.The experiment uses a fully connected tanh network with L = 4 on binary CIFAR classification.
Figure S5: Figure S5 studies final RMSE as width varies across fully connected, convolutional, and wide residual networks while changing depth and dataset size.The reported depth ranges include 1 to 16 for fully connected networks and 1 to 32 for convolutional networks.
Momentum dynamics: Momentum dynamics can be written as discrete output updates, related to continuous time by t = i√η and β̃ = (β − 1)/√η.The resulting equations admit analytic treatment for MSE loss and are compared experimentally in Figures S2, S3, and 4.
General losses: For general losses, the dynamics lack a closed-form solution but can be solved as an initial-value problem with an ODE solver.The implementation uses the dopri5 method for cross-entropy with softmax output.

D Gradient flow dynamics for training only the readout-layer

When only the readout layer is trained, wide-network outputs remain Gaussian-process distributed throughout gradient-flow training. The long-time solution matches the posterior of the neural-network Gaussian process.

The infinite-width result extends the Gaussian-process connection between Bayesian wide neural networks and training that optimizes only the readout layer.
Training only the readout layer under mean-squared loss defines a regression problem optimized by gradient flow while all other parameters remain fixed.
The readout-layer optimization yields an explicit evolution equation for the output at any test input when the empirical kernel is invertible.
As hidden-layer widths tend to infinity, network outputs throughout training converge to a Gaussian distribution over random initializations.
The trained output is a Gaussian process whose asymptotic solution equals the posterior of the neural-network Gaussian process.

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

Although NTK and standard parameterizations generally have different gradient-descent dynamics, suitable layer-dependent learning rates make them identical. Experiments show similar final performance across datasets, activations, and losses under matched normalized learning rates.

Standard parameterization networks are analyzed by transferring function-space linearization results established for NTK-parameterized networks.
NTK and standard parameterizations represent the same functions but generally produce different gradient-descent dynamics.
Choosing the prescribed layer-dependent NTK learning rates makes training dynamics identical to standard-parameterization dynamics.
Across datasets, activation functions, and loss functions, matched normalized learning rates produce similarly performing models for both parameterizations.

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

Under stated assumptions, sufficiently wide networks trained by gradient descent or gradient flow remain close to their linearized trajectories and retain stable tangent kernels. The results rely on kernel definiteness, Jacobian regularity, and width conditions, with empirical support for matching dynamics.

Assumptions: The convergence proof assumes identical hidden-layer widths, a full-rank analytic NTK, compact distinct training inputs, and a suitable activation function.
Empirical verification: Experiments report nearly identical exact and experimental output dynamics and good agreement in loss, accuracy, and test RMSE between original and linearized models.
Proof strategy: The argument uses local Lipschitzness of the Jacobian to establish convergence and tangent-kernel stability under both discrete gradient descent and gradient flow.
Gradient descent: For learning rates below ηcritical and sufficiently large width, gradient-descent trajectories remain controlled with high probability over random initialization.
Gradient flow: Gradient-flow training satisfies an analogous high-probability theorem under the same stated assumptions.

G.1 Proof of Theorem G.1

The proof controls the training trajectory inductively: concentration and Jacobian bounds establish the base event, while the learning-rate condition preserves parameter proximity and kernel stability at each step.

For sufficiently large width, initialization events satisfy the required bounds with high probability, providing the starting point for the induction.
The induction step combines prior trajectory bounds with the gradient update to control the next iterate and parameter displacement.
Convergence of the empirical NTK and the learning-rate restriction η0 < 2(λmin + λmax)^−1 preserve the required spectral conditions.
The proof completes by selecting width and radius thresholds large enough to satisfy the preceding estimates simultaneously.

G.2 Proof of Theorem G.2

The proof establishes a high-probability initialization event and shows that the trajectory remains within a controlled neighborhood of the initial parameters for all times.

The argument begins by selecting thresholds R0 and n0 so a high-probability initialization condition holds for every sufficiently wide network.
A second width threshold n1 is chosen using the same argument as Section G.1, with another high-probability bound.
The proof assumes a finite exit time t1 and shows that parameters remain in the prescribed neighborhood for every t ≤ t1.
This contradicts the definition of t1, so the trajectory never exits the neighborhood and t1 = ∞.

G.3 Proof of Lemma 1

The lemma and its extensions control network and Jacobian behavior near initialization, while subsequent bounds establish convergence of finite-width kernels toward analytic infinite-width kernels.

Proof bounds: Gaussian matrix operator-norm bounds control weight matrices throughout a neighborhood of the initialization.
Proof bounds: Inductive estimates yield local Lipschitz control of hidden representations and backpropagated quantities with respect to parameters.The estimates are obtained by starting from the first forward layer and the final backward layer.
NTK parameterization: The same framework gives a local Lipschitzness result for the Jacobian under NTK parameterization.
Kernel control: A positive smallest eigenvalue of the initial tangent kernel supports the resulting high-probability control for sufficiently wide networks.
Kernel convergence: ∥ˆΘ(n) −Θ∥F = O(1/√n) and ∥ˆK(n) −K∥F = O(1/√n) as width increases, with convergence also improving as Monte Carlo samples increase.The convergence is observed for both NNGP and tangent kernels; the figures describe one- and three-hidden-layer ReLU networks and single-hidden-layer networks.