grokking-mechanism

Grokking on modular addition — training runs

Two complete training runs of a one-layer transformer that groks modular addition: it memorizes the training set within ~100 epochs, sits near chance on held-out pairs for thousands of epochs, then abruptly generalizes.

Both runs ship every checkpoint, so the representation can be watched as it forms.

Code: 575-lab/grokking-mechanism · Live browser demo: 575-lab.github.io/grokking-mechanism

The runs

The gap between those two rows is grokking: roughly 17× and 74× longer to generalize than to memorize.

grokking-mod53 reproduces the setting of Figure 1 in Liu et al., Towards Understanding Grokking. grokking-mod113 is the Nanda-style reference run.

Architecture

One layer, four heads, no LayerNorm, no biases — the minimal architecture used in grokking replications.

d_model 128 · n_heads 4 · d_head 32 · d_mlp 512 · causal attention
tokens [a, b, "="] with "=" = token id p, answer read from the final position
AdamW, lr 1e-3, betas (0.9, 0.98), weight decay 1.0 on all parameters, full batch

What the embeddings look like

The representation the model groks is circular: embedding k lands near (cos(2πfk/p), sin(2πfk/p)) for some frequency f, which turns addition into rotation.

The circle is not in the top two principal components

Liu et al. visualize the embeddings with plain PCA. That does not reproduce here, and the reason is a property of the model rather than of the plot:

Those three panels are initialization, overfitting, and a fully grokked model at 100% test accuracy. They are visually indistinguishable.

The cause is that this architecture learns several frequencies at once, none dominant:

Four comparable frequencies means eight near-equal singular values — 6.19, 6.00, 5.63, 5.43, 5.18, 4.85, 4.73, 4.63 for p=113, i.e. four near-degenerate 2D planes. The top-2 plane can only ever hold about a quarter of the structure, and what it holds is a sum of four circles turning at unrelated rates, which is a blob.

Project onto a single frequency's plane and every circle is there:

A useful sharpness measure is in-plane purity: the fraction of the signal within a frequency's own plane that actually oscillates at that frequency. The baseline is ~50%, since the plane is chosen to maximize frequency-f content.

When the circles form

grokking-mod113 keeps all 400 checkpoints, so each circle can be followed across training:

They do not form together.

f=34 is a precursor: a clean circle by epoch 4,100, less than half way to the transition at 10,017, having accumulated variance steadily through the memorization plateau since ~epoch 100. The other three crystallize in sequence as test accuracy moves. All four first dip below the 50% baseline — f=44 to 35.9% — so early training actively disorders those planes before organizing them.

Training also prunes. Frequency 6 rose to 2.8% by epoch 9,000 and then collapsed to 0.1% by 15,000; frequency 25 was briefly a fifth active component (7% at epoch 12,000) before being pruned below threshold.

All of this is finished early. From epoch 15,000 onward nothing moves — shares and purities at 15k, 20k, 30k and 40k differ only by noise. The last 25,000 epochs, 62% of the run, change the representation not at all.

A note on reproducibility

An earlier p=113 run with a byte-identical config, seed and data split converged to a different solution: five frequencies (34, 6, 44, 21, 26) rather than four, spread thinner and slightly less pure, ending at 0.9984 test accuracy instead of 1.0000. The difference is GPU floating-point nondeterminism — reduction order varying between sessions — amplified over 40,000 steps into a different basin.

So the number and identity of the learned frequencies is not determined by the seed in this setup. The qualitative picture (several comparable frequencies, no circle in PCA, clean circles per frequency) held in both.

Layout

grokking-mod53/
  config.json          run configuration
  data_split.npz       seeded train/test index split
  metrics.npz          per-epoch train/test loss and accuracy
  checkpoints/0..199   Orbax checkpoints, one per 100 epochs
  embedding_frequency.png, embedding_pca.png, embedding_circles.png
grokking-mod113/
  config.json, data_split.npz, metrics.npz
  checkpoints/0..399   Orbax checkpoints, one per 100 epochs
  frequency_trajectory.npz    per-frequency variance share and purity over training
  training_curve.png, frequency_trajectory.png
  embedding_frequency.png, embedding_pca.png, embedding_circles.png

Checkpoint index n corresponds to epoch (n + 1) × 100.

Loading a checkpoint

git clone https://github.com/575-lab/grokking-mechanism && cd grokking-mechanism
uv sync
hf download davidnet/grokking-mechanism --local-dir runs

from pathlib import Path
import embedding_pca, train

run_dir = Path("runs/grokking-mod113")
cfg = embedding_pca.load_config(run_dir)

model = train.load_checkpoint(cfg, run_dir, 399)        # final
early = train.load_checkpoint(cfg, run_dir, 40)         # epoch 4,100

power = embedding_pca.fourier_power(embedding_pca.number_embeddings(model, cfg.p))
print(power.argsort()[::-1][:4])                        # the learned frequencies

load_checkpoint pins the target sharding to the local device, so these GPU-trained checkpoints restore on a CPU-only machine.

Reproducing

uv run python main.py train --p 113 --frac-train 0.3 --epochs 40000 --run-id grokking-mod113
uv run python main.py embedding-pca --run-dir runs/grokking-mod113 --epochs 0 1000 40000 --projection frequency
uv run python main.py circles --run-dir runs/grokking-mod113

About 21 minutes on an L4; roughly 3 hours on CPU.

References

• Liu, Kitouni, Nolte, Michaud, Tegmark, Williams. Towards Understanding Grokking: An Effective Theory of Representation Learning. NeurIPS 2022.
• Power, Burda, Edwards, Babuschkin, Misra. Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. 2022.
• Nanda, Chan, Lieberum, Smith, Steinhardt. Progress Measures for Grokking via Mechanistic Interpretability. ICLR 2023.