A new ArXiv study has mapped the precise internal mechanism behind 'grokking' — the delayed generalisation effect in neural networks — identifying a two-phase lifecycle in a geometric structure called the spectral edge that is 4,000 times more influential than random network directions at the moment generalisation occurs.

Grokking was first documented by researchers at OpenAI in 2022 and became one of the more puzzling observations in deep learning: a model trained on a task first memorises its training data, then — sometimes thousands of steps later — abruptly learns to generalise to new examples, without any obvious change in the training process. The new paper, from researchers posting to ArXiv CS.LG, asks a precise question: what is actually happening inside the network at that moment, and why?

What the Spectral Edge Actually Is

The study focuses on the Gram matrix of parameter updates — a mathematical object that captures how different parts of a network's weights are changing together during training. The dominant direction of this matrix, called the spectral edge, represents the single most influential axis along which the network is being reshaped at any given step. Think of it as the primary 'direction of travel' for the network's internal geometry.

The researchers studied two sequence tasks — Dyck-1, a formal language task involving balanced brackets, and SCAN, a compositional instruction-following benchmark — using small transformer models trained until grokking occurred.

At grokking, the spectral edge stops driving learning and becomes a compression axis — perturbation-flat yet ablation-critical.

Before grokking, they found, the spectral edge is shaped almost entirely by gradients: it reflects what the loss function is pushing the network to do. The network is actively learning. But at the moment grokking occurs, something changes. Weight decay — the regularisation technique that penalises large parameter values — aligns with the gradient component, and the spectral edge transitions from a learning driver to what the authors call a compression axis.

A Compression Axis, Not an Erasure

The distinction matters. A compression axis might sound like information is being discarded — squeezed out of the network. But the paper's findings suggest the opposite. Using nonlinear probes to test what information remains encoded after grokking, the researchers found an MLP probe achieved R² = 0.99, compared to R² = 0.86 for a linear probe applied to the same representations. In plain terms: the information is still there, but it has been re-encoded into a more compressed, nonlinear form that simpler probes cannot read.

This distinction is meaningful for interpretability research. It suggests that what looks like forgetting or compression from the outside is actually a restructuring of how knowledge is stored — the algorithm the network has learned is preserved, just represented differently.

The paper also reports that removing weight decay after grokking has occurred reverses the compression process while leaving the underlying learned algorithm intact. That finding implies weight decay is not just a training-time regulariser but an active force shaping the geometric structure of a generalised solution — and that its effects can, in principle, be disentangled from the solution itself.

Three Universality Classes and a Predictive Equation

One of the study's more ambitious claims is that spectral edges across different training runs and tasks fall into three universality classes: functional (the edge is actively driving computation), mixed (a transitional state), and compression (post-grokking). The authors propose a gap flow equation that predicts which class a given spectral edge belongs to based on the relative magnitudes of the gradient and weight-decay components.

If this classification holds up under broader testing, it would give researchers a diagnostic tool: by measuring a network's spectral edge during training, they could identify where in the grokking lifecycle it sits, without waiting for generalisation to appear in evaluation metrics.

The study is based on relatively small transformer models trained on formal language tasks, which are standard testbeds for grokking research but not representative of the scale or complexity of production AI systems. Whether the same lifecycle and universality classes appear in larger models trained on natural language or other domains is an open question the paper does not resolve.

Why Grokking Research Has Broader Stakes

Understanding grokking matters beyond academic curiosity. If neural networks routinely pass through a memorisation phase before arriving at genuine generalisation, then evaluation benchmarks that test models too early may be systematically misleading. A model that appears to have learned a task may simply not have grokked yet. Conversely, a model that has grokked may store its solutions in compressed, nonlinear representations that are harder to audit using standard interpretability tools.

The paper's finding that the spectral edge is perturbation-flat — meaning small changes to the network along that axis don't affect outputs — but ablation-critical — meaning removing it collapses performance — also has implications for model compression research. It suggests that the most geometrically prominent feature of a generalised network may be exactly the feature that is hardest to safely remove.

The results are preliminary and based on self-reported benchmarks from the authors' own experimental setup. Independent replication across model scales and task types will be needed before the gap flow equation or the three universality classes can be treated as established findings.

What This Means

This research gives scientists a more precise language — and potentially a predictive tool — for tracking when and how neural networks transition from memorisation to genuine generalisation, which could improve both training diagnostics and interpretability methods for AI systems.