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

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

Causal tracing . The transformer could be viewed as a causal graph [ 46 ] that propagates information from the input to the output through a grid of intermediate states, which allows for a variety of causal analyses on its internal computations [ 63 , 35 , 19 , 65 , 12 ] . For convenience, we will refer to a hidden state by S ⁡ [ i , a ] S[i,a] , where i i is the layer index and a a marks the role of the input token at the same position as the state (one of { h , r 1 , r 2 } \{h,r_{1},r_{2}\} ). We illustrate our method in Figure 3 , where the hidden state of interest is S ⁡ [ 4 , r 1 ] S[4,r_{1}] and the target is the model’s prediction state S ⁡ [ 8 , r 2 ] S[8,r_{2}] . There are in total three steps: 1. The normal run records the model’s hidden state activations on a regular input ( h , r 1 , r 2 ) (h,r_{1},r_{2}) . Note that since the model maintains perfect training performance throughout grokking, the final prediction is always the ground truth tail entity t t . 6 6 6 For convenience, when we refer to a state as a token, we mean the top token of the state via logit lens.

2. In the perturbed run , a slightly perturbed input is fed to the model which changes the prediction, where again the hidden state activations are recorded. For the perturbation, prior work has explored adding noise to the input [ 35 ] and replacing key tokens with semantically close ones [ 63 , 12 ] . We adopt token replacement which avoids unnecessary distribution shifts [ 74 ] . Specifically, for the hidden state of interest, we replace the input token at the same position as the state to be a random alternative of the same type (e.g., r 1 → r 1 ′ r_{1}\rightarrow r_{1}^{\prime} ) that leads to a different target prediction (e.g., t → t ′ t\rightarrow t^{\prime} ).

3. Intervention . During the normal run, we intervene the state of interest by replacing its activation with its activation in the perturbed run. We then run the remaining computations and measure if the target state (top- 1 1 token through logit lens) is altered. The ratio of such alterations (between 0 0 and 1 1 ) among the examples quantitatively characterizes the causal strength between the state of interest and the target.

The generalizing circuit . We run a set of causal tracing and logit lens experiments across different model checkpoints throughout training. The discovered generalizing circuit (i.e., the causal computational pathways after grokking) is illustrated in Figure 4 (a). Specifically, we locate a highly interpretable causal graph consisting of states in layer 0 0 , 5 5 , and 8 8 , where we have pruned away the weak nodes/connections (details in Appendix D ). Layer 5 5 splits the circuit into lower and upper layers, where 1) the lower layers retrieve the first-hop fact ( h , r 1 , b ) (h,r_{1},b) from the input h , r 1 h,r_{1} , store the bridge entity b b in S ⁡ [ 5 , r 1 ] S[5,r_{1}] , and “delay” the processing of r 2 r_{2} to S ⁡ [ 5 , r 2 ] S[5,r_{2}] ; 2) the upper layers retrieve the second-hop fact ( b , r 2 , t ) (b,r_{2},t) from S ⁡ [ 5 , r 1 ] S[5,r_{1}] and S ⁡ [ 5 , r 2 ] S[5,r_{2}] , and store the tail t t to the output state S ⁡ [ 8 , r 2 ] S[8,r_{2}] .

What happens during grokking? To understand the underlying mechanism behind grokking, we track the strengths of causal connections and results from logit lens across different model checkpoints during grokking (the “start” of grokking is the point when training performance saturates). We observe two notable amplifications (within the identified graph) that happen during grokking. The first is the causal connection between S ⁡ [ 5 , r 1 ] S[5,r_{1}] and the final prediction t t , which is very weak before grokking (Appendix D ) and grows significantly during grokking (Figure 4 (b)). The second is the r 2 r_{2} component of S ⁡ [ 5 , r 2 ] S[5,r_{2}] via logit lens, for which we plot its mean reciprocal rank (MRR) (Figure 4 (c)). Additionally, we find that the state S ⁡ [ 5 , r 1 ] S[5,r_{1}] has a large component of the bridge entity b b throughout grokking (Figure 4 (c)). These observations strongly suggest that the model is gradually forming the second hop in the upper layers ( 5 5 - 8 8 ) during grokking . This also indicates that, before grokking, the model is very likely mostly memorizing the examples in train_inferred ID {}_{\text{ID}} by directly associating ( h , r 1 , r 2 ) (h,r_{1},r_{2}) with t t , without going through the first hop.

Why does grokking happen? These observations suggest a natural explanation of why grokking happens through the lens of circuit efficiency [ 61 ] . Specifically, as illustrated above, there exist both a memorizing circuit C m ​ e ​ m C_{mem} and a generalizing circuit C g ​ e ​ n C_{gen} that can fit the training data. While C m ​ e ​ m C_{mem} is learned first (which causes training performance to saturate quickly), C g ​ e ​ n C_{gen} is relatively more efficient , in the sense that it could fit the data with a lower complexity. To see this, we can compare the amount of facts C m ​ e ​ m C_{mem} and C g ​ e ​ n C_{gen} need to store (denoted as N m ​ e ​ m N_{mem} and N g ​ e ​ n N_{gen} ) as a proxy for their complexity. 7 7 7 While the circuits also consist of other components, they pale in comparison as the number of facts scales. C m ​ e ​ m C_{mem} stores both atomic facts and inferred facts in the weights. C g ​ e ​ n C_{gen} (Figure 4 (a)) stores the atomic facts in the lower layers, and another copy of the atomic facts that appear as the second hop in the inferred facts in the upper layers. As the inferred/atomic ratio ϕ \phi increases, N m ​ e ​ m N_{mem} would increase rapidly while N g ​ e ​ n N_{gen} increases slowly and is always bounded by two times the total amount of atomic facts, and hence, the relative efficiency of C g ​ e ​ n C_{gen} increases. In the long run, the model will be incentivized to transition from C m ​ e ​ m C_{mem} to C g ​ e ​ n C_{gen} due to implicit bias of the optimization [ 53 ] and explicit regularization such as weight decay which prefers more efficient circuits, and the transition would happen faster as ϕ \phi increases. This also explains why the training data size does not affect the speed of grokking, since solely increasing the size does not change the relative efficiency of C m ​ e ​ m C_{mem} and C g ​ e ​ n C_{gen} . The explanation also implies that a larger regularization factor should accelerate grokking (and vice versa), which we confirm by varying the degree of weight decay (Appendix E.1 ).

Explaining and mitigating the deficiency in OOD generalization. The configuration of C g ​ e ​ n C_{gen} also has another important implication: while the model does acquire compositionality through grokking, it does not have any incentive to store atomic facts in the upper layers that do not appear as the second hop during training . This explains why the model fails in the OOD setting where facts are only observed in the atomic form, not in the compositional form---the OOD atomic facts are simply not stored in the upper layers when queried during the second hop. 8 8 8 We verified that in the OOD setting, S ⁡ [ 5 , r 1 ] S[5,r_{1}] and S ⁡ [ 5 , r 2 ] S[5,r_{2}] encode b b and r 2 r_{2} respectively as in the ID case. Such issue originates from the non-recurrent design of the transformer architecture which forbids memory sharing across different layers. Our study provides a mechanistic understanding of existing findings that transformers seem to reduce compositional reasoning to linearized pattern matching [ 10 ] , and also provides a potential explanation for the observations in recent findings that LLMs only show substantial positive evidence in performing the first hop reasoning but not the second [ 71 ] . 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 improve their generalization. We also show that a variant of the parameter-sharing scheme in Univeral Transformer [ 7 ] can improve OOD generalization in composition (Appendix E.2 ).

## 4 Comparison—Systematic Generalization via Parallel Circuit

We have just shown that the vanilla transformer fails to achieve OOD generalization for composition, but is the vanilla transformer generally incapable of acquiring systematic implicit reasoning skills? We show that for comparison , a task where SoTA LLMs such as GPT-4 also struggle [ 1 ] , the vanilla transformer does have the capability to acquire systematic generalization, again through grokking. On the surface, it seems that the comparison task is no different than the composition task—both require retrieving and reasoning over two pieces of facts. However, as it turns out through our analysis, the comparison task emits a “parallel circuit” that is learned by the transformer during grokking, which allows atomic facts to be stored and retrieved in the same region and enables systematicity to happen.

Setup. The comparison task involves comparing the attribute values of entities. We assume there are | ℰ | = 1000 |\mathcal{E}|=1000 entities, | 𝒜 | = 20 |\mathcal{A}|=20 attributes and | 𝒱 | = 20 |\mathcal{V}|=20 ordinal values for the attributes. Each attribute a ∈ 𝒜 a\in\mathcal{A} has a label space { a < , a = , a > } \{a_{<},a_{=},a_{>}\} , a set of relations specifying its comparative form. For example, an attribute age would have a < , a = , a > a_{<},a_{=},a_{>} to be younger, contemporary, older , respectively.

The atomic facts are (entity, attribute, value) triplets, where we assign a random value v ∈ 𝒱 v\in\mathcal{V} for each ( e , a ) ∈ ℰ × 𝒜 (e,a)\in\mathcal{E}\times\mathcal{A} . Again, we randomly partition the atomic facts into atomic ID {}_{\text{ID}} and atomic OOD {}_{\text{OOD}} ( 90 90 %: 10 10 %). The rules of comparison are: ∀ e 1 , e 2 ∈ ℰ , ∀ a ∈ 𝒜 , ∀ v 1 , v 2 ∈ 𝒱 , ( e 1 , a , v 1 ) ∧ ( e 2 , a , v 2 ) ∧ v 1 < v 2 ⟹ ( a , e 1 , e 2 , a < ) , \displaystyle\forall e_{1},e_{2}\in\mathcal{E},\forall a\in\mathcal{A},\forall v_{1},v_{2}\in\mathcal{V},(e_{1},a,v_{1})\land(e_{2},a,v_{2})\land v_{1}<v_{2}\implies(a,e_{1},e_{2},a_{<}), (2) ∀ e 1 , e 2 ∈ ℰ , ∀ a ∈ 𝒜 , ∀ v 1 , v 2 ∈ 𝒱 , ( e 1 , a , v 1 ) ∧ ( e 2 , a , v 2 ) ∧ v 1 = v 2 ⟹ ( a , e 1 , e 2 , a = ) , \displaystyle\forall e_{1},e_{2}\in\mathcal{E},\forall a\in\mathcal{A},\forall v_{1},v_{2}\in\mathcal{V},(e_{1},a,v_{1})\land(e_{2},a,v_{2})\land v_{1}=v_{2}\implies(a,e_{1},e_{2},a_{=}), ∀ e 1 , e 2 ∈ ℰ , ∀ a ∈ 𝒜 , ∀ v 1 , v 2 ∈ 𝒱 , ( e 1 , a , v 1 ) ∧ ( e 2 , a , v 2 ) ∧ v 1 > v 2 ⟹ ( a , e 1 , e 2 , a > ) . \displaystyle\forall e_{1},e_{2}\in\mathcal{E},\forall a\in\mathcal{A},\forall v_{1},v_{2}\in\mathcal{V},(e_{1},a,v_{1})\land(e_{2},a,v_{2})\land v_{1}>v_{2}\implies(a,e_{1},e_{2},a_{>}). Take the attribute age as an example, the first rule means if the age of e 1 e_{1} is smaller than the age of e 2 e_{2} , then we can infer “In terms of age , the relation between e 1 e_{1} and e 2 e_{2} is younger ”. Each entity/attribute/value/label is assigned a unique token, and training/testing is done by having the model predict the last token (attribute value for atomic facts; comparative relation for inferred facts).

Results & analysis . Figure 1 (right) shows the results for ϕ = 7.2 \phi=7.2 , and we include more results in Appendix E.3 . It can be seen that 1) the model again acquires robust generalization only through grokking; 2) surprisingly, the model also achieves systematicity in generalization, different from the case of composition.

Analyzing the model’s internals similarly as in § 3.3 (details in Appendix D ), we find the generalizing circuit for comparison illustrated in Figure 5 (a). On a separate stream, the model prepares the label space { a < , a = , a > } \{a_{<},a_{=},a_{>}\} from a a and stores it in S ⁡ [ 7 , a ] S[7,a] . In the lower layers ( 0 0 - 5 5 ), the model retrieves the two atomic facts and stores the attribute values v 1 v_{1} and v 2 v_{2} at S ⁡ [ 5 , e 1 ] S[5,e_{1}] and S ⁡ [ 5 , e 2 ] S[5,e_{2}] . Then, the upper layers ( 5 5 - 8 8 ) compare v 1 , v 2 v_{1},v_{2} and fetch the label from S ⁡ [ 7 , a ] S[7,a] based on the comparison result. Importantly, there is a major difference compared with the circuit for composition: the two atomic facts are retrieved in parallel , which suggests that the atomic facts are stored solely in the lower layers, without having separate copies across different regions as in the circuit for composition. This explains why systematicity could happen: OOD facts are now stored and accessed in the same way as ID facts. Tracking the changes in the model throughout grokking, we observe significantly strengthened causal connections from S ⁡ [ 7 , a ] S[7,a] and S ⁡ [ 5 , e 1 ] S[5,e_{1}] to the final prediction (Figure 5 (b)). We also find that throughout grokking, S ⁡ [ 7 , a ] S[7,a] always encodes the label space and S ⁡ [ 5 , e 1 ] , S ⁡ [ 5 , e 2 ] S[5,e_{1}],S[5,e_{2}] gradually encode the two attribute values (Figure 5 (c)). This confirms that a similar transition from C m ​ e ​ m C_{mem} to C g ​ e ​ n C_{gen} happens during grokking.

The findings here showcase transformer’s ability to learn parallel solutions to seemingly sequential problems, akin to the findings in Liu et al. [ 30 ] where it is shown that transformers can learn “shortcuts” to automata. The difference in the acquired generalization across the two tasks that we study also emphasizes the need for controlled and mechanistic study on understanding the transformer’s reasoning before making general claims on its limitations.

## 5 The Power of Parametric Memory for Complex Reasoning

At the high level, our study so far paves the way towards better understanding and improving transformer’s reasoning with parametric representation of knowledge and rules. But why is parametric memory practically important? Can we not simply enhance LLMs with non-parametric memory, e.g., by using their long-context modes and/or doing explicit retrieval, to solve the tasks at hand?

We believe parametric memory has its unique capability to perform deep compression and integration of information for complex reasoning. To showcase the potential of parametric memory for complex reasoning, we create a difficult reasoning task with a large search space , and show that 1) it is far out of reach even for current SoTA models (e.g., GPT-4-Turbo [ 43 ] and Gemini-Pro-1.5 [ 16 ] ) based on non-parametric memory; 2) a fully grokked transformer can solve the task with near-perfect accuracy.

Our task is a variation of the comparison task above where we use a simple way to massively expand the search space, based on an additional set of rules that are already contained within the task itself, namely, the (anti-)symmetry and transitivity of comparison: ∀ e 1 , e 2 ∈ ℰ , ∀ a ∈ 𝒜 , ( a , e 1 , e 2 , a < / a = / a > ) ⟹ \displaystyle\forall e_{1},e_{2}\in\mathcal{E},\forall a\in\mathcal{A},(a,e_{1},e_{2},a_{<}/a_{=}/a_{>})\implies ( a , e 2 , e 1 , a > / a = / a < ) , \displaystyle(a,e_{2},e_{1},a_{>}/a_{=}/a_{<}), (3) ∀ e 1 , e 2 , e 3 ∈ ℰ , ∀ a ∈ 𝒜 , ∀ y ∈ { a < , a = , a > } , ( a , e 1 , e 2 , y ) \displaystyle\forall e_{1},e_{2},e_{3}\in\mathcal{E},\forall a\in\mathcal{A},\forall y\in\{a_{<},a_{=},a_{>}\},(a,e_{1},e_{2},y) ∧ ( a , e 2 , e 3 , y ) ⟹ ( a , e 1 , e 3 , y ) . \displaystyle\land(a,e_{2},e_{3},y)\implies(a,e_{1},e_{3},y).