Introduction to Grokking

Introduction

Every machine-learning course teaches the same reflex: watch the validation curve, and stop training when it stops improving. Once a model scores 100% on its training data while doing no better than chance on new data, it has overfit — it memorized instead of learning — and more training is supposed to make things worse, not better.

In January 2022, Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin (OpenAI) and Vedant Misra published a paper that broke that reflex: "Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets." They trained small transformers on tables like $a \circ b = c$ for operations modulo a prime $p = 97$ — addition, subtraction, division, permutation composition — showing the network only a fraction of the table and holding the rest out. In their headline example, training accuracy hit 100% after roughly a thousand optimization steps. Validation accuracy sat near chance... and then, after on the order of a hundred thousand steps, shot up to nearly 100%. The network had stopped looking things up and started computing.

They borrowed a word from Robert Heinlein's 1961 novel Stranger in a Strange Land: to grok is to understand something so thoroughly that it becomes part of you. The experiment is simple enough to write in a few lines:

pairs  = all (a, b) with 0 ≤ a, b < p          # p² examples
label  = (a + b) mod p
train  = random fraction f of pairs; test = the rest
repeat for a very long time:
    full-batch step of AdamW on cross-entropy, with weight decay λ
    log train accuracy and test accuracy
# grokking: train → 100% early; test stays low, then jumps late

Grokking matters because it is a clean, reproducible case of a question at the heart of deep learning: when does a network memorize, when does it find the real rule, and what pushes it from one to the other?

Train It Live

Below is a genuine training run, not an animation. A two-layer network with 48 hidden units learns $(a+b) \bmod 23$ from part of the $23^2 = 529$ possible pairs. The blue curve is training accuracy, the red curve is accuracy on pairs it has never seen, and the dashed green curve is a progress measure explained below. The horizontal axis is logarithmic, as in the original paper.

<!-- Grokking demo: controls, live chart and readouts -->
<p class="hint">This is <b>real training</b>, live in your browser. A tiny two-layer network learns <b>(a + b) mod 23</b> from a fraction of the 529 possible pairs; the rest are held out as a test set. It has 48 hidden units with a square activation and is trained with AdamW. Nothing is precomputed. It is far smaller than the original experiments (a transformer, p = 97, up to 10⁵ steps or more), so the delay here is roughly tenfold rather than a hundredfold — but watch the red curve stay low long after the blue one hits 100%.</p>
<div class="ctrls">
  <label class="ctrl">Training data <b id="fracv">50%</b>
    <input type="range" id="frac" min="20" max="90" step="5" value="50" aria-label="Training data">
  </label>
  <label class="ctrl">Weight decay <b id="wdv">1.0</b>
    <input type="range" id="wd" min="0" max="2" step="0.1" value="1" aria-label="Weight decay">
  </label>
</div>
<div class="btns">
  <button id="go" type="button">▶ Train</button>
  <button id="reset" type="button" class="ghost">Reset</button>
  <button id="seed" type="button" class="ghost">New seed</button>
</div>
<div class="presets">
  <span class="plabel">Presets:</span>
  <button type="button" class="chip" data-f="50" data-w="1">Grokking (50%, wd 1)</button>
  <button type="button" class="chip" data-f="50" data-w="0">No weight decay</button>
  <button type="button" class="chip" data-f="70" data-w="1">More data (70%)</button>
  <button type="button" class="chip" data-f="30" data-w="1">Less data (30%)</button>
</div>
<div id="wrap"><canvas id="cv" role="img" aria-label="Train and test accuracy over training steps on a logarithmic axis"></canvas></div>
<div class="legend">
  <span><i class="sw train"></i>train accuracy</span>
  <span><i class="sw test"></i>test accuracy (unseen pairs)</span>
  <span><i class="sw four"></i>Fourier concentration</span>
  <span><i class="sw chance"></i>chance (1/23)</span>
</div>
<div class="readout">
  <span>step <b id="rs">0</b></span>
  <span>train <b id="rtr">–</b></span>
  <span>test <b id="rte">–</b></span>
  <span>Fourier <b id="rfc">–</b></span>
</div>
<div class="status" id="status"></div>

/* Layout and colors for the grokking demo */
* { box-sizing: border-box; }
body { font-family: system-ui, sans-serif; color: #222; margin: 0; background: #fff; }
.hint { font-size: .86rem; color: #444; margin: 0 0 .6rem; line-height: 1.45; }
.ctrls { display: flex; flex-wrap: wrap; gap: .4rem 1.2rem; margin-bottom: .5rem; }
.ctrl { display: flex; flex-direction: column; font-size: .85rem; font-weight: 600; color: #1d3557; min-width: 140px; flex: 1; }
.ctrl b { font-variant-numeric: tabular-nums; }
.ctrl input { width: 100%; accent-color: #1d3557; }
.btns, .presets { display: flex; gap: .4rem; flex-wrap: wrap; align-items: center; margin-bottom: .45rem; }
button { font: 600 14px system-ui, sans-serif; padding: .42rem .85rem; border: 1px solid #1d3557;
         background: #1d3557; color: #fff; border-radius: 8px; cursor: pointer; }
button.ghost { background: #fff; color: #1d3557; }
button.chip { font-size: 12px; padding: .28rem .6rem; background: #eef3f8; color: #1d3557; border-color: #cdd9e3; border-radius: 999px; }
.plabel { font-size: .8rem; color: #555; }
#wrap { width: 100%; }
canvas { display: block; width: 100%; height: 250px; border: 1px solid #e2e8f0; border-radius: 8px; background: #fbfcfe; }
.legend { display: flex; flex-wrap: wrap; gap: .25rem .9rem; font-size: .78rem; color: #555; margin-top: .4rem; }
.sw { display: inline-block; width: 14px; height: 3px; border-radius: 2px; vertical-align: middle; margin-right: .3rem; }
.sw.train { background: #2563eb; }
.sw.test { background: #dc2626; }
.sw.four { background: repeating-linear-gradient(90deg, #059669 0 4px, transparent 4px 7px); }
.sw.chance { background: repeating-linear-gradient(90deg, #9ca3af 0 2px, transparent 2px 4px); }
.readout { display: flex; flex-wrap: wrap; gap: .2rem 1rem; font-size: .85rem; margin-top: .45rem; font-variant-numeric: tabular-nums; }
.status { font-size: .92rem; font-weight: 600; margin-top: .45rem; min-height: 2.6em; line-height: 1.35; color: #1d3557; }
.status.ok { color: #0a7d33; }
.status.bad { color: #b45309; }

const i18nFmt = (s, v) => s.replace(/%(\w+)%/g, (_, k) => String(v[k] ?? ''));
// A real, tiny grokking experiment: learn (a + b) mod P from a random subset of all P*P pairs.
// Model: h = (U[:,a] + V[:,b])^2 elementwise, logits = W h. Full-batch AdamW, cross-entropy.
// Weight decay is the knob that turns memorization into generalization.
const P = 23;          // prime modulus
const H = 48;          // hidden neurons
const LR = 0.01;       // AdamW learning rate
const MAX_STEPS = 3000;
const MEM_AT = 0.99, GEN_AT = 0.90; // 'memorized' and 'generalized' accuracy thresholds

// Small seeded PRNG so every run is reproducible.
function mulberry32(a) {
  return function () {
    a |= 0; a = a + 0x6D2B79F5 | 0;
    let t = Math.imul(a ^ a >>> 15, 1 | a);
    t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t;
    return ((t ^ t >>> 14) >>> 0) / 4294967296;
  };
}

// Gaussian init for U, V (one column per input number) and W (output layer),
// then a seeded shuffle of all pairs split into train and test.
function createRun(frac, seed) {
  const rnd = mulberry32(seed);
  const gauss = () => { let u = 0; while (!u) u = rnd(); const v = rnd(); return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v); };
  const NP = 3 * H * P; // U | V | W
  const prm = new Float64Array(NP);
  for (let i = 0; i < 2 * H * P; i++) prm[i] = gauss();
  for (let i = 2 * H * P; i < NP; i++) prm[i] = gauss() / Math.sqrt(H);
  const idx = [];
  for (let i = 0; i < P * P; i++) idx.push(i); // pair (a, b) is stored as a*P + b
  for (let i = idx.length - 1; i > 0; i--) { const j = Math.floor(rnd() * (i + 1)); const t = idx[i]; idx[i] = idx[j]; idx[j] = t; }
  const n = Math.max(1, Math.round(frac * P * P));
  return {
    prm, grad: new Float64Array(NP), m: new Float64Array(NP), v: new Float64Array(NP), t: 0,
    train: Int32Array.from(idx.slice(0, n)), test: Int32Array.from(idx.slice(n)),
    z: new Float64Array(H), h: new Float64Array(H), lg: new Float64Array(P), dh: new Float64Array(H)
  };
}

// Hidden layer: add the two input columns, then square (a quadratic activation).
function hidden(S, a, b) {
  const { prm, z, h } = S; const oV = H * P;
  for (let k = 0; k < H; k++) { const s = prm[k * P + a] + prm[oV + k * P + b]; z[k] = s; h[k] = s * s; }
}

function predict(S, a, b) {
  hidden(S, a, b);
  const { prm, h } = S; const oW = 2 * H * P;
  let best = -Infinity, bi = 0;
  for (let c = 0; c < P; c++) {
    let s = 0; const o = oW + c * H;
    for (let k = 0; k < H; k++) s += prm[o + k] * h[k];
    if (s > best) { best = s; bi = c; }
  }
  return bi;
}

function accuracy(S, set) {
  let ok = 0;
  for (let i = 0; i < set.length; i++) {
    const e = set[i], a = (e / P) | 0, b = e % P;
    if (predict(S, a, b) === (a + b) % P) ok++;
  }
  return set.length ? ok / set.length : 0;
}

// One full-batch step: forward pass, softmax cross-entropy, manual backprop,
// then an AdamW update with decoupled weight decay.
function trainStep(S, lr, wd) {
  const { prm, grad, z, h, lg, dh, train } = S;
  const oV = H * P, oW = 2 * H * P, n = train.length;
  grad.fill(0);
  let loss = 0;
  for (let i = 0; i < n; i++) {
    const e = train[i], a = (e / P) | 0, b = e % P, y = (a + b) % P;
    hidden(S, a, b);
    let mx = -Infinity;
    for (let c = 0; c < P; c++) {
      let s = 0; const o = oW + c * H;
      for (let k = 0; k < H; k++) s += prm[o + k] * h[k];
      lg[c] = s; if (s > mx) mx = s;
    }
    let se = 0;
    for (let c = 0; c < P; c++) { lg[c] = Math.exp(lg[c] - mx); se += lg[c]; }
    loss -= Math.log(lg[y] / se);
    dh.fill(0);
    for (let c = 0; c < P; c++) {
      const d = (lg[c] / se - (c === y ? 1 : 0)) / n; // dLoss/dlogit = softmax - one-hot
      const o = oW + c * H;
      for (let k = 0; k < H; k++) { grad[o + k] += d * h[k]; dh[k] += d * prm[o + k]; }
    }
    for (let k = 0; k < H; k++) { const dz = 2 * dh[k] * z[k]; grad[k * P + a] += dz; grad[oV + k * P + b] += dz; }
  }
  // AdamW: Adam moments with bias correction, plus wd * weight shrinkage.
  S.t++;
  const b1 = 0.9, b2 = 0.98, c1 = 1 - Math.pow(b1, S.t), c2 = 1 - Math.pow(b2, S.t);
  const { m, v } = S;
  for (let i = 0; i < prm.length; i++) {
    const g = grad[i];
    m[i] = b1 * m[i] + (1 - b1) * g;
    v[i] = b2 * v[i] + (1 - b2) * g * g;
    prm[i] -= lr * (m[i] / c1 / (Math.sqrt(v[i] / c2) + 1e-8) + wd * prm[i]);
  }
  return loss / n;
}

// Progress measure: for each hidden neuron, share of its input-weight energy in its strongest
// nonzero frequency. Random weights give about 27%; a pure cosine wave gives 100%.
function fourierConcentration(S) {
  const { prm } = S; const F = (P - 1) / 2; let acc = 0;
  for (let k = 0; k < H; k++) {
    let tot = 0, top = 0;
    for (let f = 1; f <= F; f++) {
      let re = 0, im = 0;
      for (let a = 0; a < P; a++) { const x = prm[k * P + a]; const th = 2 * Math.PI * f * a / P; re += x * Math.cos(th); im -= x * Math.sin(th); }
      const pw = re * re + im * im; tot += pw; if (pw > top) top = pw;
    }
    acc += tot > 0 ? top / tot : 0;
  }
  return acc / H;
}

// ---------- user interface ----------
const $ = (id) => document.getElementById(id);
const cv = $('cv'), ctx = cv.getContext('2d');
const pct = (x) => Math.round(x * 100) + '%';
const LOGSET = new Set();
for (let i = 0; i <= 120; i++) LOGSET.add(Math.round(Math.pow(MAX_STEPS, i / 120)));
const shouldEval = (s) => s <= 10 || s % 25 === 0 || LOGSET.has(s);

let seed = 1, S = null, hist = [], running = false, memStep = -1, genStep = -1;
const fracVal = () => Number($('frac').value) / 100;
const wdVal = () => Number($('wd').value);

function setStatus(text, cls) { const s = $('status'); s.textContent = text; s.className = 'status ' + (cls || ''); }

function record() {
  const tr = accuracy(S, S.train), te = accuracy(S, S.test), fc = fourierConcentration(S);
  hist.push({ s: Math.max(1, S.t), tr, te, fc });
  if (memStep < 0 && tr >= MEM_AT && S.t > 0) memStep = S.t;
  if (genStep < 0 && te >= GEN_AT && S.t > 0) genStep = S.t;
  $('rs').textContent = S.t; $('rtr').textContent = pct(tr); $('rte').textContent = pct(te); $('rfc').textContent = pct(fc);
}

function reset() {
  running = false; $('go').textContent = "▶ Train";
  S = createRun(fracVal(), seed); hist = []; memStep = -1; genStep = -1;
  record();
  setStatus(i18nFmt("Ready: %n% training pairs, %m% held out. Press Train or pick a preset.", { n: S.train.length, m: S.test.length }), '');
  draw();
}

function finish() {
  running = false; $('go').textContent = "▶ Train";
  const last = hist[hist.length - 1];
  if (memStep < 0) setStatus(i18nFmt("Did not even fit the training set (train accuracy %tr%).", { tr: pct(last.tr) }), 'bad');
  else if (genStep >= 0) setStatus(i18nFmt("Grokked! Memorized at step %m%, generalized (90% test) at step %g% — about %r%× later.", { m: memStep, g: genStep, r: Math.max(1, Math.round(genStep / memStep)) }), 'ok');
  else setStatus(i18nFmt("Memorized at step %m%, but test accuracy is only %te% after %max% steps: no grokking in this window.", { m: memStep, te: pct(last.te), max: MAX_STEPS }), 'bad');
}

function frame() {
  if (!running) return;
  const t0 = performance.now();
  while (performance.now() - t0 < 14 && S.t < MAX_STEPS) {
    trainStep(S, LR, wdVal());
    if (shouldEval(S.t)) record();
  }
  draw();
  if (S.t >= MAX_STEPS) { finish(); return; }
  if (memStep >= 0 && genStep < 0) setStatus(i18nFmt("Memorized the training set at step %m% — test accuracy is still %te%. Keep watching.", { m: memStep, te: pct(hist[hist.length - 1].te) }), '');
  else setStatus(i18nFmt("Training… step %s% of %max%", { s: S.t, max: MAX_STEPS }), '');
  requestAnimationFrame(frame);
}

// Chart: accuracies against log(step), plus memorize / generalize markers.
function draw() {
  const dpr = window.devicePixelRatio || 1;
  const Wc = Math.max(260, $('wrap').clientWidth), Hc = 250;
  if (cv.width !== Math.round(Wc * dpr)) { cv.width = Math.round(Wc * dpr); cv.height = Math.round(Hc * dpr); }
  ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
  ctx.clearRect(0, 0, Wc, Hc);
  const L = 40, R = 10, T = 12, B = 34, pw = Wc - L - R, ph = Hc - T - B;
  const X = (s) => L + Math.log(Math.max(1, s)) / Math.log(MAX_STEPS) * pw;
  const Y = (v) => T + (1 - v) * ph;
  ctx.font = '11px system-ui, sans-serif'; ctx.lineWidth = 1;
  ctx.strokeStyle = '#e5e7eb'; ctx.fillStyle = '#6b7280'; ctx.textAlign = 'right'; ctx.textBaseline = 'middle';
  for (const v of [0, 0.25, 0.5, 0.75, 1]) { ctx.beginPath(); ctx.moveTo(L, Y(v)); ctx.lineTo(L + pw, Y(v)); ctx.stroke(); ctx.fillText(pct(v), L - 4, Y(v)); }
  ctx.textAlign = 'center'; ctx.textBaseline = 'top';
  for (const s of [1, 10, 100, 1000, MAX_STEPS]) { ctx.beginPath(); ctx.moveTo(X(s), T); ctx.lineTo(X(s), T + ph); ctx.stroke(); ctx.fillText(String(s), X(s), T + ph + 3); }
  ctx.fillText("training steps (log scale)", L + pw / 2, T + ph + 17);
  ctx.setLineDash([2, 3]); ctx.strokeStyle = '#9ca3af';
  ctx.beginPath(); ctx.moveTo(L, Y(1 / P)); ctx.lineTo(L + pw, Y(1 / P)); ctx.stroke();
  const marks = [[memStep, '#2563eb'], [genStep, '#dc2626']];
  for (const [s, col] of marks) if (s > 0) { ctx.strokeStyle = col; ctx.globalAlpha = 0.45; ctx.beginPath(); ctx.moveTo(X(s), T); ctx.lineTo(X(s), T + ph); ctx.stroke(); ctx.globalAlpha = 1; }
  ctx.setLineDash([]);
  const line = (key, col, dash) => {
    if (!hist.length) return;
    ctx.strokeStyle = col; ctx.lineWidth = 2; ctx.setLineDash(dash || []);
    ctx.beginPath(); hist.forEach((p, i) => { const x = X(p.s), y = Y(p[key]); i ? ctx.lineTo(x, y) : ctx.moveTo(x, y); }); ctx.stroke();
    ctx.setLineDash([]); ctx.lineWidth = 1;
  };
  line('fc', '#059669', [5, 3]);
  line('tr', '#2563eb');
  line('te', '#dc2626');
}

function syncLabels() { $('fracv').textContent = $('frac').value + '%'; $('wdv').textContent = wdVal().toFixed(1); }

$('go').addEventListener('click', () => {
  if (running) { running = false; $('go').textContent = "▶ Train"; return; }
  if (S.t >= MAX_STEPS) reset();
  running = true; $('go').textContent = "❚❚ Pause";
  requestAnimationFrame(frame);
});
$('reset').addEventListener('click', reset);
$('seed').addEventListener('click', () => { seed++; reset(); });
$('frac').addEventListener('input', () => { syncLabels(); reset(); });
$('wd').addEventListener('input', () => { syncLabels(); if (S.t > 0) reset(); });
document.querySelectorAll('.chip').forEach((b) => b.addEventListener('click', () => {
  $('frac').value = b.dataset.f; $('wd').value = b.dataset.w; syncLabels(); reset();
  running = true; $('go').textContent = "❚❚ Pause"; requestAnimationFrame(frame);
}));
window.addEventListener('resize', draw);
syncLabels();
reset();

<!-- Grokking demo: controls, live chart and readouts -->
<p class="hint">This is <b>real training</b>, live in your browser. A tiny two-layer network learns <b>(a + b) mod 23</b> from a fraction of the 529 possible pairs; the rest are held out as a test set. It has 48 hidden units with a square activation and is trained with AdamW. Nothing is precomputed. It is far smaller than the original experiments (a transformer, p = 97, up to 10⁵ steps or more), so the delay here is roughly tenfold rather than a hundredfold — but watch the red curve stay low long after the blue one hits 100%.</p>
<div class="ctrls">
  <label class="ctrl">Training data <b id="fracv">50%</b>
    <input type="range" id="frac" min="20" max="90" step="5" value="50" aria-label="Training data">
  </label>
  <label class="ctrl">Weight decay <b id="wdv">1.0</b>
    <input type="range" id="wd" min="0" max="2" step="0.1" value="1" aria-label="Weight decay">
  </label>
</div>
<div class="btns">
  <button id="go" type="button">▶ Train</button>
  <button id="reset" type="button" class="ghost">Reset</button>
  <button id="seed" type="button" class="ghost">New seed</button>
</div>
<div class="presets">
  <span class="plabel">Presets:</span>
  <button type="button" class="chip" data-f="50" data-w="1">Grokking (50%, wd 1)</button>
  <button type="button" class="chip" data-f="50" data-w="0">No weight decay</button>
  <button type="button" class="chip" data-f="70" data-w="1">More data (70%)</button>
  <button type="button" class="chip" data-f="30" data-w="1">Less data (30%)</button>
</div>
<div id="wrap"><canvas id="cv" role="img" aria-label="Train and test accuracy over training steps on a logarithmic axis"></canvas></div>
<div class="legend">
  <span><i class="sw train"></i>train accuracy</span>
  <span><i class="sw test"></i>test accuracy (unseen pairs)</span>
  <span><i class="sw four"></i>Fourier concentration</span>
  <span><i class="sw chance"></i>chance (1/23)</span>
</div>
<div class="readout">
  <span>step <b id="rs">0</b></span>
  <span>train <b id="rtr">–</b></span>
  <span>test <b id="rte">–</b></span>
  <span>Fourier <b id="rfc">–</b></span>
</div>
<div class="status" id="status"></div>

/* Layout and colors for the grokking demo */
* { box-sizing: border-box; }
body { font-family: system-ui, sans-serif; color: #222; margin: 0; background: #fff; }
.hint { font-size: .86rem; color: #444; margin: 0 0 .6rem; line-height: 1.45; }
.ctrls { display: flex; flex-wrap: wrap; gap: .4rem 1.2rem; margin-bottom: .5rem; }
.ctrl { display: flex; flex-direction: column; font-size: .85rem; font-weight: 600; color: #1d3557; min-width: 140px; flex: 1; }
.ctrl b { font-variant-numeric: tabular-nums; }
.ctrl input { width: 100%; accent-color: #1d3557; }
.btns, .presets { display: flex; gap: .4rem; flex-wrap: wrap; align-items: center; margin-bottom: .45rem; }
button { font: 600 14px system-ui, sans-serif; padding: .42rem .85rem; border: 1px solid #1d3557;
         background: #1d3557; color: #fff; border-radius: 8px; cursor: pointer; }
button.ghost { background: #fff; color: #1d3557; }
button.chip { font-size: 12px; padding: .28rem .6rem; background: #eef3f8; color: #1d3557; border-color: #cdd9e3; border-radius: 999px; }
.plabel { font-size: .8rem; color: #555; }
#wrap { width: 100%; }
canvas { display: block; width: 100%; height: 250px; border: 1px solid #e2e8f0; border-radius: 8px; background: #fbfcfe; }
.legend { display: flex; flex-wrap: wrap; gap: .25rem .9rem; font-size: .78rem; color: #555; margin-top: .4rem; }
.sw { display: inline-block; width: 14px; height: 3px; border-radius: 2px; vertical-align: middle; margin-right: .3rem; }
.sw.train { background: #2563eb; }
.sw.test { background: #dc2626; }
.sw.four { background: repeating-linear-gradient(90deg, #059669 0 4px, transparent 4px 7px); }
.sw.chance { background: repeating-linear-gradient(90deg, #9ca3af 0 2px, transparent 2px 4px); }
.readout { display: flex; flex-wrap: wrap; gap: .2rem 1rem; font-size: .85rem; margin-top: .45rem; font-variant-numeric: tabular-nums; }
.status { font-size: .92rem; font-weight: 600; margin-top: .45rem; min-height: 2.6em; line-height: 1.35; color: #1d3557; }
.status.ok { color: #0a7d33; }
.status.bad { color: #b45309; }

const i18nFmt = (s, v) => s.replace(/%(\w+)%/g, (_, k) => String(v[k] ?? ''));
// A real, tiny grokking experiment: learn (a + b) mod P from a random subset of all P*P pairs.
// Model: h = (U[:,a] + V[:,b])^2 elementwise, logits = W h. Full-batch AdamW, cross-entropy.
// Weight decay is the knob that turns memorization into generalization.
const P = 23;          // prime modulus
const H = 48;          // hidden neurons
const LR = 0.01;       // AdamW learning rate
const MAX_STEPS = 3000;
const MEM_AT = 0.99, GEN_AT = 0.90; // 'memorized' and 'generalized' accuracy thresholds

// Small seeded PRNG so every run is reproducible.
function mulberry32(a) {
  return function () {
    a |= 0; a = a + 0x6D2B79F5 | 0;
    let t = Math.imul(a ^ a >>> 15, 1 | a);
    t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t;
    return ((t ^ t >>> 14) >>> 0) / 4294967296;
  };
}

// Gaussian init for U, V (one column per input number) and W (output layer),
// then a seeded shuffle of all pairs split into train and test.
function createRun(frac, seed) {
  const rnd = mulberry32(seed);
  const gauss = () => { let u = 0; while (!u) u = rnd(); const v = rnd(); return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v); };
  const NP = 3 * H * P; // U | V | W
  const prm = new Float64Array(NP);
  for (let i = 0; i < 2 * H * P; i++) prm[i] = gauss();
  for (let i = 2 * H * P; i < NP; i++) prm[i] = gauss() / Math.sqrt(H);
  const idx = [];
  for (let i = 0; i < P * P; i++) idx.push(i); // pair (a, b) is stored as a*P + b
  for (let i = idx.length - 1; i > 0; i--) { const j = Math.floor(rnd() * (i + 1)); const t = idx[i]; idx[i] = idx[j]; idx[j] = t; }
  const n = Math.max(1, Math.round(frac * P * P));
  return {
    prm, grad: new Float64Array(NP), m: new Float64Array(NP), v: new Float64Array(NP), t: 0,
    train: Int32Array.from(idx.slice(0, n)), test: Int32Array.from(idx.slice(n)),
    z: new Float64Array(H), h: new Float64Array(H), lg: new Float64Array(P), dh: new Float64Array(H)
  };
}

// Hidden layer: add the two input columns, then square (a quadratic activation).
function hidden(S, a, b) {
  const { prm, z, h } = S; const oV = H * P;
  for (let k = 0; k < H; k++) { const s = prm[k * P + a] + prm[oV + k * P + b]; z[k] = s; h[k] = s * s; }
}

function predict(S, a, b) {
  hidden(S, a, b);
  const { prm, h } = S; const oW = 2 * H * P;
  let best = -Infinity, bi = 0;
  for (let c = 0; c < P; c++) {
    let s = 0; const o = oW + c * H;
    for (let k = 0; k < H; k++) s += prm[o + k] * h[k];
    if (s > best) { best = s; bi = c; }
  }
  return bi;
}

function accuracy(S, set) {
  let ok = 0;
  for (let i = 0; i < set.length; i++) {
    const e = set[i], a = (e / P) | 0, b = e % P;
    if (predict(S, a, b) === (a + b) % P) ok++;
  }
  return set.length ? ok / set.length : 0;
}

// One full-batch step: forward pass, softmax cross-entropy, manual backprop,
// then an AdamW update with decoupled weight decay.
function trainStep(S, lr, wd) {
  const { prm, grad, z, h, lg, dh, train } = S;
  const oV = H * P, oW = 2 * H * P, n = train.length;
  grad.fill(0);
  let loss = 0;
  for (let i = 0; i < n; i++) {
    const e = train[i], a = (e / P) | 0, b = e % P, y = (a + b) % P;
    hidden(S, a, b);
    let mx = -Infinity;
    for (let c = 0; c < P; c++) {
      let s = 0; const o = oW + c * H;
      for (let k = 0; k < H; k++) s += prm[o + k] * h[k];
      lg[c] = s; if (s > mx) mx = s;
    }
    let se = 0;
    for (let c = 0; c < P; c++) { lg[c] = Math.exp(lg[c] - mx); se += lg[c]; }
    loss -= Math.log(lg[y] / se);
    dh.fill(0);
    for (let c = 0; c < P; c++) {
      const d = (lg[c] / se - (c === y ? 1 : 0)) / n; // dLoss/dlogit = softmax - one-hot
      const o = oW + c * H;
      for (let k = 0; k < H; k++) { grad[o + k] += d * h[k]; dh[k] += d * prm[o + k]; }
    }
    for (let k = 0; k < H; k++) { const dz = 2 * dh[k] * z[k]; grad[k * P + a] += dz; grad[oV + k * P + b] += dz; }
  }
  // AdamW: Adam moments with bias correction, plus wd * weight shrinkage.
  S.t++;
  const b1 = 0.9, b2 = 0.98, c1 = 1 - Math.pow(b1, S.t), c2 = 1 - Math.pow(b2, S.t);
  const { m, v } = S;
  for (let i = 0; i < prm.length; i++) {
    const g = grad[i];
    m[i] = b1 * m[i] + (1 - b1) * g;
    v[i] = b2 * v[i] + (1 - b2) * g * g;
    prm[i] -= lr * (m[i] / c1 / (Math.sqrt(v[i] / c2) + 1e-8) + wd * prm[i]);
  }
  return loss / n;
}

// Progress measure: for each hidden neuron, share of its input-weight energy in its strongest
// nonzero frequency. Random weights give about 27%; a pure cosine wave gives 100%.
function fourierConcentration(S) {
  const { prm } = S; const F = (P - 1) / 2; let acc = 0;
  for (let k = 0; k < H; k++) {
    let tot = 0, top = 0;
    for (let f = 1; f <= F; f++) {
      let re = 0, im = 0;
      for (let a = 0; a < P; a++) { const x = prm[k * P + a]; const th = 2 * Math.PI * f * a / P; re += x * Math.cos(th); im -= x * Math.sin(th); }
      const pw = re * re + im * im; tot += pw; if (pw > top) top = pw;
    }
    acc += tot > 0 ? top / tot : 0;
  }
  return acc / H;
}

// ---------- user interface ----------
const $ = (id) => document.getElementById(id);
const cv = $('cv'), ctx = cv.getContext('2d');
const pct = (x) => Math.round(x * 100) + '%';
const LOGSET = new Set();
for (let i = 0; i <= 120; i++) LOGSET.add(Math.round(Math.pow(MAX_STEPS, i / 120)));
const shouldEval = (s) => s <= 10 || s % 25 === 0 || LOGSET.has(s);

let seed = 1, S = null, hist = [], running = false, memStep = -1, genStep = -1;
const fracVal = () => Number($('frac').value) / 100;
const wdVal = () => Number($('wd').value);

function setStatus(text, cls) { const s = $('status'); s.textContent = text; s.className = 'status ' + (cls || ''); }

function record() {
  const tr = accuracy(S, S.train), te = accuracy(S, S.test), fc = fourierConcentration(S);
  hist.push({ s: Math.max(1, S.t), tr, te, fc });
  if (memStep < 0 && tr >= MEM_AT && S.t > 0) memStep = S.t;
  if (genStep < 0 && te >= GEN_AT && S.t > 0) genStep = S.t;
  $('rs').textContent = S.t; $('rtr').textContent = pct(tr); $('rte').textContent = pct(te); $('rfc').textContent = pct(fc);
}

function reset() {
  running = false; $('go').textContent = "▶ Train";
  S = createRun(fracVal(), seed); hist = []; memStep = -1; genStep = -1;
  record();
  setStatus(i18nFmt("Ready: %n% training pairs, %m% held out. Press Train or pick a preset.", { n: S.train.length, m: S.test.length }), '');
  draw();
}

function finish() {
  running = false; $('go').textContent = "▶ Train";
  const last = hist[hist.length - 1];
  if (memStep < 0) setStatus(i18nFmt("Did not even fit the training set (train accuracy %tr%).", { tr: pct(last.tr) }), 'bad');
  else if (genStep >= 0) setStatus(i18nFmt("Grokked! Memorized at step %m%, generalized (90% test) at step %g% — about %r%× later.", { m: memStep, g: genStep, r: Math.max(1, Math.round(genStep / memStep)) }), 'ok');
  else setStatus(i18nFmt("Memorized at step %m%, but test accuracy is only %te% after %max% steps: no grokking in this window.", { m: memStep, te: pct(last.te), max: MAX_STEPS }), 'bad');
}

function frame() {
  if (!running) return;
  const t0 = performance.now();
  while (performance.now() - t0 < 14 && S.t < MAX_STEPS) {
    trainStep(S, LR, wdVal());
    if (shouldEval(S.t)) record();
  }
  draw();
  if (S.t >= MAX_STEPS) { finish(); return; }
  if (memStep >= 0 && genStep < 0) setStatus(i18nFmt("Memorized the training set at step %m% — test accuracy is still %te%. Keep watching.", { m: memStep, te: pct(hist[hist.length - 1].te) }), '');
  else setStatus(i18nFmt("Training… step %s% of %max%", { s: S.t, max: MAX_STEPS }), '');
  requestAnimationFrame(frame);
}

// Chart: accuracies against log(step), plus memorize / generalize markers.
function draw() {
  const dpr = window.devicePixelRatio || 1;
  const Wc = Math.max(260, $('wrap').clientWidth), Hc = 250;
  if (cv.width !== Math.round(Wc * dpr)) { cv.width = Math.round(Wc * dpr); cv.height = Math.round(Hc * dpr); }
  ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
  ctx.clearRect(0, 0, Wc, Hc);
  const L = 40, R = 10, T = 12, B = 34, pw = Wc - L - R, ph = Hc - T - B;
  const X = (s) => L + Math.log(Math.max(1, s)) / Math.log(MAX_STEPS) * pw;
  const Y = (v) => T + (1 - v) * ph;
  ctx.font = '11px system-ui, sans-serif'; ctx.lineWidth = 1;
  ctx.strokeStyle = '#e5e7eb'; ctx.fillStyle = '#6b7280'; ctx.textAlign = 'right'; ctx.textBaseline = 'middle';
  for (const v of [0, 0.25, 0.5, 0.75, 1]) { ctx.beginPath(); ctx.moveTo(L, Y(v)); ctx.lineTo(L + pw, Y(v)); ctx.stroke(); ctx.fillText(pct(v), L - 4, Y(v)); }
  ctx.textAlign = 'center'; ctx.textBaseline = 'top';
  for (const s of [1, 10, 100, 1000, MAX_STEPS]) { ctx.beginPath(); ctx.moveTo(X(s), T); ctx.lineTo(X(s), T + ph); ctx.stroke(); ctx.fillText(String(s), X(s), T + ph + 3); }
  ctx.fillText("training steps (log scale)", L + pw / 2, T + ph + 17);
  ctx.setLineDash([2, 3]); ctx.strokeStyle = '#9ca3af';
  ctx.beginPath(); ctx.moveTo(L, Y(1 / P)); ctx.lineTo(L + pw, Y(1 / P)); ctx.stroke();
  const marks = [[memStep, '#2563eb'], [genStep, '#dc2626']];
  for (const [s, col] of marks) if (s > 0) { ctx.strokeStyle = col; ctx.globalAlpha = 0.45; ctx.beginPath(); ctx.moveTo(X(s), T); ctx.lineTo(X(s), T + ph); ctx.stroke(); ctx.globalAlpha = 1; }
  ctx.setLineDash([]);
  const line = (key, col, dash) => {
    if (!hist.length) return;
    ctx.strokeStyle = col; ctx.lineWidth = 2; ctx.setLineDash(dash || []);
    ctx.beginPath(); hist.forEach((p, i) => { const x = X(p.s), y = Y(p[key]); i ? ctx.lineTo(x, y) : ctx.moveTo(x, y); }); ctx.stroke();
    ctx.setLineDash([]); ctx.lineWidth = 1;
  };
  line('fc', '#059669', [5, 3]);
  line('tr', '#2563eb');
  line('te', '#dc2626');
}

function syncLabels() { $('fracv').textContent = $('frac').value + '%'; $('wdv').textContent = wdVal().toFixed(1); }

$('go').addEventListener('click', () => {
  if (running) { running = false; $('go').textContent = "▶ Train"; return; }
  if (S.t >= MAX_STEPS) reset();
  running = true; $('go').textContent = "❚❚ Pause";
  requestAnimationFrame(frame);
});
$('reset').addEventListener('click', reset);
$('seed').addEventListener('click', () => { seed++; reset(); });
$('frac').addEventListener('input', () => { syncLabels(); reset(); });
$('wd').addEventListener('input', () => { syncLabels(); if (S.t > 0) reset(); });
document.querySelectorAll('.chip').forEach((b) => b.addEventListener('click', () => {
  $('frac').value = b.dataset.f; $('wd').value = b.dataset.w; syncLabels(); reset();
  running = true; $('go').textContent = "❚❚ Pause"; requestAnimationFrame(frame);
}));
window.addEventListener('resize', draw);
syncLabels();
reset();

Result

Console

Things to try:

Grokking (50%, wd 1). Blue reaches 100% within roughly a hundred steps. Red stays near chance for a while, then climbs to ~99%. In this tiny model the gap is about tenfold; in the paper's transformers it was closer to a hundredfold.
No weight decay. The network memorizes just as fast — and then nothing happens. Test accuracy stays at or even below chance: a pure lookup table carries no information about unseen pairs.
More data / less data. With 70% of the table, generalization arrives much sooner. With 30%, it does not arrive within this 3,000-step window. Power et al. saw the same trend: the smaller the training fraction, the longer the wait.
Watch the green curve. For each hidden neuron, it measures how much of its input-weight energy sits in one single frequency of a discrete Fourier transform. Random weights score about 27%. As the network groks, it climbs towards 90%: every neuron turns into a clean wave. The red curve follows it, not the other way around.

The demo is deliberately small so it runs in seconds, so the numbers differ from the paper's — but every point on the chart is computed live from real gradients.

What Is Really Going On

The objective. With training set $D$ and weights $\theta$ , the network minimizes cross-entropy with a pull towards small weights:

\mathcal{L}(\theta) = -\frac{1}{|D|}\sum_{(a,b)\in D} \log \operatorname{softmax}\big(f_\theta(a,b)\big)_{(a+b) \bmod p} \;+\; \frac{\lambda}{2}\,\lVert\theta\rVert^2

(AdamW applies the second term as a separate "decoupled" shrink, $\theta \leftarrow \theta - \eta\lambda\theta$ , but the effect is the same: weights that are not needed slowly decay.) Each full-batch step in the demo costs $O(|D|\cdot p \cdot H)$ operations for $H$ hidden units.

Two solutions, one race. The training data can be fit in two ways. A memorizing solution stores each of the $|D|$ answers separately — a lookup table. A generalizing solution implements the actual rule. Memorization is easier to find, so gradient descent gets there first. But the lookup table needs weight for every example, while the algorithm reuses the same weights for all of them. Varma et al. (DeepMind, 2023) made this precise as circuit efficiency: the generalizing circuit produces larger correct logits per unit of parameter norm. Weight decay keeps charging for norm, so it slowly starves the table and feeds the algorithm — and since a table gets more expensive the more examples it must store, there is a critical dataset size below which memorization wins outright. The same paper reported "ungrokking" (going back to memorizing when the data are shrunk) and "semi-grokking" (delayed but partial generalization).

What the generalizing circuit looks like. Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith and Jacob Steinhardt (ICLR 2023) fully reverse-engineered a one-layer transformer that had grokked addition mod $p = 113$ . With $\omega_k = 2\pi k / p$ for a handful of key frequencies $k$ , the network:

embeds each input as waves: $a \mapsto \big(\cos \omega_k a,\ \sin \omega_k a\big)$ ;
combines them with the angle-addition identities, $\cos \omega_k(a+b) = \cos \omega_k a \cos \omega_k b - \sin \omega_k a \sin \omega_k b$ (and the same for $\sin$ );
scores each candidate answer $c$ by $\ell_c \propto \sum_k \cos\big(\omega_k (a + b - c)\big)$ , which peaks exactly when $c \equiv a + b \pmod p$ .

Addition becomes rotation around a circle — a discrete Fourier transform learned from scratch. The tiny network in the demo does the same with a square activation: if $u_a = \cos(\omega a + \phi)$ and $v_b = \cos(\omega b + \psi)$ , then $(u_a + v_b)^2$ contains the cross term $2u_a v_b = \cos\big(\omega(a+b) + \phi + \psi\big) + \cos\big(\omega(a-b) + \phi - \psi\big)$ , which already depends on $a + b$ . Andrey Gromov (2023) wrote down such solutions for two-layer networks analytically.

Not so sudden after all. Using this understanding, Nanda et al. defined progress measures and found three phases: memorization, circuit formation (the Fourier circuit grows while the table still dominates), and cleanup (weight decay removes the memorizing components). The jump in test accuracy happens at cleanup. The underlying change is gradual; the accuracy metric just makes it look like a cliff. The green curve in the demo is a home-made progress measure of the same kind.

Other pieces of the puzzle.

Weight norm (Omnigrok). Ziming Liu, Eric Michaud and Max Tegmark (ICLR 2023) showed that training loss against weight norm is L-shaped while test loss is U-shaped: good solutions live in a "Goldilocks zone" of norm. Start with large weights and the network first fits the data far outside that zone, then slowly decays into it. With a large initialization they induced grokking on MNIST, text and molecule data too.
Optimizer effects. Thilak et al. (Apple, 2022) described a "Slingshot" mechanism — cyclic spikes in the norm of the last layer under adaptive optimizers — and observed grokking without explicit weight decay around those spikes.
Numerical limits. Prieto, Barsbey, Mediano and Birdal (ICLR 2025) argued that without regularization the softmax runs into floating-point "Softmax Collapse", which blocks grokking; avoiding it (their StableMax) brings grokking back without weight decay.
Weight decay is a strong lever, not a law. Gromov's two-layer networks grok with plain gradient descent on a squared-error loss and no regularization at all, so the delay is about which solution the dynamics reach, not about one specific trick.

What is still open. There is no general theory that predicts when a given network will grok, or how large the delay will be. The clean picture comes from toy algorithmic tasks; evidence that large language models grok specific skills during pretraining exists but is still being debated.

Where It Matters

Grokking was found on toy tables, but it touches several everyday questions:

Early stopping can stop too early. If the validation curve can stay flat for thousands of steps before improving, "stop when it stalls" may throw away the generalizing solution. Grokking is the extreme case of the epoch-wise double descent Nakkiran et al. had observed; Davies, Langosco and Krueger (2023) proposed one account covering both.
Regularization, re-explained. In the classical bias-variance picture, weight decay just trims variance. Grokking shows a sharper role: it decides which of several perfect training fits gradient descent ends up in.
A test bed for interpretability. Modular addition became a standard example of mechanistic interpretability because the learned algorithm could be read off the weights and checked. Progress measures built from that understanding show the change coming long before the accuracy jumps.
Emergence and scaling. Sudden jumps in capability as models grow are often compared to grokking. Neural scaling laws describe smooth average loss; grokking is a reminder that individual skills can still switch on abruptly — and that the abruptness may partly be an effect of the metric chosen.
Data efficiency. The critical-data-size results put numbers on an old intuition: below some amount of data memorization is the cheapest answer, above it the rule is.

Conclusion

Grokking turns a textbook rule on its head: a network can reach perfect training accuracy, look hopelessly overfit, and still go on to find the real rule — if training continues and something, usually weight decay, keeps penalizing the expensive lookup table.

The best-understood case is beautifully concrete. A network that learns addition modulo a prime rediscovers waves and trigonometric identities and computes $a + b$ as a rotation on a circle. The "sudden" moment of understanding is the visible end of a slow, measurable process: a general circuit grows quietly beside the memorized one until the memorized one is cleaned away.

Much is still open — above all, predicting when grokking will happen in large models. But the lesson already stands: a flat validation curve does not always mean nothing is being learned.

Comments

Loading comments...

https://www.kipuhub.com/en/article/grokking/Content licensed under CC BY-NC 4.0.

Cookies & privacy

Grokking

When a network suddenly understands — long after it memorized