ML Mathematics: NTK & Optimization¶

Job: define the neural tangent kernel (Jacot–Gabriel–Hongler) governing the infinite-width lazy-training dynamics of gradient descent, contrast the lazy / NTK regime with the feature-learning / mean-field regime (Chizat–Bach, Mei–Montanari–Mignacco), and state the over-parameterized convergence theorems (Du–Zhai–Poczos–Singh, Allen-Zhu–Li–Song) and the loss-landscape results that explain why gradient descent reaches a global minimum on a non-convex NN landscape.

Core¶

The neural tangent kernel¶

Let $f_\theta(x)$ be a fully-connected network with parameters $\theta$ (all weights and biases). The empirical NTK at initialization $\theta_0$ is

\[ K_{\theta_0}(x,x')=\langle\nabla_\theta f_{\theta_0}(x),\nabla_\theta f_{\theta_0}(x')\rangle, \]

a positive-semidefinite kernel on the input space. The NTK in the infinite-width limit is the deterministic limit

\[ K(x,x')=\lim_{m\to\infty}K_{\theta_0}^{(m)}(x,x'),\qquad\theta_0\sim\Theta_0\text{ (random init)}, \]

where $m$ is the width and the expectation is over the initialization distribution $\Theta_0$. $K$ depends on the activation $\sigma$, the depth, and the initialization scheme, but not on a specific draw — it is a universal kernel for the architecture.

Theorem (Jacot–Gabriel–Hongler 2018). In the NTK regime — width $m\to\infty$, standard initialization (e.g. $W_{ij}\sim\mathcal{N}(0,1/m)$), learning rate $\eta=O(1/m)$ so the weights barely move (the lazy regime) — the function $f_{\theta_t}(x)$ evolves under gradient descent on the empirical loss as $$ \partial_t f_t(x)=-\sum_{i=1}^n K(x,x_i)\bigl(f_t(x_i)-y_i\bigr), $$ i.e. the function values at the data points follow kernel regression with the NTK $K$. If $K$ is positive definite on the data $\{x_i\}$ (the kernel matrix $K_n=(K(x_i,x_j))\succ 0$), the training loss converges to zero exponentially at rate $e^{-\lambda_{\min}(K_n)t}$, where $\lambda_{\min}(K_n)$ is the smallest eigenvalue.

The key point: in the lazy regime the function evolves as a kernel method, with the kernel fixed at initialization. The NN has degenerated to a kernel machine, and gradient descent on the non-convex parameter landscape reduces to convex kernel regression in function space.

Lazy training vs feature learning (Chizat–Bach)¶

Theorem / Definition (Chizat–Bach 2019). In the scaling limit where the initialization scale $1/\sqrt{m}$ is replaced by $1/m$ (a larger initialization amplitude — the "mean-field" or feature-learning scaling), the weights move by $O(1)$ during training, and the infinite-width limit is a Wasserstein gradient flow on the measure over neuron parameters: $$ \partial_t \rho_t = -\nabla_\theta\cdot\bigl(\rho_t\,\nabla_\theta \Psi(\theta;\rho_t)\bigr),\qquad \Psi(\theta;\rho)=\mathbb{E}{(x,y)}\bigl[(f\rho(x)-y)\,\sigma(\theta\cdot(x,1))\bigr], $$ where $f_\rho(x)=\int\sigma(\theta\cdot(x,1))\,d\rho(\theta)$ is the population-net function. This is the feature-learning regime — the measure $\rho_t$ evolves non-trivially, the neurons learn features (their directions in parameter space move), and the dynamics is not a fixed-kernel regression.

The two regimes are distinguished by whether the weights move: lazy / NTK — weights move by $O(1/m)$, dynamics $=$ kernel regression, no feature learning, the NN is a kernel method; feature-learning / mean-field — weights move by $O(1)$, dynamics $=$ Wasserstein gradient flow, features are learned, the NN can outperform kernel methods (the Barron-class advantage of Approximation Theory is realized only here). The NTK regime is the degenerate limit where the expressivity advantage of depth is lost; the feature-learning regime is where the deep-learning advantage lives.

Regime	Scaling	Weight motion	Limit dynamics	Convergence	Feature learning
NTK / lazy (Jacot et al.)	$W_{ij}\sim\mathcal{N}(0,1/m)$, $\eta=O(1/m)$	$O(1/m)$	kernel regression with fixed $K$	exponential, $\lambda_{\min}(K_n)$	none (kernel method)
Mean-field / feature-learning (Chizat–Bach)	$W_{ij}\sim\mathcal{N}(0,1/m)$, scaled loss by $m$	$O(1)$	Wasserstein GF on neuron measure	global (under conditions); partial	yes (directions move)
Over-parameterized finite width (Du et al., Allen-Zhu et al.)	width $m\gg n$	$O(1/\sqrt{m})$	NTK approximation; lazy at large $m$	GD reaches global min, $m=\Omega(n^2/\lambda_{\min}^2)$	limited (lazy at large $m$)

Over-parameterized convergence¶

Theorem (Du–Zhai–Poczos–Singh 2018; Allen-Zhu–Li–Song 2018, over-parameterized GD). For a depth-$L$ ReLU net with width $m\ge\mathrm{poly}(n,\lambda_{\min}(K_n)^{-1},L)$ and random initialization, gradient descent on the empirical loss with step size $\eta=O(1/m)$ converges to a global minimum (zero training loss) at a linear rate, with high probability. The proof uses the NTK approximation: at large width the empirical NTK $K_{\theta_t}$ stays close to $K_{\theta_0}=K$ throughout training, so the dynamics stays in the kernel-regression regime, and $\lambda_{\min}(K_n)\succ 0$ ensures convergence.

The width requirement $m=\Omega(n^2/\lambda_{\min}^2)$ (or similar) is polynomial in the data size and the conditioning of the NTK; it is a sufficient condition for GD to reach a global minimum in the lazy regime. The result explains the empirical observation that GD trains over-parameterized nets to zero loss — the landscape has no spurious local minima that trap GD in the NTK regime.

Loss landscape geometry¶

Result	Statement
NTK regime (Jacot et al.)	GD on over-parameterized nets $=$ kernel regression; global minimum reached exponentially
Over-parameterized (Du et al., Allen-Zhu et al.)	width $\gg n$ $\Rightarrow$ GD reaches global min, polynomial width sufficient
Interpolation threshold	at width $\approx n$ (or rank $\approx n$), the landscape has spurious local minima / saddle plateaus; the NTK analysis breaks down
No bad local minima (under conditions)	for 1-hidden-layer nets with squared loss and certain activations, all local minima are global (Safran–Shamir); fails for general depth / loss
Lottery ticket (Frankle–Carbin)	a sparse sub-network ("lottery ticket") of the initialized net trains to comparable accuracy; theoretical analysis via sparse / implicit regularization
Mean-field (Chizat–Bach)	the feature-learning dynamics is a Wasserstein GF; global convergence under conditions (e.g. for 2-layer nets with squared loss, the GF reaches the global optimum)

The landscape is benign in the over-parameterized regime because the NTK stays positive-definite throughout training; at the interpolation threshold (width $\approx n$) the landscape develops spurious minima / saddle plateaus, and GD can stall — the regime where the loss landscape is genuinely non-convex and the NTK analysis is insufficient.

See it¶

The NTK of a single-hidden-layer ReLU net is the arc-cosine kernel (Cho–Saul). For a single-hidden-layer net $f_\theta(x)=\sum_{i=1}^m a_i\sigma(w_i\cdot x)$ with $w_i\sim\mathcal{N}(0,I_d/m)$, $a_i\sim\mathcal{N}(0,1)$ at initialization, compute the NTK

\[ K(x,x')=\mathbb{E}_{w\sim\mathcal{N}(0,I_d/m),\,a\sim\mathcal{N}(0,1)}\bigl[\partial_{w_i}f_\theta(x)\cdot\partial_{w_i}f_\theta(x') + \partial_{a_i}f_\theta(x)\cdot\partial_{a_i}f_\theta(x')\bigr]. \]

The $\partial_{a_i}$ term gives $\mathbb{E}[\sigma(w\cdot x)\sigma(w\cdot x')]=\kappa_1(\langle x,x'\rangle)$, and the $\partial_{w_i}$ term gives $\mathbb{E}[a^2\sigma'(w\cdot x)\sigma'(w\cdot x')\,x x']=\kappa_0(\langle x,x'\rangle)\,x x'$, where $\kappa_0,\kappa_1$ are the Cho–Saul arc-cosine kernels of degree 0 and 1. Concretely, for normalized $x,x'\in S^{d-1}$ with $\theta=\arccos(x\cdot x')$, the ReLU arc-cosine kernel of degree $1$ is

\[ \kappa_1(\theta)=\frac{1}{\pi}(\pi-\theta)\cos\theta+\frac{1}{\pi}\sin\theta, \]

and the full NTK is $K(x,x')=\kappa_0(x,x')+\kappa_1(x,x')$ with $\kappa_0(\theta)=\frac{1}{\pi}(\pi-\theta)\frac{\sin\theta}{\|x\|\|x'\|}$ after normalization — the arc-cosine / angular kernel measuring the alignment of $x$ and $x'$ through the angle $\theta$. The kernel is positive definite on $S^{d-1}$, so $\lambda_{\min}(K_n)\succ 0$ for distinct data, and the Jacot et al. theorem gives exponential kernel-regression convergence at rate $e^{-\lambda_{\min}(K_n)t}$. This single computation — "the infinite-width ReLU NTK is the arc-cosine kernel, a function of the angle between inputs" — is the operational content of the NTK theorem: the training dynamics of a wide ReLU net is a deterministic kernel method with a closed-form kernel.

Check¶

Why is the NTK / lazy regime degenerate from the perspective of the deep-learning advantage, and which structural property of the feature-learning regime recovers the advantage?

In the lazy regime the weights move by $O(1/m)$ — the neurons stay near their initialization directions, so the features (the functions $x\mapsto\sigma(w_i\cdot x)$) are fixed at initialization and the net is a kernel machine in the span of those fixed features. The expressivity / approximation advantage of deep nets (depth separation, Barron dimension-free rates — see Approximation Theory) is lost: a kernel method cannot adapt its features to the target, so it cannot realize the depth-separation gap or the Barron rate for structured targets (the NTK has a fixed RKHS, independent of the data, and is limited to functions in that RKHS — a much smaller class than the Barron space). The feature-learning regime recovers the advantage because the neuron directions move (the Wasserstein GF transports $\rho_t$), so the net adapts its features to the target — this is the regime where a 2-layer net can realize the Barron $O(1/\sqrt{m})$ approximation for a Barron-bounded target by learning the directions $w_i$ rather than sampling them at random. The structural property is the non-trivial transport of the neuron measure $\rho_t$ (the weights move by $O(1)$, not $O(1/m)$), which is absent in the lazy regime.

Compute the convergence rate of the NTK kernel-regression dynamics and identify the spectral quantity controlling it, then state the regime in which the rate is exponential.

The dynamics is $\partial_t f_t(x)=-\sum_i K(x,x_i)(f_t(x_i)-y_i)$, i.e. on the data values $u_t=(f_t(x_i))\in\mathbb{R}^n$, $\partial_t u_t=-K_n(u_t-y)$ where $K_n=(K(x_i,x_j))_{i,j}\in\mathbb{R}^{n\times n}$ is the NTK Gram matrix. The solution is $u_t-y=e^{-K_n t}(u_0-y)$, so $\|u_t-y\|\le e^{-\lambda_{\min}(K_n)t}\|u_0-y\|$, exponential at rate $\lambda_{\min}(K_n)$ — the smallest eigenvalue of the NTK Gram matrix. The regime is: width $m$ large enough that the empirical NTK $K_{\theta_t}$ stays within $O(\lambda_{\min}/2)$ of $K$ throughout training (requires $m=\Omega(n^2/\lambda_{\min}^2)$ for the Du et al. bound), and $K_n\succ 0$ strictly (the data is distinct, the kernel is PD). If $\lambda_{\min}(K_n)=0$ (e.g. data lies on a degenerate subspace where $K$ is degenerate, or the kernel has a non-trivial null-space on the data), the convergence is only on the image of $K_n$ and a residual persists — the rate degenerates. Hence the spectral quantity is $\lambda_{\min}(K_n)$, and the regime is "lazy / NTK with positive-definite Gram matrix."

Pitfalls¶

NTK / lazy is a degenerate regime: the lazy regime is where the net is a kernel method, not where the deep-learning advantage lives. Treating NTK convergence as "the explanation of deep learning" overstates the result — it explains optimization, not the feature-learning advantage.
Feature-learning / mean-field is not fully solved: the Wasserstein GF global-convergence results are partial (2-layer nets, squared loss, conditions); the general deep feature-learning convergence is open. Do not claim "feature-learning always converges to the global minimum."
Width requirements are polynomial but large: the Du et al. / Allen-Zhu et al. width $m=\Omega(n^2/\lambda_{\min}^2)$ can be much larger than $n$; the over-parameterized regime is sufficient but not necessary.
At the interpolation threshold the NTK analysis breaks: the lazy-regime convergence theorem requires $m\gg n$; at $m\approx n$ the empirical NTK drifts, spurious minima appear, and the GD dynamics is genuinely non-convex (see Generalization for the double-descent peak at this threshold).
Lottery ticket is a hypothesis with partial theory: the sparse-sub-network observation (Frankle–Carbin) is empirical; the rigorous analysis (implicit regularization, sparse-recovery / compressed-sensing analogies) is active and partly conjectural — do not state it as a theorem in full generality.

Refs¶

Jacot–Gabriel–Hongler, Neural tangent kernel: convergence and generalization in neural networks (NeurIPS 2018)
Chizat–Bach, On the global convergence of gradient descent for over-parameterized models using optimal transport (NeurIPS 2018); Implicit regularization of deep learning via lazy training (2019)
Du–Zhai–Poczos–Singh, Gradient descent provably optimizes over-parameterized neural networks (ICLR 2019)
Allen-Zhu–Li–Song, A convergence theory for deep learning via over-parameterization (ICML 2019)
Mei–Montanari–Mignacco, The mean-field analysis of neural networks (2020); Cho–Saul, Kernel methods for deep learning (NeurIPS 2009)

Next¶

Approximation Theory · Generalization · Information-Theoretic Limits · ML Mathematics index · Optimization · Functional analysis · Probability · Numerical analysis