TLDR: Eigenanalysis of the empirical Neural Tangent Kernel (eNTK) can identify features learned by trained neural networks in toy models. The research demonstrates that eNTK exhibits sharp spectral cliffs whose top eigenspaces align with ground-truth features in “Toy Models of Superposition” and Fourier features in modular arithmetic. This method also provides a diagnostic for detecting grokking phase transitions and localizing features to specific network layers.
Understanding how neural networks process information and represent learned features is a fundamental challenge in mechanistic interpretability. Researchers are constantly seeking methods to “reverse-compile” these complex models into human-understandable components, revealing the effective features they use and how these interact to produce outputs.
A recent research paper, “Feature Identification via the Empirical NTK”, by Jennifer Lin, explores a promising avenue: using eigenanalysis of the empirical Neural Tangent Kernel (eNTK) to uncover these hidden features. The Neural Tangent Kernel (NTK) is a theoretical tool that provides a closed-form formula for the function learned by a neural network under specific conditions. While realistic models often operate beyond these strict conditions, the empirical NTK (eNTK), which evaluates the NTK at the end of training, has been conjectured and empirically verified to be a good approximation in many situations.
The core hypothesis of this work is that the eigenvectors of the eNTK may align with significant features learned by a network. The study provides compelling evidence for this in two standard toy models used in mechanistic interpretability: Toy Models of Superposition (TMS) and a one-layer Multi-Layer Perceptron (MLP) trained on modular addition.
Insights from Toy Models of Superposition (TMS)
In the TMS setup, a simple autoencoder learns to reconstruct synthetic features. The research found that the eNTK eigenspectrum exhibits distinct “cliffs,” which are sharp drops in the ordered eigenvalues. Crucially, the leading eigenspace—the set of eigenvectors corresponding to these top eigenvalues—aligns almost perfectly with the ground-truth features the model is designed to learn. This alignment was observed in both sparse (high superposition) and dense (low superposition) regimes, indicating the robustness of the method. The study also showed that by adjusting an “importance rescaling” parameter, the eNTK could reveal both high-importance and lower-importance features.
Discoveries in Modular Arithmetic (Grokking)
The second experiment involved an MLP trained on modular arithmetic, a task known to exhibit “grokking”—a phenomenon where a model suddenly generalizes to unseen data long after achieving perfect training accuracy. Here, the eNTK spectrum developed two distinct cliffs. The first cliff was found to localize to the first layer of the model and its span matched the Fourier features learned by that layer. Through a technique called “graph smoothness,” the researchers were able to rotate the eNTK eigenvector basis within this subspace to precisely align with the expected Fourier features.
Even more remarkably, the second cliff in the eNTK spectrum appeared precisely at the grokking phase transition. This cliff aligned with “sum” and “difference” Fourier features. This suggests that eNTK spectral analysis can serve as a practical diagnostic tool, not only for localizing features to specific layers but also for detecting significant phase changes in a model’s learning process. Interestingly, while the eigenvalues associated with this second cliff jumped at grokking, the eigenvectors themselves evolved smoothly, indicating a continuous formation of these features.
The eNTK Pipeline for Feature Discovery
The methodology involves eigendecomposing the eNTK and its variants. A “cliff” is defined as a significant drop in eigenvalues, and the “eNTK feature subspace” is the span of eigenvectors corresponding to these eigenvalues. To assess alignment, the cosine similarity between eNTK eigenvectors and ground-truth feature vectors is computed. Layerwise NTK variants are used to attribute features to specific layers. For the modular arithmetic task, a “graph smoothness” algorithm was employed to disentangle and identify specific Fourier modes within the eNTK eigenspaces.
Also Read:
- Understanding Neural Network Performance: New Insights into Scaling Laws and Weight Spectra
- Unlocking the Mathematical Foundations of Vector-Valued Neural Networks and Operators
Future Directions
The findings suggest that eNTK analysis offers a practical approach for feature discovery and diagnosing phase changes in smaller models. Future work aims to extend this pipeline to more complex models like Transformers, including those trained on realistic datasets. Understanding the analytical properties of the weights and data that lead to these clear mechanistic signals is another interesting direction. Furthermore, exploring ways to identify features without relying on explicit structural priors, perhaps by learning such priors directly from the kernel, could enhance the method’s generalizability.
