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

TL;DR

This paper identifies a new phase called anti-grokking, where neural networks lose their generalization after initially succeeding, and introduces spectral diagnostics using WeightWatcher to detect this phenomenon.

Contribution

The study uncovers anti-grokking as a late-stage collapse of generalization and proposes spectral eigenvalue analysis with WeightWatcher as a diagnostic tool.

Findings

Anti-grokking occurs after successful generalization, leading to test accuracy collapse.

Correlation Traps, large eigenvalues in weight spectra, signal anti-grokking.

Large-scale LLMs exhibit similar anti-grokking pathologies.

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 $\alpha$ 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.