Grokking in Transformers: The Decoder Bottleneck and the Influence of Numerical Representation

The phenomenon of "grokking" in transformer models, characterized by a long delay between fitting the training data and abrupt generalization, represents one of the most intriguing challenges in machine learning. Understanding the root causes of this delay is crucial for optimizing the development and deployment of Large Language Models (LLMs) and other transformer-based systems. Recent research sheds new light on this enigma, suggesting that the problem lies not in the model's acquisition of structure, but rather in the decoder's ability to effectively access and leverage it.

The study focused on encoder-decoder arithmetic models, specifically in one-step Collatz sequence prediction. Results indicate that the encoder is capable of organizing parity and residue structure within a few thousand training steps. However, output accuracy remains near chance for tens of thousands of additional steps, highlighting a disconnect between the encoder's internal learning and the model's ability to produce correct results. This observation led to the "decoder bottleneck hypothesis."

The Critical Role of the Decoder in Generalization

To test the decoder bottleneck hypothesis, researchers conducted a series of causal interventions. The results were significant: transplanting an already trained encoder into a new model accelerated the grokking process by 2.75 times. Conversely, transplanting a trained decoder actively hurt performance, suggesting that the decoder is the limiting factor for generalization.

A particularly revealing experiment involved freezing a converged encoder and subsequently Fine-tuning only the decoder. This strategy completely eliminated the learning plateau, leading to an accuracy of 97.6%. This figure contrasts sharply with the 86.1% accuracy obtained with joint training of both the encoder and decoder. These results reinforce the idea that the decoder is the primary obstacle to generalization, and that once the encoder has learned the structure, the decoder requires a targeted learning process to fully exploit it.

The Influence of Numerical Representation

The research also explored how numerical representation influences the decoder's ability to perform its task. Through the analysis of 15 different numeral bases, it emerged that those whose factorization aligns with the Collatz map's arithmetic (e.g., base 24) allowed models to achieve 99.8% accuracy. This suggests that some representations inherently facilitate the decoder's task.

In contrast, binary representation showed complete failure, due to the collapse of its internal representations, which never recovered. The choice of numeral base thus acts as an inductive bias, controlling the amount of local digit structure the decoder can exploit. This leads to substantial differences in the model's "learnability," even if the underlying task remains identical.

Implications for LLM Development and On-Premise Deployment

These findings have significant implications for the design and optimization of LLMs, especially for those evaluating self-hosted or on-premise deployments. Understanding that generalization capability can be limited by a "bottleneck" in the decoder, and that data representation can act as a powerful inductive bias, offers new avenues for improving model efficiency.

For infrastructure architects and DevOps leads, a model that "groks" faster or requires fewer resources to achieve high accuracy directly translates into a lower TCO. Optimizing Fine-tuning strategies, perhaps by focusing on the decoder once the encoder has learned essential features, or exploring more efficient data representations, could reduce VRAM requirements, the computational power needed for Inference, and training times. This is particularly relevant in air-gapped environments or those with resource constraints, where every optimization contributes to maximizing the value of hardware and infrastructure investment. AI-RADAR offers analytical frameworks on /llm-onpremise to evaluate these trade-offs and support informed deployment decisions.