Technical Summary: Detecting Overfitting in Neural Networks during Long-Horizon Grokking using Random Matrix Theory

Problem Statement

Training deep neural networks (NNs) without overfitting is challenging, and detecting overfitting is equally difficult, particularly when access to training data, test data, optimizer states, or training histories is unavailable. This is a critical issue for open-weight foundation models, where users may possess model weights and cards but lack the context to determine if a checkpoint is robustly trained or has overfit to specific data correlations.

The paper focuses on the phenomenon of grokking, where a model's training accuracy reaches near-perfection while test accuracy remains near chance for an extended period before abruptly improving. While previous research has focused on the transition from memorization to generalization, this work investigates the long-horizon regime after grokking. In this phase, termed "anti-grokking," a model that has successfully generalized can lose its held-out performance and revert to overfitting its training data, despite maintaining near-perfect training accuracy. Distinguishing this harmful "anti-grokking" phase from the initial "pre-grokking" phase (where the model has not yet found a generalizing solution) is difficult using standard train/test accuracy curves alone, as both regimes exhibit high training accuracy and poor test accuracy.

Methodology

The authors propose a novel diagnostic method based on Random Matrix Theory (RMT) that operates solely on the model's weight matrices without requiring access to data.

Core Concept: Correlation Traps

The method relies on the principle of self-averaging. In a well-trained, generalizing model, the weight matrices should exhibit statistical properties consistent with random matrices when their internal structure is randomized.

1. Randomization: For each layer weight matrix $W$, the entries are shuffled element-wise to create a randomized matrix $W_{rand}$.
2. Spectral Analysis: The empirical spectral distribution (ESD) of the covariance matrix $X_{rand} = \frac{1}{N} W_{rand}^\top W_{rand}$ is computed.
3. Marchenko-Pastur (MP) Baseline: The ESD is fitted to the Marchenko-Pastur distribution, which serves as the theoretical baseline for a self-averaging, uncorrelated random matrix.
4. Trap Detection: The authors identify Correlation Traps as eigenvalue outliers ($\lambda_{trap}$) that lie significantly beyond the right edge of the fitted MP distribution (accounting for Tracy-Widom fluctuations).
   • Definition: A Correlation Trap is an eigenvalue $\lambda_{trap} > \lambda_{MP}^+ + \Delta_{TW}$.
   • Significance: These outliers indicate a violation of self-averaging, suggesting the presence of structured, non-random correlations within the layer that persist even after entry-wise shuffling.

Distinguishing Harmful vs. Benign Traps

Since the presence of traps indicates non-self-averaging but does not inherently prove harmful overfitting, the authors introduce a data-free JSD (Jensen-Shannon Divergence) diagnostic ablation test:

• Procedure: A specific trap (eigenvector) is replaced with a random vector. The model is then run on probe inputs (random tokens for transformers, Gaussian inputs for MLPs).
• Metric: The JSD between the output logits of the original model and the ablated model is calculated.
• Classification:
  • Harmful Trap: Removing the trap significantly changes the model's output distribution (high JSD) and alters test accuracy.
  • Benign Trap: Removing the trap has a negligible effect on outputs and accuracy.

Key Contributions

1. Definition of Correlation Traps: The authors define Correlation Traps as spectral outliers in shuffled layer weight matrices that violate the Marchenko-Pastur baseline, serving as a signature of non-self-averaging structure.
2. Identification of Anti-Grokking: The paper characterizes "anti-grokking" as a distinct phase following successful generalization where test accuracy degrades while training accuracy remains high. This phase is structurally distinct from pre-grokking.
3. Correlation Traps as a Diagnostic: The study demonstrates that Correlation Traps are absent during pre-grokking and grokking but emerge and grow in number and scale specifically during the onset of anti-grokking.
4. Mechanistic Link to Overfitting: The authors link these traps to localization (where variance is concentrated on a small set of coordinates) and spectral condensation (where a single dominant mode carries macroscopic variance), both of which prevent the generalization error from concentrating.
5. Empirical Validation: The method is validated across three distinct architectures and tasks:
   • A 3-layer MLP on MNIST.
   • A small Transformer on a Modular Addition (MA) task.
   • A GPT2-style Transformer on a synthetic composition task.
6. Application to Foundation Models: The diagnostic was applied to open-weight checkpoints of OpenAI's gpt-oss-20b and gpt-oss-120b, revealing the presence of Correlation Traps, suggesting potential overfitting in frontier-scale models.

Results

• Phase Separation: In all three benchmarks (MLP, MA, GPT2), the number of Correlation Traps remains effectively zero during pre-grokking and grokking. A sharp increase in trap count correlates precisely with the decline in test accuracy during the anti-grokking phase.
• Weight Decay Control: Experiments with weight decay (WD=0.01) showed that while anti-grokking still occurs, the growth of Correlation Traps and the severity of test accuracy degradation are significantly suppressed compared to the no-weight-decay setting (WD=0).
• Perturbation Robustness: Perturbation studies confirmed that trap counts are not merely a proxy for the global scale (norm) of the weights. Even when global weight norms were altered, traps remained absent in pre-grokking/grokking and appeared only in anti-grokking.
• Mechanistic Insight (MLP): In the MLP case study, localized trap directions were found to evolve into "prototype-like" structures in pixel space. These structures caused the model to rely on global input magnitude (e.g., $\|x\|^2$) rather than semantic digit patterns, leading to brittle behavior on shuffled or Gaussian inputs while maintaining perfect training accuracy.
• JSD Ablation: The JSD diagnostic successfully distinguished between traps that, when removed, degraded model performance (harmful) and those that had little effect (benign).

Significance and Claims

The paper claims that Correlation Traps provide a practical, weight-only diagnostic for detecting signatures of potentially harmful overfitting in trained models without access to training or test data.

• Structural Distinction: The primary significance is the ability to structurally distinguish between a model that has not yet learned (pre-grokking) and a model that has learned and then unlearned/generalized poorly (anti-grokking). Both states look identical via accuracy metrics but differ fundamentally in their weight spectra.
• Foundation Model Screening: The detection of Correlation Traps in large open-weight LLMs suggests that foundation-scale models may inadvertently overfit their training data in ways that are detectable via spectral analysis.
• Regulatory and Practical Utility: The authors suggest this method could serve as an unbiased tool for regulatory bodies or users to gauge model safety and robustness, particularly when proprietary data is unavailable.
• Limitations: The authors modestly note that their evidence comes from controlled long-horizon grokking settings. They do not claim that all overfitting produces Correlation Traps, nor do they assert that the presence of traps in frontier models definitively proves harmful behavior without further task-based evaluation. The trap count is a warning signal, not a universal quality score.