# Neural Tangent Kernel

In the infinite-width limit, gradient descent on a network behaves like kernel regression with a fixed kernel determined at initialisation.

**Rests on:** Forward Propagation, Gradient Descent

*Deep Learning · Deep Learning*


## Why learn this

Understanding the Neural Tangent Kernel helps explain a puzzling observation in deep learning: why extremely large, complex models often work better than traditional math suggests they should. Normally, we expect a model with millions of adjustable parts to memorize its training data and fail on new data. However, when these networks are made vast—far wider than necessary—they become surprisingly stable and easy to analyze. This concept explains that stability. It shows how a complex network can be viewed as a simpler, predictable system during training, which is crucial for debugging issues like models that refuse to improve or become unstable.

For example, if you double the width of the network while keeping other factors constant, the prediction changes in a way that can be calculated precisely using this kernel, rather than requiring complex simulation.


## The intuition

Imagine you are trying to fit a smooth curve through a set of scattered points on a graph. You could use a rigid straight line (which misses most points), a gentle parabola, or a wildly wiggly high-degree polynomial that connects every single point exactly. In standard machine learning, we don't choose the shape by hand; we let an algorithm adjust millions of numbers (weights) to find a shape that fits best.

Here is the puzzle: if you have millions of adjustable weights but only a few thousand data points, why doesn't the model just memorize those specific points and fail everywhere else? The Neural Tangent Kernel explains part of the answer by looking at what happens when you make the network extremely wide—so wide it has an almost unimaginable number of neurons in each layer.

In this extreme limit, the learning process simplifies dramatically. Instead of changing its fundamental structure to learn new patterns (which is hard to track), the network behaves as if it is making tiny, fixed adjustments to a pre-existing set of features. The 'kernel' acts like a lens that reveals the smooth, predictable motion behind the apparent chaos. For example, if you double the width of the network while keeping other factors constant, the prediction changes in a way that can be calculated precisely using this kernel, rather than requiring complex simulation.


## An analogy

Think of a massive orchestra with tens of thousands of musicians, each holding a piece of sheet music (the initial random weights). When they start playing, the sound is chaotic noise. Now, imagine a conductor who wants to shape this noise into a specific melody by making tiny adjustments to each musician's volume or timing (gradient descent).

In a normal-sized orchestra, as the conductor adjusts one violinist, their playing changes in a way that affects how they interact with others—a complex web of influence. But now, scale this up: imagine an infinitely large orchestra where every individual change is so infinitesimally small relative to the whole mass that no musician actually learns or improvises. Their sheet music never truly changes; only their volume fluctuates slightly around a fixed plan. The resulting sound becomes entirely predictable and linear because the underlying structure (the interaction between musicians, defined by the initial random sheet music) never shifts. This predictable, linear behavior in the infinite limit is what the Neural Tangent Kernel captures.

Caveat: real networks are finite, so musicians *do* change their roles over time, which is where this analogy diverges from reality.


## Definition

The Neural Tangent Kernel is a mathematical function that predicts how the output of an infinitely wide neural network changes when its parameters are updated using gradient descent; in this infinite limit, the network’s training behaves exactly like a simpler, static linear model known as kernel regression.


## How it connects

This concept sits beside the Neural Tangent Kernel in understanding why extremely wide networks are stable. It also connects to Double Descent, which describes how prediction error often decreases, then briefly increases (overfitting), and then decreases again as network size grows—a phenomenon that the infinite-width limit helps explain.


## The mathematics

The mathematics of the Neural Tangent Kernel (NTK) rests on two ideas you already know: how steep a curve is at one point (a derivative), and how to combine two lists of numbers into a single measure of similarity (a dot product). We will build the definition step by step.

Step 1: The network as a function
Imagine a neural network that takes an input, such as a list of numbers representing an image, and produces one output number, such as a score for 'cat-ness'. Let us call this network f. The network has many internal settings called weights. We collect all these weights into one long list and call it theta (θ). So, the output is a function of the input x and the weights θ:

Output = f(x, θ)

Step 2: Sensitivity to each weight
In training, we want to change the weights to make the output better. To know which way to change a specific weight, we need to know how sensitive the output is to that weight. This sensitivity is called the derivative.

If you nudge one weight by a tiny amount, how much does the output change? That number is the derivative of f with respect to that weight.

Since there are many weights, there are many derivatives. We collect all these derivatives into a single list (a vector). Let us call this list the gradient vector for input x. We write it as:

Gradient(x) = [derivative w.r.t. weight 1, derivative w.r.t. weight 2, ...]

This vector tells you: 'If I change each weight by a tiny bit, here is exactly how the output will move.'

Step 3: The Neural Tangent Kernel
Now, imagine we freeze the weights at their very first random values (before any training happens). We call these initial weights θ₀.

We pick two different inputs, x and x'. For example, x might be a picture of a cat, and x' might be a picture of a dog.

1. Compute the gradient vector for x using the frozen weights θ₀. Call this G(x).
2. Compute the gradient vector for x' using the same frozen weights θ₀. Call this G(x').
3. Calculate the dot product of these two vectors.

The result is a single number. This number is the Neural Tangent Kernel between x and x'. We write it as K(x, x').

What does this number mean? The dot product measures how much the two gradient vectors point in the same direction. If G(x) and G(x') point in similar directions, the kernel value is large and positive. This means that changing the weights to improve the prediction for x will also tend to improve the prediction for x'. If they point in opposite directions, the kernel is negative. If they are perpendicular, the kernel is zero, meaning training on x tells us nothing about how to train for x'.

Step 4: Why 'Tangent' and 'Kernel'?
The word 'kernel' comes from mathematics and simply means a function that measures similarity between pairs of items.
The word 'tangent' refers to the fact that we are looking at the network's behavior as if it were a straight line (the tangent) near its starting point. In very wide networks, this straight-line approximation is so good that the kernel K(x, x') stays constant even as the weights change during training. This allows us to predict how the network will learn using simple linear algebra.

Symbol glossary:
- f: The neural network function. It takes an input and weights and returns a number.
- x: One specific input (e.g., one image).
- x': A different specific input.
- θ: The list of all weights in the network.
- θ₀: The initial random values of the weights, before training starts.
- Gradient(x): The list of derivatives of the output with respect to each weight, calculated at input x and weights θ₀.
- K(x, x'): The Neural Tangent Kernel value between inputs x and x'.


## A worked example

Let's build a tiny example. Imagine a network with just two inputs and two weights. We'll ignore the bias for simplicity. The network output is simply weight1 * input1 + weight2 * input2 (this is a 'linear' neuron, no activation function). We'll pick two inputs: Input A = [1, 2] and Input B = [3, 1]. Initial weights: w = [0.5, -0.5].

Step 1: Derivative of output w.r.t. weight1 at input A: input1 = 1   # The derivative is just the input value itself, because the output is weight * input.
Step 2: Derivative of output w.r.t. weight2 at input A: input2 = 2   # Same reasoning.
Step 3: Sensitivity vector for Input A: [1, 2]   # This is the list of derivatives, one per weight.
Step 4: Sensitivity vector for Input B: [3, 1]   # Because the inputs are [3, 1].
Step 5: Neural Tangent Kernel (A, B) = 1*3 + 2*1 = 3 + 2 = 5   # This is a dot product: multiply the first components together (1*3), multiply the second components together (2*1), and add them up.

Result: The kernel value between Input A and Input B is 5. If the kernel were zero, the two inputs would be 'independent' in the eyes of the network — changing one wouldn't affect predictions on the other. A high positive value means they are strongly related — training on Input A also improves predictions on Input B. This is the core of the NTK's power: it tells you how training on one example 'transfers' to another.


## From scratch, in Python

```python
import math

# We'll build a super tiny network: 2 inputs -> 2 hidden neurons -> 1 output
# Activation: tanh (we'll use math.tanh)

# Initial weights and biases (random, but let's pick specific numbers for clarity)
w1_1 = 0.5   # from input 1 to hidden neuron 1
w2_1 = -0.3  # from input 2 to hidden neuron 1
b1 = 0.1
w1_2 = -0.7  # from input 1 to hidden neuron 2
w2_2 = 0.4   # from input 2 to hidden neuron 2
b2 = -0.2
v1 = 0.8     # from hidden neuron 1 to output
v2 = 0.3     # from hidden neuron 2 to output
b_out = 0.0

def forward(x1, x2):
    """Return the output of the network for input (x1, x2)."""
    # Hidden layer
    h1 = math.tanh(w1_1 * x1 + w2_1 * x2 + b1)
    h2 = math.tanh(w1_2 * x1 + w2_2 * x2 + b2)
    # Output layer
    out = v1 * h1 + v2 * h2 + b_out
    return out

def derivative_w1_1(x1, x2):
    """Derivative of output with respect to w1_1."""
    # Chain rule: d(out)/d(w1_1) = v1 * d(h1)/d(w1_1)
    # d(h1)/d(w1_1) = (1 - h1^2) * x1  (derivative of tanh times the input)
    h1 = math.tanh(w1_1 * x1 + w2_1 * x2 + b1)
    return v1 * (1 - h1**2) * x1

def derivative_w2_1(x1, x2):
    h1 = math.tanh(w1_1 * x1 + w2_1 * x2 + b1)
    return v1 * (1 - h1**2) * x2

def derivative_w1_2(x1, x2):
    h2 = math.tanh(w1_2 * x1 + w2_2 * x2 + b2)
    return v2 * (1 - h2**2) * x1

def derivative_w2_2(x1, x2):
    h2 = math.tanh(w1_2 * x1 + w2_2 * x2 + b2)
    return v2 * (1 - h2**2) * x2

# Now compute the NTK between two inputs
def ntk(x1, x2, x1_prime, x2_prime):
    """K between (x1,x2) and (x1_prime,x2_prime)."""
    # Compute sensitivity vector for first input
    sens1 = [
        derivative_w1_1(x1, x2),
        derivative_w2_1(x1, x2),
        derivative_w1_2(x1, x2),
        derivative_w2_2(x1, x2)
    ]
    # For second input
    sens2 = [
        derivative_w1_1(x1_prime, x2_prime),
        derivative_w2_1(x1_prime, x2_prime),
        derivative_w1_2(x1_prime, x2_prime),
        derivative_w2_2(x1_prime, x2_prime)
    ]
    # Dot product
    dot = 0
    for i in range(4):
        dot += sens1[i] * sens2[i]
    return dot

# Test with the numerical example from the math section
print(ntk(1, 2, 3, 1))
```


## The same thing, with a framework

```python
import numpy as np
import torch
import torch.nn as nn

# In PyTorch, you'd typically use autograd to compute the gradients.
# Here's the same network using torch, and computing the NTK via autograd.

# Define the network (same as above)
net = nn.Sequential(
    nn.Linear(2, 2),  # 2 inputs to 2 hidden neurons
    nn.Tanh(),
    nn.Linear(2, 1)   # 2 hidden to 1 output, no activation
)

# (Normally you would copy the initial weights to a variable so you can reuse them)

# Function to compute the gradient of the output w.r.t. all parameters
def get_grad(x):
    x_tensor = torch.tensor(x, dtype=torch.float32, requires_grad=False)
    net.zero_grad()
    out = net(x_tensor)  # shape [1]
    grad_list = []
    for param in net.parameters():
        if param.grad is not None:
            grad_list.append(param.grad.detach().flatten())
    # This is not a full implementation, but the key is: after calling out.backward(),
    # each param.grad holds the derivative.
    return torch.cat(grad_list)

# For a real implementation, you would construct the full Jacobian.
# This is just to show the idea.
```


## What the code does

The from-scratch code is deliberately tiny: a 2-2-1 network with tanh activation. The key lines are the derivatives. For a linear layer, the derivative of the output with respect to a weight is just the input value. For a hidden neuron with tanh, we use the chain rule: the derivative of tanh is (1 - h²), where h is the neuron's output before activation. Then we multiply by the connection weight (v1 or v2). The ntk function computes the sensitivity vector (the list of all derivatives) for each input, then takes the dot product. In the framework version, we let PyTorch's autograd do the derivative work for us — it computes the same thing, but it's a black box that follows the computational graph. The important thing is that both compute the same quantity: the dot product of two gradient vectors.


## Common misconceptions

**NTK is just a fancy name for a regular kernel** — I already know what a kernel is from SVM — the NTK is just another kernel like a Gaussian RBF.

An RBF kernel is hand-crafted, but the NTK is defined by the network structure and initial random weights. Two different networks (e.g., tanh vs ReLU) give different NTKs. In fact: The NTK is a kernel that is automatically derived from the architecture and initialization. You don't pick it — the network defines it.

**NTK predicts the behaviour of all networks** — Since the NTK explains wide networks, it must also explain narrow ones, just with a bit of approximation error.

For narrow networks, the kernel changes drastically during training — the 'tangent' is no longer fixed. So any prediction made using the initial NTK quickly becomes invalid. In fact: The NTK is a limit statement: it is exactly correct only as width goes to infinity. Real finite networks deviate, and that deviation is what enables feature learning.

**The NTK says networks don't learn features** — So the NTK says that neural networks never learn features? That would mean all that 'deep learning learns hierarchies' is a myth.

In the infinite-width limit, the weights barely move, so no features are learned. But real networks do learn features — that's the whole point of deep learning. The NTK is a useful theoretical toy, not a description of real networks. In fact: The NTK describes the lazy training regime. Real networks (especially narrow ones) are in the rich regime where features do emerge, and that's where deep learning's power comes from.

**NTK requires infinite width to be useful** — If the NTK is only exact for infinite width, I can't use it for any practical network.

Many practical networks are wide enough that the NTK is a decent approximation, especially in the early stages of training. The 'finite-width' corrections are well-studied (e.g., in papers by Lee et al.). In fact: The NTK is a limit, but the approximation is often good enough for predicting training speed, setting learning rates, or diagnosing divergence.

**The NTK is just theory with no connection to code** — This is pure math — I can't compute it for a real network.

Modern autograd systems compute exact Jacobians and the NTK in a few lines. It's not just theory; it's a computational object. In fact: You can compute the NTK numerically with `torch.func` or the `neural-tangents` library. It's a concrete, computable matrix (size N×N for N data points).


## In short

The neural tangent kernel is a way to make the seeming chaos of training a neural network look simple. It says: if you make the network extremely wide, then training it is mathematically identical to classical kernel regression with a fixed kernel defined by the initial random weights. It gives you a predictive model for training dynamics, but it intentionally ignores feature learning. It's a bridge between deep learning theory and practice — useful for setting learning rates, understanding generalization, and diagnosing divergences — but it is a limit statement, not a complete description of real networks.


## Where this is used

The neural tangent kernel is not something a practitioner runs on a laptop every day, but it is the theoretical foundation for understanding when deep learning works and when it fails. Here's the honest picture:

- **Large language models (LLMs):** The NTK is used to explain why very large transformers (like GPT-4) can memorize training data and still generalize. The kernel viewpoint says that in the infinite-width limit, the model's predictions are a linear function of the training labels — this 'kernel ridge regression' regime is what happens when gradient descent converges to the minimum norm solution. Researchers at OpenAI and DeepMind have used the NTK to analyze the training dynamics of large models, especially the phase transitions where loss suddenly drops. The key insight is that the kernel at initialization determines how quickly different components of the data are learned. In practice, this shows up in the 'critical batch size' and the learning rate schedule — you can compute the NTK's 'eigenvalues' to predict which frequencies in the data are learned first and which take longer.

- **Small models on devices:** When you compress a large model into a small one (distillation), the NTK helps explain why the student can sometimes outperform the teacher. The kernel of the student is different from the teacher's, and the training dynamics can be analyzed as a kernel regression with a different kernel. This is used in research on model compression and in practical tools like TensorFlow Lite's model optimization. A concrete use case: after training a small keyword-spotter model for a smartphone, you can compute its NTK to verify that it won't 'forget' the wake word when it's fine-tuned on user data.

- **Retrieval-augmented generation (RAG):** RAG systems embed documents and queries into vectors and retrieve the nearest ones. The NTK doesn't directly apply to the retrieval step, but it does apply to the fine-tuning of the retriever. When you fine-tune a small cross-encoder to the task of relevance ranking, the NTK tells you how quickly it will adapt to new domains — a large 'kernel bandwidth' means fast adaptation. Some recent work (e.g., 'Fine-Tuning the Retrieval Model' at ICML 2023) has used the NTK to explain why fine-tuning with a tiny learning rate and only a few steps on a new dataset can improve retrieval without catastrophic forgetting. The mechanism is that the NTK's fixed kernel ensures small, stable updates.

- **Agentic systems:** Agents that plan, call tools, and act in loops are built on top of LLMs. The internal critic or reward model is often a small transformer. When you fine-tune that critic via reinforcement learning, the NTK provides a theoretical guarantee that if the network is wide enough, the policy update will be smooth — not randomly jumping. This is crucial for stability in RL with neural networks. In practice, you'd use the NTK's conditioning (the smallest eigenvalue of the kernel matrix) to set the learning rate and trust region. Without it, the agent's training may diverge.

- **Everyday applied ML:** In fraud detection or recommendation systems, you often train a gradient-boosted decision tree or a small neural net. If you switch from a small network to a wider one, the NTK explains the 'double descent' curve: test error goes down, then up, then down again as you increase width. Knowing this prevents you from being confused when a larger model performs worse than a smaller one. A developer in a fintech might use the NTK-based learning rate finder: compute the NTK of a small sample of data, estimate its largest eigenvalue, and set the learning rate to be less than 2/lambda_max. This is done in practice using tools like Google's 'Neural Tangent Kernel' library (https://github.com/google/neural-tangents).

**What developers actually write:** They usually don't write the NTK from scratch — they use libraries. The `neural-tangents` package gives you functions to compute the NTK of a neural network in JAX. PyTorch has `torch.func` (`torch.autograd.functional.jacobian`) to compute the Jacobian of a network's output with respect to its parameters. The code is usually:

```python
from torch.autograd.functional import jacobian
import torch
# Generate 10 random inputs
x = torch.randn(10, 2)
# Compute Jacobian of all outputs w.r.t. all parameters (flattened)
j = jacobian(lambda w: net(w), x)  # shape (10, 2, num_params) roughly
# Then compute the NTK by flattening and taking dot products
```

What breaks when done badly: if you don't check the conditioning of the NTK matrix, you might set the learning rate too high, causing divergence. If you treat the NTK as a fixed object when the network is narrow, you'll be wrong — the kernel changes during training, and the NTK analysis is invalid. Also, using the NTK to predict the exact training dynamics will fail for real networks because they learn features, so you'll be inaccurate if you ignore the 'feature learning' correction.


---

Read this on the web: https://basavapurushottam.com/learn-ai/neural-tangent-kernel
