Fast Generalized Neural Tangent Kernel Statistics via Trace Estimation

James Hazelden; Balaaji Reddy Nagireddy; Eric Shea-Brown

cs.LGNov 13, 2025

Fast Generalized Neural Tangent Kernel Statistics via Trace Estimation

Authors: James Hazelden, Balaaji Reddy Nagireddy, Eric Shea-Brown

Organizations: Applied Mathematics, University of Washington & The Allen Institute, Seattle, WA.

Abstract

The empirical state-space Neural Tangent Kernel (NTK) describes the local learning geometry of a finite-width neural network, but computing it explicitly is almost always impractical in terms of computation and memory costs. Here, we show that many useful NTK statistics that characterize, for example, the dimensionality of learned updates or how two models or learning rules relate, can instead be efficiently approximated to very high accuracy via matrix-free products using randomized trace estimation. Namely, we use Hutch++ to estimate the NTK trace, Frobenius norm, effective rank, and alignment. Furthermore, we show that the positive-semidefinite structure of the NTK yields one-sided estimators that require only forward- or reverse-mode automatic differentiation. We validate these estimators across MLPs, recurrent GRUs, and a natural-language Transformer with up to 410 million parameters, in which the state-space contains high-dimensional four-tensors. We demonstrate orders-of-magnitude speedups, with the fastest estimator in a given application depending on the ratio of parameter and state dimensions. Equipped with these estimators, we examine rich and lazy RNN training using hidden-state NTK alignment and use NTK alignment as a regularizer for data-scarce knowledge distillation. We find that this regularization can modestly improve generalization, especially in very data-scarce settings. Together, these results suggest state-space NTK diagnostics are practical even at large scales.

Figures & tables

Figure 1: Paper Overview: The empirical state-space NTK and trace estimation. The NTK acts on the full model state, which can form a high-dimensional tensor, and is accessed here only through randomized matrix-vector products (matvecs). For representative MLP and GRU examples (top right, see [§ 5 ]), the fastest of the three estimators discussed achieves 99% accuracy about $80$ and $10{,}000$ times faster than exact trace evaluation, respectively. Curves show medians over 50 runs with interquartile ranges.

	Pros	Cons
Hutch++	$O(1/m^{2})$ variance guarantee and robust	May be bottlenecked by either forward- or backward-mode AD
One-Sided	Only either forward- or backward-AD	Worse $O(1/m)$ variance guarantee

Table 2

Figure 2: Estimating

\text{tr}(\texttt{NTK})

when

h(\theta)

tracks the hidden unit activations of a GRU neural network at every evaluation time. A illustrates the trace estimates for Hutch++ and the one-sided Hutch estimator (Algorithms 1 and 2 , see Figure 7 for an analogous result for MLPs). B illustrates relative error

\frac{|\text{tr}(\texttt{NTK})-t_{m}|}{|\text{tr}(\texttt{NTK})|}

of each estimate,

t_{m}

versus runtime in seconds. In this case, exactly computing the trace with

n

matvecs took 1557.72 seconds (around 26 minutes). C shows the low-runtime region of B in more detail with linearly scaled runtime. In contrast to Fig. 7 , FHutch performed fastest and both one-sided estimators outperformed Hutch++ for error above

2\cdot 10^{-3}

. Hutch++ was fastest for accuracy above 99.9%.

Figure 3: Matrix-free estimation of

\text{tr}(\texttt{NTK})

for the hidden-state representations of released Pythia Transformers. The state

h(\theta)

contains 393,216 entries, and its Jacobian is taken with respect to 353,822,720 parameters. A shows trace estimates for Hutch++, RHutch, and FHutch against randomized sample count

m

. B shows relative error

\frac{|t_{m}-\text{tr}(\texttt{NTK})|}{|\text{tr}(\texttt{NTK})|}

versus runtime. C shows the corresponding speedup (we plot the fastest of the 3 methods in the graph) relative to exact state-basis trace evaluation. Curves show medians over 20 estimates, with shaded regions indicating the

25^{\mathrm{th}}

and

75^{\mathrm{th}}

percentiles.

Figure 4: Relative runtime versus error for an MLP trained on MNIST, estimating the Frobenius norm, kernel alignment, and the effective rank, respectively. Here,

\texttt{NTK}_{f}

denotes the NTK of the model post-training while

\texttt{NTK}_{0}

is the NTK pre-training. A–C plot the relative error versus the exact value, computed with exact trace evaluation. We plot these versus the percentage of total runtime required to compute the exact expression in each case (range

0\%

100\%

). Curves correspond to the median over 50 re-evaluations with Hutch++ and the shaded regions illustrate the

25^{\text{th}}

and

75^{\text{th}}

percentiles.

Figure 5: Initialization-dependent rich and lazy reorganization of recurrent representations and NTK during learning on sequential MNIST. As in ( Liu et al., 2024 ) , we swept the rank of the recurrent weights

W_{0}

before training, with all models achieving above 95% testing accuracy. A Quantifying drift (rich learning) in the RNN hidden state and hidden-state NTK before and after training per run. Each dot represents a distinct training run. B Hidden-state NTK alignment over training as test accuracy increases for four representative RNN initializations. C Corresponding NTK effective rank, showing continued expansion of tangent geometry during learning. Curves and error bars correspond to the standard error of the mean (SEM) over 3 seeds per run.

NTK Space	$\lambda$	$N$	$\Delta$ test acc. (pp)	Paired wins
Logits	0.7	250	$+0.62\pm 0.15$	$5/6$
Logits	0.7	500	$+0.77\pm 0.12$	$6/6$
Logits	0.7	1000	$+0.53\pm 0.16$	$6/6$
Hidden	0.5	250	$+0.98\pm 0.22$	$3/3$
Hidden	0.5	500	$+0.49\pm 0.23$	$3/3$
Hidden	0.5	1000	$-0.21\pm 0.16$	$1/3$

Table 7

Figure 6: Teacher-student NTK alignment as a regularizer for data-limited Fashion-MNIST students. A student is trained on a pretrained teacher’s output with limited training data and an NTK regularization term with factor

\lambda

. A Final test-accuracy gain relative to matched cross-entropy-only students as a function for training set size

N\in\{250,500,1000\}

. Positive values indicate improved generalization B Teacher-NTK alignment during training for cross-entropy-only and NTK-regularized students (

\lambda=0.4

), showing rapid and sustained alignment with the teacher geometry. C Test-accuracy gain relative to cross-entropy-only training over a longer training horizon for

N=250

and

N=500

. Shaded regions in A and C correspond to SEM over reruns.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Estimating

\text{tr}(\texttt{NTK})

when

h(\theta)

tracks the output of an MLP neural network. A illustrates the trace estimates for the three approaches in [§ 4 ]: Hutch++ (Algorithm 1 ) and the one-sided estimator (Algorithm 2 ) with forward- and backward-AD (FHutch and RHutch, respectively). B shows the relative error

\frac{|\text{tr}(\texttt{NTK})-t_{m}|}{|\text{tr}(\texttt{NTK})|}

of each estimate,

t_{m}

versus runtime in seconds. For reference, exactly computing the trace took 1.96 seconds. C zooms into the low-runtime region of B and displays results with linearly scaled runtime, showing that FHutch is the slowest for this setup. Note Hutch++ obtained relative error below

10^{-6}

in about 0.06 seconds (about

33\times

faster than the exact trace).

	Pythia-14M	Pythia-410M-deduped
Blocks $L$	6	24
Hidden width $d$	128	1,024
Attention heads	4	16
Feed-forward width	512	4,096
Sequences $B$ / token positions $T$	$1\,/\,16$	$1\,/\,16$
State dimension $n=LBTd$	12,288	393,216

Appendix

Table 1: Dimensions of the numerical validation experiment where all blocks and 16 tokens are used. The parameter count excludes the vocabulary output projection.

Figure 8: Matrix-free estimation of

\text{tr}(\texttt{NTK})

for the hidden-state representations of released Pythia Transformers. For each model,

h(\theta)

contains the hidden representations of all 16 tokens at the output of every Transformer block A shows trace estimates against randomized sample count,

m

. B shows relative error

|t_{m}-\text{tr}(\texttt{NTK})|/|\text{tr}(\texttt{NTK})|

against runtime. C shows speedup of randomized trace estimation relative to the exact state-basis evaluation. At each target accuracy and repeat, we record the fastest measured Hutch++, RHutch, or FHutch estimate satisfying the target tolerance and divide the exact-computation runtime by the estimator runtime. The most stringent-accuracy tail points should be interpreted cautiously when only a small number of estimates satisfy the target tolerance. Curves show medians over 20 estimates, and shaded regions show the

25^{\mathrm{th}}

and

75^{\mathrm{th}}

percentiles.

Explore similar work

The Global Empirical NTK: Self-Referential Bias and Dimensionality of Gradient Descent Learning
May 9, 2026James Hazelden, Laura Driscoll, Eli Shlizerman +1Neural Tangent Kernel Gradient Descent
Applied Mathematics, University of Washington; Allen Institute · Allen Institute · Neurobiology & Biophysics, University of Washington +2
Label-NTK Alignments and A Tighter Convergence Bound in the NTK Regime
May 24, 2026Ruchirinkil Marreddy, Chaoyue LiuNeural Tangent Kernel Overparameterization
Elmore Family School of Electrical and Computer Engineering Purdue University
A Function-Space Dichotomy for Compositional Learning: Exponential Sub-Optimality of the Neural Tangent Kernel
Jul 7, 2026Arkaprabha Ganguli, Emil ConstantinescuNeural Tangent Kernel Rectified Linear Unit Networks
Mathematics & Computer Science Division, Argonne National Laboratory, Lemont, IL, USA

May 9, 2026, cs.LGJ/K move · Enter open · S save

The Global Empirical NTK: Self-Referential Bias and Dimensionality of Gradient Descent Learning

James Hazelden, Laura Driscoll, Eli Shlizerman, Eric Shea-Brown

Applied Mathematics, University of Washington; Allen Institute · Allen Institute · Neurobiology & Biophysics, University of Washington+2

Neural Tangent Kernel Gradient Descent Recurrent Model Self-Referential Bias

Open paper AskPDF arXiv

In training a neural network with gradient descent (GD), each iteration induces a linear operator that governs first-order updates to a model's internal state variables. We define this operator as the Global Empirical Neural Tangent Kernel (NTK). In finite-width networks, the NTK is typically intractable to form, leading prior work to focus on restrictive settings such as tracking outputs only or taking infinite-width limits. Here, we study the structure of the NTK for a range of models. Formulating the model state as the solution to a single global implicit constraint, we derive the NTK as a product of two operators: K, accounting for immediate parameter-to-state interactions, and P, describing internal state-to-state dependencies. For a broad class of weight-based models, including RNNs and transformers, we prove a universal Kronecker-core theorem showing that K admits an exact, computable form given by the Gram matrix of weight-site variables. This core structure reveals that the NTK is structurally bottlenecked, constraining its effective rank and giving rise to a self-referential bias whereby GD preferentially learns within dominant modes of joint hidden and input activity. For recurrent models, we examine the spectrum of the NTK and show when it is biased and low-rank in space or time under the proposed decomposition. We further demonstrate that model dynamics at initialization bias the NTK, restricting learning and preventing task components from being learned effectively. Finally, we show that the NTK associated with a self-attention transformer is likewise structurally constrained to be low-rank. Overall, we show that the NTK possesses tractable structure that explains GD bias toward task solutions and the emergence of low-rank representations. To enable use of the NTK as a practical metric, we build kpflow, a library relying on randomized matrix-free numerical linear algebra.

Abstract

Explore similar work

The Global Empirical NTK: Self-Referential Bias and Dimensionality of Gradient Descent Learning

Label-NTK Alignments and A Tighter Convergence Bound in the NTK Regime

A Function-Space Dichotomy for Compositional Learning: Exponential Sub-Optimality of the Neural Tangent Kernel