Neural Tangent Kernel

Authors: Arthur Jacot, Franck Gabriel, Clément Hongler
Published: 2018
Venue: NeurIPS
Link: arXiv

Summary

This groundbreaking paper shows that infinite-width neural networks are equivalent to kernel methods with a specific kernel called the Neural Tangent Kernel (NTK). This provides a theoretical framework for understanding deep learning through the lens of kernel methods.

Key Contributions

1. Neural Tangent Kernel Definition

For a neural network $f(x; \theta)$ with parameters $\theta$ , the NTK is:

$\Theta(x, x') = \left\langle \frac{\partial f(x; \theta)}{\partial \theta}, \frac{\partial f(x'; \theta)}{\partial \theta} \right\rangle$

2. Infinite-Width Limit

As the width of each layer $\to \infty$ :

Network output at initialization converges to a Gaussian process
NTK converges to a deterministic kernel $\Theta^*$
NTK remains constant during training (lazy training)

3. Training Dynamics

Under gradient flow on squared loss: $\frac{df(x_i)}{dt} = -\eta \sum_{j=1}^n \Theta(x_i, x_j)(f(x_j) - y_j)$

In matrix form: $\frac{d\mathbf{f}}{dt} = -\eta \mathbf{\Theta}(\mathbf{f} - \mathbf{y})$

4. Closed-Form Solution

For infinite-width networks: $f(x; t) = \Theta(x, X)\Theta(X, X)^{-1}(I - e^{-\eta \Theta(X, X) t})\mathbf{y}$

At convergence ( $t \to \infty$ ): $f^*(x) = \Theta(x, X)\Theta(X, X)^{-1}\mathbf{y}$

This is kernel regression with the NTK!

Mathematical Framework

Recursive Kernel Computation

For fully-connected network with depth $L$ :

Layer 0 (input): $\Sigma^{(0)}(x, x') = x^\top x'$

Layer $l$ (recursive): $\Sigma^{(l)}(x, x') = \sigma_w^2 \mathbb{E}_{f \sim \mathcal{N}(0, \Sigma^{(l-1)})}[\phi(f(x))\phi(f(x'))] + \sigma_b^2$

NTK (recursive): $\Theta^{(L)}(x, x') = \Sigma^{(L)}(x, x') + \Theta^{(L-1)}(x, x') \cdot \dot{\Sigma}^{(L)}(x, x')$

where $\dot{\Sigma}^{(l)}$ involves derivatives of the activation function.

Specific Architectures

ReLU activation: $\dot{\Sigma}^{(l)}(x, x') = \frac{1}{2\pi}\left(\pi - \arccos\left(\frac{\Sigma^{(l-1)}(x, x')}{\sqrt{\Sigma^{(l-1)}(x, x)\Sigma^{(l-1)}(x', x')}}\right)\right)$

Erf activation: Analytic formulas available.

Implementation

import numpy as np

def ntk_relu_kernel(X1, X2, depth, sigma_w=1.0, sigma_b=0.0):
    """
    Compute NTK for fully-connected ReLU network.
    
    Parameters:
    -----------
    X1, X2 : arrays of shape (n1, d) and (n2, d)
    depth : number of layers
    sigma_w : weight variance
    sigma_b : bias variance
    """
    n1, n2 = len(X1), len(X2)
    
    # Initialize Sigma^(0)
    Sigma = X1 @ X2.T
    
    # Diagonal elements for normalization
    diag1 = np.diag(X1 @ X1.T)
    diag2 = np.diag(X2 @ X2.T)
    
    # Initialize NTK
    Theta = Sigma.copy()
    
    for l in range(depth):
        # Normalize for angle computation
        norm_matrix = np.sqrt(np.outer(diag1, diag2))
        cos_theta = np.clip(Sigma / (norm_matrix + 1e-10), -1, 1)
        theta = np.arccos(cos_theta)
        
        # Derivative kernel (for ReLU)
        Sigma_dot = (np.pi - theta) / (2 * np.pi)
        
        # Update Sigma
        Sigma_new = sigma_w**2 * (
            (np.pi - theta) * Sigma + 
            norm_matrix * np.sin(theta)
        ) / (2 * np.pi) + sigma_b**2
        
        # Update NTK
        Theta = Sigma_new + Theta * Sigma_dot
        
        # Update for next layer
        Sigma = Sigma_new
        diag1 = sigma_w**2 * diag1 / 2 + sigma_b**2
        diag2 = sigma_w**2 * diag2 / 2 + sigma_b**2
    
    return Theta

def ntk_regression(X_train, y_train, X_test, depth, reg=1e-6):
    """
    Kernel regression with NTK.
    """
    K_train = ntk_relu_kernel(X_train, X_train, depth)
    K_test = ntk_relu_kernel(X_test, X_train, depth)
    
    # Solve (K + reg*I) alpha = y
    alpha = np.linalg.solve(K_train + reg * np.eye(len(K_train)), y_train)
    
    return K_test @ alpha

Theoretical Implications

1. Generalization

NTK framework enables:

Bounds on generalization error via kernel theory
Understanding implicit regularization
Analyzing double descent phenomena

2. Training Guarantees

For sufficiently wide networks:

Global convergence of gradient descent
Linear convergence rate
Deterministic training dynamics

3. Feature Learning Perspective

NTK describes the lazy training regime:

Features (weights) barely change
Only last layer learns

This contrasts with feature learning where early layers adapt.

Limitations and Extensions

Limitations of NTK Theory

Requires infinite width: Finite networks may behave differently
Lazy training: Doesn't capture feature learning
Architecture-specific: Different NTK for each architecture

Beyond Lazy Training

Mean-field theory: Feature learning regime
Tensor programs: Systematic analysis of width limits
Catapult dynamics: Phase transitions in learning

Connection to Other Kernels

Kernel	Network Type
NTK	Fully-connected
CNTK	Convolutional
Graph NTK	Graph neural networks
Attention NTK	Transformers

Impact on Deep Learning Theory

This paper has:

Unified kernel methods and deep learning
Enabled theoretical analysis of neural networks
Inspired architecture design (infinite-depth limits)
Explained phenomena like double descent

Citation

@inproceedings{jacot2018neural,
  title={Neural tangent kernel: Convergence and generalization in neural networks},
  author={Jacot, Arthur and Gabriel, Franck and Hongler, Cl{\'e}ment},
  booktitle={Advances in Neural Information Processing Systems},
  volume={31},
  year={2018}
}

Summary​

Key Contributions​

1. Neural Tangent Kernel Definition​

2. Infinite-Width Limit​

3. Training Dynamics​

4. Closed-Form Solution​

Mathematical Framework​

Recursive Kernel Computation​

Specific Architectures​

Implementation​

Theoretical Implications​

1. Generalization​

2. Training Guarantees​

3. Feature Learning Perspective​

Limitations and Extensions​

Limitations of NTK Theory​

Beyond Lazy Training​

Connection to Other Kernels​

Impact on Deep Learning Theory​

Citation​

Further Reading​