Fast Neural Tangent Kernel Alignment, Norm and Effective Rank via Trace Estimation

TL;DR

This paper introduces fast, matrix-free methods for analyzing the Neural Tangent Kernel (NTK), enabling rapid computation of key properties like trace, norm, and rank, which accelerates NTK research and applications.

Contribution

The paper presents a novel trace estimation approach for the NTK that is faster and more efficient, especially for large models, using only forward or reverse automatic differentiation.

Findings

Trace estimation methods outperform traditional approaches in speed.

One-sided estimators are effective in low-sample regimes.

Matrix-free approaches enable orders-of-magnitude speedups.

Abstract

The Neural Tangent Kernel (NTK) characterizes how a model's state evolves over Gradient Descent. Computing the full NTK matrix is often infeasible, especially for recurrent architectures. Here, we introduce a matrix-free perspective, using trace estimation to rapidly analyze the empirical, finite-width NTK. This enables fast computation of the NTK's trace, Frobenius norm, effective rank, and alignment. We provide numerical recipes based on the Hutch++ trace estimator with provably fast convergence guarantees. In addition, we show that, due to the structure of the NTK, one can compute the trace using only forward- or reverse-mode automatic differentiation, not requiring both modes. We show these so-called one-sided estimators can outperform Hutch++ in the low-sample regime, especially when the gap between the model state and parameter count is large. In total, our results demonstrate that matrix-free randomized approaches can yield speedups of many orders of magnitude, leading to faster analysis and applications of the NTK.