Finite Neural Tangent Kernel

Finite Neural Tangent Kernel

The Neural Tangent Kernel (NTK) describes how a neural network evolves during training via gradient descent. For infinite-width networks, the NTK remains constant, and training behaves like kernel regression. However, in finite-width networks, this is not exactly true. The tangent kernel is defined as:

K(x,y) = \nabla_θ f(x) · \nabla_θ f(y)

This module explores a small, exactly finite network, illustrating the discrepancy between the actual updated function value and the first-order Taylor expansion (linearization) provided by the NTK. Specifically, from a single training point update:

f(x_(test)) ≈ f_(initial)(x_(test)) - η K(x_(test), x_(train)) (f(x_(train)) - y_(target))

Experiment

You can interactively explore the NTK in this small network. Open the Finite Neural Tangent Kernel module.

By selecting different "Network" presets, you can observe how different weight configurations ("Mixed", "Both active", "Dead", "Kink") impact the training dynamics.

1. Change the Network preset to "Dead" and notice the gradient becomes identically zero.
2. Set the Network to "Mixed", and change the Step size to 0.5.
3. Observe that the error of the first-order linear approximation increases as the step size gets larger, because the actual parameters update bilinearly for even this simple ReLU network.

Worked Example

Consider the "Mixed" preset network where parameters are a = [1, 2], w = [1, -1], and b = [0.5, 0.5]. Let the train point be x = 1, the test point be x = 0, the target be y = 0, and the step size be η = 0.1.

Initial Evaluation at Train x=1:
z_0 = 1 * 1 + 0.5 = 1.5, r_0 = 1.5
z_1 = -1 * 1 + 0.5 = -0.5, r_1 = 0
f(1) = (1 * 1.5 + 2 * 0) / sqrt(2) ≈ 1.0606601717798212
Train residual = 1.0606601717798212 - 0 = 1.0606601717798212

Initial Evaluation at Test x=0:
z_0 = 1 * 0 + 0.5 = 0.5, r_0 = 0.5
z_1 = -1 * 0 + 0.5 = 0.5, r_1 = 0.5
f(0) = (1 * 0.5 + 2 * 0.5) / sqrt(2) ≈ 1.0606601717798212

Gradient at Train x=1:
g_0 = 1, g_1 = 0
da0 = 1.5 / sqrt(2), da1 = 0
dw0 = 1 / sqrt(2), dw1 = 0
db0 = 1 / sqrt(2), db1 = 0

Gradient at Test x=0:
g_0 = 1, g_1 = 1
da0 = 0.5 / sqrt(2), da1 = 0.5 / sqrt(2)
dw0 = 0, dw1 = 0
db0 = 1 / sqrt(2), db1 = 2 / sqrt(2)

Tangent Kernel K(0, 1):
K = (1.5 * 0.5 + 0 + 0 + 0 + 1 * 1 + 0) / 2 = 1.75 / 2 = 0.875

Linearized Test Output:
f_lin(0) = 1.06066 - 0.1 * 0.875 * 1.06066 = 0.9678524067490869

The actual updated parameters would yield an output of approximately 0.9738186202153484, giving a linearization error of 0.005966213466261472.

Limitations

This module considers only an extremely simple 2-unit finite network for explanatory purposes. In practical applications, deep neural networks contain millions of parameters, and calculating the exact NTK is computationally prohibitive. The linearization approximation also breaks down rapidly with larger learning rates and real-world batch training procedures.

References

The following primary references provide the foundational understanding and mathematical background for the Neural Tangent Kernel concept.

• Jacot, A., Gabriel, F., & Hongler, C. (2018). Neural Tangent Kernel: Convergence and Generalization in Neural Networks. arXiv:1806.07572