PaperlayerView exact revision on arXiv

Source-linked AI summary

Tensor Programs II: Neural Tangent Kernel for Any Architecture

Greg Yang

arXiv:2006.14548v4stat.MLcond-mat.dis-nncs.LGcs.NE

TL;DR

The paper addresses whether modern neural networks have computable infinite-width NTKs and develops a tensor-program framework for answering that question. It proves deterministic convergence for standard architectures, characterizes when GIA is valid through Simple GIA Check, and shows that GIA can fail outside that condition.

  • Problem

    It is unclear whether modern architectures have valid, computable infinite-width NTKs and whether the Gradient Independence Assumption gives correct gradient statistics.

  • Method

    The paper extends tensor programs to forward and backward computations, decomposes NTKs into computable intermediate kernels, and uses Simple GIA Check to justify GIA when applicable.

  • Results

    The NTK of any randomly initialized standard-architecture network converges almost surely to a deterministic infinite-width limit, with concrete algorithms agreeing with simulations for batchnorm-ReLU MLPs, transformers, and RNNs.

  • Takeaways & Limitations

    The universal NTK correspondence provides a route to studying the inductive biases of neural networks with varied architectures under SGD.

  • Takeaways & Limitations

    The convergence theorem requires the BP-like condition and does not cover non-BP-like cases, which the authors defer to future work.

Abstract

from arXiv · show

We prove that a randomly initialized neural network of *any architecture* has its Tangent Kernel (NTK) converge to a deterministic limit, as the network widths tend to infinity. We demonstrate how to calculate this limit. In prior literature, the heuristic study of neural network gradients often assumes every weight matrix used in forward propagation is independent from its transpose used in backpropagation (Schoenholz et al. 2017). This is known as the *gradient independence assumption (GIA)*. We identify a commonly satisfied condition, which we call *Simple GIA Check*, such that the NTK limit calculation based on GIA is correct. Conversely, when Simple GIA Check fails, we show GIA can result in wrong answers. Our material here presents the NTK results of Yang (2019a) in a friendly manner and showcases the *tensor programs* technique for understanding wide neural networks. We provide reference implementations of infinite-width NTKs of recurrent neural network, transformer, and batch normalization at https://github.com/thegregyang/NTK4A.

1 Introduction

The paper asks whether modern neural networks have computable infinite-width NTKs and proves deterministic convergence for standard architectures. It also identifies when gradient-independence calculations are valid and provides computation procedures and implementations.

  • The paper asks whether every modern neural network has an infinite-width NTK and whether that kernel can be computed.
  • The NTK of any randomly initialized neural network with standard architecture converges almost surely to a deterministic limit as widths tend to infinity.Standard architectures include compositions involving MLPs, RNNs, skip connections, convolutions, pooling, normalization, and attention.
  • The general result applies when forward and backpropagation can be expressed through nonlinearities and matrix multiplication.
  • Simple GIA Check rigorously justifies the Gradient Independence Assumption, while calculations based on GIA can be incorrect when the condition fails.
  • Concrete infinite-width NTK algorithms for batchnorm-ReLU MLPs, transformers, and RNNs agree with simulations across widths 2^6 through 2^13.The plotted uncertainty is a 95% confidence interval over 100 random seeds.
  • NETSOR⊤ extends tensor programs to backpropagation, enabling analysis of gradients and NTKs for architectures beyond forward-only programs.
  • The resulting universal NTK correspondence opens a way to study inductive biases associated with architectures trained under SGD.

2 Background

The background introduces NTK as the kernel of neural-network parameter gradients and reviews its infinite-width training interpretation. In that limit, networks can follow fixed-kernel dynamics and fit training data under stated conditions.

  • For randomly initialized networks, the finite-width NTK converges to a deterministic kernel as widths grow to infinity.
  • At large width, gradient-descent training converges to dynamics governed by the limiting NTK for any fixed training-time interval.
  • The resulting network evolution is equivalent to kernel gradient descent with a fixed kernel, and square loss yields a linear differential equation.
  • Because the limiting NTK is generally nonsingular, a sufficiently wide network can fit any training data under this description.
  • The infinite-width NTK reflects an implicit prior determined by gradient descent, architecture, and initialization, with its spectrum indicating which functions can be learned quickly and generalize well.

3 Related Works

Related work connects NTKs to Gaussian-process limits, signal propagation, and analyses of optimization and generalization. Prior studies also examined specialized architectures and finite-width or higher-order effects.

  • The NTK characterizes training the full network by gradient descent, whereas the infinite-width neural-network Gaussian process characterizes training only the last layer.
  • Earlier work established architectural universality for NTK and NN-GP limits, followed by simplified presentations and specialized studies of particular architectures.
  • Signal-propagation research studies pathological behavior in deep random networks, with forward propagation linked to Gaussian-process limits and backward propagation linked to NTK limits.
  • NTKs have been used to study how neural networks train despite nonconvexity, including convergence questions and spectral behavior.
  • Kernel regression with corresponding NTKs can perform better than neural networks in low-data regimes, including settings such as medical data with few samples.

4 Warmup: Neural Tangent Kernel for a Multi-Layer Perceptron

The MLP warmup decomposes the NTK into forward and backward kernel factors, then uses wide-network Gaussian heuristics and gradient independence to derive recursive limits. The resulting expression recovers the established MLP NTK formula.

  • MLP setup: The MLP is expressed through recursively defined affine transformations and nonlinearities, with NTK-scaled factorized weight matrices.
  • Decomposing NTK: The NTK decomposes into contributions from weights and biases, allowing convergence to be established by analyzing each summand.
  • Forward limits: Forward quantities are treated as approximately Gaussian and coordinate-wise independent in the infinite-width limit, producing deterministic covariance kernels C_l.
  • Backward limits: Under GIA, backward quantities are likewise approximated as Gaussian and independent across coordinates, yielding recursive derivative kernels D_l.
  • Backward limits: The derivative-kernel recursion multiplies the next-layer kernel by the expected product of activation derivatives.
  • NTK limit: Combining the recurrences recovers the NTK limit formula previously derived for MLPs.

5 NTK for Any Architecture? The Issues and the Proposal

The paper addresses whether NTK analysis can extend to modern architectures, whether tied weights invalidate gradient independence, and whether these cases admit a unified treatment. It answers these questions with a decomposition into convergent terms and a checkable condition for GIA.

  • Issues: The paper asks whether tied weights preserve a finite, convergent NTK decomposition, whether GIA remains valid, and whether modern architectures can be handled systematically.These questions concern weight-tying correlations, transpose dependence in backpropagation, and architectural diversity.
  • Proposal: The NTK decomposes into sums of products of inner products that converge to deterministic finite constants as width grows.This generalizes the MLP decomposition and implies NTK convergence.
  • Proposal: Simple GIA Check requires output-layer weights to be independently sampled, zero-mean, and unused elsewhere in the network.Under this condition, GIA is justified even when hidden-to-hidden weights are correlated through reuse.
  • Caveat: When the condition fails, GIA can produce incorrect calculations because forward weights and their transposes may interact through unblocked paths.The paper gives a counterexample where violating the condition makes GIA fail.
  • Proposal: NETSOR⊤ expresses forward and backward computations using matrix multiplication and nonlinearities, enabling recursive computation of the infinite-width NTK under suitable conditions.The paper states this applies to standard architectures and yields provable convergence through Corollary 7.3.

6 Strategy for Computing the Infinite-Width NTK

The paper computes infinite-width NTKs by decomposing gradient contributions into forward and backward kernel limits and recursively evaluating them with Gaussian rules. It also demonstrates why the GIA-based rules require a structural condition.

  • General strategy: For general architectures, the NTK is computed by summing contributions from every matrix multiplication used in each network computation.The decomposition includes factorized weights and biases and compares two inputs.
  • General strategy: Forward and backward inner products converge almost surely to deterministic limits, whose products determine the limiting NTK.The limits are denoted by D and C terms for gradients and activations.
  • Kernel rules: At large width, coordinatewise nonlinearities and matrix multiplications induce recursive Gaussian rules for computing the intermediate kernels C and D.The rules treat coordinates as approximately iid while retaining correlations across vectors.
  • RNN example: For an RNN, forward and backward recursions over states and derivatives combine with the NTK decomposition to compute the infinite-width kernel.The construction applies the rules to recurrent weight sharing and input sequences.
  • GIA failure: Without the structural condition, treating W^⊤ as independent from W can predict a zero-mean backward quantity when the actual mean is nonzero.The two-hidden-layer example shows that GIA can fail when the output averages the final embedding.

7 NETSOR⊤

NETSOR⊤ is a language for expressing both forward and backward neural-network computations. Its Master Theorem rigorously establishes Gaussian-limit behavior that supports almost-sure NTK convergence for standard architectures under stated regularity conditions.

  • Language: NETSOR⊤ expresses neural-network forward and backward propagation through nonlinearities and matrix multiplication, including transposed matrix operations.It extends NETSOR, which represented forward propagation alone.
  • Master Theorem: The Master Theorem applies to BP-like programs with Gaussian matrix initialization and polynomially bounded nonlinearities.It provides the regularity framework needed for rigorous infinite-width limits.
  • Master Theorem: For polynomially bounded test functions, empirical coordinate averages converge almost surely to expectations under the Gaussian variables defined by the program.This result formalizes the recursive intuitions used to compute kernels.
  • NTK convergence: Expressing a network in NETSOR⊤, checking BP-like structure, and checking polynomially bounded activations and derivatives suffices to prove NTK convergence.The resulting limit is computed using the NTK decomposition and the kernel rules.
  • NTK convergence: Corollary 7.3 states that standard feedforward or recurrent networks satisfying Condition 1 have NTKs converging almost surely to deterministic kernels at infinite width.The result assumes polynomially bounded weak derivatives and Gaussian NTK initialization.
  • Scope: The theorem does not cover valid non-BP-like programs, whose treatment is deferred to future work requiring additional machinery.The paper also extends the framework to variable dimensions and discusses stronger languages for architectural blocks such as attention.

C Programs with Variable Dimensions

The variable-dimension extension organizes program variables into Common Dimension Classes and lets dimensions grow with finite, nonzero pairwise ratios. Within each class, empirical averages converge to Gaussian expectations, while cross-class convergence requires care.

  • Common Dimension Classes: Common Dimension Classes group variables whose dimensions are constrained to be equal by program-generation rules or input equivalence.Different classes may have different dimensions.
  • Scaling limit: The variable-dimension limit sends all dimensions to infinity while their pairwise ratios converge to finite, nonzero values.This generalizes the equal-width setting.
  • Convergence theorem: For each Common Dimension Class, polynomially bounded functions of its variables converge almost surely to expectations under a corresponding Gaussian distribution.The theorem accommodates Gaussian initial vectors with class-specific means and covariances.
  • Limitation: A naive convergence statement across different Common Dimension Classes is unavailable because their dimensions need not match and no common empirical average can be formed.The paper heuristically treats such vectors as approximately independent because their main randomness comes from different matrix variables.
  • Extensions: The variable-dimension Master Theorem extends the same convergence principle to parameter-controlled functions and NETSOR⊤+ programs.The extension also requires the stated rank-stability and parameter-control conditions.

D Writing Backpropagation of Standard Architectures in NETSOR⊤

NETSOR⊤ extends NETSOR with matrix transposes so forward and backward computations of standard architectures can be represented and analyzed. The framework expresses convolutions, pooling, batch normalization, recurrent networks, layer normalization, and attention, enabling systematic NTK calculations.

  • Language extension: NETSOR⊤ extends NETSOR to express backpropagation alongside forward propagation, allowing gradient and NTK analysis in the infinite-width limit.NETSOR handles forward propagation, whereas NETSOR⊤ also represents matrix transposes used during backpropagation.
  • Convolution: Convolution and graph convolution can be represented as sums of weight-shared dense matrix multiplications, with corresponding forward and backward NETSOR⊤ programs.The construction supports common padding choices and generalizes to strides, dilations, and graph convolutions.
  • Batch normalization: Batch normalization is represented by coordinatewise batchwise nonlinearities, with an analogous coordinatewise backpropagation operation.The batch size is fixed while width tends to infinity.
  • Recurrent networks: GRU and LSTM computations fit NETSOR⊤ because they combine recurrent dense matrix multiplications with coordinatewise nonlinearities.The GRU forward equations and gradient propagation use only MatMul and Nonlin operations.
  • Layer normalization and attention: Layer normalization and attention require the more expressive NETSOR⊤+ extension to represent scalar means, variances, and related operations.The extension is introduced to express architectural blocks not covered by basic NETSOR⊤.
  • NTK calculation tools: The V-transform maps covariance matrices through multivariate or coordinatewise nonlinearities and supplies reusable formulas for NTK calculations.For a Gaussian input z with covariance K, VΦ is defined through E[Φ(z)Φ(z)⊤].

E.1 MLP

The appendix applies the NETSOR⊤ Master Theorem to compute infinite-width NTKs, first illustrating the calculation for an MLP and then extending the same framework to recurrent and convolutional architectures. The resulting limits are obtained from deterministic forward and backward recursions.

  • MLP NTK: The formal MLP calculation computes the infinite-width NTK by evaluating limits of the squared norms that compose the kernel.Theorem A.6 supplies the limits needed for these norm terms.
  • MLP NTK: The MLP example uses a ReLU network with equal hidden widths and applies the theorem because ReLU and its derivative are polynomially bounded.The argument does not require the nonlinearities to be smooth.
  • MLP NTK: Theorem 7.2 recovers the MLP kernel calculation rigorously, including terms involving matrix transposes that cannot be handled by the forward-only NETSOR theorem.The transpose-dependent limits arise in the gradient contributions to the NTK.
  • Recurrent NTK: For recurrent networks, the NTK is computed from forward and backward recursions, with distinct initial conditions for last-state and average-state outputs.Last-state outputs simplify the double sum because most backward variables vanish by induction, whereas average pooling retains more terms.
  • Convolutional NTK: For convolutional networks, deterministic covariance and gradient scalars satisfy layerwise recurrences derived from the NETSOR⊤ Master Theorem.These recurrences support both global average pooling and vectorized outputs.
  • General result: The NETSOR⊤+ construction yields a well-defined almost-sure NTK limit for supported architectures, summarized by vectorized formulas.The program’s nonlinearities are parameter-controlled and satisfy the required assumption.

G Proof of Main Theorem

The proof establishes a BP-like NETSOR⊤ Master Theorem: polynomially bounded functions of program variables converge almost surely to expectations under a corresponding Gaussian limit. Its induction requires rank and zero stability to control covariance degeneracies and pseudoinverses.

  • Theorem: For a BP-like NETSOR⊤ program with polynomially bounded nonlinearities, polynomially bounded observables converge almost surely to Gaussian expectations determined by their limiting mean and covariance.This is the content of Theorem A.6, equivalent to Theorem 7.2.
  • Proof organization: The proof proceeds by simultaneous induction on the moment statement and a CoreSet statement rather than by induction on program variables alone.The inductive organization yields Moments(M), which is Theorem A.6.
  • CoreSet: CoreSet provides basis and null-avoidance properties that ensure asymptotic linear independence and prevent concentration on Lebesgue measure zero sets.These properties support the rank and conditioning arguments used in the induction.
  • Rank stability: Rank stability is needed because convergence of pseudoinverses is guaranteed when the covariance matrices retain fixed rank in the limit.The proof obtains rank stability through Lemma G.6 and uses it to establish pseudoinverse convergence.
  • Induction: The input-variable base cases follow from the law of large numbers, while the inductive case uses Gaussian conditioning and the established stability properties.The base CoreSet can be chosen from a subset preserving the covariance rank.

G.2 Inductive Case: Setup

The inductive case analyzes a new variable formed by multiplying a random matrix with a function of earlier variables, while accounting for both the matrix and its transpose. Conditioning produces projection terms and requires rank-stable pseudoinverses.

  • Inductive setup: The new G-variable is formed as g = Ah, while earlier uses of A and A⊤ produce constrained variables ˆG and ˇG.The proof conditions on all preceding G-vars and tracks both forward and transpose-related constraints.
  • Conditional representation: Conditioning on previous variables constrains A through ˆG = AˆH and ˇG = A⊤ˇH, introducing projections onto the relevant column spaces.An independent copy of A and projection operators are used in the conditional representation.
  • Transpose dependence: Unlike the forward-only NETSOR theorem, transpose operations prevent treating the new variable as having iid Gaussian coordinates because of an orthogonal projection term.This correlation is a central source of additional proof complexity.
  • BP-like condition: The BP-like assumption makes the conditional correction approximately ˆGˆH+, a relation that generally fails without BP-like structure.This identifies why the theorem’s inductive argument depends on the BP-like condition.
  • Stability: Induction supplies convergence of the covariance quantities, conditional parameters, and pseudoinverses required for the next moment calculation.Fixed-rank convergence gives ˆΛ+ → ˆΛ̊+ and ˇΛ+ → ˇΛ̊+ almost surely.

G.3 Rank Stability and Zero Stability

This section establishes rank and zero stability for limiting kernel matrices. Under the stated basis and NullAvoid properties, kernels, images, and ranks agree almost surely for all sufficiently large widths.

  • Lemma G.6 establishes rank stability for kernel matrices generated by polynomially bounded functions.
  • Almost surely, ker K = ker ˚ K, im K = im ˚ K, and rank K = rank ˚ K for sufficiently large n.
  • Zero Stability follows by applying NullAvoid to the measure-zero set where a nonnegative function vanishes, preserving the required core-set properties.
  • The proof uses a basis of the limiting kernel's null space and compares quadratic forms of finite-width and limiting matrices.
  • Proposition G.8 shows that the core-set covariance submatrix has full rank and that the core variables span the limiting Gaussian variables.
  • When the limiting residual variance vanishes, the corresponding function is almost surely a linear combination of the other functions.

G.4.2 If ˚σ > 0.

For positive limiting conditional variance, the argument preserves NullAvoid while controlling Gaussian-conditioned quantities. Gaussian averaging and moment bounds then yield the required almost-sure convergence.

  • The proof enlarges the core set with g when it cannot asymptotically lie in the span of the existing core variables.
  • Positive residual variance gives the conditional distribution of gm a density, allowing measure-zero slice arguments to establish NullAvoid.
  • The Sherman–Morrison argument and eigenvalue control show the relevant covariance submatrices remain nonsingular.
  • The proof decomposes the target error into terms A, B, and C, showing each converges almost surely to zero.
  • Conditioning expresses the final Gaussian coordinate as a linear prediction plus independent Gaussian noise with variance ˚σ.
  • Amenable sequences and moment bounds provide almost-sure uniform control needed to apply the induction hypothesis after conditioning.

G.5.3 B Converges Almost Surely to 0

This section proves that term B converges almost surely to zero. The proof combines a partition by projection diagonal size with Gaussian-smoothing bounds and polynomial-growth control.

  • The argument requires polynomially bounded nonlinearities so their compositions remain integrable against the Gaussian measure.
  • The proof partitions indices into U, where dα < 1/2, and its complement V, with U uniformly bounded in size.
  • The contribution from Eq. (49) converges almost surely to zero by bounding its terms separately.
  • Gaussian averaging makes Ψα smooth, enabling Eq. (50) to be shown almost surely o(1/n).
  • The assumption ˚σ > 0 keeps inverse-variance factors uniformly bounded as σ converges almost surely to ˚σ.
  • Together, the bounds on the decomposed terms finish the proof that B converges almost surely to zero.
Loading 2006.14548v4…