Demystifying the Neural Tangent Kernel (NTK): When Neural Networks Become Linear
Demystifying the Neural Tangent Kernel (NTK): When Neural Networks Become Linear
Training a deep neural network often feels like alchemy. We navigate a high-dimensional, non-convex loss landscape, hoping that Stochastic Gradient Descent (SGD) finds a global minimum despite the theoretical nightmare of saddle points and local minima.
But what if I told you that as a neural network gets wider, it actually becomes simpler?
In this post, we dive into the Neural Tangent Kernel (NTK), a theoretical breakthrough that proves that in the "infinite-width limit," neural networks behave like linear models governed by a static kernel.
The Core Intuition: From Parameters to Functions
Traditionally, we view training as updating a vector of parameters $\theta$ in a massive parameter space. The NTK framework shifts this perspective. Instead of asking "How do the weights change?", it asks "How does the output function $f(x)$ change?"
The "Infinite-Width" Magic
The central thesis of NTK is that as the number of hidden units $n \to \infty$, the network's parameters move very little from their initialization, yet the network can still fit any data perfectly. In this limit, the network behaves as a linear model in a high-dimensional feature space.
The training dynamics are then governed by a static kernel—the Neural Tangent Kernel—which defines the direction and speed of convergence in function space.
The Mathematical Foundation
To understand NTK, we need to look at how the network output $f_\theta(x)$ evolves during gradient descent.
1. Defining the NTK
The NTK is essentially the inner product of the gradients of the network output with respect to its parameters for two different inputs $x$ and $x'$:
$$\text{NTK: } \Theta(\theta)(x, x') = \sum_{p=1}^{P} \frac{\partial f_\theta(x)}{\partial \theta_p} \frac{\partial f_\theta(x')}{\partial \theta_p}$$
2. Evolution in Function Space
If we use gradient descent to minimize a cost function $C$, the evolution of the network's output over time $t$ is described by:
$$\frac{df_\theta(t, x)}{dt} = -\int \Theta(\theta(t))(x, x') \nabla_{f(x')} C(f(t)) dx'$$
The Breakthrough: As the width of the layers tends toward infinity, $\Theta(\theta)$ becomes constant ($\Theta_\infty$). The complex, non-linear training of a neural network simplifies into a linear differential equation.
Architecture & Workflow
The following diagram illustrates the two paths to the same destination: training via standard SGD (Parameter Space) versus predicting via the NTK closed-form solution (Function Space).
Implementation: Proving the Theory with PyTorch
Can we actually see this in code? Yes. By creating a very wide network, we can compute the empirical NTK at initialization and use it to predict the final output of gradient descent without ever actually running the training loop.
The Implementation
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_regression
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error
class NTKNetwork(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim=1):
super(NTKNetwork, self).__init__()
self.layers = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, output_dim)
)
# NTK scaling: Crucial for stability as width -> infinity
with torch.no_grad():
nn.init.normal_(self.layers[0].weight, std=1.0 / np.sqrt(hidden_dim))
nn.init.zeros_(self.layers[0].bias)
nn.init.normal_(self.layers[2].weight, std=1.0 / np.sqrt(hidden_dim))
nn.init.zeros_(self.layers[2].bias)
def forward(self, x):
return self.layers(x)
def compute_ntk_matrix(model, x):
n_samples = x.shape[0]
ntk_matrix = torch.zeros((n_samples, n_samples))
for i in range(n_samples):
x_i = x[i].unsqueeze(0).requires_grad_(False)
out_i = model(x_i)
grads_i = []
for param in model.parameters():
grad = torch.autograd.grad(out_i, param, retain_graph=True)[0]
grads_i.append(grad.view(-1))
flat_grad_i = torch.cat(grads_i)
for j in range(i, n_samples):
x_j = x[j].unsqueeze(0).requires_grad_(False)
out_j = model(x_j)
grads_j = []
for param in model.parameters():
grad = torch.autograd.grad(out_j, param, retain_graph=True)[0]
grads_j.append(grad.view(-1))
flat_grad_j = torch.cat(grads_j)
val = torch.dot(flat_grad_i, flat_grad_j)
ntk_matrix[i, j] = val
ntk_matrix[j, i] = val
return ntk_matrix
def solve_ntk_regression(x, y, kernel_matrix):
# Closed form solution: f(x) = f(x, 0) + K(x, X) * K(X, X)^-1 * (y - f(X, 0))
reg = 1e-4 * torch.eye(kernel_matrix.shape[0])
alpha = torch.linalg.solve(kernel_matrix + reg, y)
return torch.matmul(kernel_matrix, alpha)
# --- Execution ---
X_raw, y_raw = make_regression(n_samples=50, n_features=10, noise=0.1, random_state=42)
scaler = StandardScaler()
X = torch.tensor(scaler.fit_transform(X_raw), dtype=torch.float32)
y = torch.tensor(y_raw, dtype=torch.float32).view(-1, 1)
# Wide network to approximate NTK limit
model = NTKNetwork(10, 500, 1)
K = compute_ntk_matrix(model, X)
# Path A: Gradient Descent
optimizer = optim.SGD(model.parameters(), lr=0.1)
criterion = nn.MSELoss()
for epoch in range(200):
optimizer.zero_grad()
loss = criterion(model(X), y)
loss.backward()
optimizer.step()
# Path B: NTK Closed Form
final_pred_gd = model(X).detach()
final_pred_ntk = solve_ntk_regression(X, y, K)
print(f"MSE (Gradient Descent): {mean_squared_error(y, final_pred_gd):.4f}")
print(f"MSE (NTK Closed Form): {mean_squared_error(y, final_pred_ntk.detach().numpy()):.4f}")Key Takeaways
1. Why does this matter?
The NTK provides a theoretical bridge between deep learning and kernel methods (like Support Vector Machines). It explains why over-parameterized networks generalize well: they are essentially performing a form of kernel regression.
2. The Trade-off
While the NTK is a powerful analytical tool, it describes the "Lazy Training" regime. In real-world networks, we often want the kernel to evolve (feature learning), which happens when networks are not infinitely wide or when we use different learning rates.
3. Summary Table
| Feature | Parameter Space (SGD) | Function Space (NTK) |
|---|---|---|
| Perspective | Updating weights $\theta$ | Updating function $f(x)$ |
| Complexity | Non-convex, high-dimensional | Linear, convex |
| Computation | Iterative (Epochs) | Closed-form (Matrix Inverse) |
| Limit | Finite width | Infinite width ($n \to \infty$) |
By understanding the NTK, we move one step closer to turning the "black box" of deep learning into a transparent, predictable mathematical system.