A pre-registered predictive theory of the grokking step, and its falsification on unseen task families

# A pre-registered predictive theory of the grokking step, and its falsification

1. What was committed, and when

The point of this paper is the ordering of events, so it is stated first.

The freeze command refuses to run if any held-out trace already exists on disk, so the second lock is only obtainable while the answers are still unknown. Fitting-set and held-out configurations were checked for collision programmatically: zero overlap.

A theory frozen after seeing its test is worth nothing, and the only defence against that is a timestamped commitment made when the answer is not yet available. We also report a case where this discipline cost us: the tolerance ladder committed at stage 1 turned out to contain an unreachable headline, for reasons given in §5, and we report the failure rather than moving the goalposts.

2. The clock, which is derived and which survived

Under AdamW the parameter update is

theta <- (1 - eta*lambda) theta - eta m_hat / (sqrt(v_hat) + eps)

Wherever sqrt(v_hat) >> eps, the second term has modulus approximately 1 per coordinate whatever the gradient magnitude is: Adam normalises it away. The decoupled factor (1 - eta*lambda) is the only term in the update that is not normalised. Therefore, after the training loss has collapsed, eta*lambda is the only rate constant the optimiser leaves behind**, and any post-memorisation timescale must be proportional to 1/(eta*lambda).

This is a constraint, not a fit: it forbids the rate from depending on the gradient norm, the loss scale, the parameter count or the logit scale. Measured within each training-fraction group, the slope of log10 t_grok against log10(1/(eta*lambda)) is

-0.98 (frac 0.7), -1.11 (frac 0.6), -0.73 (frac 0.5), -0.96 (frac 0.4)

so the exponent is -1 and the problem reduces to the dimensionless prefactor A = t_grok*eta*lambda.

Two consequences of the same argument were used as structural gates and both held on held-out data. Coupled L2 (adam_l2) puts the decay inside the preconditioner, where sqrt(v_hat) divides it out; plain SGD and SGD with momentum have no preconditioner at all, so the modulus-1 argument fails. The predictor therefore refuses to emit a number for any optimiser other than AdamW. On the held-out set that gate fired 9 times and was correct 9 times; not one refusal landed on a run that generalised. In the fitting set, adam_l2 never reaches even 99% training accuracy at any weight decay tried — final train accuracy 0.025 to 0.033, i.e. chance — so "delayed generalisation" is not defined there.

3. The frozen law

t_grok = A0 U^w / ( D eta lambda C^(w+1) (f - f_c)^v ), f_c = U ln(U) / N

with U the number of unknowns (symbol classes summed over argument roles), C the number of target classes, N the number of table rows, f the training fraction, and

D — the architecture term, with zero fitted constants. D is the number of weight matrices on the shortest token-to-logit path: 2 for a one-hot MLP, 2 + n_layers for an embedding MLP, and 2 for a transformer regardless of depth, because the residual stream leaves an embed-to-unembed path open. The mechanism: under uniform decay a degree-d path relaxes at d*eta*lambda, the logits are a sum over paths, and a sum of exponentials is governed by its slowest — hence shortest — term.

Six constants were fitted (A0, v, w, plus a scatter law); a seventh, the exponent on N, was fixed at 0 by argument and declared in the registry as unidentified by the fitting set, because every fitting task happened to satisfy N = C^2. That declaration was made in advance precisely because the held-out set contains tasks where it is false.

The predictor is probe-free: it reads the configuration dictionary and nothing else. It never inspects a trace of the run it is predicting. Earlier rounds of the search produced only predictors that leaned on an early-time probe window; forcing zero-shot cost almost nothing in accuracy (0.61 to 0.58 in-sample) while making the claim much stronger.

4. Result: falsified

80 held-out configurations, 61 generalised, 19 did not.

The primary falsification condition failed. We report it as such.

4.1 The failure is structured, not diffuse

Per family, median signed error (positive = predicted too late): modadd +0.030, modmul −0.087, moddiv +0.670, modadd3 +1.508, modmax +2.103, parity +2.677. Sixty-eight percent of unseen-family errors are one-sided late.

4.2 What genuinely worked

Zero-shot extrapolation across scale. The canonical configurations — the ones the grokking literature actually quotes — were in the held-out set and were never seen in any form:

Three of four within a factor of 1.7, from the configuration dictionary alone, at moduli 97 and 113 when fitting never exceeded 59, at a training fraction of 0.3 when fitting never went below 0.35, and on a two-layer model.

Cross-task generalisation, where the family is structurally similar. modmul was never fitted and came out at median error −0.087.

The optimiser gate, 9/9 as described in §2.

4.3 The diagnosis

f_c = U ln U / N is a coupon-collector threshold: it measures how much of the table must be observed before the entries determine the answer. What actually sets the time is the difficulty of the rule. max(a,b) is a one-line comparison and groks in 28 steps; the law demanded 7700 because the table has 961 cells. Sparse parity groks in 307; the law demanded 164778. Table size and rule difficulty are strongly correlated inside modular arithmetic and decouple completely outside it.

The gamma = 0 risk declared in the registry did fire as predicted on modadd3 (+1.508, the one held-out family with N != C^2), but it is not the whole story: modmax satisfies N = C^2 and is still wrong by +2.10. The defect is the table-size premise itself, not merely its exponent.

5. A bound we discovered too late, and the tolerance we could not reach

The stage-1 lock committed to placing 80% of held-out configurations within ±0.15. That target was unreachable when it was set, and we did not know it.

Grokking time is intrinsically stochastic. Across 85 runs identical in every field except the initialisation seed, sd(log10 t_grok) is 0.038 for attention-free architectures and 0.218 for transformers, with disjoint per-cell ranges and a worst case of 15.2x within one cell. Against sigma = 0.218, a predictor that returns the exact conditional median places at most 51% of transformer runs within ±0.15. The held-out set is 77/80 transformers.

So the headline rung was impossible before any theory existed. The ±0.30 rung — which was on the same committed ladder — has a ceiling of 83%, and the theory reached 0.550 against a forecast of 0.61. That is a real failure, not a ceiling effect. We report both, and we report that the ceiling was discovered after the commitment rather than pretending it informed it.

6. Cycle two: the diagnosed repair does not work either

Having identified rule complexity as the missing quantity, we attacked it directly: a module of seven complexity measures computed from the task definition alone (no trace, no trained network), and eight new task families constructed so that the table stays at 961 cells while the rule ranges from firstarg (the output is an input token) through cmp, modmin, modgcd, modexp, polymix to modaffine. The fast families grok in 32 to 113 steps against modular addition's 24276 at low eta*lambda — a 760x dynamic range, which is what a complexity term needs to be identifiable.

The measures carry real signal. Across 19 task families the correlation of median log10 A with linear_top1_share is +0.588 and with linear_erank is −0.551, and the ordering is right: the families the frozen law over-predicted by two decades are exactly the ones the measures rank as easy.

Three of the four repairs failed. Under leave-one-task-family-out (excluding the degenerate modadd fold, which removes 290 of 396 runs), on 106 scored rows:

The fourth is a real advance and we report it as such. Its new term is a rank exponent

q = log( effective rank of A ) / log( n_values ), A[i,j] = P(target unchanged | token swapped i->j)

averaged over argument positions: the number of genuinely distinct behaviours a position exhibits, normalised by its own vocabulary and therefore invariant in p by construction. That invariance is the whole point — it makes the confound above unrepresentable, because no direction in the fitted parameter space can alter the modulus dependence. The channel exponent is a literal −1 in the source rather than a fitted parameter, and after fitting the measured slope over p = 19…113 is −1.01 with a worst per-prime deviation of 0.230 dex.

It repairs exactly what cycle 1 broke. modmax has the identical 961-cell table as modadd but q = 0.73 against 1.00. Sparse parity, which the frozen law missed by +2.677 dex, comes out at median error 0.030. Per family under leave-one-family-out: parity, modcube and perm 1.000 inside ±0.30; cmp 0.833; modsq 0.750; modmul 0.714. Still failing: modadd3 (1.769 dex), moddiv (1.257), modmax (0.912), modquad and modexp.

Two caveats we insist on. First, on grokked rows alone the figure is 0.535 — so the headline is not inflated by censored rows, but 35 of the 106 pooled rows are censored and scored correct merely by predicting past the budget. Second, and more seriously, the feature set, the order-load exponent, the family weighting and the channel definition were all chosen by inspecting this same leave-one-family-out score. That is model selection on the evaluation metric, and 0.538 is therefore optimistic for a genuinely unseen family. We say so because it is the reason a second held-out set exists (§6.1) rather than a caveat discovered afterwards.

The clock transfers across task families. Whether the prefactor does is now an open, pre-registered question rather than a settled negative.

6.1 A second pre-registration, frozen before this paper was written

The cycle-2 law, its five fitted constants and a binding forecast are hash-committed in registry/heldout2.lock.json, together with 59 configurations drawn from ten task families that appear in no fitting sweep and in no cycle-1 held-out run — modinv, modsqsum, modhalf, modsum2, modtriple, digitsum, absdiff, modand, modor, modcmp3 — verified programmatically against every family ever executed (overlap: zero). Predictions span 186 to 45,942 steps, a 247× range.

The forecast is discounted for the selection bias above: 0.45 inside ±0.30, against the 0.538 the selected metric reports. Three failures are named in advance rather than explained afterwards:

• modsum2 and modtriple are argument-asymmetric abelian rules, pushed to ~40,000 steps by the

order-load term. The one asymmetric abelian family in fitting (modsub) already fails +1.05 dex. We expect these to fail late by about a decade; if they do, the order-load term is wrong rather than mis-tuned.

• modhalf is (a+b)/2 mod p — abelian and symmetric, so it must behave like modadd. If modhalf

fails, the abelian core is wrong and the cycle-2 result collapses.

• The lowest-q families (modcmp3, digitsum, modand) are where the repair should pay off. If

they fail, the repair did not work and cycle 2 falls for the same reason cycle 1 did.

Those runs are executing as this paper is submitted and their outcome is not known to us. It will be reported in full whatever it shows.

One structural warning for anyone attempting this: within the modular family a larger modulus is faster (more target classes, more routes to a generalising circuit), while across families higher rule complexity is slower. Spectral measures such as effective rank equal the modulus for modular addition, so a single column carries both effects with opposite signs. Fitting one column inverts the other — we did exactly this and watched the channel-count coefficient collapse to 0.09, destroying the term that had survived cycle 1. The best repair in the table above is the only one that pinned that exponent instead of fitting it.

7. Methodological findings that outlive the theory

In-sample fit is not evidence, and the gap is enormous. Of twenty candidate theories scored on 262 fitting runs, the best in-sample was a 26-parameter empirical curve fit at 0.61 — with a leave-one-axis-out worst-case of exactly 0.00. It gets every prediction wrong on at least one held-out axis. The mechanistic four-constant laws scored lower in-sample and held up under cross-validation. We deliberately made the empirical baseline strong so that it could not flatter a weak theory, and it collapsed the moment it had to extrapolate.

An order parameter that creeps is useless as a clock. Weight-norm transport — "grokking happens when the norm reaches a critical value" — scored 0.13, last of thirteen. The reason is structural rather than incidental: the norm approaches its fixed point asymptotically, so a small error in the threshold becomes a large error in the step. We predicted this from the traces before benchmarking and the ranking confirmed it. This is independent of whether the mechanism is real.

Kramers escape cannot be the mechanism in the full-batch setting, on the arithmetic rather than on evidence: full-batch training is a deterministic map with exactly zero gradient-noise covariance, so an escape rate driven by sampling noise is identically zero while the runs demonstrably grok.

A leak we found in our own scorer, before it mattered. The truncated run handed to a predictor originally carried censored_at, and the training loop stops a successful run at ceil(1.35*t_grok)+10, so that field had a median ratio of 1.308 to the answer. A predictor returning budget/1.308 would have scored near-perfectly with no theory at all. It was caught by an adversarial audit before the held-out runs and fixed; the exploit now scores 0.060. We report it because a scorer with that leak would have produced a spectacular and entirely false result.

8. Limitations

One optimiser family carries all the positive results (AdamW; the others do not memorise at these settings). Architectures are one-layer and two-layer transformers, embedding MLPs and one-hot MLPs; no convolutions, no normalisation-heavy models, no realistic scale. Tasks are algorithmic. The held-out set is 77/80 transformers, so the scatter law was effectively untested. Several held-out families have only three runs, so per-family medians are thin. The cycle-2 leave-one-family-out figures exclude the modadd fold as degenerate, and we report that exclusion rather than hiding it.

9. Conclusion

We do not yet have a predictive theory of grokking. We have a derived clock that transfers across task families, a zero-shot prefactor that worked inside its fitted family and failed outside it by up to two and a half decades, a diagnosis of why, and a repair that fixes the diagnosed defect on the metric it was selected against — with a second pre-registration already frozen to find out whether that survives contact with families nobody has run. The pre-registration is what makes this reportable: without it, the within-family result and the canonical-configuration predictions in §4.2 would have made a persuasive paper, and it would have been wrong in exactly the way the held-out set exposed.

Everything is released: the harness, 396 runs with full traces, both hash-committed registries, the scoring code with its probe-window enforcement, and all twenty candidate theories including the failures.