Elamur Research⌕
☰
home / papers / 2022 / Towards Understanding Grokking: An Effective Theory of Representation Learning

Towards Understanding Grokking: An Effective Theory of Representation Learning

Paper · 2022 · Ziming Liu, Ouail Kitouni, Niklas Nolte et al.

Summary

Grokking is the behavior where a neural network fits its training data early but only begins to generalize much later. The authors ask three questions: how such models generalize at all, why the time to generalize grows sharply as the training set shrinks toward a threshold, and when the delay appears. Their main testbed is a simple model. Two input symbols get learnable embeddings, which are summed and fed to an MLP decoder trained on addition, as either regression or classification. They also run a transformer on modular addition and a classifier on MNIST. To measure structure, they define a representation quality index (RQI), which counts how many arithmetically consistent parallelograms the embeddings form. They also build an effective theory in which embeddings evolve under a simplified loss that ignores the decoder. This theory predicts a critical training fraction near 0.4 and matches how embeddings actually evolve. Grid searches over hyperparameters reveal four regimes: comprehension, grokking, memorization and confusion. Good representations arise only in an intermediate "Goldilocks" band. The work explains where grokking comes from and shows that hyperparameter tuning can eliminate the delay.

Key points

  • Generalization coincides with structured embeddings, such as parallelograms for addition and a circle for modular addition. Accuracy predicted from RQI closely tracks measured accuracy.
  • The effective theory and the experiments both show a phase transition around a training data fraction of r_c ≈ 0.4. Below this fraction, a linear representation is unlikely to form.
  • The four learning phases depend on decoder capacity and learning speed, which is shaped by factors like learning rate and weight decay. In transformers, the grokking phase lies nearer to memorization than to comprehension.
  • In the transformer, weight decay on the decoder speeds up generalization. Heavy dropout brings generalization below 10³ steps and makes grokking disappear.
  • The authors report the first observation of grokking on a standard benchmark (MNIST) and control it via optimization hyperparameters. Code is at github.com/ejmichaud/grokking-squared (NeurIPS 2022).

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, topics and concepts

📚 Shelves 2The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 works

Cite

Cite as Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark and Mike Williams. 2022. Towards Understanding Grokking: An Effective Theory of Representation Learning. arXiv:2205.10343.

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

@misc{liu2022understanding,
  title = {{Towards Understanding Grokking: An Effective Theory of Representation Learning}},
  author = {Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark and Mike Williams},
  year = {2022},
  doi = {10.48550/arxiv.2205.10343},
  eprint = {2205.10343},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2205.10343}
}

References

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

  1. 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.
  2. Neel Nanda and Tom Lieberum. A mechanistic interpretability analysis of grokking, 2022. URL https://www.alignmentforum.org/posts/N6WM6hs7RQMKDhYjB/ a-mechanistic-interpretability-analysis-of-grokking. ↗
  3. Beren Millidge. Grokking ’grokking’. 2022-01-11-Grokking-Grokking/, 2022. https://beren.io/Shah. Alignment Newsletter \#159. https: ↗
  4. Rohin //www.alignmentforum.org/posts/zvWqPmQasssaAWkrj/ an-159-building-agents-that-know-how-to-experiment-by\#DEEP_LEARNING_, 2021. find
  5. Yedid Hoshen and Shmuel Peleg. Visual learning of arithmetic operation. In AAAI, 2016. find
  6. Yang-Hui He. Machine-learning mathematical structures. arXiv preprint arXiv:2101.06317, 2021. — [arXiv:2101.06317](https://arxiv.org/abs/2101.06317) find
  7. Sergei Gukov, James Halverson, Fabian Ruehle, and Piotr Sułkowski. Learning to unknot. Machine Learning: Science and Technology, 2(2):025035, 2021. find
  8. Alex Davies, Petar Veličković, Lars Buesing, Sam Blackwell, Daniel Zheng, Nenad Tomašev, Richard Tanburn, Peter Battaglia, Charles Blundell, András Juhász, et al. Advancing mathematics by guiding human intuition with ai. Nature, 600(7887):70–74, 2021. find
  9. Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechanics: Theory and Experiment, 2021(12):124003, 2021. find
  10. Preetum Nakkiran, Prayaag Venkat, Sham Kakade, and Tengyu Ma. Optimal regularization can mitigate double descent. arXiv preprint arXiv:2003.01897, 2020. — [arXiv:2003.01897](https://arxiv.org/abs/2003.01897) find
  11. Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8): 1798–1828, 2013.
  12. Yassine Ouali, Céline Hudelot, and Myriam Tami. An overview of deep semi-supervised learning. arXiv preprint arXiv:2006.05278, 2020. — [arXiv:2006.05278](https://arxiv.org/abs/2006.05278) find
  13. Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in Neural Information Processing S… find
  14. Phuc H Le-Khac, Graham Healy, and Alan F Smeaton. Contrastive representation learning: A framework and review. IEEE Access, 8:193907–193934, 2020. find
  15. James Halverson, Anindita Maiti, and Keegan Stoner. Neural networks and quantum field theory. Machine Learning: Science and Technology, 2(3):035002, 2021. find
  16. Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint arXiv:2106.10165, 2021. — [arXiv:2106.10165](https://arxiv.org/abs/2106.10165) find
  17. Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.04728, 2020. — [arXiv:2012.04728](https://arxiv.org/abs/2012.04728) find
  18. Yansong Gao and Pratik Chaudhari. A free-energy principle for representation learning. In International Conference on Machine Learning, pages 3367–3376. PMLR, 2020. find
  19. Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc Mézard, and Lenka Zdeborová. Generalisation error in learning with random features and the hidden manifold model. In International Conference on Machine Learning, pages 3452–3462. PMLR, 2020. find
  20. Mohammad Pezeshki, Amartya Mitra, Yoshua Bengio, and Guillaume Lajoie. Multi-scale feature learning dynamics: Insights for double descent. In International Conference on Machine Learning, pages 17669–17690. PMLR, 2022. find
  21. Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mezard, and Lenka Zdeborova. The gaussian equivalence of generative models for learning with shallow neural networks. In Joan Bruna, Jan Hesthaven, and Lenka Zdeborova, editors, Proceedings of the 2nd Mathematical and Scientific M… find
  22. R Kuhn and S Bos. Statistical mechanics for neural networks with continuous-time dynamics. Journal of Physics A: Mathematical and General, 26(4):831, 1993. find
  23. 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… — In-context Learning and Induction Heads
  24. Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory. Cambridge University Press, 2022. https://deeplearningtheory.com. ↗
  25. Deep Ganguli, Danny Hernandez, Liane Lovitt, Nova DasSarma, Tom Henighan, Andy Jones, Nicholas Joseph, Jackson Kernion, Ben Mann, Amanda Askell, et al. Predictability and surprise in large generative models. arXiv preprint arXiv:2202.07785, 2022.
  26. Jacob Steinhardt. Future ML Systems Will Be Qualitatively Different. https://www.lesswrong.com/s/4aARF2ZoBpFZAhbbe/p/pZaPhGg2hmmPwByHc, 2022. ↗
  27. Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017. find
  28. Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117 (40):24652–24663, 2020. Thomson problem — Wikipedia, the free encyclope-[29] Wikipedia contributors. dia. https://en.wiki… ↗
  29. Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15750–15758, 2021. find
  30. Zhi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. In International Conference on Neural Information Processing, pages 264–274. Springer, 2019. find
  31. Yaoyu Zhang, Zhi-Qin John Xu, Tao Luo, and Zheng Ma. A type of generalization error induced by initialization in deep neural networks. In Mathematical and Scientific Machine Learning, pages 144–164. PMLR, 2020. find
  32. Ziming Liu, Eric J. Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data, 2022.
  33. Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 o… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/liu-2022-towards-understanding-grokking-effective-theory/