Grokking at the Edge of Numerical Stability

Lucas Prieto · Melih Barsbey · Pedro Mediano^* · Tolga Birdal^*

Imperial College London

Under Review

This is the official implementation of our paper Grokking at the Edge of Numerical Stability (arXiv: 2501.04697). Here you can find guidance to reproduce the main results of the paper.

Replicating Our Results

To replicate the main figures from our paper, use the run_main_experiments.sh script. This script generates and logs all necessary metrics for the primary figures (Figures 1, 2, 4, and 6), and uses cuda:0 by default. Once the relevant metrics have been saved, you can generate the figures from the paper using paper_plots.ipynb.

Usage

./run_main_experiments.sh [--figures FIGURE_LIST] [--device DEVICE]

Requirements

1. Install PyTorch:

   Visit the PyTorch Get Started page to choose the appropriate installation command for your system.

2. Install Python Packages: Install the remaining dependencies using pip:

   pip install pandas==2.2.3
   pip install matplotlib==3.10.0
   

Citation

If you found this code or paper useful, please consider citing:

@article{prieto2025grokking,
  title={Grokking at the Edge of Numerical Stability},
  author={Prieto, Lucas and Barsbey, Melih and Mediano, Pedro and Birdal, Tolga},
  year = {2025},
  eprint={2501.04697},
  archivePrefix={arXiv},
  primaryClass={cs.CV}
}