This is the code repo for the paper: "Omnigrok: Grokking Beyond Algorithmic Data", accpeted in ICLR 2023 as spotlight. We elucidate the grokking phenomenon from the perspective of loss landscapes, and show that grokking can not only happen for algorithmic datasets and toy teacher-students setups, but also for standard machine learning datasets (e.g., MNIST, IMDb movie reviews, QM9 molecules).
The examples used in this paper are relatively small-scale. We make our codes as minimal as possible: each example is self-consistent, kept in a single folder.
| Examples | Figure in paper | Folder |
|---|---|---|
| Teacher-student | Figure 2 | ./teacher-student |
| MNIST handwritten digits | Figure 3 | ./mnist |
| IMDb Movie Reviews | Figure 4 | ./imdb |
| QM9 Molecule properties | Figure 5 | ./qm9 |
| Modular addition | Figure 6 & 8 | ./mod-addition |
| MNIST Representation | Figure 7 | ./mnist-repr |
For each example, we conduct two kinds of experiments:
- (1) reduced landscape analysis: the weight norm is fixed during training.
- (2) grokking experiments: the weight norm is not fixed during training (standard training).
Each folder (except for MNIST representation) contains two subfolders, for (1) "landscape" and (2) "grokking".