Paper: Neural Tangent Kernel: Convergence and Generalization in Neural Networks" (Jacot et al., 2018)

 

Neural Tangent Kernel (NTK): Key Ideas

Paper: "Neural Tangent Kernel: Convergence and Generalization in Neural Networks" (Jacot et al., 2018)
Core Contribution: A theoretical framework to analyze the training dynamics of infinitely wide neural networks using kernel methods.

---

1. Intuition Behind NTK

• At infinite width, a neural network behaves like a linear model in parameter space around its initialization.

• The NTK is a kernel function that describes how small changes in parameters affect the network’s output during gradient descent.

• Key Insight: Training dynamics of wide networks simplify to kernel regression with the NTK.

---

2. Mathematical Definition

For a neural network $f_{\theta }(x)$ with parameters $\theta$:

• The NTK $\mathrm{\Theta }(x,x^{\mathrm{\prime }})$ is defined as:

  $$\mathrm{\Theta }(x,x^{\mathrm{\prime }})={\mathbb{E}}_{\theta \sim \text{init}}⟨{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x),{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x^{\mathrm{\prime }})⟩$$

  where ${\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)$ is the gradient (Jacobian) of the network’s output w.r.t. parameters.

• At Infinite Width:

  • The NTK becomes deterministic and stays constant during training (lazy training regime).

  • Training reduces to solving:

    $$f_t(x)\approx f_0(x)+\mathrm{\Theta }(x,\cdot {)}^T{\mathrm{\Theta }}^{-1}(y-f_0(X))$$

    where $f_0$ is the network at initialization.

---

3. Key Results

1. Convergence Guarantees:

   • Infinitely wide networks trained with gradient descent converge to global minima if the NTK is positive definite.

2. Generalization:

   • NTK theory explains why wide networks generalize despite overparameterization.

3. Linearized Training Dynamics:

   • The network’s evolution can be approximated by:

     $$\frac{d{f}_t(x)}{dt}\approx -\eta \sum _{i=1}^n\mathrm{\Theta }(x,x_i)(f_t(x_i)-y_i)$$

     where $\eta$ is the learning rate.

---

4. Practical Implications

• NTK Regime:

  • Networks behave like linear models when width ≫ depth (e.g., wide ResNets, MLPs).

  • Explains success of random feature models and shallow networks.

• Beyond NTK:

  • Finite-width networks deviate from NTK predictions (feature learning becomes important).

  • Modern architectures (Transformers, GNNs) can also be analyzed with NTK extensions.

---

5. Limitations

• Finite-Width Networks: NTK assumptions break down when width is not extreme.

• Feature Learning: NTK ignores non-linear feature adaptation (critical in deep narrow networks).

• Kernel Computability: Exact NTK is expensive to compute for large architectures.

---

6. Code Example (NTK Approximation)

python

Copy

import jax.numpy as jnp
from neural_tangents import stax

# Define an infinite-width network (linearized)
init_fn, apply_fn, kernel_fn = stax.serial(
    stax.Dense(512), stax.Relu(),
    stax.Dense(1)
# Compute NTK matrix
ntk = kernel_fn(X_train, X_train, 'ntk')

---

7. Follow-Up Work

• NTK for Deep Networks: Arora et al., 2019

• NTK and Generalization: Cao & Gu, 2019

• Beyond NTK (Feature Learning): Yang & Hu, 2021

---

Why NTK Matters

• Connects neural networks to classical kernel methods.

• Provides theoretical guarantees for wide networks.

• Inspires new optimization techniques (e.g., NTK-aware initialization).

For deeper analysis, see the original paper or the Neural Tangents library. 

Link to paper: https://arxiv.org/pdf/1806.07572