Skip to content
Sagan

Paper

Feature Repulsion and Spectral Lock-in: An Empirical Study of Two-Layer Network Grokking

Unreadunread

AI summary

The paper empirically tests a theoretical "feature repulsion" mechanism from Tian (2025) that predicts similar learned features should repel each other during the grokking phase (delayed generalization) of two-layer network training. The authors confirm the predicted sign structure holds robustly across modular addition tasks, but find that the spectral signature in weight updates—whether repulsion shows up as a rank-2 eigenvalue gap—depends entirely on the activation function. With squared activations (x²), a clear rank-2 spectrum emerges; with ReLU, the spectrum stays rank-1 and repulsion is invisible in the weight updates, even though the sign structure is still correct.

Main takeaways:

  • Tian's feature-repulsion sign rule (similar features have negative off-diagonal entries in the B matrix) holds empirically with high accuracy in grokking setups.
  • The spectral signature of repulsion—a detectable rank-2 eigenvalue gap in weight updates—only appears with x² activation, not ReLU.
  • With x², a simple eigengap detector fires reliably at the grokking transition (epoch ~174) in all 15 seeds, and never in non-grokking controls.
  • With ReLU, the detector never fires and the spectrum remains rank-1, despite identical sign-structure validity.
  • This activation-dependent dissociation aligns with Tian's distinction between "focused" (power-law) and "spreading" (ReLU) memorization dynamics.