On this page
Maximilian Nicholson · University of Bath
Read the full paper (PDF, 49 pages)
Suggested viewing before diving into the paper: The most complex model we actually understand — Welch Labs.
For a written introduction, read Understanding grokking in terms of representation learning dynamics — Eric J. Michaud. It explores how structured representations emerge during training, with visualisations of learned embeddings and a representation-learning perspective on grokking.
Where and when does a learned algorithm become visible in a network’s parameter-gradient geometry? Mechanistic analyses of modular addition locate the trigonometric computation in the MLP, so one might expect the MLP empirical neural tangent kernel (eNTK) to show the clearest task structure. After independently normalizing each parameter block, we find the opposite: in a 20-seed, p = 113 experiment, final unembedding alignment completely separates seeds grokked by a 7,000 epoch landmark from right-censored seeds (with Mann–Whitney p = 0.0002, r_rb = 1.00), tracking final accuracy (ρ = 0.928). A prospectively locked GELU with pre-LayerNorm replication reproduces this normalized result.
The apparent contradiction dissolves through the readout identity derived by Lin [2025]: Kₜ⁽ᵘⁿᵉᵐᵇᵉᵈ⁾ = pGₜ in our scalarized convention, making the normalized readout kernel exactly the geometry of the final residual representation. That visibility does not imply universal additive dominance. The unembedding holds a median 90.8% of endpoint target mass in the primary ReLU/no-LayerNorm model, whereas the GELU+pre-LayerNorm condition assigns median shares of 66.9% to embeddings, 28.1% to attention, and 2.85% to the unembedding. Activation Grams then reveal where the visible representation arises: addition alignment is negligible through attention, rises to 0.524 at the MLP output, and remains high downstream.
Finally, in the ten grokked primary models, immediate removal of a rank-5 eNTK-guided sum-orbit subspace reduces mean accuracy from 0.9987 to about 0.10, while retaining only it preserves about 0.25. This supports necessity for these evaluated solutions and partial immediate sufficiency, not unique mechanistic ownership. The readout is therefore a window onto an upstream representation, while visibility, additive attribution, pathway emergence, and functional relevance remain distinct.