Technical Summary: Two Speeds of Learning: A Representation-Readout Decomposition of Grokking and Double Descent

Problem Statement

Standard signals for monitoring deep neural network training, such as training loss and accuracy, often fail to capture the complex dynamics of generalization. Two prominent phenomena complicate this picture:

1. Grokking: A model achieves near-perfect training performance while test accuracy remains at chance levels for a prolonged period, followed by an abrupt emergence of generalization.
2. Epoch-wise Double Descent: Training loss decreases monotonically, while test loss or error exhibits a non-monotonic rise and fall.

Existing explanations are frequently task-specific (e.g., relying on algebraic structures in modular arithmetic) or rely on binary regime transitions (e.g., "lazy" to "rich" learning). There is a lack of a unified, task-agnostic framework capable of diagnosing these phenomena across diverse architectures and tasks, distinguishing genuine generalization from spurious learning artifacts induced by non-standard training recipes.

Methodology: Representation-Readout Decomposition

The authors propose a framework that decomposes a neural network $f(x)$ into two distinct components:

1. Encoder ($\phi_\theta$): A nonlinear map from input space to an $N$-dimensional representation space.
2. Readout ($W$): A linear map from the representation space to output logits.

This decomposition allows for the separate tracking of representation learning (changes in $\phi_\theta$) and readout calibration (changes in $W$). To analyze the interplay between these components, the authors employ four task-agnostic diagnostic measures derived from representational geometry and Neural Tangent Kernel (NTK) theory:

1. Critical Dimension ($N_{crit}$): Quantifies the robustness of linear separability in the representation space. It is defined as the expected smallest projection dimension required to preserve separability with probability $\geq 0.5$. A lower $N_{crit}$ indicates better "untangled" manifolds and higher generalizability.
2. Task-Relevant Geometric Measures (GLUE): A set of four metrics ($D, R, \rho_c, \rho_a$) derived from the Geometry Linked to Untangling Efficiency (GLUE) theory. These characterize the effective manifold dimension, radius, and the alignment of centers and axes between label-conditioned manifolds.
3. Linear Probe Accuracy: A regularized linear classifier trained on frozen representations. The gap between probe accuracy and model accuracy diagnoses readout misalignment (overfitting or underfitting relative to the encoder's capabilities).
4. Last-Layer NTK-Label Alignment: Measures the alignment between the feature Gram matrix (last-layer NTK) and the label kernel. This quantifies how well the current feature geometry supports gradient flow toward label-discriminative solutions, evaluated separately on train and test sets.

Key Contributions and Results

1. Two Speeds of Learning in Grokking

Applying the decomposition to canonical grokking tasks (modular addition, permutation composition, sparse parity) across MLP and Transformer architectures, the authors refute the strict "lazy-to-rich" account, which posits that representation learning is dormant until the onset of grokking.

• Continuous Representation Learning: $N_{crit}$ decreases gradually and continuously throughout training, including the pre-grokking phase. Representation learning is slow but not absent; it proceeds continuously rather than as a sudden phase transition.
• Train-Biased Readout: Before grokking onset, the readout learning direction is systematically biased toward training data. The NTK alignment gap (train minus test) is large and positive, indicating that the readout overfits the training set while the representations are not yet sufficiently robust to support generalization.
• Mechanism: Grokking arises not from a sudden switch in learning regimes, but from the gradual equalization of the readout direction as representation quality ($N_{crit}$) eventually catches up to the speed of readout calibration.

2. Architectural Differences in Geometric Mechanisms

The framework reveals that different architectures achieve generalization through distinct geometric pathways.

• MLPs vs. Transformers: In modular addition and permutation composition tasks, MLPs and Transformers exhibit opposite trends in center alignment ($\rho_c$) and axis alignment ($\rho_a$). While both architectures reduce manifold dimension and radius, they manipulate inter-class geometry differently to achieve the same surface-level grokking behavior.

3. Diagnosing Spurious Learning

The framework successfully distinguishes genuine generalization from artifacts caused by non-standard training recipes.

• MNIST Grokking Artifact: A previously reported grokking example on MNIST (induced by high initial weight norms and small training sets) is diagnosed as spurious. The framework reveals representation degradation (increasing $N_{crit}$ and decreasing linear probe accuracy) and readout overfitting during the "pre-grokking" phase. The apparent delayed generalization is an artifact of the readout eventually correcting for a degrading representation, rather than a genuine improvement in feature learning.
• Epoch-wise Double Descent: In double descent scenarios induced by label noise, the framework identifies a two-stage process: an initial phase of representation deterioration followed by a phase where readout updates are dominated by spurious signals (noisy labels). Removing the non-standard ingredients (e.g., label noise) eliminates both the double descent pattern and the diagnostic signatures.

4. Top-Down Analysis Framework

The authors position this decomposition as a "top-down" complement to "bottom-up" mechanistic interpretability (e.g., circuit discovery). By characterizing the global structure of representations (manifold geometry, separability, alignment), the framework provides consistent observables across train/test splits and diverse tasks, offering a reliable signal for monitoring learning dynamics where standard loss curves fail.

Significance

The paper establishes the representation-readout decomposition as a unified, task-agnostic framework for understanding learning dynamics. Its primary significance lies in:

• Refining Theoretical Accounts: It moves beyond the binary lazy-to-rich dichotomy, demonstrating that representation learning is continuous and that grokking is driven by the relative speeds of representation improvement versus readout calibration.
• Diagnostic Utility: It provides specific signatures to identify spurious learning dynamics (e.g., representation degradation or readout misalignment) that may be mistaken for genuine learning phenomena in non-standard training setups.
• Generalizability: The framework applies consistently across algorithmic tasks, real-world datasets (MNIST), and architectures (MLP, Transformer, ResNet), offering a principled tool for interpretability research and safety-critical model monitoring.

The authors conclude that while the framework is currently applied at the final linear layer, it offers a robust foundation for future investigations into intermediate layers, task-specific structures, and continual learning scenarios.