Fast Finite Width Neural Tangent Kernel (2206.08720v1)

Abstract: The Neural Tangent Kernel (NTK), defined as $\Theta_\theta^f(x_1, x_2) = \left[\partial f(\theta, x_1)\big/\partial \theta\right] \left[\partial f(\theta, x_2)\big/\partial \theta\right]^T$ where $\left[\partial f(\theta, \cdot)\big/\partial \theta\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.

• The paper introduces innovative AD-based algorithms that significantly reduce computational complexity and memory requirements in NTK computation.
• It exploits neural network structure using Jacobian-vector and vector-Jacobian products to avoid instantiating large matrices.
• The open-source integration in the Neural Tangents package paves the way for enhanced initialization, architecture search, and meta-learning.

Fast Finite Width Neural Tangent Kernel

The paper "Fast Finite Width Neural Tangent Kernel" addresses the computational challenges inherent in calculating the Neural Tangent Kernel (NTK) for finite width neural networks. The NTK has become crucial in understanding the dynamics and generalization properties of neural networks, and its computation for finite widths offers advantages such as improving initialization, enabling architecture search, and facilitating meta-learning. However, its computational costs have traditionally been prohibitive, especially in large-scale applications.

Key Contributions

The authors present several innovations to efficiently compute the finite width NTK. They employ two novel algorithms that significantly reduce both computational and memory requirements, making NTK computation viable for practical applications in deep learning. Specifically, the algorithms they put forward alter the exponent of the computational complexity, offering notable performance enhancements.

1. Exploiting Neural Network Structure: By leveraging the inherent structure of neural networks, the authors propose methods that improve the NTK computation's time complexity. They do this without fundamentally altering the networks' architectures or requiring any modifications that could affect the networks' outputs or training dynamics.
2. Algorithms for Efficient NTK Computation: The two presented algorithms utilize automatic differentiation (AD) techniques supported by the JAX library, streamlining NTK computation for any differentiable function. These methods demonstrate flexibility and can be applied in a black-box manner to diverse neural architectures.
3. Implementation and Open Source: The practical implementations of these algorithms are integrated into the Neural Tangents package, offering researchers access to efficient NTK calculations without the need for extensive computational resources.

Methodological Advancements

The advancements are rooted in automatic differentiation and the strategic exploitation of the neural network's layer-wise structure. By focusing on Jacobian-vector products (JVP) and vector-Jacobian products (VJP), the authors manage to compute the NTK without instantiating large matrices, which is typically computationally intense.

• JVP and VJP Usage: These fundamental components of AD underpin both proposed algorithms, enabling the authors to compute derivatives effectively and employ novel contraction techniques to handle large parameter spaces efficiently.
• Structured Derivatives: By identifying and taking advantage of structured derivatives in neural networks, the authors realize further computational savings. This reduces the operational complexity associated with NTK computation and prevents exponential growth in resource usage with network size.

Implications and Speculations

The implications of this work extend across several domains of machine learning and AI. Efficient NTK computation unlocks new avenues for research in model initialization, architecture search, and the exploration of generalization properties in neural networks. These improvements foster greater model accuracy and training efficiency, which are critical in deploying large-scale, real-world applications.

In future developments, the principles underlying these algorithms could be extrapolated to enhance other computationally intensive tasks in AI research, such as computing the Fisher Information matrix and advancing kernel methods in neural networks. The integration of efficient NTK calculation into mainstream machine learning libraries could potentially standardize its use in model evaluation, leading to widespread adoption beyond theoretical studies.

Conclusion

This paper's contribution lies not only in the theoretical underpinnings of faster NTK computation but also in its practical implications and application potential. By transforming NTK computation from an intractable resource-heavy task to a manageable operation, the authors pave the way for enhanced neural network analysis and deployment. This work stands to significantly impact how NTKs are utilized across machine learning research and practice, contributing meaningfully to the predictive power and optimization of neural networks.