Is Grokking Worthwhile? Functional Analysis and Transferability of Generalization Circuits in Transformers

TL;DR

The paper investigates whether grokking yields genuinely transferable compositional reasoning in transformers and whether the associated computational cost is justified. Using mechanistic analysis with the logit lens on a parameter-sharing Tiny_recursive_model across ID/OOD two-hop tasks, it shows that grokked and non-grokked models follow the same reasoning path and that the Generalization Circuit often forms early and is not the sole driver of improved transfer. It also demonstrates that behavioral grokking can occur without the intended mechanistic circuit and that transfer of the circuit to new facts remains limited, requiring additional grokking or substantial finetuning. These findings challenge the view that grokking equates to fully transferable, human-like compositional mastery and highlight a data-compute trade-off: with sufficient supervised data, models can achieve equivalent performance without waiting for grokking, while sparse supervision may necessitate grokking-time investments.

Abstract

While Large Language Models (LLMs) excel at factual retrieval, they often struggle with the "curse of two-hop reasoning" in compositional tasks. Recent research suggests that parameter-sharing transformers can bridge this gap by forming a "Generalization Circuit" during a prolonged "grokking" phase. A fundamental question arises: Is a grokked model superior to its non-grokked counterparts on downstream tasks? Furthermore, is the extensive computational cost of waiting for the grokking phase worthwhile? In this work, we conduct a mechanistic study to evaluate the Generalization Circuit's role in knowledge assimilation and transfer. We demonstrate that: (i) The inference paths established by non-grokked and grokked models for in-distribution compositional queries are identical. This suggests that the "Generalization Circuit" does not represent the sudden acquisition of a new reasoning paradigm. Instead, we argue that grokking is the process of integrating memorized atomic facts into an naturally established reasoning path. (ii) Achieving high accuracy on unseen cases after prolonged training and the formation of a certain reasoning path are not bound; they can occur independently under specific data regimes. (iii) Even a mature circuit exhibits limited transferability when integrating new knowledge, suggesting that "grokked" Transformers do not achieve a full mastery of compositional logic.

Figures (4)

• Figure 1: Mechanistic Comparison Between traditional transformer and parameter-sharing transformer. (Left) Standard Transformers struggle with representational mismatch across independent layers and are unable to generalize to compositional reasoning. (Right) Transformer with shared parameter resolves the bridge entity $b$ and is able to use shallow parameters to get the final answer.
• Figure 2: Transformer model "Grokking" under different data regime.
• Figure 3: "Grokked" Transformer learning new knowledge, detailed visualization on all settings can be seen in Figure \ref{['fig:Finetune plot 3_4']} in Appendix \ref{['Appendix:figure']}
• Figure 4: "Grokked" transformer finetuned on new facts.