Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization — part 3/4

Cleaned-text code points [31098, 46688). Bibliography omitted here; source.md includes it.

In the original setting (§ 4 ) of the task, for the OOD test set, one can simply retrieve the two OOD facts and compare the attribute values, which requires no further search. We change the task setting via the following. For each attribute, 1) we do not add the OOD atomic facts into training, and 2) we add a random portion of the comparisons between ID entities and OOD entities into training. We test the models on queries consisting of derivable (from all training facts) comparisons between OOD entities where any possible proof would involve rules from both Eqs.( 2 ) and Eqs.( 3 ). Consequently, answering a test query would require the model to successfully locate two ID bridge entities which can connect the two query entities into a proof (Figure 6 ). We select a balanced (by a < , a = , a > a_{<},a_{=},a_{>} ) subset from these queries for evaluation. More details are included in Appendix F .

The difficulty of such a task is two-fold. First, the search space is large. For example, on average, each query entity connects with more than 50 50 facts, and each bridge entity in the ground truth proof connects with more than 900 900 facts. Second, there are no surface form clues to exploit and bias the search towards the ground truth proof, unlike most conventional QA benchmarks where the proof steps are transparent from the query.

To test LLMs based on non-parametric memory, we translate the facts into natural language by simple templates (Appendix F ). Facts/queries for each attribute are grouped/tested separately. 9 9 9 This could also be thought of as performing a retrieval step into the memory based on the attribute. We test both the vanilla setup where all facts ( 28.2 28.2 K on average) are loaded into the LLM context, and the retrieval-augmented setup ( 5.4 5.4 K facts retrieved on average) where the two-hop neighborhoods of the two query entities are retrieved, which includes enough facts to deduce the answer. We also try both standard prompting where the model answers directly, and chain-of-thought (CoT) prompting where the model is prompted to verbalize the reasoning. We test GPT-4-Turbo and Gemini-Pro-1.5, where for GPT-4-Turbo we only test the retrieval-augmented setup due to context length limit.

Results . As shown in Table 1 , all models based on non-parametric memory fail badly, where the only setting that surpasses random guess ( 33.3 33.3 %) is the retrieval-augmented setting with Gemini-Pro-1.5 and direct answer prediction. Intriguingly, LLMs perform worse (especially Gemini) when prompted to reason verbally. Through closer examinations, we find that with CoT, 70.7 70.7 % of Gemini’s responses conclude that the answer cannot be decided (which we treat as wrong since the answer can be decided) after a series of search steps. 10 10 10 The ratio drops a bit to 58.7 58.7 % when augmenting with retrieval. More shockingly, most of the CoT rationales that achieve the correct final answer are actually wrong due to either hallucinating underivable facts or logical errors . This illustrates the current models’ inability to reason deeply with non-parametric memory. On the other hand, a grokked transformer, which is trained extensively on the given facts to compress and integrate the information to the extreme, could achieve near-perfect accuracy. By examining the model, we find that it acquires the same generalizing circuit as in Figure 5 (a), and remarkably, even though not explicitly encouraged/trained to do this, the model successfully infers most of the OOD entities’ attribute values by integrating the observed training facts (Appendix F ).

## 6 Related Work

Knowledge and reasoning in language models . Numerous work finds that transformer language models, even SoTA ones such as GPT-4, struggle in implicit reasoning over their parametric knowledge [ 56 , 23 , 52 , 51 , 48 , 1 , 71 ] , suggesting their limitations in inducing structured and compressed representations of facts and rules during training. A series of efforts try to understand transformer’s knowledge and reasoning through controlled experiments [ 24 , 49 , 10 , 66 ] , which is also our focus. We find that transformers can learn implicit reasoning over knowledge through grokking, and characterize the connection between the acquired systematicity level and the inductive bias of transformer.

“Chain-of-Thought” and verbalized reasoning . A series of studies prompt/fine-tune language models to verbalize (i.e., generate) the intermediate knowledge and reasoning steps [ 67 , 64 , 73 , 55 , 31 , 72 ] during inference, which has been shown to improve performance especially for large models with strong generation capabilities. There are also theoretical results showing the benefits of such verbalizations [ 11 , 27 ] . Our focus here on implicit reasoning is orthogonal, and it is an interesting open problem to have principled understandings of the role of such verbalizations in reasoning problems, and also develop methods that can decide the appropriate balance between implicit and explicit reasoning to handle challenging problems with large intrinsic complexity. Relatedly, recent work also finds that explicit verbalizations could be a useful medium for teaching models to reason implicitly via distillation or curriculum [ 9 , 8 ] .

Grokking is first discovered by Power et al. [ 47 ] on a set of small algorithmic reasoning tasks. The intriguing phenomenon inspired follow-up works proposing different explanations and expanding the set of tasks where grokking is observed [ 59 , 32 , 6 , 41 , 39 , 61 , 36 , 38 , 33 , 78 , 21 ] . To our knowledge, we are the first work to observe grokking in the domain of knowledge-based reasoning, and our controlled experiments suggest potential corrections of prior hypotheses based on critical data size. Our formulation of rule induction from atomic and inferred facts is general, which we hope could inspire future work on understanding grokking and generalization in deep learning.

Analyzing the inner workings of neural models . Recent work tries to open up the “black box” of neural models through a wide range of techniques; see survey in Ferrando et al. [ 13 ] . We apply causal tracing [ 63 , 35 , 19 , 65 , 12 ] and logit lens [ 40 , 15 , 71 ] to discover interpretable circuits in the model to understand the grokking process and how/why generalization happens.

Parametric and non-parametric memory . Our focus in this work is on parametric memory in language models, and an orthogonal direction is to enhance models with non-parametric memory, such as extending the effective context length [ 70 , 43 , 14 , 16 ] and augmenting with retrieval [ 18 , 26 , 3 , 58 , 75 , 37 ] . The two types of memory are largely complementary to each other—parametric memory has its unique ability to compress and integrate information but is also inevitably lossy and subject to hallucination, while non-parametric memory is lossless and could also provide attribution. Similarly for humans—a human acquires expertise in a domain by acquiring and structuring knowledge in the brain (parametric), but he/she also wouldn’t memorize all pieces of details and could refer to the source when necessary (non-parametric). How to decide the tradeoff between parametric and non-parametric memory (or, how to define the objective for such a tradeoff) is another interesting open problem for future work.

## 7 Conclusion

We find that transformers are capable of learning to implicitly reason over parametric knowledge, however, such a skill is only robustly acquired through extended training far beyond the point of overfitting, or grokking . Mechanistic analysis into the model’s internals reveals the configuration and gradual formation of the generalizing circuit, and also explains the different levels of systematicity the model acquires across tasks. These findings guide data and training setup to better induce implicit reasoning, and suggest potential improvements to the transformer architecture to further unlock its generalization. We conclude by showcasing the unique power of parametric memory on a challenging reasoning task with a large search space.

## Limitations

Scope of our task formulation . As mentioned in § 1 , we formulate the implicit reasoning problem as induction and application of inference rules from a mixture of atomic and inferred facts. This may not apply to the full spectrum of reasoning which has a range of different types and meanings [ 20 ] . Still, our formulation could capture the nature of a wide range of reasoning problems, and crucially, we believe that it is a good conceptualization of certain aspects of language model (pre-)training, where the model needs to both memorize the “atomic” world knowledge and also induce generalizable rules from the massive amount of records of human activities, which allow the model to connect and reason over knowledge, and ultimately help with humans.

Abstract nature and connection with practice . Our study has an abstract nature, which is required for the complete solidity and rigor of our experiments and evaluation. We also strive to make sure that the results are robust to different setups closer to practice through additional experiments (Appendix B , C , E ). Still, there are certain distances from our settings to those in practice. However, we believe that it is far more important to build solid understandings, even having distances with practice, than to draw conclusions or make claims that are closer to practice but questionable due to insufficient control over data and evaluations.

## Acknowledgement

The authors would like to thank colleagues from the OSU NLP group and CMU NeuLab for their thoughtful comments. This research was supported in part by NSF CAREER #1942980, ARL W911NF2220144, NSF OAC 2112606, and Ohio Supercomputer Center [ 5 ] . The views and conclusions contained herein are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the U.S. government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notice herein.

## Appendix A Training Details

All implementations were done with PyTorch [ 45 ] and Huggingface Transformers [ 68 ] . All model training runs are done on NVIDIA A6000 and A100 GPUs and last 96 96 hours at maximum.

## Appendix B Effect of Model Scale

We run the experiments on composition with larger model scales with | ℰ | = 2000 |\mathcal{E}|=2000 and ϕ ∈ { 5.4 , 9.0 , 18.0 } \phi\in\{5.4,9.0,18.0\} . The results are shown in Figure 7 . Overall, it could be seen that scaling up the model won’t qualitatively change the model’s generalization behaviors, and the main pattern is that larger models converge in fewer optimization steps, which shares with prior findings [ 60 , 28 ] .

## Appendix C Effect of Tokenizations

In the experiments in our main content, tokenization is done by having a unique token for each entity. 11 11 11 Preliminary experiments show that tokenization of the relations does not exhibit notable impacts, which is expected since relations are always explicitly given. This is different from how real-world entities are typically tokenized—in practice, entities are usually multi-token, and different entities could share tokens at different positions. We investigate the effect of tokenizations on the composition task by having two tokens for each entity (resembling the first name and last name of a person) in the setting with | ℰ | = 2000 |\mathcal{E}|=2000 and ϕ = 12.6 \phi=12.6 . We generate a set of unique first names and a set of unique last names with equal size from which the two tokens for each entity are randomly chosen (we make sure each entity gets a unique ordered pair of tokens). We define token multiplicity to be the number of entities that share the same first/last name. For example, when the size of the set of first/last names is 50 50 , the token multiplicity would be 2000 / 50 = 40 2000/50=40 .

Figure 8 (a) shows the ID test results, where the training and OOD results are the same from earlier (training performance saturates quickly, OOD result remains zero). It can be seen that a larger token multiplicity would delay the generalization, which is expected to a certain degree since the scale of the model is effectively smaller due to having fewer tokens in the vocabulary. Nevertheless, ID generalization always happens. We also run linear probing on S ⁡ [ 5 , r 1 ] S[5,r_{1}] throughout training to predict the second token of the bridge entity b b in the setting with token multiplicity 40 40 , 12 12 12 Recall that S ⁡ [ 5 , r 1 ] S[5,r_{1}] is the state that encodes the bridge entity in the generalizing circuit (Figure 4 (a)) , where the results are shown in Figure 8 (b). It can be seen that the second token of b b can be perfectly decoded from S ⁡ [ 5 , r 1 ] S[5,r_{1}] after grokking, and the decodability improves throughout grokking. This suggests that for the multi-token case, the model is additionally storing the second token of b b into S ⁡ [ 5 , r 1 ] S[5,r_{1}] throughout grokking, which may be another factor that further delays the speed of grokking. These results also share with recent findings that in many cases, tokens beyond the immediate next token are linearly encoded in the hidden states [ 69 , 44 , 2 , 4 ] .

In summary, different tokenizations affect the results in rather expected ways, and do not influence our main findings and conclusions.

## Appendix D More Details on Circuit Analysis

### D.1 Composition

We run causal tracing on hidden states in layer 1 1 - 7 7 and every position, where the target is the final prediction state S ⁡ [ 8 , r 2 ] S[8,r_{2}] . The changes in the strengths are monotone and smooth, and we show in Figure 9 the strengths for the model checkpoint at the start and end of grokking, and also their difference (same as Figure 4 (b)). We also find that after grokking, the state S ⁡ [ 5 , r 2 ] S[5,r_{2}] (which encodes r 2 r_{2} ) is not affected by perturbing the input h h or r 1 r_{1} .

Deriving the generalizing circuit . Starting from the 9 9 states in layers 0 , 5 , 8 0,5,8 we can directly eliminate S ⁡ [ 8 , h ] S[8,h] and S ⁡ [ 8 , r 1 ] S[8,r_{1}] since they have no computational paths connecting to S ⁡ [ 8 , r 2 ] S[8,r_{2}] . S ⁡ [ 5 , h ] S[5,h] can be eliminated as could be seen by Figure 9 (c). The connections from S ⁡ [ 0 , h ] S[0,h] and S ⁡ [ 0 , r 1 ] S[0,r_{1}] to S ⁡ [ 5 , r 2 ] S[5,r_{2}] could be eliminated as mentioned earlier.

### D.2 Comparison

Figure 10 includes the causal tracing results where the target is the prediction state S ⁡ [ 8 , e 2 ] S[8,e_{2}] , and Figure 11 includes the results with the target state S ⁡ [ 5 , e 2 ] S[5,e_{2}] . It can be seen that after grokking, S ⁡ [ 5 , e 2 ] S[5,e_{2}] does not depend on e 1 e_{1} , which gives the generalization circuit in Figure 5 (a).

Figure 12 shows the rank (via logit lens) of the three relations { a < , a = , a > } \{a_{<},a_{=},a_{>}\} in the label space at state S ⁡ [ 7 , a ] S[7,a] , where we use Recall@3 as the measure. It can be seen that S ⁡ [ 7 , a ] S[7,a] encodes the label space throughout grokking.

## Appendix E Additional Results

### E.1 Weight decay

Figure 13 shows the ID generalization performance when varying the degree of weight decay ( | ℰ | = 2000 |\mathcal{E}|=2000 and ϕ = 9.0 \phi=9.0 ). It can be seen that a larger weight decay can improve the speed of grokking, and vice versa.

### E.2 Transformer with parameter sharing