Title: Geometric Compression in Grokking: The Three-Stage Modular Dynamics of Transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Grokking, Geometric Deep Learning, Transformers

Abstract: A central mystery in deep learning is how generalizable algorithms emerge from the complex dynamics of training. The phenomenon of grokking serves as a canonical example of this puzzle. While mechanistic reverse engineering has successfully identified the final algorithms networks discover, the dynamic process of their formation remains largely uncharacterized. Progress in understanding these dynamics is hindered by complexity metrics that are either monolithic or presuppose non-universal architectural properties like piecewise linearity. We approach this challenge from a geometric perspective, introducing the Geometric Coherence Score (GCS), which quantifies how consistently networks transform local geometric structures across inputs, revealing the hidden geometric evolution of algorithmic learning. Applying GCS to Transformer grokking on modular arithmetic, we discover a universal three-stage construct-then-compress dynamic with precise modular division of labor: (I) Early Geometric Convergence attention achieves high geometric coherence through simple memorization patterns; (II) Geometric Restructuring attention actively decreases coherence by constructing complex structured representations necessary for generalization; (III) System wide Consolidation all computational flows coordinately increase coherence, stabilizing the generalizable algorithm. We substantiate this discovery through multiple lines of evidence: the dynamic persists across various activation functions (ReLU, GeLU, SiLU); it distinguishes successful grokking from overfitting (where geometric restructuring fails); and GCS dynamics directly correspond to the evolution of attention patterns from uniform simplicity to algorithmic sophistication. Our work reframes grokking as sophisticated modular geometric reorganization, providing the first direct geometric evidence for construct then compress mechanisms in neural networks and offering a principled diagnostic tool for interpreting emergent algorithms.

TL;DR: We show that grokking in Transformers is not monotonic simplification, but a "construct-then-compress" algorithm where the Self-Attention module must first increase its geometric complexity to enable a subsequent, rapid compression in the FFN.