Grokking

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:

• openai/grok

Paper:

• Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets