arXiv · 2511.10796
Fast Neural Tangent Kernel Alignment, Norm and Effective Rank via Trace Estimation
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.
Explore related subjects
Keep this discovery
James Hazelden. 2025-11-13. Fast Neural Tangent Kernel Alignment, Norm and Effective Rank via Trace Estimation. https://arxiv.org/abs/2511.10796
Cite the original work for its findings. Save a collection to share your selection of sources.