The Long Delay to Arithmetic Generalization: When Learned Representations Outrun Behavior

A new research paper by Laura Gomezjurado Gonzalez, 'The Long Delay to Arithmetic Generalization: When Learned Representations Outrun Behavior,' tackles the mysterious phenomenon of 'grokking' in AI. Grokking describes when a model trained on algorithmic tasks, like arithmetic, shows a long period of memorizing training data with near-chance accuracy before suddenly and sharply generalizing. The study argues this delay isn't due to a failure to learn the underlying structure but rather a bottleneck in accessing that knowledge. Using a one-step Collatz prediction task, the researchers found that the model's encoder organizes the necessary mathematical concepts (like parity and residue) within the first few thousand training steps. However, the decoder struggles to translate this internal representation into correct outputs for tens of thousands of steps longer.

Causal interventions strongly support this 'decoder bottleneck' hypothesis. Transplanting a fully trained encoder into a fresh, untrained model accelerated the grokking process by 2.75 times. Conversely, transplanting a trained decoder actively hurt performance. Perhaps most strikingly, freezing a converged encoder and retraining only the decoder eliminated the accuracy plateau entirely, achieving 97.6% accuracy compared to 86.1% for standard joint training. The research also revealed that the choice of how numbers are represented (the numeral base) acts as a critical inductive bias. For the Collatz task, bases whose factorization aligns with the map's arithmetic, like base 24, achieved 99.8% accuracy, while binary representation failed completely because its representations collapsed and never recovered.