Feature Identification via the Empirical NTK

TL;DR

This work investigates whether eigenvectors of the empirical NTK at the end of training ($e$NTK) reveal the features learned by neural networks. By analyzing two toy benchmarks—Toy Models of Superposition (TMS) and a 1-layer MLP trained on modular addition—the authors show sharp spectral cliffs in the $e$NTK spectrum whose top eigenspaces align with ground-truth features, with layerwise $e$NTK pinpointing feature localization by layer. They also demonstrate that the evolution of the $e$NTK spectrum tracks grokking and can diagnose phase transitions in small models, using techniques like a two-stage graph-smoothness rotation to surface Fourier features. Across these results, the paper proposes a practical pipeline for feature discovery via $e$NTK eigenanalysis, with potential extensions to larger architectures like Transformers for mechanistic interpretability. This approach offers a concrete link between kernel spectra and interpretable representations learned by neural networks, and suggests new diagnostic tools for training dynamics and representation learning.

Abstract

We provide evidence that eigenanalysis of the empirical neural tangent kernel (eNTK) can surface the features used by trained neural networks. Across two standard toy models for mechanistic interpretability, Toy Models of Superposition (TMS) and a 1-layer MLP trained on modular addition, we find that the eNTK exhibits sharp spectral cliffs whose top eigenspaces align with ground-truth features. In TMS, the eNTK recovers the ground-truth features in both the sparse (high superposition) and dense regimes. In modular arithmetic, the eNTK can be used to recover Fourier feature families. Moreover, we provide evidence that a layerwise eNTK localizes features to specific layers and that the evolution of the eNTK spectrum can be used to diagnose the grokking phase transition. These results suggest that eNTK analysis may provide a practical handle for feature discovery and for detecting phase changes in small models.

Figures (5)

• Figure 1: Eigenvalue spectrum of the flattened eNTK (top row) and column-normalized, importance-rescaled flattened eNTK eigenvector-feature indicator heatmaps for TMS models with $n = 50$ ground-truth features, importance $I_i = 0.8^{i}$, (hidden-layer size, sparsity) $m = 10, S= 0.3$ (left column); $m=10, S = 0.9$ (middle column) or $m=40, S = 0.9$ (right column); and importance rescaling parameters $\beta = 1$ (middle row) and $\beta = 0.3$ (bottom row) used as described in the main text. The center column is an example of a TMS model in the sparse regime (high superposition), while the left and right columns are two different ways to push the model to the dense regime (low superposition).
• Figure 2: Full and layerwise eNTK spectrum of the modular arithmetic model after training to convergence (500 epochs). The left column displays the spectrum of the full eNTK, the middle column the spectrum of the layerwise eNTK wrt the parameters \ref{['l1']} in layer 1, and the right column the spectrum of the layerwise eNTK wrt the parameters \ref{['l2']} in layer 2.
• Figure 3: Spectral change at the grokking phase transition. Left: Train and test accuracy shows evidence for a grokking phase transition (sudden onset of generalization) around epoch 90. Right: at the same training time, the eNTK spectrum develops a kink at the base of the second cliff.
• Figure 4: Results from applying the two-stage per-axis graph smoothness algorithm described in Appendix \ref{['sA']} to rotate the eNTK eigenvector basis in the Layer 1 cliff. Left: heatmap showing the squared norm of the inner product of eNTK eigenvectors in the cliff with the $(\cos a, \sin a)$ or $(\cos b, \sin b)$ feature families. Right: heatmap showing the squared norm of the inner product of basis vectors in the Layer 1 cliff found from applying the two-stage graph smoothness algorithm described in Appendix \ref{['sA']} with the $(\cos a, \sin a)$ or $(\cos b, \sin b)$ feature families.
• Figure 5: Results from applying the two-stage graph smoothness algorithm described in Appendix \ref{['sA']} wrt the sum and difference Laplacians \ref{['e131']}, \ref{['e141']} to rotate the eNTK eigenvectors in the second eigenvalue cliff of the modular arithmetic experiment at different epochs. Each plot is a heatmap showing the squared norm of the inner product of rotated basis vectors in the second eigenvalue cliff with the $(\cos a+b, \sin a+b)$ or $(\cos a-b, \sin a-b)$ feature families.