Elamur Research
☰
home / papers / 2023 / Progress measures for grokking via mechanistic interpretability

Progress measures for grokking via mechanistic interpretability

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive licence to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2301.05217

Summary

Deep networks sometimes gain new abilities abruptly, which makes their behavior hard to anticipate. This work asks whether such jumps hide smooth underlying progress, and whether interpretability can reveal it. The test case is grokking: small transformers learning addition modulo a prime first memorize the training set and only generalize much later. The authors fully reverse-engineer a one-layer transformer trained on this task. The model embeds each input as sines and cosines at a few key frequencies, combines them with trigonometric identities in its attention and MLP layers, and reads out the answer by making waves interfere constructively at the correct sum. They call this algorithm Fourier multiplication. They support it with periodic structure in weights and activations, approximations of individual neurons, and ablations in Fourier space. Building on this mechanism, they define two progress measures, restricted loss and excluded loss. Both improve steadily well before test accuracy jumps. This shows that grokking is gradual circuit formation followed by weight-decay-driven removal of memorization, rather than a sudden discovery.

Key points

  • Setup: P = 113, a one-layer ReLU transformer (d = 128, 4 heads of size 32, 512 MLP units, no LayerNorm), trained on 30% of all 113·113 input pairs. Training used full-batch AdamW (lr 0.001, weight decay 1) for 40,000 epochs, and generalization appeared around epoch 10,000.
  • The neuron-logit map W_L is roughly rank 10, covering the sine and cosine of 5 key frequencies. Most neurons and attention heads fit well as degree-2 trigonometric polynomials of a single frequency.
  • Removing the key frequencies drops accuracy to chance level. Removing the remaining ~95% of frequencies slightly improves performance.
  • Training splits into three phases: memorization, circuit formation, and cleanup. The sharp rise in test accuracy happens during cleanup, and grokking did not appear without regularization.
  • Code and interactive figures are available at neelnanda.io/grokking-paper.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Mathematics & statistics

Cited by

156 papers of our library cite it (of the 140673 whose bibliographies we read); 3 of them have a page here, 153 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2301.05217 · arXiv 2301.05217

@misc{nanda2023progress,
  title = {{Progress measures for grokking via mechanistic interpretability}},
  author = {Nanda, Neel and Chan, Lawrence and Lieberum, Tom and Smith, Jess and Steinhardt, Jacob},
  year = {2023},
  doi = {10.48550/arxiv.2301.05217},
  eprint = {2301.05217},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2301.05217}
}

References

26 references, 4 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprint arXiv:2207.08799, 2022. find
  2. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. find
  3. Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: Circuits. Distill, 2020. doi: 10.23915/distill.00024. https://distill.pub/2020/circuits. — [doi:10.23915/distill.00024](https://doi.org/10.23915/distill.00024) ↗
  4. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Ja… find
  5. Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:1803.03635, 2018. find
  6. Deep Ganguli, Danny Hernandez, Liane Lovitt, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova Dassarma, Dawn Drain, Nelson Elhage, et al. Predictability and surprise in large generative models. In 2022 ACM Conference on Fairness, Accountability, and Transparency, pp. 1747–1764, 2022. find
  7. Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Rı́o, Mar… find our card
  8. Niall Hurley and Scott Rickard. Comparing measures of sparsity. IEEE Transactions on Information Theory, 55(10):4723–4741, 2009. find
  9. Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. arXiv preprint arXiv:2205.10343, 2022. find
  10. Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017. find
  11. Decoupled weight decay regularization. Thomas McGrath, Andrei Kapishnikov, Nenad Tomašev, Adam Pearce, Demis Hassabis, Been Kim, Ulrich Paquet, and Vladimir Kramnik. Acquisition of chess knowledge in alphazero. arXiv preprint arXiv:2111.09259, 2021. — [arXiv:2111.09259](https://arxiv.org/abs/2111.09259) find
  12. Beren Millidge. Grokking ’grokking’, 2022. 2022-01-11-Grokking-Grokking/. find
  13. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amode… find
  14. Alexander Pan, Kush Bhatia, and Jacob Steinhardt. The effects of reward misspecification: Mapping and mitigating misaligned models. arXiv preprint arXiv:2201.03544, 2022. — [arXiv:2201.03544](https://arxiv.org/abs/2201.03544) find
  15. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. find our card
  16. Plotly Technologies Inc. Collaborative data science, 2015. URL https://plot.ly. arXiv preprint URL https://www.beren.io/ ↗
  17. Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022. find
  18. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  19. Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein-like notation. In International Conference on Learning Representations, 2022. URL https://openreview. net/forum?id=oapKSVM2bcj. find
  20. Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929–1958, 2014. find
  21. Jacob Steinhardt. More is different for ai, Feb 2022. URL https://bounded-regret. ghost.io/more-is-different-for-ai/. find
  22. Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206.04817, 2022. — [arXiv:2206.04817](https://arxiv.org/abs/2206.04817) find
  23. Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593, 2022. find
  24. Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682, 2022a. find our card
  25. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022b. find our card
  26. Wes McKinney. Data Structures for Statistical Computing in Python. In Stéfan van der Walt and Jarrod Millman (eds.), Proceedings of the 9th Python in Science Conference, pp. 56 – 61, 2010. doi: 10.25080/Majora-92bf1922-00a. find