Source-linked AI summary
Finite Depth and Width Corrections to the Neural Tangent Kernel
Boris Hanin, Mihai Nica
TL;DR
The paper addresses whether the NTK remains deterministic and frozen when network depth and width grow together. It derives finite-depth and finite-width scaling for NTK moments and its first SGD update in randomly initialized fully connected ReLU networks. The results show persistent stochasticity and non-trivial kernel evolution in deep-and-wide regimes, while the scope is limited by stability and architecture assumptions.
Problem
The joint statistical effects of depth and width on finite-size NTKs, including whether they remain deterministic and frozen during training, remain unclear.
Method
The paper analyzes moments of the NTK and its first SGD update for fully connected ReLU networks at finite depth and width using a sum-over-path approach.
Results
When d/n is bounded away from 0, the NTK is not approximately deterministic at initialization, and under square loss its mean first SGD update can be comparable in scale to the kernel itself.
Takeaways & Limitations
Deep-and-wide ReLU networks may learn data-dependent features through stochastic, evolving NTKs even in the lazy-training regime.
Takeaways & Limitations
The results are established for fully connected ReLU networks, and their extension to other nonlinearities and architectures remains an open question.
Abstract
from arXiv · showhide
We prove the precise scaling, at finite depth and width, for the mean and variance of the neural tangent kernel (NTK) in a randomly initialized ReLU network. The standard deviation is exponential in the ratio of network depth to width. Thus, even in the limit of infinite overparameterization, the NTK is not deterministic if depth and width simultaneously tend to infinity. Moreover, we prove that for such deep and wide networks, the NTK has a non-trivial evolution during training by showing that the mean of its first SGD update is also exponential in the ratio of network depth to width. This is sharp contrast to the regime where depth is fixed and network width is very large. Our results suggest that, unlike relatively shallow and wide networks, deep and wide ReLU networks are capable of learning data-dependent features even in the so-called lazy training regime.
1. Introduction
The paper studies finite-depth and finite-width corrections to the NTK in fully connected ReLU networks, focusing on how depth and width jointly affect initialization and training. It contrasts deep-and-wide networks with the fixed-depth, infinite-width regime, where the kernel becomes deterministic and remains frozen.
- Infinite-width baseline: Fixed depth with width tending to infinity makes the NTK converge to the infinite-width limit of its mean and remain frozen during training.In this regime, the limiting kernel depends on network depth and nonlinearity but not on the dataset.
- Motivation: Finite-width NTK fluctuations matter because the kernel can evolve during training and potentially learn data-dependent features absent from the infinite-width mean kernel.The paper motivates finite-width analysis partly because finite but wide networks can outperform infinite-width NTK predictions under initially large learning rates.
- Scope and approach: The paper analyzes the simultaneous effects of depth and width for fully connected ReLU networks at initialization, including finite but large widths and arbitrary depth.Its results are formulated for variable layer widths, with d/n used as the exposition-level depth-to-width parameter.
- Initialization: When d/n is bounded away from 0, the NTK is not approximately deterministic at initialization, even when depth and width are both large.For fixed input x, the standard deviation of K_N(x, x) is at least as large as its mean, so its distribution is not close to a delta function.
- Training dynamics: Under square loss, the mean first SGD update of K_N(x, x) can remain comparable in scale to the kernel itself when depth and width grow together.The stated scaling implies non-trivial, potentially data-dependent kernel evolution, unlike the fixed-depth, wide-network regime.
- Implications and future work: The paper identifies numerical stability as a boundary: existing results indicate gradient-based training is not stable for fully connected ReLU networks unless d/n is relatively small.The authors therefore conjecture a weak feature-learning NTK regime with 0 < d/n ≪ 1, while broader architectures and nonlinearities remain future work.
2. Formal Statement of Results
The paper formalizes finite-depth and finite-width NTK behavior for randomly initialized ReLU networks, including initialization statistics and the first SGD update. In the deep-and-wide scaling regime, the NTK can remain non-deterministic and evolve during training, unlike the fixed-depth, very-wide regime.
- Network and initialization: The model is a depth-d ReLU network with input dimension n0, hidden widths n1,...,n_{d−1}, scalar linear output, and zero-initialized biases.The biases remain trainable despite being initialized to zero.
- Proof strategy: The proofs use path-and-edge notation and sum-over-path expressions to compute moments of the NTK and its first time derivative.The article introduces notation, outlines the proofs, and then develops detailed arguments in later sections.
- Network and initialization: The results apply at finite depth and width, with universal implicit constants and variable layer widths.The formal setup assumes a fixed probability measure for initialization with a density and standard moment conditions.
- Initialization statistics: E[KN(x,x)]^2 ≃ exp(5β)(1 + O(β/n)), where β = d/n, for equal hidden widths ni = n.Here ≃ denotes upper and lower bounds by universal constant factors.
- Initialization statistics: In the deep-and-wide double-scaling limit, the NTK does not converge to a constant in probability.This contrasts with the regime of fixed depth and widths tending to infinity.
- Training evolution: The paper analyzes whether the on-diagonal NTK changes after one size-1 SGD update under square loss and learning rate λ.The update is denoted ΔKN(x,x) for a batch containing x.
- Training evolution: The mean first SGD update has a theorem-level scaling with multiplicative finite-width error, and for equal hidden widths its behavior is governed by depth-to-width scaling.The formal result is stated for the mean time derivative of the on-diagonal NTK at initialization.
- Training evolution: When d is fixed and n→∞, the update prefactor scales like 1/n; when β is bounded away from 0 and ∞ and n0 is proportional to d, the average update matches the NTK's order.The zero-mean conclusion for d^2/(nn0)→0 is stated only for the size-1 batch and on-diagonal setting; larger batches and off-diagonal arguments remain unclear.
3. Notation
The notation formalizes the network as a layered computational graph and represents paths, their weights, and the edge multisets they generate. These objects support later path-based moment calculations.
- Computational graph: The computational graph is a directed multipartite graph whose vertices are neurons in successive layers and whose edges connect every adjacent-layer pair.Layer ℓ contains vertices indexed by [nℓ].
- Paths: A path from layer ℓ1 to layer ℓ2 is a sequence of neuron indices, with fixed endpoints and one neuron selected in every intermediate layer.The path is written γ = (γ(ℓ1), …, γ(ℓ2)).
- Path weights: For a fixed input, a path weight combines the normalized edge weights along the path with the event that all neurons on the path are open.The associated edge weights and openness event are defined for paths ending at the output.
- Edge multisets: Collections of paths induce unordered multisets of edges, whose left and right endpoint multisets record the neurons incident to those edges.These multisets organize which edges are covered and how often in path-moment calculations.
- Even coverings: The notation also distinguishes edge multisets in which every edge appears an even number of times, a condition used in the moment expansions.The corresponding path collections are denoted with the even qualifier.
4. Overview of Proof of Theorems 1 and 2
The proof analyzes NTK moments and the first SGD update through sums over paths in the network. Path interactions, counted as loops, determine the depth–width dependence of these quantities.
- Target quantities: The proofs estimate moments of the NTK and its first SGD update, including weight, bias, and mixed bias-weight contributions.The weight-only terms are the most technically involved, while bias terms are treated similarly with paths starting inside the network.
- Path expansion: The network output and its parameter derivatives are expressed as sums over paths, with path weights containing edge-weight products and neuron-openness indicators.Almost surely, differentiating the openness indicators contributes zero.
- Moment structure: Mean calculations retain only path collections in which every edge is traversed an even number of times, because the edge weights have mean zero.For the first NTK moment, this restricts the contributing two-path collections substantially.
- Path interactions: For four-path terms, a loop occurs when all four paths pass through the same neuron in a layer, and the counting Jacobians scale with the number of loops.The proof changes variables from path collections to unordered edge multisets and then to pairs of paths.
- Expectation estimates: The layerwise normalization turns sums over path pairs into expectations for two independent paths that sample neurons uniformly in each layer.These expectations are bounded using a proposition for products of independent events.
5. Proof of Proposition 3
The proof of Proposition 3 derives the weight contributions to the NTK and its first update by enumerating four-path configurations and estimating their loop-dependent expectations. The resulting bounds are then evaluated in depth–width regimes.
- Path representation: The weight contributions to KN and ΔKN admit almost-sure sum-over-paths representations involving pairs of paths and shared edges.These representations provide the starting point for estimating their moments.
- Loop combinatorics: Loop counts organize the change of variables from four paths to unordered edge multisets and then to pairs of paths.A loop in a pair of paths occurs when the number of occupied vertices changes from one layer to the next as specified by the loop count.
- Bounding expectations: The layerwise interaction factors are bounded using independent-event estimates, yielding upper and lower bounds for the relevant expectations.The proof applies Proposition 10 to collision and path-coincidence events across layers.
- Fixed-depth behavior: For fixed depth and increasing widths, the normalized expectation of Δww relative to Kw converges to zero because the update expression retains additional 1/nℓ factors.This distinguishes the first-update scaling from the corresponding kernel moment in the fixed-depth regime.
6. Proof of Proposition 4
Proposition 4 estimates the bias contributions to the NTK and its first SGD update by adapting the weight-based path analysis. The argument reduces the needed expectations to simpler path sums.
- Bias path expansion: The bias analysis estimates E[Kb], E[Kb^2], and the corresponding first-update contributions using a sum-over-paths expression.The proof is presented as an analog of the weight-contribution argument.
- Bias-specific structure: Bias paths begin at neurons in the interior of the network rather than at the input layer, reflecting the location of the differentiated bias parameter.The number of paths from the starting neuron to the output enters the expectation.
- Four-path calculation: The key four-path calculation fixes two neurons and uses delta constraints to force paired paths through the same neuron at the relevant layer.Symmetry then removes dependence on the particular neuron index.
- Expectation reduction: The resulting inner sums reduce to expectations over pairs of paths whose subsequent neurons are chosen independently and uniformly across layers.This is the same type of path expectation used in the weight analysis.
- Conclusion: Combining the bias lemmas yields the stated bounds in Proposition 4.The final step applies the preceding path-expectation estimates.
7. Proof of Proposition 5
The proof evaluates the expectation of the NTK product by organizing path tuples and conditioning on shared neurons or edges. It then evaluates the first-update term using analogous path-counting and collision arguments, completing Proposition 5.
- Expectation of KbKw: The proof represents E[KbKw] through four-tuples of computational-graph paths with specified starting neurons.The path family Γ4(Z,Z,a) contains four paths, with γ1, γ2 starting from Z and γ3, γ4 starting from a1 and a2.
- Path counting: The path-counting terms track whether γ3 and γ4 share an edge at a layer and whether edge multisets induce collisions between layers.The indicator T_i_3,4 records shared edges, while C(E,i,j) denotes a collision between layers i and j.
- Expectation of KbKw: Conditioning on the common neuron forces γ3 and γ4 through the same neuron at layer ℓ(Z), enabling bounds on E[KbKw].The proof uses delta functions to impose the shared-neuron condition before applying the path-counting argument.
- First-update term: The proof evaluates E[Δwb] using a decomposition with a mean-zero term and analogous counting arguments, then completes Proposition 5.The derivation proceeds as in earlier lemmas and the proof of Proposition 3 before invoking the resulting expression.