Late-Stage Generalization Collapse in Grokking: Detecting anti-grokking with Weightwatcher

TL;DR

The paper investigates late-stage generalization collapse, anti-grokking, in neural networks by applying WeightWatcher with HTSR/SETOL theory to detect Correlation Traps and track the tail exponent $\alpha$ across training. It analyzes a 3-layer MLP on a MNIST subset and a transformer on modular addition, revealing three phases: pre-grokking, grokking, and anti-grokking, with anti-grokking marked by numerous traps and $\alpha<2$ in some layers. Correlation Traps emerge as robust indicators of anti-grokking, even without access to training or test labels, while $\alpha$ provides a secondary signal. The results highlight distinct memorization regimes (prototype memorization in MLP, rule-based memorization in MA) and suggest that long-term spectral diagnostics can diagnose and potentially prevent catastrophic forgetting in large models.

Abstract

\emph{Memorization} in neural networks lacks a precise operational definition and is often inferred from the grokking regime, where training accuracy saturates while test accuracy remains very low. We identify a previously unreported third phase of grokking in this training regime: \emph{anti-grokking}, a late-stage collapse of generalization. We revisit two canonical grokking setups: a 3-layer MLP trained on a subset of MNIST and a transformer trained on modular addition, but extended training far beyond standard. In both cases, after models transition from pre-grokking to successful generalization, test accuracy collapses back to chance while training accuracy remains perfect, indicating a distinct post-generalization failure mode. To diagnose anti-grokking, we use the open-source \texttt{WeightWatcher} tool based on HTSR/SETOL theory. The primary signal is the emergence of \emph{Correlation Traps}: anomalously large eigenvalues beyond the Marchenko--Pastur bulk in the empirical spectral density of shuffled weight matrices, which are predicted to impair generalization. As a secondary signal, anti-grokking corresponds to the average HTSR layer quality metric $α$ deviating from $2.0$. Neither metric requires access to the test or training data. We compare these signals to alternative grokking diagnostic, including $\ell_2$ norms, Activation Sparsity, Absolute Weight Entropy, and Local Circuit Complexity. These track pre-grokking and grokking but fail to identify anti-grokking. Finally, we show that Correlation Traps can induce catastrophic forgetting and/or prototype memorization, and observe similar pathologies in large-scale LLMs, like OSS GPT 20/120B.

Figures (18)

• Figure 1: The three phases of grokking. Training curves for a depth-3, width-200 MLP on MNIST. The initial pre-grokking phase (grey): training accuracy (red line) surges at $10^2$ steps, saturating between $10^4-10^5$ steps, while test accuracy (purple line) remains low; the grokking phase (yellow): with test accuracy rapidly increasing after $\sim 10^5$ steps, and reaching a maximum at $10^6$ steps; and the newly revealed late-stage anti-grokking phase (green): test accuracy collapses (to $0.5$).
• Figure 2: Left: Example of the ESD derived from a well-correlated $\mathbf{W}$ (blue) with a power-law tail fit (red), on a log–log plot. Right: Example of the ESD of $\mathbf{W}^{\mathrm{rand}}$ (light purple) with a Marchenko–Pastur (MP) fit (red), shown on a log–linear plot.
• Figure 3: Examples of Correlation Traps. ESDs of $\mathbf{W}^{rand}$ (light purple) of Layer 2 compared to an MP fit (red). Correlation traps $\lambda_{trap}$ appear as spikes to the right of the MP bulk (log x-axis). Left: Right Before Collapse ($\sigma_{mp}\approx0.9879$; KS $p\approx4\times10^{-13}$). A single dominant trap at $\lambda_{trap}\approx10^{6.5}$. Right: Final Generalization Collapse (KS $p\approx1.9\times10^{-5}$). Multiple traps $\lambda_{trap}\in[10^{2.x},10^{6.5}]$.
• Figure 4: Average Correlation Traps across layers during optimization steps. The increase in the number of traps coincides with the "anti-grokking" performance drop seen in Fig. \ref{['fig:training_curves']} at $1M$ steps. This is our primary signal for anti-grokking,
• Figure 5: Average $\alpha$ across layers during optimization. Note the drop below the critical threshold $\alpha=2$, coinciding with the "anti-grokking" performance drop seen in Fig. \ref{['fig:training_curves']} after $\sim 1M$ steps. This is our secondary signal for anti-grokking. From Figure \ref{['fig:replicate-alpha']}.