Elamur Research⌕
☰
home / papers / 2022 / Omnigrok: Grokking Beyond Algorithmic Data

Omnigrok: Grokking Beyond Algorithmic Data

Paper · 2022 · Ziming Liu, Eric J. Michaud and Max Tegmark

Summary

Grokking is when a network fits its training data early but only begins to generalize much later. It had been documented mainly on algorithmic tasks, and its cause was unclear. Liu, Michaud and Tegmark study it through loss-landscape geometry. They use a "reduced" loss, obtained by optimizing over weight directions while holding the L2 weight norm fixed. Plotted against weight norm, the training loss falls and then stays flat (an "L"), while the test loss is low only within a band of norms called the Goldilocks zone (a "U"). They name this the LU mechanism. A model initialized with an overly large norm overfits quickly, then drifts slowly toward the generalizing zone under weak regularization such as weight decay. Guided by this picture, the authors enlarge initialization scale to induce grokking on MNIST, IMDb reviews and QM9. Conversely, they nearly remove grokking on algorithmic data by constraining weight norm. They attribute the especially abrupt grokking on algorithmic tasks to representation learning. The work offers a simple, unifying account of delayed generalization.

Key points

  • Toy teacher-student setup.
    • Architecture: a 5-100-100-5 tanh MLP.
    • Data: 100 training and 100 test samples.
    • Training: Adam (lr 3×10⁻⁴) for 10⁵ steps.
    • Students start at scale α = 2.0 or 0.5, combined with weight decay γ ∈ {0, 0.03, 1}.
  • Effect of initialization and weight decay.
    • With large initialization (α = 2.0):
      • No regularization never generalizes.
      • Small γ produces grokking.
      • Large γ generalizes fast.
    • With small initialization (α = 0.5), the model generalizes quickly regardless of γ.
  • Timing.
    • Steps to reach 95% train accuracy do not depend on γ.
    • Steps to reach 95% test accuracy scale roughly as 1/γ, matching the predicted delay of t ≈ ln(w₀/w_c)/γ.
  • MNIST with cross-entropy loss.
    • Cross-entropy gives a wider Goldilocks zone than MSE, so a much larger scale (α = 100, 200 training samples) is needed.
    • Test accuracy plateaus near 30–40%, then later climbs to about 75% while train accuracy stays at 100%.
  • Algorithmic data and code.
    • Algorithmic tasks grok sharply because a bad representation has a larger w_c than a good one. For base-p addition, the decoder must handle O(p²) versus O(p) cases.
    • Code: github.com/KindXiaoming/Omnigrok.

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

 Institutions

🎓Massachusetts Institute of TechnologyCambridge · United States274 works · 82,013 citations🔬The NSF AI Institute for Artificial Intelligence and Fundamental InteractionsCambridge · United States3 works · 159 citations

 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, Eric J. Michaud and Max Tegmark. 2022. Omnigrok: Grokking Beyond Algorithmic Data. arXiv:2210.01117.

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

@misc{liu2022omnigrok,
  title = {{Omnigrok: Grokking Beyond Algorithmic Data}},
  author = {Ziming Liu, Eric J. Michaud and Max Tegmark},
  year = {2022},
  doi = {10.48550/arxiv.2210.01117},
  eprint = {2210.01117},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2210.01117}
}

References

30 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. 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.
  3. 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
  4. 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
  5. Stanislav Fort and Adam Scherlis. The goldilocks zone: Towards better understanding of neural network loss landscapes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3574–3581, 2019. find
  6. 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
  7. Andrew Ng and Tengyu Ma. Cs229 lecture notes. spring2022/main_notes.pdf, page 115, 2022. https://cs229.stanford.edu/lectures- ↗
  8. Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint arXiv:1611.01232, 2016. — [arXiv:1611.01232](https://arxiv.org/abs/1611.01232) find
  9. Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. Advances in neural information processing systems, 30, 2017. find
  10. Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141–142, 2012. find
  11. Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735–1780, 1997. find
  12. Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142–150, Portland, Oregon, USA,…
  13. Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014. find
  14. Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
  15. Rohin Shah. Alignment Newsletter \#159. https://www.alignmentforum.org/posts/zvWqPmQasssaAWkrj/ an-159-building-agents-that-know-how-to-experiment-by\#DEEP_ LEARNING_, 2021. ↗
  16. Beren Millidge. Grokking ’grokking’. 2022-01-11-Grokking-Grokking/, 2022. https://beren.io/ ↗
  17. Fatih Furkan Yilmaz and Reinhard Heckel. Regularization-wise double descent: Why it occurs and how to eliminate it. In 2022 IEEE International Symposium on Information Theory (ISIT), pages 426–431. find
  18. Lin Chen, Yifei Min, Mikhail Belkin, and Amin Karbasi. Multiple descent: Design your own generalization curve. Advances in Neural Information Processing Systems, 34:8898–8912, 2021. find
  19. Stéphane d’Ascoli, Levent Sagun, and Giulio Biroli. Triple descent and the two kinds of overfitting: Where & why do they appear? Advances in Neural Information Processing Systems, 33:3058–3069, 2020. find
  20. Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249–256. JMLR Workshop and Conference Proceedings, 2010. find
  21. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026–1034, 2015. find
  22. Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam S Schoenholz, Jascha Sohl-Dickstein, and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics, 11(1), 2020. find
  23. Li Jing, Yichen Shen, Tena Dubcek, John Peurifoy, Scott Skirlo, Yann LeCun, Max Tegmark, and Marin Soljačić. Tunable efficient unitary neural networks (eunn) and their application to rnns. In International Conference on Machine Learning, pages 1733–1741. PMLR, 2017. find
  24. Maciej Skorski, Alessandro Temperoni, and Martin Theobald. Revisiting initialization of neural networks. arXiv preprint arXiv:2004.09506, 2020. — [arXiv:2004.09506](https://arxiv.org/abs/2004.09506) find
  25. 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
  26. 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
  27. Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. arXiv preprint arXiv:1810.12281, 2018. — [arXiv:1810.12281](https://arxiv.org/abs/1810.12281) find
  28. Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.05350, 2017. — [arXiv:1706.05350](https://arxiv.org/abs/1706.05350) find
  29. Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with l_2 regularization. Advances in Neural Information Processing Systems, 33:4790–4799, 2020. find
  30. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/liu-2022-omnigrok-grokking-algorithmic-data/