Fast Finite Width Neural Tangent Kernel
This work addresses a bottleneck for researchers and practitioners using NTKs in tasks like initialization, architecture search, and meta-learning, by making finite-width NTK computation more practical, though it is incremental in improving existing methods.
The paper tackles the high computational cost of computing the Neural Tangent Kernel (NTK) for finite-width neural networks, proposing two novel algorithms that significantly improve efficiency by changing the exponent of compute and memory requirements, and open-sourcing implementations in the Neural Tangents package.
The Neural Tangent Kernel (NTK), defined as $Θ_θ^f(x_1, x_2) = \left[\partial f(θ, x_1)\big/\partial θ\right] \left[\partial f(θ, x_2)\big/\partial θ\right]^T$ where $\left[\partial f(θ, \cdot)\big/\partial θ\right]$ is a neural network (NN) Jacobian, has emerged as a central object of study in deep learning. In the infinite width limit, the NTK can sometimes be computed analytically and is useful for understanding training and generalization of NN architectures. At finite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning. Unfortunately, the finite width NTK is notoriously expensive to compute, which severely limits its practical utility. We perform the first in-depth analysis of the compute and memory requirements for NTK computation in finite width networks. Leveraging the structure of neural networks, we further propose two novel algorithms that change the exponent of the compute and memory requirements of the finite width NTK, dramatically improving efficiency. Our algorithms can be applied in a black box fashion to any differentiable function, including those implementing neural networks. We open-source our implementations within the Neural Tangents package (arXiv:1912.02803) at https://github.com/google/neural-tangents.