🧠 Infinite-Width Limits: The Lazy/Rich Dichotomy

Table of Contents

• 1. Setup
• 2. The Neural Tangent Kernel
  • 2.1 Definition and Gradient Flow
  • 2.2 Infinite-Width Limit of the NTK
• 3. The Lazy Regime
  • 3.1 Why Features Freeze
  • 3.2 Frozen NTK and Linear Dynamics
  • 3.3 Equivalence to Kernel Regression
• 4. The Rich Regime
  • 4.1 Mean-Field Scaling and Breakdown of Linearization
  • 4.2 The McKean-Vlasov Limit
  • 4.3 Feature Learning as Evolution of the Weight Distribution
• 5. The Output Multiplier as Control Parameter
• 6. The Maximal Update Parameterization
  • 6.1 Deep Networks and the Tensor Programs Framework
  • 6.2 The Two Non-Trivial Limits
  • 6.3 Hyperparameter Transfer
• 7. Summary Comparison
• References

1. Setup 🎯

Consider a shallow neural network of width \(N\): \[f(x;\theta) = \frac{\alpha}{\sqrt{N}} \sum_{i=1}^N a_i \,\sigma(w_i^\top x), \qquad \theta = \{a_i, w_i\}_{i=1}^N\]

where \(x \in \mathbb{R}^d\), \(a_i \in \mathbb{R}\), \(w_i \in \mathbb{R}^d\), \(\sigma\) is a pointwise nonlinearity (e.g. ReLU), and \(\alpha > 0\) is an output multiplier — the key parameter controlling the lazy/rich dichotomy.

Initialize: \[a_i \overset{\text{iid}}{\sim} \mathcal{N}(0,1), \qquad w_i \overset{\text{iid}}{\sim} \mathcal{N}(0, I_d).\]

Train by gradient flow on a dataset \(\{(x_\mu, y_\mu)\}_{\mu=1}^n\): \[\dot{\theta}(t) = -\nabla_\theta L(\theta(t)), \qquad L(\theta) = \frac{1}{n}\sum_{\mu=1}^n \ell\!\left(f(x_\mu;\theta),\, y_\mu\right).\]

Goal. Understand the \(N \to \infty\) limit of these dynamics, and how \(\alpha\) governs qualitatively distinct dynamical regimes.

2. The Neural Tangent Kernel 🔑

2.1 Definition and Gradient Flow

Definition (Neural Tangent Kernel). For a network \(f(x;\theta)\), define: \[K_\theta(x,x') \;:=\; \bigl\langle \nabla_\theta f(x;\theta),\, \nabla_\theta f(x';\theta) \bigr\rangle \;=\; \sum_k \frac{\partial f(x;\theta)}{\partial \theta_k} \frac{\partial f(x';\theta)}{\partial \theta_k}.\]

Under gradient flow, the output at any input \(x\) evolves as: \[\dot{f}(x;\theta(t)) = \nabla_\theta f(x;\theta(t))^\top \dot{\theta}(t) = -\sum_\mu K_{\theta(t)}(x,x_\mu)\,\frac{\partial L}{\partial f(x_\mu)}.\]

This is a coupled ODE in \(f\)-space, driven by the Gram matrix \(K_{\theta(t)}\) evaluated on the training set. The fundamental question is: does \(K_{\theta(t)}\) change during training, or is it frozen?

2.2 Infinite-Width Limit of the NTK

For our shallow network, compute the partial derivatives: \[\nabla_{a_i} f(x;\theta) = \frac{\alpha}{\sqrt{N}} \sigma(w_i^\top x), \qquad \nabla_{w_i} f(x;\theta) = \frac{\alpha}{\sqrt{N}} a_i\, \sigma'(w_i^\top x)\, x.\]

The NTK is therefore: \[K_\theta(x,x') = \frac{\alpha^2}{N} \sum_{i=1}^N \Bigl[ a_i^2\, \sigma'(w_i^\top x)\,\sigma'(w_i^\top x')\,\langle x,x'\rangle + \sigma(w_i^\top x)\,\sigma(w_i^\top x') \Bigr].\]

Proposition (NTK convergence). As \(N \to \infty\), by the law of large numbers: \[K_\theta(x,x') \xrightarrow{\;\mathrm{a.s.}\;} \alpha^2\, K^\infty(x,x'),\]

where \(K^\infty\) is the deterministic limiting kernel: \[K^\infty(x,x') = \underbrace{\mathbb{E}_{w}\!\left[\sigma'(w^\top x)\,\sigma'(w^\top x')\right]\langle x,x'\rangle}_{\text{hidden-layer contribution}} + \underbrace{\mathbb{E}_{w}\!\left[\sigma(w^\top x)\,\sigma(w^\top x')\right]}_{\text{output-layer contribution}}\]

with \(w \sim \mathcal{N}(0,I_d)\) and \(a \sim \mathcal{N}(0,1)\) (note \(\mathbb{E}[a^2]=1\)).

3. The Lazy Regime 😴

3.1 Why Features Freeze

The lazy regime corresponds to \(\alpha = O(1)\) — i.e., standard LeCun initialization with no extra output scaling.

Compute the gradient flow equation for the feature parameters \(w_i\): \[\dot{w}_i(t) = -\nabla_{w_i} L = -\frac{\alpha}{\sqrt{N}} a_i(t) \sum_\mu \frac{\partial \ell}{\partial f(x_\mu)} \sigma'(w_i(t)^\top x_\mu)\, x_\mu.\]

Since \(a_i = O(1)\), \(\frac{\partial \ell}{\partial f} = O(1)\), and \(\sigma' = O(1)\), we have \(\|\dot{w}_i\|_2 = O(\alpha/\sqrt{N})\).

After \(T\) time units: \[\|w_i(T) - w_i(0)\|_2 \;\leq\; \int_0^T \|\dot{w}_i(t)\|_2\,dt \;=\; O\!\left(\frac{\alpha T}{\sqrt{N}}\right) \;\xrightarrow{N\to\infty}\; 0.\]

The same calculation holds for \(a_i\). In particular, the change in hidden activation at any input \(x\): \[|\Delta h_i(x)| \;:=\; |\sigma(w_i(T)^\top x) - \sigma(w_i(0)^\top x)| \;\leq\; \|\sigma'\|_\infty\, \|\Delta w_i\|\,\|x\| \;=\; O\!\left(\frac{\alpha T}{\sqrt{N}}\right).\]

In the \(N \to \infty\) limit, hidden representations are frozen at their initialization values. This is the defining property of the lazy regime.

3.2 Frozen NTK and Linear Dynamics

Since \(w_i(t) \approx w_i(0)\) and \(a_i(t) \approx a_i(0)\) for all \(t \leq T\), the NTK barely changes: \[K_{\theta(t)}(x,x') = \alpha^2 K^\infty(x,x') + O\!\left(\frac{1}{\sqrt{N}}\right).\]

Substituting into the gradient flow equation for \(f\): \[\dot{f}(x;\theta(t)) \approx -\alpha^2 \sum_\mu K^\infty(x,x_\mu)\,\frac{\partial L}{\partial f(x_\mu)}.\]

For squared loss \(\ell(f,y) = \frac{1}{2}(f-y)^2\), define the residual vector \(r(t) \in \mathbb{R}^n\) with \((r)_\mu = f(x_\mu;t) - y_\mu\). Then: \[\dot{r}(t) = -\alpha^2 \mathbf{K}\, r(t), \qquad \mathbf{K}_{\mu\nu} = K^\infty(x_\mu, x_\nu).\]

This is a linear ODE with solution: \[r(t) = e^{-\alpha^2 \mathbf{K} t}\, r(0).\]

The loss decays exponentially at rate determined by the smallest eigenvalue \(\lambda_{\min}(\mathbf{K})\): \[L(\theta(t)) = \frac{1}{2}\|r(t)\|^2 \sim e^{-2\alpha^2 \lambda_{\min}(\mathbf{K})\, t} \cdot L(\theta(0)).\]

3.3 Equivalence to Kernel Regression

As \(t \to \infty\) under the lazy dynamics, assuming \(\mathbf{K} \succ 0\): \[f(x;\theta(\infty)) = K^\infty(x, \mathbf{x})\,[\mathbf{K}]^{-1} y,\]

where \(K^\infty(x,\mathbf{x}) \in \mathbb{R}^n\) has entries \(K^\infty(x,x_\mu)\). This is exactly kernel regression with kernel \(K^\infty\).

The NTK regime thus reduces the problem of training a neural network to a kernel method with a fixed architecture-dependent kernel. This is exact in the \(N \to \infty\) limit and explains why the lazy regime cannot exhibit feature learning: the network’s effective kernel \(K^\infty\) is determined entirely by the initialization and architecture, and never adapts to the data.

4. The Rich Regime 🌊

4.1 Mean-Field Scaling and Breakdown of Linearization

The rich regime arises when we downscale the output by an extra factor of \(1/\sqrt{N}\), changing the parameterization to: \[f(x;\theta) = \frac{1}{N}\sum_{i=1}^N \Phi(\theta_i, x), \qquad \Phi(\theta_i, x) := a_i\,\sigma(w_i^\top x).\]

Compare with the lazy case: the prefactor is \(1/N\) instead of \(1/\sqrt{N}\), so \(f(x;\theta_0) = O(1/\sqrt{N}) \to 0\) at initialization — the output starts near zero and must grow during training.

The NTK for this parameterization: \[K^{MF}_\theta(x,x') = \frac{1}{N^2}\sum_{i=1}^N \bigl\langle \nabla_{\theta_i}\Phi(\theta_i,x),\, \nabla_{\theta_i}\Phi(\theta_i,x') \bigr\rangle = O\!\left(\frac{1}{N}\right) \xrightarrow{N\to\infty} 0.\]

Critically, the NTK vanishes in the \(N \to \infty\) limit. This means the linearization breaks down — the network cannot be approximated by its first-order Taylor expansion. One must deal with the full nonlinear dynamics, which is why feature learning can occur.

4.2 The McKean-Vlasov Limit

The correct \(N \to \infty\) limit for the mean-field network is a distributional description. Define the empirical measure over the \(N\) particles: \[\mu^N_t \;:=\; \frac{1}{N}\sum_{i=1}^N \delta_{\theta_i(t)} \;\in\; \mathcal{P}(\mathbb{R}^{d+1}).\]

The mean-field output is: \[f(x;\mu^N_t) = \int \Phi(\theta, x)\,d\mu^N_t(\theta).\]

Each particle evolves under gradient flow: \[\dot{\theta}_i(t) = -\nabla_{\theta_i} L = -\mathbb{E}_{x,y}\!\left[\frac{\partial \ell}{\partial f}\bigl(f(x;\mu^N_t), y\bigr)\,\nabla_{\theta_i}\Phi(\theta_i,x)\right].\]

Theorem (Mean-Field Limit — Mei et al. 2018, Chizat & Bach 2018). As \(N \to \infty\), \(\mu^N_t \rightharpoonup \mu_t\) weakly, where \(\mu_t\) satisfies the McKean-Vlasov PDE: \[\partial_t \mu_t + \nabla_\theta\cdot\!\left(\mu_t\,\mathbf{v}_t\right) = 0,\]

\[\mathbf{v}_t(\theta) \;:=\; -\,\mathbb{E}_{(x,y)\sim\mathcal{D}}\!\left[\frac{\partial \ell}{\partial f}\bigl(f(x;\mu_t),\,y\bigr)\,\nabla_\theta\Phi(\theta,x)\right],\]

\[f(x;\mu_t) \;:=\; \int \Phi(\theta,x)\,d\mu_t(\theta).\]

This is a continuity equation — mass is conserved — with a velocity field \(\mathbf{v}_t\) that depends nonlinearly on \(\mu_t\) itself (through the output \(f(x;\mu_t)\)). The self-referential dependence is the “McKean” part.

4.3 Feature Learning as Evolution of the Weight Distribution

📐 The key quantity distinguishing the two regimes is the marginal distribution over first-layer features: \[\rho_t(w) \;:=\; \int_{\mathbb{R}} \mu_t(a, w)\,da.\]

This is the precise mathematical definition of feature learning: the marginal weight distribution \(\rho_t\) changes non-trivially during training.

5. The Output Multiplier as Control Parameter 🎛️

The lazy and rich regimes are unified by the output multiplier \(\alpha\). Write: \[f(x;\theta) = \frac{\alpha}{\sqrt{N}}\sum_{i=1}^N a_i\,\sigma(w_i^\top x).\]

The gradient of \(L\) w.r.t. \(w_i\) is proportional to \(\alpha/\sqrt{N}\), so cumulative feature change after time \(T\): \[\|w_i(T) - w_i(0)\|_2 = O\!\left(\frac{\alpha T}{\sqrt{N}}\right).\]

Proposition (Chizat et al. 2019). In the \(N \to \infty\) limit: - Lazy (\(\alpha = O(1)\)): the NTK is frozen, training dynamics are linear, and the network converges to a kernel predictor. - Rich (\(\alpha = \Theta(\sqrt{N})\)): the NTK changes during training, linearization fails, and the network exhibits feature learning described by the McKean-Vlasov PDE.

Proof sketch. In the lazy case, the gradient satisfies \(\|\nabla_\theta L\| = O(\alpha/\sqrt{N})\), so parameter changes \(\|\Delta\theta\| = O(\alpha T/\sqrt{N}) \to 0\). Since \(K_\theta\) is a smooth function of \(\theta\), \(\|K_{\theta(t)} - K_{\theta(0)}\| \to 0\) uniformly, and the NTK is frozen. Linearizing around \(\theta(0)\) then introduces only \(O(\|\Delta\theta\|^2) = O(\alpha^2 T^2/N) \to 0\) error.

In the rich case, write \(\alpha = \beta\sqrt{N}\). Then \(f = \frac{\beta}{N}\sum_i a_i \sigma(w_i^\top x)\), which is precisely the mean-field parameterization, and the McKean-Vlasov limit applies. \(\square\)

6. The Maximal Update Parameterization 📐

6.1 Deep Networks and the Tensor Programs Framework

For deep networks, one cannot simply count \(1/\sqrt{N}\) vs. \(1/N\) prefactors — the multi-layer structure creates a web of interdependencies between initialization variances and learning rates. The Tensor Programs framework (Yang & Hu 2021) resolves this systematically.

Consider a depth-\(L\) MLP with hidden width \(N\): \[h^0 = x, \quad z^\ell = W^\ell h^{\ell-1}, \quad h^\ell = \sigma(z^\ell), \quad f = v^\top h^L,\]

where \(W^\ell \in \mathbb{R}^{N \times N}\) and \(v \in \mathbb{R}^N\). Write the hyperparameters as: \[W^\ell_{ij} \sim \mathcal{N}(0,\, \sigma^2_\ell / N), \qquad v_i \sim \mathcal{N}(0, \sigma^2_{\text{out}}/N), \qquad \eta_\ell = \eta \cdot N^{c_\ell},\]

where \(\sigma^2_\ell\) and \(c_\ell\) are per-layer constants to be chosen.

The question: For which exponents \(c_\ell\) does the network exhibit non-trivial, non-explosive training dynamics as \(N \to \infty\)?

The Tensor Programs answer this by tracking the \(N\)-scaling of every intermediate quantity (pre-activations, activations, backpropagated signals, weight updates) through the forward and backward passes, using a set of master equations called tensor programs.

6.2 The Two Non-Trivial Limits

The framework shows that — up to degenerate cases — there are exactly two consistent, non-explosive parameterizations that survive the \(N \to \infty\) limit with non-trivial dynamics:

The name “Maximal Update Parameterization” refers to the fact that \(\mu P\) is the unique parameterization where every layer achieves the maximum possible contribution to the output change without causing divergence.

6.3 Hyperparameter Transfer

💡 The most practically important consequence of μP is hyperparameter transfer across widths.

Claim: In μP, the optimal learning rate \(\eta^*\) is approximately independent of width \(N\) for large enough \(N\).

Why? Under μP, the per-step output change is \(O(\eta)\) regardless of \(N\) (by construction — this is what “maximal updates” means). So the effective learning rate as experienced by the loss is \(O(\eta)\) at all widths, meaning the optimal \(\eta\) doesn’t need to be re-tuned when \(N\) changes.

Under NTP, the per-step output change is \(O(\eta/\sqrt{N})\), so the optimal \(\eta\) must scale as \(\sqrt{N}\) — it is width-dependent.

Practical consequence: Tune \(\eta\) on a small proxy model of width \(N_0 \ll N_\text{target}\) (cheap), then deploy the same \(\eta\) on the full model. This reduces the cost of hyperparameter search by orders of magnitude for large models.

7. Summary Comparison 📊

References