Technical Summary: Is Grokking Worthwhile?

Problem Statement

Large Language Models (LLMs) frequently exhibit a "curse of two-hop reasoning," where they fail to compose separately learned facts into a coherent chain of inference. For instance, a model may recall "Barack's wife is Michelle" and "Michelle was born in 1964" but fail to infer "Barack's wife was born in 1964" through internal computation. This failure is attributed to a representational mismatch: standard Transformers retrieve atomic facts in shallow layers, but deeper layers are not trained to utilize these specific latent representations as inputs for subsequent retrieval steps.

Recent literature suggests that parameter-sharing Transformers can overcome this by undergoing a prolonged training phase known as "grokking," during which a "Generalization Circuit" emerges. This circuit allows the model to internalize reasoning chains by explicitly recovering intermediate (bridge) entities in middle layers. However, fundamental questions remain: Is a grokked model truly superior to its non-grokked counterpart on downstream tasks? Is the extensive computational cost of waiting for the grokking phase justified? Does the reasoning capability acquired through grokking effectively transfer to novel facts?

Methodology

The authors employ a mechanistic study using a state-of-the-art parameter-sharing Transformer architecture on a synthetic compositional reasoning task.

• Dataset Construction: A synthetic dataset comprising 2,000 entities and 200 unique relations was constructed, totaling 40,000 atomic facts. The data is partitioned into In-Distribution (ID) and Out-of-Distribution (OOD) subsets. Models are trained on all atomic queries but only on compositional queries constructed from ID facts. Evaluation occurs on compositional queries involving OOD facts to test for genuine generalization rather than memorization.
• Mechanistic Analysis: The study utilizes the "logit lens" technique to decode hidden states across all layers, visualizing the maturation of Generalization Circuits. An inference trace is categorized as "Bridge found" if the decoded hidden state matches the ground-truth bridge entity.
• Experimental Variations:
  1. Natural Grokking: Training on standard ID-only compositional data to observe the natural emergence of grokking.
  2. Full Supervision Augmentation: Providing compositional supervision on OOD facts (First Hop Full Augmentation and Second Hop Full Augmentation) to accelerate learning and test if the reasoning path changes.
  3. Transferability Evaluation: Fine-tuning pre-trained models (both naturally grokked and "fake grokked") on a mix of retained ID facts and 2,000 entirely new atomic facts to assess the ability to assimilate new knowledge and compose it.

Key Findings

1. Identical Reasoning Paths

Contrary to the hypothesis that grokking represents the sudden acquisition of a new reasoning paradigm, the study finds that non-grokked and grokked models follow identical inference paths for in-distribution compositional queries.

• The "Generalization Circuit" is not a new structure formed only after grokking; rather, the reasoning path is established early in training for ID facts.
• Grokking is characterized as the process of integrating sparsely supervised OOD facts into this naturally established circuit.
• Models trained with full compositional supervision on OOD facts achieve the same reasoning path and transfer performance as naturally grokked models but in significantly fewer steps.

2. Dissociation of Behavior and Mechanism

The paper demonstrates that behavioral grokking (high accuracy on unseen cases) is not intrinsically linked to the formation of a Generalization Circuit.

• Under specific data regimes (e.g., Second Hop Full Augmentation), models can achieve high accuracy on OOD compositions without ever correctly recovering the bridge entity in intermediate layers.
• These "fake grokked" models lack the underlying mechanism of the Generalization Circuit and, crucially, fail to transfer this capability to new facts, unlike models that possess the true circuit.

3. Limited Transferability

Even when a mature Generalization Circuit is present, its transferability to new knowledge is limited.

• While the circuit allows models to recover bridge entities for new facts, it does not guarantee the ability to complete the second hop of reasoning if the second hop involves a novel fact.
• To establish a reliable circuit for new knowledge, models still require either an additional grokking phase or substantially more fine-tuning data covering the new facts and their compositions.
• "Fake grokked" models fail to adapt to compositional reasoning mixing new and old facts, merely memorizing new atomic facts without learning to reason over them.

Significance and Claims

The paper challenges the prevailing view that grokking represents the genuine acquisition of a generalized reasoning paradigm. Instead, it posits that grokking is a mechanism for integrating sparse data into an existing structure.

The authors conclude that the decision to pursue grokking involves a trade-off between data scarcity and computational resources:

• Under sparse supervision: Training until grokking may be necessary to force the integration of facts into the reasoning circuit.
• Under sufficient supervision: If sufficient compositional data is available, models can stop at training saturation and achieve equivalent performance without the computational overhead of the prolonged grokking phase.

Ultimately, the study suggests that "grokked" Transformers do not achieve a full mastery of compositional logic comparable to human-like reasoning, as their ability to generalize reasoning patterns across domains remains constrained.

Limitations

The authors acknowledge two primary limitations:

1. Computational Constraints: The investigation into "fake grokking" is not exhaustive due to the extreme training cycles required to observe grokking phenomena, leaving the precise boundary conditions for circuit collapse as an open question.
2. Architectural Scope: The findings are specific to parameter-sharing Transformers. Standard Transformers with independent layer weights inherently struggle with this form of implicit generalization, and the results may not extend to those architectures.