The Grokked Illusion: True Equilibrium Mitigates Catastrophic Forgetting

Technical Summary: The Grokked Illusion: True Equilibrium Mitigates Catastrophic Forgetting

Problem Statement
While neural network (NN) performance is traditionally evaluated via training and test metrics, these do not fully capture the robustness of a learned representation. A model may generalize perfectly to a test set yet remain fragile; subsequent parameter updates required to learn new information can lead to catastrophic forgetting, where previously acquired knowledge collapses. This disconnect between generalization and robustness raises a fundamental question: what properties of a neural network determine its resilience to interference? Recent work has established a "high-entropy advantage," where solutions occupying larger volumes in parameter space (quantified by Boltzmann entropy) exhibit superior generalizability compared to conventionally optimized solutions. This paper investigates whether this high-entropy advantage extends beyond generalization to robustness against catastrophic forgetting.

Methodology
The authors utilize modular arithmetic ($x + y \mod 67$) as a controlled experimental setting, specifically leveraging the "grokking" phenomenon where models eventually achieve perfect generalization. To isolate the effect of entropy from generalization performance, the study compares two types of models that achieve identical perfect test accuracy (100%) and similar training losses under the same weight norm constraint ($\|w\|=30$):

1. AdamW-goldilocks: A standard Transformer trained with the AdamW optimizer.
2. WLMD-equilibrium: A Transformer sampled from the Boltzmann entropy landscape using Wang-Landau Molecular Dynamics (WLMD), representing a high-entropy state.

Noise Injection Experiment:
To evaluate robustness, the authors design a noise injection protocol. The vocabulary of the trained models is expanded, and the models are fine-tuned on a mixed dataset containing the original training data and new noisy data with random or structured labels. The goal is to force the models to fully memorize the new data (achieving >99.8% training accuracy on the combined set) while monitoring the retention of the original task's test accuracy. Three noise types are tested:

• Random noise.
• Structured noise with partial similarity ($x^2 + y \mod 37$).
• Structured noise with high similarity ($x + y \mod 37$).

Mechanistic Analysis:
The study employs Singular Value Decomposition (SVD) on the weight matrices of attention and MLP layers to analyze the Effective Rank (ER). ER quantifies the number of significant singular directions, serving as a measure of representational richness. Additionally, cosine similarity is calculated between parameter vectors before and after noise injection to measure global parameter displacement.

Key Results

1. The "Grokked Illusion": Despite achieving perfect generalization initially, AdamW-trained models suffer from catastrophic forgetting when forced to memorize new data. In the most challenging condition (random noise), their original task test accuracy drops from 100% to below 75%. In contrast, WLMD-equilibrium (high-entropy) models maintain approximately 95% test accuracy under the same conditions.
2. Entropy and Robustness: The robustness gap persists across all noise types but narrows as the new data becomes structurally more similar to the original task. High-entropy models consistently outperform AdamW models in retaining prior knowledge.
3. Effective Rank as a Buffer: SVD analysis reveals that WLMD-equilibrium models possess significantly higher effective rank in attention and MLP layers both before and after noise injection compared to AdamW models. While both models experience a reduction in effective rank after learning new data, the high-entropy models retain a much richer feature representation. This suggests that the abundance of representational capacity in high-entropy states acts as a buffer against forgetting.
4. Parameter Displacement: The cosine similarity between pre- and post-injection parameters is lower for WLMD models, indicating larger parameter displacements. However, because these models reside in a high-volume, high-entropy basin, they remain within the same functional region despite the displacement, whereas AdamW models, residing in narrower low-entropy basins, fall out of the functional region.

Significance and Claims
The paper claims to identify a phenomenon termed the "grokked illusion," where solutions that appear to have "grokked" (achieved saturated test performance) via conventional training can still be highly fragile to new information. The primary contribution is demonstrating that perfect generalization does not imply equal robustness.

The authors argue that the statistical physics framework, specifically the concept of Boltzmann entropy, provides a principled lens for understanding model resilience. They posit that high-entropy states, by occupying larger volumes in parameter space, inherently possess a robustness advantage. Furthermore, the study links this robustness to higher effective rank, suggesting that richer, more distributed feature representations serve as a buffer against interference. This finding challenges the prevailing view that low-rank structures are always associated with robustness (as seen in adversarial training contexts), proposing instead a third regime where high-entropy sampling yields high-rank, intrinsically robust representations.

The paper concludes that these findings offer a new perspective on catastrophic forgetting in continual learning, suggesting that entropy-guided strategies could position models in regions of parameter space that are intrinsically more robust to forgetting, rather than merely regularizing updates.