Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation

Greg Yang

Source-linked AI summary

Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation

Greg Yang

arXiv:1902.04760v3Feb 13, 2019cs.NEcond-mat.dis-nncs.LGmath-phstat.ML

TL;DR

Wide random neural networks motivate scaling-limit analyses across Gaussian processes, signal propagation, and training dynamics. This paper introduces tensor programs to express most neural-network computations and characterize their large-width limits, yielding Gaussian-process and Neural Tangent Kernel convergence results alongside corrected gradient-dynamics analyses. The framework is broad but does not fully cover every deep-learning layer and has important scope boundaries around weight sharing and batch normalization.

Problem
Scaling-limit analyses of wide random neural networks lacked a unified treatment across modern architectures and forward, backward, and training-related computations.
Method
The paper introduces straightline tensor programs that express most neural-network computations and characterizes their scaling limits for large randomized tensors.
Results
The framework gives Gaussian-process convergence for standard architectures, identifies when gradient independence is correct, and proves Neural Tangent Kernel convergence without batch normalization.
Takeaways & Limitations
The results support analyzing Gaussian-process behavior, signal propagation, Jacobian singular values, and training dynamics within one framework.
Takeaways & Limitations
The framework does not fully cover some layers, and weight sharing involving matrix transposes has scaling behavior distinct from sharing without transposes.

Abstract

from arXiv · show

Several recent trends in machine learning theory and practice, from the design of state-of-the-art Gaussian Process to the convergence analysis of deep neural nets (DNNs) under stochastic gradient descent (SGD), have found it fruitful to study wide random neural networks. Central to these approaches are certain scaling limits of such networks. We unify these results by introducing a notion of a straightline \emph{tensor program} that can express most neural network computations, and we characterize its scaling limit when its tensors are large and randomized. From our framework follows (1) the convergence of random neural networks to Gaussian processes for architectures such as recurrent neural networks, convolutional neural networks, residual networks, attention, and any combination thereof, with or without batch normalization; (2) conditions under which the \emph{gradient independence assumption} -- that weights in backpropagation can be assumed to be independent from weights in the forward pass -- leads to correct computation of gradient dynamics, and corrections when it does not; (3) the convergence of the Neural Tangent Kernel, a recently proposed kernel used to predict training dynamics of neural networks under gradient descent, at initialization for all architectures in (1) without batch normalization. Mathematically, our framework is general enough to rederive classical random matrix results such as the semicircle and the Marchenko-Pastur laws, as well as recent results in neural network Jacobian singular values. We hope our work opens a way toward design of even stronger Gaussian Processes, initialization schemes to avoid gradient explosion/vanishing, and deeper understanding of SGD dynamics in modern architectures.

1. Introduction

The paper unifies scaling-limit analyses of wide random neural networks through tensor programs that express standard deep-learning computations. Its scope covers compositions of MLPs, recurrent networks, skip connections, attention, convolution, and batch normalization.

Tensor programs provide a common language for analyzing scaling limits as network width increases under Glorot initialization.The framework targets scenarios corresponding roughly to inference and backpropagation.
Standard architectures include compositions of MLPs, RNNs, skip connections, attention, convolution, and batch normalization.A readout layer is any linear layer converting hidden states to an output vector.
Most corollaries are stated for standard architectures, although the results are typically more general.

2. Related Works and Our Corollaries

The paper unifies scaling-limit analyses of wide neural networks through tensor programs, extending Gaussian-process and Neural Tangent Kernel results across standard architectures. It also establishes when gradient-independence approximations are correct, identifies corrections for weight sharing and readout choices, and connects the framework to classical random-matrix laws.

Gaussian Behavior of Wide Neural Networks: The paper generalizes the DNN-GP correspondence to standard architectures and very general nonlinearities under stated growth and randomization conditions.The informal corollary covers fixed standard architectures with linear readouts, bounded nonlinearities, finite input sets, and independently sampled readout weights.
Signal Propagation in Neural Networks: The paper distinguishes significant scaling behavior between sharing matrix transposes and sharing no matrix transposes.This distinction is presented as a consequence of weight tying and motivates corrections to analyses that treat weight sharing as straightforward.
Signal Propagation in Neural Networks: The gradient independence assumption is correct for MLPs but can fail when the readout is global mean pooling or when weights are tied.For tied and untied weights, the paper reports agreement for simple RNNs but disagreement for general gated RNNs, while global mean pooling requires corrected equations.
Signal Propagation in Neural Networks: Theoretical LSTM gradient-norm predictions are already very precise for networks with 1000 neurons.The scaling-limit prediction closely matches empirical simulation in the reported LSTM experiment.
Neural Tangent Kernel: The NTK converges almost surely on finite input sets as widths grow together for fixed standard architectures without batch normalization.The result assumes a linear readout and nonlinearities with polynomially bounded weak derivatives; the paper contrasts simultaneous-width convergence with prior sequential limits.
Other Works: The framework is broad enough to rederive the semicircle and Marchenko-Pastur laws and to justify related empirical Fisher and NTK spectral computations.The paper also connects its tensor-program analysis to Approximate Message Passing and State Evolution.

3. Tensor Programs

Tensor programs provide a straightline representation for many neural-network computations and define randomized wide-width scaling regimes through dimension classes and input sampling.

Tensor programs consist of typed assignment lines for vector inputs, matrix inputs, transposes, matrix multiplication, linear combinations, and coordinatewise nonlinear functions.
Examples: These programs express MLP, recurrent, residual, convolutional, batch-normalized, and combined neural-network computations, while layer normalization is not directly captured.
Dimension constraints: Common dimension classes encode equality constraints among variables and represent varying widths across practical architectures, including residual-network blocks.
Randomization and limits: A randomized program is specified by width sequences, matrix variances, input means, covariances, and independent Gaussian input coordinates for each dimension class.
Randomization and limits: The scaling analysis sends every dimension class to infinity while requiring convergent width ratios, matrix scales, input moments, and eventual covariance ranks.
Scaling behavior: Under suitable conditions, G-vars converge to Gaussians with specific means and covariances, although general tensor programs need not have Gaussian G-var limits.

4. Programs with No Transposes

For tensor programs without transposes, the framework establishes Gaussian scaling limits under controlled nonlinearities and rank-convergence assumptions, while clarifying a fixed-depth tied-weight paradox.

Assumptions: Almost sure rank convergence is assumed for collections of variables produced from each shared matrix, with linear combinations accounting for persistent dependencies.
Theorem 4.3 analyzes transpose-free skeletons whose nonlinear functions satisfy an α-controlled growth condition for α < 2.
Limit structure: G-vars created from the same matrix can remain correlated, while other variables are asymptotically independent apart from linear-combination structure.
Limit structure: Under these conditions, coordinate-level limits are Gaussian with specified mean and covariance and are iid across coordinates.
Tied-weight example: For tied and untied deep linear networks, fixed-depth α-controlled moments match asymptotically even though intuition suggests tied powers may amplify eigenvector directions.
Tied-weight example: The apparent contradiction is resolved because the theorem fixes depth before width grows; the depth required for exponentiation increases with width.

5. Backprop with Zero Mean Gradients

The extended-syntax construction represents backpropagation without transposes and supports Gaussian scaling limits when gradient-related inputs are zero-mean and independent of forward inputs.

An extended program appends transposed matrix inputs, new vector inputs, and constrained odd functions to represent backpropagation after a transpose-free forward skeleton.
Gradient construction: The added vector inputs represent neural-network outputs and are sampled with zero mean independently from the forward input variables.
Scaling conditions: The extended scaling recursion assigns zero means and zero cross-covariances between newly introduced variables and earlier forward variables.
Scaling conditions: The theorem requires polynomially bounded functions and almost sure rank convergence, then gives convergence for polynomially bounded test functions.
Boundary: The result does not apply to batch normalization because its Jacobian is singular on a one-dimensional affine subspace, including the origin.

6. General Tensor Programs

General tensor programs require detransposition to track correlations and linear-combination components induced by transposed shared matrices, yielding correct scaling limits beyond zero-mean-gradient settings.

Motivation: When gradient-input variables lack zero mean, the simpler theorem can compute incorrect expectations, as shown by a quadratic-activation one-hidden-layer MLP.
Extended syntax: Extended syntax permits nonlinear functions to consume H-vars and is semantically equivalent to original syntax through expanded definitions.
Detransposition: Detransposition converts transposed matrix operations into an extended transpose-free program while maintaining a mapping between original and transformed variables.
Detransposition: For matrix reuse, detransposition records both Gaussian components and coefficients of linear combinations induced by correlations with earlier variables.
Applications: The resulting framework rederives random-matrix laws, including the semicircle and Marchenko–Pastur laws, and applies to random MLP examples.
Limit structure: The general intuition is that transposed shared-matrix operations produce a Gaussian term plus a linear combination of earlier variables, so the limit need not be Gaussian.

7. Proof Techniques

The proof strategy uses Gaussian conditioning to characterize tensor-program limits, with moment and correlation bounds extending beyond earlier AMP analyses.

Gaussian conditioning represents a conditioned Gaussian matrix as a mean plus projected independent Gaussian residual.
Compared with earlier AMP analysis, the framework assumes input G-vars have all finite moments rather than only a bounded number.This removes a smoothness assumption on the functions corresponding to tensor-program nonlinearities.
Linear-programming bounds control off-diagonal projection correlations, while Hermite expansions bound moments of projected nonlinear transforms.
The approach can handle nonsmooth functions such as the step-function gradient of ReLU, which a more direct earlier argument would not support.
When covariance is nonsingular, distributional derivatives can be interpreted against the Gaussian density; singular covariance is handled by restricting to a maximal-rank subset.

8. Discussion

The framework unifies scaling-limit analyses across modern neural-network computations, while discussion identifies unsupported layers, fixed-skeleton scope, and possible automation.

The tensor program framework expresses almost all modern neural-network computations and characterizes their scaling limits, yielding DNN-GP, signal-propagation, and NTK results.
The authors expect Gaussian-sampling results to extend to other “nice” distributions with a few finite moments.
The presented framework does not cover a few layers completely, with dropout and layernorm described as extendable and gradient reasoning for batchnorm left constrained.
Scaling-limit results apply only to fixed tensor-program skeletons, leaving joint limits for dataset size and network width as future work.
The framework could support software that automatically computes corresponding µc and Kc from PyTorch tapes or TensorFlow computation graphs.
Examples encode shared weights across multiple inputs, residual-block merge styles, recurrent time steps, batch normalization, and convolutional pixel preactivations.
Notation distinguishes straightline-program variables from arbitrary variables, with superscripts, line indices, dimensions, and node collections specifying program structure.

D.1. Warmup: MLP

The MLP warmup derives Gaussian-process, gradient-covariance, and NTK limits using tensor programs, then illustrates the framework on classical random-matrix laws.

The warmup constructs tensor programs for fully connected feedforward MLPs and uses them to study GP correspondence, gradient dynamics, and NTK convergence.
The NTK is defined as the inner product of parameter gradients, and for finite input sets it converges almost surely under width growth and suitable randomization.
The NTK convergence theorem assumes nonlinearities with polynomially bounded weak derivatives, excluding batch normalization in this formulation.
The MLP analysis formalizes computations previously used for GP correspondence, signal propagation, and NTK derivations.
In the infinite-width limit, the randomized MLP output converges almost surely to a Gaussian process whose kernel is determined by the limiting covariance.
Backpropagation is represented by transposed matrix multiplications and coordinatewise derivatives, producing gradient covariances Πl alongside forward covariances.
For global mean pooling with an odd final nonlinearity, the Gradient Independence Assumption gives the correct gradient covariance and NTK computation.
The tensor-program method rederives the semicircle law for the Gaussian Orthogonal Ensemble using moment calculations and Catalan-number limits.

D.4. DNN-GP correspondence

The tensor-program framework shows that randomized wide networks converge to Gaussian processes across batch, sequence, and standard architectural settings.

D.4. DNN-GP correspondence: A tensor program represents the network body and its output representations using G- and H-vars.The framework applies when the computation F(z; θ) is represented by a tensor program π.
D.4. DNN-GP correspondence: Gaussian input embeddings arise when independently sampled Gaussian matrices multiply the inputs, with covariance determined by pairwise input products.This provides the jointly Gaussian starting point for the scaling-limit analysis.
D.4. DNN-GP correspondence: In the infinite-width limit, normalized inner products of equal-sized representations converge almost surely to entries of a positive semidefinite kernel matrix C.The kernel is computable through the framework’s main scaling-limit theorem.
D.4. DNN-GP correspondence: Batch-parallel networks with shared readout weights converge in distribution to a Gaussian process.This is the batch-input specialization of the DNN-GP correspondence.
D.4. DNN-GP correspondence: Sequence-to-sequence RNNs processing batches of sequences converge in distribution to a multivariate Gaussian process.The result covers parallel processing of sequences with shared recurrent computation.
D.4. DNN-GP correspondence: Similar Gaussian-process convergence results hold for any standard architecture.The stated architecture class includes compositions of common neural-network components.

D.5. Gradient Independence Assumption

The framework determines when forward and backward weights may be treated as independent, and extends singular-value analyses to tied and arbitrary architectures.

D.5. Gradient Independence Assumption: The gradient-independence assumption is correct for MLPs under the theorem’s stated conditions.The paper explicitly establishes correctness for MLPs.
D.5. Gradient Independence Assumption: For standard architectures without batch normalization and without transposed forward weight sharing, gradient dynamics can be computed while treating forward G-vars as independent of backward-pass weights.The result assumes zero-mean readout weights and polynomially bounded weak derivatives.
D.5. Gradient Independence Assumption: The relevant forward-program variables are asymptotically free from the corresponding backward-program variables in the mixed-moment analysis.This asymptotic freeness underlies independence-style calculations in the wide limit.
D.5. Gradient Independence Assumption: The resulting corollary justifies Stieltjes-transform calculations and almost-sure convergence of singular-value distributions.It connects the tensor-program analysis to prior MLP Jacobian results.
D.5. Gradient Independence Assumption: Tensor programs can compute neural-network Jacobian singular-value distributions even with weight tying and arbitrary architecture.The same moment-computation machinery handles these broader settings.

D.6. Signal Propagation

The paper uses tensor programs to rigorously recover wide-network signal-propagation results, including weight-tied autoencoders and RNNs under weaker activation assumptions.

D.6. Signal Propagation: The signal-propagation analysis covers simple RNNs viewed as MLPs with parameters tied across layers.The comparison uses identically sampled weights and biases for the RNN and MLP.
D.6. Signal Propagation: A weight-tied autoencoder is analyzed by letting every layer width diverge while neighboring width ratios converge to positive constants.The limit considers nl →∞ with nl/nl−1 →αl ∈(0, ∞).
D.6. Signal Propagation: The framework reproduces the forward signal-propagation theorem for the weight-tied autoencoder without requiring smooth nonlinearities or test functions.The result permits polynomially bounded activation functions.
D.6. Signal Propagation: In the wide limit, each autoencoder hidden state is distributed as a Gaussian vector with covariance proportional to the identity.The stated result is xl ∼= N(0, ¯τlInl).
D.6. Signal Propagation: The allowed test functions are coordinatewise and polynomially bounded, while the resulting convergence is almost sure rather than in probability.Removing smoothness for noncoordinatewise test functions remains an open direction.

D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS

The framework rigorously justifies prior semirigorous equations for MLPs, residual networks, convolutional networks, and batch-normalized networks, while excluding batchnorm gradient dynamics.

D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS: Tensor-program analysis rigorously recovers signal-propagation equations previously derived heuristically.The theorem is applied across the examples presented in this section.
D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS: For residual and convolutional networks, the framework expresses the architecture and its randomized parameters in tensor-program form.The examples specify residual parameters and convolutional channel dimensions.
D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS: The framework justifies prior convolutional signal-propagation calculations involving circular cross-correlation.The convolutional example uses ⋆ for 2D circular cross correlation.
D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS: Backward signal-propagation recurrences can likewise be justified rigorously using the gradient-analysis theorem.The equations include reversed decoder weights in the recurrence.
D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS: For batch-normalized networks, the framework gives infinite-width results over finite sets of input batches under α-controlled nonlinearities.The batch size and input dimension remain fixed while hidden widths diverge.
D.6.1. JUSTIFYING SEMIRIGOROUS EQUATIONS: The theorems do not rigorously justify batch-normalization gradient dynamics because the batchnorm Jacobian has singularities.The limitation is specifically tied to the singularity of the Jacobian.

D.6.2. A TASTE OF WEIGHT-TYING

Weight tying changes the scaling behavior of residual networks, requiring coupled forward and backward analyses. The tensor-program framework derives limiting kernels and NTKs for such architectures, including a CNN example.

Weight-Tied Residual Networks: Weight-tied residual networks provide a minimal setting where sharing weights across layers can produce behavior different from untied networks.The section uses this architecture to study forward norms, paired-vector dynamics, and gradients.
Forward Dynamics: The norm of the hidden state grows much faster for tied-weight residual networks than for untied-weight residual networks.This comparison is stated directly for the weight-tied residual example.
Forward Dynamics: The simultaneous evolution of two vectors is characterized through limiting quantities K(l, m) and C(l, m), with symmetry and zero-boundary identities.The derivation expands these quantities recursively through the residual dynamics.
Backward Dynamics: Backward dynamics follow by defining h_t as the loss gradient with respect to h_t and g_t as W^T h_t.Boundary conditions are then used to analyze gradients through the tied residual network.
Neural Tangent Kernel: The framework computes NTK contributions by expanding network outputs and gradients into a single tensor program under rank-convergence and derivative-growth conditions.For the CNN example, the resulting NTK converges almost surely for polynomially bounded weak derivatives.
Neural Tangent Kernel: The CNN example gives an explicit limiting NTK, with separate contributions from layers and the linear readout under the stated activation condition.The section presents this as an example of an NTK not previously appearing in the literature.

D.8. Approximate Message Passing

This section connects tensor-program limits to Approximate Message Passing and related Gaussian conditioning arguments. It shows that smoothness assumptions can be weakened in a Gaussian initialization setting while imposing a stronger distributional assumption.

AMP Background: Approximate Message Passing reconstructs a sparse vector from undersampled noisy linear observations using iterative nonlinear updates.Its asymptotic behavior is described by state evolution as n, N →∞ with n/N →δ.
Tensor-Program Corollary: The tensor-program theorem yields an AMP corollary that forgoes smoothness assumptions on η_t when each component of x_0 is iid Gaussian.The generalized formulation uses sequences of functions f_t and g_t.
Tensor-Program Corollary: The AMP recursion is translated into a tensor program and analyzed under Gaussian matrices with entries A_ij ∼N(0, 1/n) and n/N →δ ∈(0, ∞).The resulting corollary applies to sequences whose empirical distributions converge weakly to Gaussian limits.
Comparison with Prior AMP Theory: Compared with Bayati and Montanari, the result allows nonlinearities and test functions without smoothness assumptions but requires Gaussian empirical limits for x_0 and w.The Gaussian-limit requirement is explicitly identified as a stronger assumption than finite bounded moments.
Technical Tools: The appendix develops Gaussian-conditioning and pseudoinverse tools for singular matrices, including projections onto column and null spaces.These lemmas support conditioning random Gaussian matrices on linear constraints.

F. Proof of Main Theorems

The proofs establish tensor-program scaling limits by induction, conditioning Gaussian matrix multiplications on prior variables and tracking convergent covariance and projection quantities. Rank convergence and growth conditions control the almost-sure limits.

Main Induction: The induction is straightforward for linear combinations and nonlinearities, while matrix multiplication requires conditioning on all previously generated Gaussian variables.The conditioned matrix is decomposed using an independent copy and projections onto prior spans.
Conditional Gaussian Limit: Conditioned coordinates of a new matrix product become asymptotically Gaussian with conditional means and variances determined by covariance and projection terms.The variance limit is K_c(g, g) − K_c(g, G)K_c(G, G)+K_c(G, g).
Limit Control: The proof establishes convergence of empirical moments by applying induction hypotheses, rank convergence, and strong laws to polynomially or α-controlled functions.Uniform boundedness arguments permit the required limit exchanges and almost-sure convergence.
Assumptions: Transpose-containing programs require stronger polynomial-boundedness assumptions because compositions of generally α-controlled functions need not be Gaussian-integrable.This is explicitly presented as a more stringent condition than in Theorem 4.3.
Main Induction: Theorem 6.3 proves almost-sure convergence for polynomially bounded test functions when the detransposed program satisfies almost-sure rank convergence.The proof proceeds by induction over program lines after sampling inputs according to the prescribed scheme.