Late-Stage Generalization Collapse in Grokking: Detecting anti-grokking with Weightwatcher
Hari K Prakash, Charles H Martin
cs.LG
2026-02-03
Grokking setups trained to 10^7 steps later lose test accuracy while train stays 100%. WeightWatcher flags anti-grokking from Correlation Traps in shuffled weights, with no data.
Grokking is the delayed-generalization pattern: train accuracy hits 100% early, test accuracy sits near chance for a long time, then jumps. After Power et al. (2022) made it a named phenomenon on algorithmic tasks, most explanations leaned on weight decay, shrinking ℓ2 norms, or a fight between memorization and generalization circuits.
Those accounts usually stop the clock once test accuracy rises. Prakash (UC San Diego) and Martin (Calculation Consulting) rerun two canonical setups out to 10^7 steps and find a third phase. Train accuracy stays perfect. Test accuracy falls off its peak. They call it anti-grokking. It is not Varma et al.'s ungrokking, which retrains an already grokked net on a smaller dataset. Anti-grokking happens on the original data, and the headline run uses WD=0.
The diagnostic is the open-source WeightWatcher tool (v0.7.5.5), built on HTSR (Heavy-Tailed Self-Regularization) and the newer SETOL (Semi-Empirical Theory of Learning). It never looks at train or test examples. Each layer weight matrix W is converted to a correlation matrix X = (1/N) W^T W, then to an empirical spectral density (ESD).
Well-trained layers typically grow a power-law tail ρ(λ) λ^{-α}. HTSR reads α as layer quality:
The primary alarm is a Correlation Trap. Shuffle the entries of W to get Wrand. Shuffling should destroy correlations, so the spectrum should collapse to a Marchenko-Pastur (MP) bulk. Eigenvalues that still sit far past the MP right edge are traps: magnitude-driven pathologies that SETOL predicts will hurt generalization.
Two long runs:
Baselines include the ℓ2 weight norm and Golechha's Activation Sparsity, Absolute Weight Entropy, and Approximate Local Circuit Complexity.
On the MLP with WD=0, train accuracy rises around 10^2 steps and saturates by 10^4–10^5 while test stays low. Test then climbs after 10^5 steps and peaks near 10^6. After that it falls to about 0.5 while train remains 100%. The figure caption names 0.5 as the collapse floor; the abstract says "chance." Chance for 10-class MNIST is 0.1. The main text never reports a point estimate of peak test accuracy.
Layer-averaged α and trap counts line up with the three phases:
| Phase | Step scale | Mean α | FC1 traps |
| pre-grokking | 10^5 | 4.0 ± 0.6 | 0 |
| grokking | 10^6 | 2.9 ± 0.2 | 0 |
| anti-grokking | 10^7 | 1.1 ± 0.3 | 7.5 ± 5.6 |
FC2 α falls to 1.3 ± 0.3 in the last phase, with about one trap. With WD=0.01, test accuracy groks, dips slightly, then plateaus near the peak, and mean α sits around 2. Traps still appear, but FC1 only has 2.0 ± 0.0 of them and the collapse is much milder. Activation Sparsity and related progress measures track the first two phases, then keep moving in the same direction, so they cannot tell grokking from the later crash.
On modular addition, train accuracy is 1.0 ± 0.0 in all three phases. Test goes from 0.40 ± 0.28 (pre) to 0.97 ± 0.02 (grok) to 0.68 ± 0.11 (anti). Mean α tightens from 4.10 ± 0.83 to 2.02 ± 0.44, then rebounds to 3.89 ± 0.68. Trap counts stay at 0 until the last phase, then average 4.88 ± 1.55 per layer.
The two models memorize differently. In the MLP, pre-grokking is weakly correlated: some layers still have α ≥ 5. During collapse, the leading right singular vector of the first layer shrinks from a smooth global template into a digit-shaped blob that looks like an "8", and rows of W1 become literal digit prototypes. The paper calls this prototype overfitting. The modular-addition network is very-heavy-tailed across the whole ESD in every phase. At peak grokking many layers have only one or a few giant eigenvalues (rule-based memorization). When it collapses, α rises rather than falls, closer to catastrophic forgetting.
A closing scan of gpt-oss-20b and gpt-oss-120b reports an unusually large share of layers with α < 2, plus traps that are not plotted. There is no intervention and no task-level correlation.
For anyone reproducing grokking on small models, "train long enough and it will generalize" is incomplete. Generalization can peak and then die. Weight decay here is concrete: WD=0.01 does not remove traps, but it turns the crash into a plateau. If grokking is supposed to be a generalization mechanism, the clock needs to run to 10 million steps before that claim is safe.
The more portable pitch is the data-free probe. α and traps depend only on weight matrices, so they could in principle watch overfitting when a validation set is missing. The catch is that the alarm is glued to WeightWatcher / SETOL, and Martin is the tool's author. This paper is a substantial extension of an ICML workshop paper. Treating it as a production monitor is premature.
The gpt-oss histograms are a teaser, not evidence that those models have already anti-grokked.
The authors say α is a phenomenological binning. Values near 2 often line up with good generalization, but not always. Some layers sit below 2 and still generalize; some well-trained layers sit well above 2. In modular addition, anti-grokking spectra are already atypical, with rank collapse and multimodal power laws, so the α fit itself is shaky.
The experiments are tiny: a 3-layer MLP on 1,000 MNIST images, and a LayerNorm-free toy transformer on modular addition. The main run uses 10^7 steps, float64, and initialization scaled by 8.0, far from modern LLM training. Appendix F gives a BBP-style sufficient condition for a single oversized entry to survive shuffling as a trap. There is no ablation that removes the trap singular directions and checks whether test accuracy returns. The gpt-oss figure has no layer counts and no cleaner comparison model.
Figure 1 marks the collapse floor at 0.5; the abstract says chance. Those are not the same number for 10-way MNIST. Peak test accuracy is also missing from the main tables. An appendix perturbation study reports grok-phase test 0.8879 ± 0.0051 and anti-grok 0.6184 ± 0.1086, under a different phase split. Read the main figure, the abstract, and the appendix as three separate calibrations.