Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets

Problem

Framing

Small algorithmic datasets expose a gap between interpolation and true rule learning: networks can memorize every seen equation yet fail on held-out ones for orders of magnitude more steps. The paper identifies a delayed phase transition, "grokking," where the same model later jumps from chance-level validation accuracy to near-perfect generalization.

Currently Used Methods

Foundational

• @vaswaniAttentionAllNeed2017 — decoder-only Transformers for symbolic sequence modeling.
  • Limitation in context: point estimates hide late generalization dynamics after memorization.
• @frankleLotteryTicket2019 — sparse subnetworks can train as well as dense networks.
  • Limitation in context: does not explain delayed transition from memorization to rule learning.
• @srivastavaDropout2014 — dropout regularizes networks with stochastic masking.
  • Limitation in context: weaker than weight decay for data-efficient generalization here.
• @loshchilovAdamW2019 — decoupled weight decay for adaptive optimization.
  • Limitation in context: regularization was not studied as a trigger for grokking dynamics.

Proposed Method

Architecture

The paper uses a decoder-only Transformer on equations of the form $a \circ b = c$, with each symbol tokenized separately. The network has 2 layers, width 128, and 4 attention heads, for about $4 \cdot 10^5$ non-embedding parameters.

Loss / Objective

Training minimizes next-token cross-entropy on the answer token only.

Algorithm

For noisy-update ablations, the paper perturbs Adam updates with Gaussian noise.

Training Procedure

• Model: decoder-only Transformer, 2 layers, width 128, 4 heads.
• Parameters: about $4 \cdot 10^5$ non-embedding parameters.
• Main Section 3.1 runs: Adam, no weight decay.
• Optimization budget: $10^6$ steps for Section 3.1.
• Learning-time curves: budget increased to $5 \cdot 10^5$ steps.
• Repeats: 3 random seeds; 7 seeds for Section 3.1.1.

Evaluation

Datasets

• Binary operation tables over discrete symbols.
• Modular arithmetic tasks, including $x+y \bmod p$, $x-y \bmod p$, and $x/y \bmod p$.
• Polynomial-style operations over finite fields.
• Composition and product tasks in the abstract group $S_5$.
• Random train/validation splits over all equations.

Metrics

• Train accuracy.
• Validation accuracy.
• Train loss.
• Validation loss.
• Steps to first reach $99\%$ validation accuracy.

Headline results

• Modular division, $50\%$ train data: train accuracy reaches $100\%$ by about $10^3$ steps.
• Modular division, $50\%$ train data: validation stays near chance until about $10^5$ steps.
• Modular division, $50\%$ train data: validation then rises to about $100\%$ by about $10^6$ steps.
• $S_5$ product, $25$-$30\%$ train data: a $1\%$ data drop increases median time-to-$99\%$ by $40$-$50\%$.
• $S_5$ product, $10^5$-step budget: weight decay more than halves samples needed versus most interventions.

Ablations

• Optimizer choice: full-batch, minibatch, and noisy variants all generalize differently under fixed compute.
• Weight decay: strongest data-efficiency gain; more than halves required samples on $S_5$.
• Training-data fraction: less data preserves final accuracy, but sharply delays generalization onset.
• Task structure: less symmetric or more complex operations need higher train fractions to generalize.

Method Strengths and Weaknesses

Strengths

• Isolates memorization versus generalization with exact train/validation equation splits.
• Shows a stark temporal separation: perfect fit by $10^3$, generalization near $10^6$ steps.
• Demonstrates strong regularization effects; weight decay sharply improves data efficiency.
• Uses multiple algebraic tasks, showing grokking is not unique to one dataset.

Weaknesses

• Evidence is centered on tiny synthetic operation tables, not naturalistic tasks.
• Main mechanism remains descriptive; no predictive theory for grokking onset.
• Results depend on very long optimization horizons, which may be impractical.
• Architecture sweep is narrow: one small Transformer family dominates experiments.

Suggestions from the authors

• Test whether grokking appears on datasets beyond small algorithmic tables.
• Evaluate whether sharpness and related minima measures predict grokking onset.
• Study compute–data tradeoffs behind rapidly increasing time to generalization.
• Investigate why SGD noise and flatter solutions correlate with better generalization.

Links

Prior Papers

• @frankleLotteryTicket2019 — adjacent perspective on implicit structure in overparameterized networks and which subnetworks support generalization.
• @vaswaniAttentionAllNeed2017 — provides the Transformer architecture used for the symbolic equation model.
• @srivastavaDropout2014 — dropout is one of the regularization baselines compared in the optimization ablations.
• @loshchilovAdamW2019 — AdamW weight decay is a central intervention in the paper's generalization experiments.

Further Papers

No vault papers identified as further work yet.