Grokking

May 2, 2026 ยท View on GitHub

An implementation of the OpenAI 'Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets' paper in PyTorch.

Installation

  • Clone the repo and cd into it: 
    git clone https://github.com/danielmamay/grokking.git
cd grokking
  • Install dependencies: 
    uv sync

Usage

The project uses Weights & Biases to keep track of experiments. Run uv run wandb login to use the online dashboard, or uv run wandb offline to store the data on your local machine.

  • To run a single experiment using the CLI:

    uv run python grokking/cli.py

  • To run a grid search using W&B Sweeps:

    uv run wandb sweep sweep.yaml
uv run wandb agent {entity}/grokking/{sweep_id}

Development

  • Type checking: 
    uv run pyright grokking/
  • Linting: 
    uv run ruff check grokking/

References

Code:

Paper: