Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization — part 1/4
Cleaned-text code points [0, 15667). Bibliography omitted here; source.md includes it.
# Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization
###### Abstract
We study whether transformers can learn to implicitly reason over parametric knowledge, a skill that even the most capable language models struggle with. Focusing on two representative reasoning types, composition and comparison, we consistently find that transformers can learn implicit reasoning, but only through grokking , i.e., extended training far beyond overfitting. The levels of generalization also vary across reasoning types: when faced with out-of-distribution examples, transformers fail to systematically generalize for composition but succeed for comparison. We delve into the model’s internals throughout training, conducting analytical experiments that reveal: 1) the mechanism behind grokking, such as the formation of the generalizing circuit and its relation to the relative efficiency of generalizing and memorizing circuits, and 2) the connection between systematicity and the configuration of the generalizing circuit. Our findings guide data and training setup to better induce implicit reasoning and suggest potential improvements to the transformer architecture, such as encouraging cross-layer knowledge sharing. Furthermore, we demonstrate that for a challenging reasoning task with a large search space, GPT-4-Turbo and Gemini-1.5-Pro based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation, while a fully grokked transformer can achieve near-perfect accuracy, showcasing the power of parametric memory for complex reasoning. 1 1 1 Code and data: https://github.com/OSU-NLP-Group/GrokkedTransformer .
## 1 Introduction
Large language models (LLMs) have been shown deficient in implicit reasoning with their parametric memory of knowledge and rules. For example, a range of LLMs are found to be incapable of robustly composing internalized facts [ 48 , 71 ] , and even GPT-4 [ 42 ] cannot adequately compare entities’ attributes despite knowing them [ 1 ] .
Deficiency in implicit reasoning has profound implications. It implies the models’ limitations in inducing structured and compressed representations of facts and rules, which lead to redundant knowledge storage and difficulty in propagating knowledge updates [ 76 ] , and importantly, fundamentally impede the model from systematic generalization over knowledge [ 25 ] . While explicit verbalizations of reasoning steps (e.g., chain-of-thought rationales) can improve task performance [ 67 , 64 , 73 , 55 , 31 ] , they are not available during large-scale (pre-)training where the model’s core capabilities are acquired [ 77 , 29 ] .
Is implicit reasoning doomed given that even the most capable models struggle? Can it be resolved by further scaling data and compute, or are there fundamental limitations of the transformer [ 62 ] that prohibit robust acquisition of this skill?
In this paper, we rigorously study these questions by constructing synthetic training and evaluation datasets, training transformers from scratch, and examining their generalization. We conceptualize reasoning as the induction and application of inference rules , and expose the model to a mixture of ‘‘atomic facts’’ and ‘‘inferred facts’’ (which are deduced from the atomic facts via a set of latent rules), resembling ‘‘axioms’’ and ‘‘theorems’’ in a formal system. To evaluate how well the model learns the rules, we test its ability to make novel deductions (i.e., completing unseen inferred facts) in both in-distribution (ID) and out-of-distribution (OOD) scenarios. 2 2 2 Definitions of ID/OOD are introduced in § 2 . This approach allows us to control the training data and perform clean evaluations, which would be challenging when studying existing LLMs trained on uncontrolled data.
Our experiments reveal that transformers can learn to perform implicit reasoning, but this skill is only robustly acquired through extended training far beyond overfitting (Figure 1 ), a phenomenon known as grokking [ 47 ] . We find that the speed of improvement in generalization correlates with the ratio between inferred and atomic facts in training, and depends little on the absolute size of the training data (Figure 2 ). This suggests a correction of prior explanations of grokking based on critical data size [ 33 , 61 , 78 , 21 ] , in that it should instead be the critical data distribution that decides the characteristics of grokking. Our findings extend prior observations of the grokking phenomenon primarily in algorithmic and linguistic tasks [ 47 , 38 ] to the domain of knowledge-based reasoning, and deepen our understanding of the grokking phenomenon.
Moreover, we find that the transformer exhibits different levels of systematicity across reasoning types. While ID generalization is consistently observed, in the OOD setting, the model fails to systematically generalize for composition but succeeds in comparison (Figure 1 ). To understand why this happens, we conduct mechanistic analysis of the internal mechanisms of the model. The analysis uncovers the gradual formation of the generalizing circuit throughout grokking and establishes the connection between systematicity and its configuration, specifically, the way atomic knowledge and rules are stored and applied within the circuit. Our findings imply that proper cross-layer memory-sharing mechanisms for transformers such as memory-augmentation [ 54 , 17 ] and explicit recurrence [ 7 , 22 , 57 ] are needed to further unlock transformer’s generalization.
Finally, to demonstrate the power and potential of parametric memory for complex reasoning, we show that for a reasoning task with a large search space, a fully grokked transformer can achieve near-perfect accuracy, while state-of-the-art LLMs like GPT-4-Turbo [ 43 ] and Gemini-1.5-Pro [ 16 ] based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation.
## 2 General Setup
Training data & ID/OOD evaluation. As stated in § 1 , we are interested in whether transformers can induce and apply latent rules over knowledge implicitly in a generalizable way. We create a data-generating process consisting of 1) sampling a set of basic atomic facts , and 2) using the atomic facts and latent rules to deduce inferred facts . To better characterize the level of generalization acquired by the model, we evaluate the model’s in-distribution (ID) and out-of-distribution (OOD) performance. We prepare two separate sets of atomic facts: atomic ID {}_{\text{ID}} and atomic OOD {}_{\text{OOD}} . Our training set includes all the atomic facts and a uniformly random portion of the inferred facts deduced from atomic ID {}_{\text{ID}} , which we call train_inferred ID {}_{\text{ID}} . For evaluation, (1) ID generalization aims to evaluate whether the model learns the latent rules correctly, by testing its ability to complete unseen inferred facts also deduced from atomic ID {}_{\text{ID}} , which we denote by test_inferred ID {}_{\text{ID}} . (2) OOD generalization aims to evaluate the systematicity [ 25 ] acquired by the model, namely, the ability to apply rules over knowledge regardless of its distribution. To do this, we test the model on the facts deduced from atomic OOD {}_{\text{OOD}} , denoted by test_inferred OOD {}_{\text{OOD}} .
Model & optimization. We use a standard decoder-only transformer model as in GPT-2 [ 50 ] with 8 8 layers, 768 768 hidden dimensions and 12 12 attention heads (we explore the impact of different model scales in Appendix B ). Optimization is done by AdamW [ 34 ] with learning rate 10 − 4 10^{-4} , batch size 512 512 , weight decay 0.1 0.1 and 2000 2000 warm-up steps. Notably, models are trained for a large number of epochs/steps beyond the point where training performance saturates. More details are in Appendix A .
## 3 Composition—Delayed Generalization without Systematicity
We begin our investigation with composition , where a model needs to “chain” different pieces of facts, e.g., “Barack’s wife is Michelle” and “Michelle is born in 1964” , to successfully complete a compositional sentence, e.g., “Barack’s wife is born in [1964]” . Prior work extensively studied whether transformer-based language models can perform implicit composition, and negative results are consistently reported [ 48 , 1 , 71 ] . Specifically, there exists a “compositionality gap” [ 48 ] , i.e., the frequency at which the model knows all the underlying basic facts but fails to compose them, which is considerable across different LLMs and does not decrease as models scale. Are transformers doomed to fail on such kind of reasoning, and if so, why?
### 3.1 Setup
We focus on two-hop composition in this work. For atomic facts, we generate a random knowledge graph 𝒢 \mathcal{G} consisting of | ℰ | |\mathcal{E}| entities and | ℛ | = 200 |\mathcal{R}|=200 relations, where each entity (as the subject) has 20 20 random distinct relations that each connects to another random entity (as the object). The atomic facts are then the edges, i.e., (subject, relation, object) triplets in 𝒢 \mathcal{G} , which we partition disjointly into atomic ID {}_{\text{ID}} and atomic OOD {}_{\text{OOD}} ( 95 95 %: 5 5 %). The rule of (two-hop) composition is ∀ h , b , t ∈ ℰ , ∀ r 1 , r 2 ∈ ℛ , ( h , r 1 , b ) ∧ ( b , r 2 , t ) ⟹ ( h , r 1 , r 2 , t ) , \forall h,b,t\in\mathcal{E},\forall r_{1},r_{2}\in\mathcal{R},(h,r_{1},b)\land(b,r_{2},t)\implies(h,r_{1},r_{2},t), (1)
which is used to deduce the ID and OOD inferred facts from atomic ID {}_{\text{ID}} and atomic OOD {}_{\text{OOD}} , respectively. For convenience, in the above rule, we will call h h the head entity, b b the bridge entity, and t t the tail entity. For both atomic and inferred facts, training/testing is done by having the model predict the final tail entity. We assign a unique token to each relation/entity by default, and also find that the results are robust to different tokenizations (details in Appendix C ).
We study the influence of the following two aspects on the model’s learned behaviors: • Ratio between inferred and atomic facts. By varying the amount of inferred facts included in training, we study the effect of the ratio ϕ = | train_inferred ID | / | atomic ID | \phi=|\text{{{train\_inferred${}_{\text{ID}}$}}}|/|\text{{{atomic${}_{\text{ID}}$}}}| on the model.
• Training data size. We study the impact of the training data size by varying | ℰ | |\mathcal{E}| , the total number of entities, while controlling the ratio ϕ \phi . Note that the size of training data (both atomic/inferred facts) scales linearly with | ℰ | |\mathcal{E}| .
### 3.2 Results
Grokking observed in ID generalization but not in OOD generalization . Figure 1 (left) shows the model’s accuracy on the train and test facts throughout optimization, with | ℰ | = 2000 |\mathcal{E}|=2000 and ϕ = 7.2 \phi=7.2 . We find that the model can generalize to ID test examples, but high performance is only achieved through extended training far beyond overfitting, a phenomenon called grokking [ 47 ] . Specifically, the training performance saturates (over 99 99 % accuracy on both atomic and inferred facts) at around 14K optimization steps, before which the highest ID generalization accuracy is merely 9.2 9.2 %. However, generalization keeps improving by simply training for longer, and approaches almost perfect accuracy after extended optimization lasting around 50 50 times the steps taken to fit the training data. On the other hand, OOD generalization is never observed. We extend the training to 2 2 million optimization steps, and there is still no sign of OOD generalization.
Inferred/atomic ratio ϕ \phi correlates with generalization speed . Figure 2 (a) shows the ID test accuracy across different ϕ \phi . We omit the other splits since for all settings, the training performance saturates quickly and the OOD test accuracy remains at zero as earlier. 3 3 3 The training performances of all settings saturate within 25 25 K steps, where larger ϕ \phi takes more steps. It could be seen that the ratio ϕ \phi strongly correlates with the speed of generalization. A very large ratio can push generalization to improve at a similar pace as the model fits the training data, reducing the need for extended training. 4 4 4 When ϕ = 18.0 \phi=18.0 , the model achieves 96.7 % 96.7\% accuracy before training performance saturates.
Training data distribution , instead of training data size, qualitatively influences generalization behavior . When ϕ \phi increases and | ℰ | |\mathcal{E}| holds constant, the size of training data also gets larger. Prior studies hypothesize that training data size plays a central role in order for grokking to happen. In particular, previous work connects grokking with the notion of critical data size (CDS) [ 33 , 61 , 78 , 21 ] , where it is hypothesized that CDS marks the shift from memorization to generalization (via grokking), and the speed of generalization improves as the training data further scales. However, results from our controlled experiments seem to contradict such a hypothesis. Figure 2 (b) shows the results of varying | ℰ | |\mathcal{E}| with a fixed ϕ = 9.0 \phi=9.0 , where we change the horizontal axis from optimization step to epoch for better visualization. 5 5 5 The optimization steps for each epoch scale linearly with the training size since we use a fixed batch size. When fixing the ratio ϕ \phi , the training data size does not qualitatively affect the model’s generalization. Specifically, scaling the data affects neither the relative speed of ID generalization and training improvement (as seen by the rather constant “gap” between train_inferred ID {}_{\text{ID}} and test_inferred ID {}_{\text{ID}} curves), nor the systematicity level (OOD performance stays zero). We also run the experiments across different ϕ \phi and find the results to be consistent. This suggests that critical data “distribution”, not size, may be the actual deciding factor behind grokking and generalization. In addition, we find that scaling up the model size also does not qualitatively change the generalization behaviors observed here (Appendix B ), and the main pattern is that larger models converge in fewer optimization steps, which shares with prior findings [ 60 , 28 ] .
Summary. We have shown that transformers are capable of acquiring the rule of composition through grokking, with controlled experiments suggesting the crucial factor of data distribution (e.g., the inferred/atomic ratio ϕ \phi ) in characterizing the model’s generalization. However, important questions still remain: what happens during grokking, why does it happen, and why do transformers struggle with OOD examples? Answering these questions requires a deeper understanding of (the changes in) the model’s inner workings, which we investigate next.
### 3.3 Analyzing the inner workings of the model throughout grokking
We analyze the internal mechanisms within the model via a combination of two prevalent approaches: logit lens and causal tracing. We apply our analysis to the setting with | ℰ | = 2000 , ϕ = 9.0 |\mathcal{E}|=2000,\phi=9.0 on 300 300 random examples from train_inferred ID {}_{\text{ID}} .
Logit lens . We interpret individual hidden states via logit lens [ 40 , 15 , 71 ] , where the activation is converted into a set of logits for each vocabulary token by multiplying with the output embedding matrix. We follow the recent practice [ 71 ] where the activation first goes through the transformer’s final normalization layer before multiplying with the output embedding (Figure 3 , top right).