Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization
Authors: Boshi Wang
Xiang Yue
Yu Su
Huan Sun
Venue: NeurIPS 2024 poster (2024)
Keywords: Reasoning
Grokking
Systematic Generalization
Mechanistic Interpretability
Transformer
Primary area: interpretability_and_explainability
TL;DR: We discover and analyze Transformer's grokking phenomenon on the task of implicit reasoning.

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.

Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization
Boshi Wang* Xiang Yue +
Yu Su
Huan Sun
*The Ohio State University [wang. . 13930, yue. 149, su . 809, sun sun.397] mosso.sueeee .ddu
Carnegie Mellon University
Abstract
We study whether transformers can learn to implicitly reason over parametric knowl- edge, skill that even the most capable language models struggle with. Focusing 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 outttf-distributiin examples, transformers fail to systematically generalize for composition but succeed for comparison. We delve into the model's 's internals throughout training, conducting analytical exper- iments that reveal: 1) the mechanism behind grokking, such as the formation of the generalizing circuit and its relation 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 challenging reasoning task with large search space, GPT-4-Turbo and Gemini- 1.5-Pr based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation, while fully grokked transformer can achieve near- cear-perfect accuracy, showcasing the power parametric memory for complex reasoning
Introduction
Large language models (LLMs) have been shown deficient implicit reasoning with their parametric memory of knowledge and rules. For example, a range of LLMSs are found to be incapable of robustly composing internalized facts [48 71, and even GPT-4 cannot adequately compare entities'
attributes despite knowing them
Deficiency in implicit reasoning has profound implications. It implies the models' limitations in induc- ing 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
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?
Project started when at OSU.
Code and data: http:////.....aaooo/////////////
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Composition
Comparison
Sarnck
wife
Michhlle
A
OEchello
Jurn In
1964
Biden
wilre
(1964]
Frung
Biden
1.0
0.9
Accuracy
Accuracy
0.6
0.8
0.4
0.7
Train (ID)
Train (ID)
0.2
Test (ID)
0.6
Test (ID)
(OOD)
(OOD)
0.0
0.5
10
10a
10
10a
Optimization Step (Log Scale)
Optimization Step (Log Scale)
Figure 1: We find that transformers can learn to reason implicitly, but this skill is only robustly acquired through grokking, i.c., an extended period of training far beyond overfitting. Moreover, the transformer fails to systematically generalize for composition, yet succeeds for comparison. We conduct mechanistic study into the model internals throughout grokking, which reveals distinct generalizing circuits across the tasks (Figure 4 5
that explains the variations in systematicity.
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 mixture of "atomic facts" and "inferred facts" (which are deduced from the atomic facts via 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-distribttion (OOD) scenarios 3 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 [.) a phenomenon known grokking 47. We find that the speed of improvement in generalizatron 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 B3
61 78 (1]. in that it should instead be the critical data distribution that decides the characteristies of grokking. Our findings extend prior observations of the grokking phenomenon primarily in algorithmic and linguistic tasks [7 38 to the domain of knowledge-baid -baed 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 consistently observed, in the OOD setting, the model fails to systematically generalize for composition but succeeds in comparison (Figure IT. 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-augmentarte [54] 17
and explicit
recurrence 22 57
are needed to further unlock transformer's generalization.
Finally, demonstrate the power and potential of parametric memory for complex reasoning, show that for reasoning task with large search space. fully grokked transformer can achieve near-perfect accuracy, while state-of-the-art LLMs like GPT-4-Turbo 4K] and Gemini-1.5-Pie [16]
based on parammemtt memory fail badly regardless of prompting styles or retrieval augmentation.
Definitioos of ID/OOD are introduced in 2

---

2 General Setup
Training data & ID/OOD evaluation. As stated in T 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 I) sampling 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 uut-of-ddistriuuionn (OOD) performance. We prepare two separate sets of atomic facts: at atoiic Cop and atomi Coop. Our training set includes all the atomic facts and uniformly random portion of the inferred facts deduced from at atomi CiD. which we call train in inferredp. 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 a t omi CTD, which we denote by test_ festtiiffeeee in (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 at omi CooD. denoted by test_inferredoon
Model & optimization. We use a standard decoder-only transformer model as in GPT-2 50 with 8 layers, 768 hidden dimensions and 12 attention heads (we explore the impact of different model scales in Appendix B). Optimization is done by Adam W (34) with learning rate 10 batch size 512, weight decay 0.1 and 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-Delayeee Generalization without Systematicity
We begin our investigation with composition, where model needs to "chain" different pieces of facts, e.g. "Barack's wife is Michelle" and "Michelle is born in 1964", to successfully complete compositional sentence, e.g.. "Barack's 's wife is born in [1964]". Prior work extensively studied whether transformer-basee language models can perform implicit composition, and negative results are consistently reported [48 71]. Specifically, there exists a "compositionallty gap" [48), i.c.. 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 & consisting of |g| entities and [ = 200 relations, where each entity (as the subject) has 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 G. which we partition disjointly into atomic. and at omi Cow (95%: 5%). The rule of (two-hop) composition is
Yh, b, E C,Vri,r r2 R., (h, b) A (b, r). = (h,r1,r,,t).
(1)
which is used to deduce the ID and OOD inferred facts from at omicrp and at oitomicoo CooD. respectively. For convenience, in the above rule, we will call h the head entity, b the bridge entity, and 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 unique token to each relation/entity by default, and also find that the results are robust to different tokenizations (details in Appendix
We study the influence of the following two aspects on the model's '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 ø = ra iniinnerreeeeeee
/a//mmiomoiio cra| on the model. Training data size. We study the impact of the training data size by varying (&), the total number of entities, while controlling the ratio é. Note that the size of training data (both atomic/inferred facts) scales linearly with |
3.2 Results
Grokking observed in ID generalization but not in OOD generalization. Figure Ileft shows the model's accuracy the train and test facts throughout optimization, with |E| = 2000 and ø 7.2.
3

---

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 477. Specifically, the training performance saturates (over 99% accuracy on both atomic and inferred facts) at around 14K optimization steps, before which the highest ID generalization accuracy is merely 9.2%. However, generalization keeps improving by simply training for longer, and approaches almost perfect accuracy after extended optimization lasting around 50 times the steps taken to fit the training data. On the other hand, OOD generalization is never observed. We extend the training to million optimization steps, and there is still no sign of OOD generalization.
Inferred/atomic ratio & correlates with generalization speed. Figure 2[a) shows the ID test accuracy across different 6. We omit the other splits since for all settings, the training performance saturates quickly and the OOD test accuracy remains at zero as earlier It could be seen that the ratio strongly correlates with the speed of generalization. A very large ratio can push generalization to improve at similar pace as the model fits the training data, reducing the need for extended training
Ratio (Inferred/Atomic)
le]: 10K
Train (ID)
3.8
7.2
12.6
[e]: 5K
Test (ID)
5.4
9.0
18.0
[e]: 2K
Test (COD)
1.0
1.0
0.8
0.8
Accuracy
0.6
Accuracy
0.6
0.4
0.4
0.2
0.2
0.0
0.0
10"
10"
50
100
150
Optimization Step (Log Scale)
# Epoch
(a) Effect of changing ratio & ((E) = 2000).
(b) Effect of changing |8] (Ø = 9.0).
Figure 2: The speed of grokking on the in-distribution (ID) test performance (a) correlates with the ratio between inferred and atomic facts, and (b) is not influenced by the size of training data.
Training data distribution, instead of training data size, qualitatively influences generalization behavior. When & increases and |&] 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) B3 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 (bb) shows the results of varying |g| with fixed é 9.0, where we change the horizontal axis from optimization step to epoch for better visualization When fixing the ratio ø, the training data size does not qualitatively affect the model's 's generalization. Specifically, scaling the data affects neither the relative speed of ID generalization and training improvement (as seen by the rather constant "gap" curves), nor the systematicity level (OOD performance stays zero). We also run the experiments across different & 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
between ra n_ infereed and rest_inferree
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 Φ) in characterizing the model's generalization. However, important questions
"The training performances of all settings saturate within 25K steps, where larger takes more steps.
'When & = 18.0, the model achieves 96.7% accuracy before training performance saturates. 'The optimization steps for each epoch scale linearly with the training size since we fixed batch size.

---

still remain: what happens during grokking, why does it happen, and why do transformers struggle with OOD examples? Answering these questions requires 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 combination of two prevalent approaches: logit lens and causal tracing. We apply our analysis to the setting with |8] = 2000, ø 9.0 on 300 random examples from train inferred...
(2) Perturbed run
Normal run
In It Cg 65
L
Layer
Input & Output Embosiddngs
(3) Intervention
Pinal LayerNorm
Layer
5
Encoding
Activations
Normal
O
States
Layer
by Inservention
Tape Tama.
t'(# t)
Measure change
Figure 3: Illustaaioon of our circuit analysis approach (on the composition task). We use logit lens interpret individual states, and use causal tracing measure the strength of connections between states. Details are in the main content.
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 top right).
Logit lens. We interpret individual hidden states via logit lens 40 15
Causal tracing. The transformer could be viewed as a causal graph [66] that propagates information from the input to the output through grid of intermediate states, which allows for a variety causal will refer hidden state by Sli, a). where is the layer index and a marks the role the input token at the same position the state (one of {h, T1,
analyses its internal computations [63] 35 19 65 [1]. For convenience,
method in Figure 3 where the hidden state of interest S r1] and the target is the model's prediction state S(8, r2). There are in total three steps:
(2,72) )2)). We illustrate
1. The normal run records the model's hidden state activations on regular input (h, T1, r2). Note that since the model maintains perfect training performance throughout grokking, the final prediction is always the ground truth tail entity {
model which changes the prediction, where again the hidden state activations 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
the perturbed run, slightly perturbed input fed
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., r1 > ri) that leads to different target prediction (e.g., t'). 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 the target state (top-1 token through logit lens) is altered. The ratio of such alterations (between 0 and 1) among the examples quantitatively characterizes the causal strength between the state of interest and the target.
For convenience, when refer state token, the top token of the state via logit lens.
5

---

by directly Grokking starts
1:00
1.0
Layer
Layer 3-
0.25
0.6
5
Layer 4
0.00
0.4
MARLD) st: AT
Layer
-0.25
Test (ID)
L
-0.50
Text (DOD)
0.0
Layer 7-
1.0
2,0
2.5
3.0
Optimization (1e5)
(a)
(b)
(c)
Figure 4: The (evolution of) generalizing circuit for composition. (a) The generalizing circuit. (b) The change in causal strengths during grokking, where the target is the prediction state. (c) Mean reciprocal rank (via logit lens) of the bridge entity at S15, r1] and second relation r2 at r2].
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 com- putational pathways after grokking) is illustrated in Figure 4 a). Specifically, we locate highly interpretable causal graph consisting of states in layer 0, 5, and 8, where we have pruned away the weak nodes/connections (details Appendix Dj. Layer splits the circuit into lower and upper layers, where 1) the lower layers retrieve the first-hop fact (h. r1, from the input h, r1, store the bridge entity b in S(5, and "delay" the processing of r2 to S[5, r2]: ; 2) the upper layers retrieve the second-hop fact (6, 72, t) from S15, r1] and S[5, r2]. and store the tail & to the output state S(8,
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 the causal connection between S|5, rı] and the final prediction t. which is very weak before grokking (Appendix D and grows significantly during grokking (Figure 4 ()). The second is the r2 component of S|5, r2] via logit lens, for which we plot its mean reciprocal rank (MRR) (Figure 4c)). Additionally, we find that the state S15, r1] has large component of the bridge entity b throughout grokking (Figure (c)). These observations strongly suggest that the model is gradually forming the second hop in the upper layers (5-8) during grokking. This also indicates that, before grokking, the model is very likely mostly memorizing the examples in t associating (h, T2) with I, 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 Cinem and a generalizing circuit Cgen that can fit the training data. While Cmem is learned first (which causes training performance to saturate quickly), Cgen is relatively more efficient, the sense that it could fit the data with a lower complexity. To see this, we can compare the amount of facts Cmmm and C Gen need to store (denoted as Nmem and N Ngen) as proxy for their complexity Cmem stores both atomic facts and inferred facts in the weights. Cgen (Figure (()) 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 ø increases, Neem would increase rapidly while Ngen increases slowly and is always bounded by two times the total amount of atomic facts, and hence, the relative efficiency of Cgen increases. In the long run, the model will be incentivized to transition from Cwem to Cgen due to implicit bias the optimization (53) and explicit regularization such weight decay which prefers more efficient circuits, and the transition would happen faster as & 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 Cinem and Cgen- 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.
Whiie the circuits also consist of other components, they pale in comparison the number of facts scales.

---

Explaining and mitigating the deficiency in OOD generalization. The configuration of Cgen 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 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 [TO], and also provides 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 77.] Our findings imply that proper cross-layer memory-sharing mechanisms for transformers such as memory-augmentatioon [54] 17 and explicit recurrence 7 [2 57) are needed improve their generalization. We also show that variant of the parameter-sharing scheme in Univeral Transformer can improve OOD generalization in composition (Appendix E.2)
4 Comparisn 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, task where SoTA LLMs such GPT-4 also struggle M. 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 "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 entities, |A| = 20 attributes and | = 20 ordinal values for the attributes. Each attribute E A has a label space fac, a=, a> }, a set of relations specifying its comparative form. For example, an attribute age would have az, a=, ay to be younger, contemporary, older, respectively.
The atomic facts are (entity, attribute. value) triplets, where we assign a random value E V for each (e, a) E & x A. Again, we randomly partition the atomic facts into a attmiciy Ctp and a atomi Coon (90%: 10%) The rules of comparison are:
Vei, e2 E EE, Va E A, Vu U2 E V, (ei, v1) A (e2, 0,)) Avj < 02 =
e2,
Vel, E E, Va E Yu U2 E V, (e1, a, 01) A (e2, a, 02) A v1 = U2 = (0,e), 02,
(2)
Vel, E E, Va E A. U2 E V. (e1, a, v1) A (e2, a, 02) A > 02 = e2,
Take the attribute age as an example, the first rule means if the age of e₁ is smaller than the age of e2, then we can infer "In terms of age, the relation between ei and e2 is younger". Each emittylatribuuee/llll//bbb 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 rightt shows the results for ø = 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 33.3 (details in Appendix [D). we find the generalizing circuit for comparison illustrated in Figure 5la). a). On a separate stream, the model prepares the label space {ac,, a,, 0>} } from a and stores in S a]. In the lower layers (0-5), the model retrieves the two atomic facts and stores the attribute values U and at S(5,e1) and $15, e2). Then, the upper layers (5-8) compare U1, U2 and fetch the label from S17, a] based on the comparison result. Importantly, there is 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 copics across different regions as in the circuit for composition. This explains why sy systematicity could happen: OOD facts are now stored and accessed in the same way as ID)
We verified that in the OOD setting, S(5, rı] and S15, r2] encode b and r2 respectively in the ID case.
7

---

facts. Tracking the changes in the model throughout grokking, we observe significantly strengthened
causal connections from S[7, a] and S|5, ei] to the final prediction (Figure b))). We also find that
throughout grokking, S a] always encodes the label space and S|5, ei S e2 gradually encode the two attribute values (Figure ())).. This confirms that a similar transition from Crem Cgen happens during grokking.
a
Grokking
1.00
0.75
Layer
Layer
0.50
Layer
Layer
0.00
MR) s[5
Layer?
-0.25
MAR(v)) SES, ex)
Layer
Layer
-0.56
Test (OOD)
Test
-0.75
Layer
-1.00
Optimization stop (1e5)
(b)
Figure 5: The (evolution of) generalizing circuit for comparison. (a) The generalizing circuit. (b) The change in causal strengths during grokking, where the target the prediction state. (c) Mean reciprocal rank (via logit lens) of the two attribute values (V1, (2) at S|5, e1] and S[5,
The findings here showease transformer's ability learn parallel solutions to seemingly sequential problems, akin to the findings in Liu et al. (30) where it is shown that transformers can learn "shorteuts" 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 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 is far out of reach even for current SoTTA models (e.g., GPT-4-Turbo 43 and Gemini-Pro 1.5 (16) based on non non-paramette pa-paarmmttrcc memory; 2) a fully grokked transformer can solve the task with near rear-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 additional of rules that already contained within the task itself, namely, the (anti- monti-)symmetry and transitivity of comparison:
Vel,e2 e2 E, Va A, (a, e1, e2, a2,a</a_/a>> - (a, e2, el, e1,
(3)
Yei, C2 63 E & E V. E (ac,a=, (a, e1, 22,9) A (a, e2, 63,8) = (a, (0,e1))3.1). v)-
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) add 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. Consequently, answering test query would require the model to successfully locate two ID bridge entities which can connect the two query entities into proof (Figure 6). We select a balanced (by az. a_. subset from these queries for evaluation. More details included in Appendix
The difficulty such task is two-fold. First, the search space large. For example, on average, each query entity connects with more than 50 facts, and each bridge entity in the ground truth proof

---

Bob
5
Leo
younger
older
Eve
youngey
Query Entity
Joe
Mia
older
John
younger
10
Rick
Bridge Eutity
ulder
younger
Sam
Gold Proof
Ben
Max
Sumee age AS
Tom
Figure 6: Illustration of the complex reasoning task, which involves comparing the attributes of two query entities based set of facts encompassing large search space.
connects with more than 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 10 We test both the vanilla setup where all facts (28.2K on average) are loaded into the LLM context, and the cerreeal-auugmented setup (5.4K facts retrieved on average) where the two-hop neighhorhoods 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 tfhuuugh (CoT) prompting where the model is prompted to verbalize the reasoning. We test GPT-4-Turbo and Remini-Pro--i where for GPT-4-Turbo we only test the retrieval-aus -auevvl---geeeeedd setup due to context length limit.
the
Table Results on the complex reasoning task. Direct/CoT: predict the answer directly/verbalize the reasoning steps. "+R": retrieval augmentation.
GPT-4-Turbo
Gemini-Pro-1.5
Grokked Transformer
Direct+R
CoT+R
Direct
CoT
Direct+
CoT+R
Accuracy (%)
33.3
31.3
28.7
11.3
37.3
12.0
99.3
Results. As shown in Table all models based non-parametric memory fail badly, where the only setting that surpasses random guess (33.3%) is the retrieval-augmented setting with Gemini-Pro-1 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% of Gemini's responses conclude that the answer cannot be decided (which treat as wrong since the answer can be decided) after a series of search steps] [1 More shockingly. most of the CoT rationales that achieve the correct final answer are actually wrong due either hallucinating underivable facts or logical errors. This illustrates the current models' inability to reason deeply with non- non-parametric memory. On the other hand, grokked transformer, which is trained extensively on the given facts compress and integrate the information to the extreme, could achieve near- peer-ee accuracy. By examining the model, we find that acquires the same generalizing circuit in Figure 5[a). and remarkably, even though not explicitly encouraged//rained 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-44 struggle in implicit reasoning over their parametric knowl-
48 771. suggesting their limitations in inducing structured and compressed representations of facts and rules during training. A series of efforts try to understand transformer's
edge [56 23
52
51
knowledge and reasoning through controlled experiments 24 49
66 which is also our focus. We
10This could also thought of as performing a retrieval step into the memory based on the attribute.
TTe ratio drops bit to 58.7% when augmenting with retrieval.

---

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-Thoueee Thought" and verbalized reasoning. A series of studies prompt/fine-um language models to verbalize (i.e., generate) the intermediate knowledge and reasoning steps (67 64 73 55
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 I1 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 useful medium for teaching models to reason implicitly via distillation or curriculum
Grokking is first discovered by Power et al. 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
36 38 [3 78, 21). To our knowledge. 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.
6 41 39 61
we are the first work observe grokking in the domain of knowledge-basee
Analyzing the inner workings of neural models. Recent work tries 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| [2] 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 parametric memory in language models, and an orthogonal direction to enhance models with non non-parametric memory, such as extending the effective context length [70 [4 [6] and augmenting with retrieval [8] 26 3 [5] 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 domain by acquiring and structuring knowledge in the brain (parametric), but he/she also wouldn't 't memorize all pieces of details and could refer to the source when necessary non-ppaameetric.. How to decide the tradeoff between parametric and paramaareeticc 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, 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 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 challenging reasoning task with a large search space.
Limitations
Scope of our task formulation. As mentioned in $ we formulate the implicit reasoning problem as induction and application of inference rules from 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 wide range of reasoning problems, and crucially, we believe that it a good conceptualization of certain aspects of language model (pre-) ttrininaii 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.
10

---

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 [ C 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 thank colleagues from the OSU NLP group and CMU NeuLab for their thoughtful comments. This research was supported in part by NSF CAREER #1942980, NSF OAC 21 12606, 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 implied, of the U.S. government. The U.S. Government authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notice herein.
ARL W91
References
[1] Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipula- tion. In arXiv preprint: 0e.22009.1402. 2023.
[2] Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens. In arXiv preprint: abs/2303.08112 2023.
[3] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean- -aptittt Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on muchine learning, pages 2206-2240. PMLR, 2022.
[4] Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads, In arXiv preprint: abx/2401. 107744 2024.
[5] Ohio Supercomputer Center. Ohio supercomputer center, 1987.
[6] Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. In NeurlPS ML Safety Workshop, 2022.
[7] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. In International Conference on Learning Representations, 2019.
[8] Yuntian Deng, Yejin Choi, and Stuart Shieber. From explicit cot to implicit cot: Learning to internalize cot step by step. In arXiv preprint: abs/2405.14838 2024.
[9] Yuntian Deng, Kiran Prasad, Roland Fernandez, Paul Smolensky, Vishrav Chaudhary, and Stuart Shieber. Implicit chain of thought reasoning via knowledge distillation. In arXiv preprint: ab01233 2023.
[10] Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck. Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers compositionality. In Thirly-sevnnth Conference Neural Information Processing Systems, 2023.
[11] Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. In shirty- Conference Neural Information Processing Systems, 2023.

---

[12] Jiahai Feng and Jacob Steinhardt. How do language models bind entities in context? In The Twelfth International Conference on Learning Representations, 2024.
[13] Javier Ferrando, Gabriele Sarti, Arianna Bisazza, and Marta R. Costa-jussä, A primer on the inner workings of transformer-bes based language models. In arXiv preprint: abs/2405.00208,, 2024.
[14] Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context, 2024.
[15] Mor Geva, Avi Caciularu, Kevin Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pages 30-45, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[16] Google, Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. 2024.
[17] Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwióska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626):471447777 2016.
[18) Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pages 3929-3938, PMLR, 2020.
[19] Michael Hanna, Ollie Liu, and Alexandre Variengien. How does GPT-2 compute greater- than?: Interpreting mathematical abilities in pre-trained language model. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[20] Jie Huang and Kevin Chen-Chuan Chang. Towards reasoning in large language models: A survey. Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Findings of the Association for Computational Linguistics: ACL 2023, pages 1049-100 1065, Toronto, Canada, July 2023. Association for Computational Linguistics.
[21] Yufei Huang. Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. Unified view of grokking, double descent and emergent abilities: A perspective from circuits competition. In arXiv preprint: abs/2402.15175 2024.
[22] DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- recurrent transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[23] Nora Kassner, Benno Krojer, and Hinrich Schütze. Are pretrained language models symbolic reasoners over knowledge? In Proceedings of the 24th Conference on Computational Natural Language Learning, pages 552-564, Online, November 2020. Association for Computational Linguistics.
[24] Nora Kassner, Benno Krojer, and Hinrich Schütze. Are pretrained language models symbolic reasoners over knowledge? In Proceedings of the 24th Conference on Computational Natural Language Learning, pages 552-564, Online, November 2020. Association for Computational Linguistics.
[25] Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of soquence-to-sequeer recurrent networks. In International conference on machine learning, pages 2873-2883 PMLR, 2018.
[26] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmente generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 14:4559-9474 2020.
12

---

[27] Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers solve inherently serial problems. In The Twelfth International Conference on Learning Representations, 2024.
[28] Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gon- zalez. Train big, then compress: Rethinking model size for efficient training and inference of transformers. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5958-5968. PMLR, 13-18 Jul 2020.
[29] Bill Yuchen Lin, Abhilasha Ravichander, Ximing Lu, Nouha Dziri, Melanie Sclar, Khyathi Chandu, Chandra Bhagavatula, and Yejin Choi. The unlocking spell on base LLMs: Rethinking aligmment via in-context learning. In The Twelfth International Conference on Learning Representations, 2024.
[30] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Trans- formers learn shorteuts to automata. In The Eleventh International Conference on Learning Representations, 2023.
[31] Jiacheng Liu, Ramakanth Pasunuru, Hannaneh Hajishirzi, Yejin Choi, and Asli Celikyilmaz. Crystal: Introspective reasoners reinforced with self-feedback. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pages 11557-11572, Singapore, December 2023. Association for Computational Linguistics.
[32] Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[33] Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference on Learning Representations, 2023.
[34] Uya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[35] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[36] William Merrill, Nikolaos Tsilivis, and Aman Shukla. tale of two circuits: Grokking as competition of sparse and dense subnetworks. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
[37] Sewon Min, Weijia Shi, Mike Lewis, Xilun Chen, Wen-tau Yih, Hannanch Hajishirzi, and Luke Zettlemoyer. Nonparametric masked language modeling- In Anna Rogers, Jordan Boyd- Graber, and Naoaki Okazaki, editors, Findings of the Association for Computational Linguisties: ACL 2023, pages 2097-2118, Toronto, Canada, July 2023. Association for Computational Linguistics.
[38] Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. Grokking of hierarchical structure in vanilla transformers. In Anna Rogers, Jordan Goyd-Graber and Naoaki Okazaki, editors, Proceedings of the 61 Anmual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 439-448 448, Toronto, Canada, July 2023. Association for Computational Linguistics.
[39] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
[40] nostalgebraist. interpreting gpt: the logit lens. 2020.
13

---

[41] Pascal Jr. Tikeng Notsawo, Hattie Zhou, Mohammad Pezeshki, Irina Rish, and Guillaume Dumas. Predicting grokking long before happens: A look into the loss landscape of models which grok. In arXiv preprint: 0bs/2306.13253, 2023.
[42] OpenAI. Gpt-4 technical report. In arXiv preprint: abs/2303.08774, 2023.
[43] OpenAl. New models and developer products announced at devday. 2023.
[44] Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: Antici- pating subsequent tokens from single hidden state. In Jing Jiang, David Reitter, and Shumin Deng, editors, Proceedings of the 27th Conference on Computational Natural Language Learn- ing (CoNLL), pages 548-560, Singapore, December 2023. Association for Computational Linguistics.
[45] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performande deep learning library. Advances in neural information processing systems, 32, 2019.
[46] Judea Pearl. Causality. Cambridge university press, 2009.
[47] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. In arXiv preprint: abs/2201.02177 2022.
[48] Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah Smith, and Mike Lewis. Mcasuring and narrowing the compositionality gap in language models. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Findings of the Association for Computational Linguistics: EMNLP 2023, pages 5687-57111 Singapore, December 2023, Association for Computational Linguistics.
[49] Ben Prystawski, Michael Y. Li, and Noah Goodman. Why think step by step? reasoning emerges from the locality of experience. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[50] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[51] Jack W Rac, Sebastian Borgeaud, Trevor Cai, Katic Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446, 2021.
[52] Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in BERTology: What we know about how BERT works. Transactions of the Association for Computational Linguistics, 8:842-866 2020.
[53] Daniel Soudry, Elad Hoffer, and Nathan Srebro. The implicit bias of gradient descent on separable data. In International Conference on Learning Representations, 2018.
[54] Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. Advances in neural information processing systems, 28, 2015,
[55] Zhiqing Sun, Xuezhi Wang. Yi Tay, Yiming Yang, and Denny Zhou. language models. In The Eleventh International Conference on Learning Representations, 2023,
[56] Alon Talmor, Yanai Elazar, Yoav Goldberg, and Jonathan Berant. olmpics-on what language model pre-training captures. Transactions of the Association for Computational Linguistics, 8:743-758, 2020.
[57] Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse universal transformer. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 169-179, Singapore, December 2023. Association for Computational Linguistics.
14

---

[58] Yi Tay, Vinh Q. Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, and Donald Metzler. Transformer memory differentiable search index. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022.
[59] Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. In arXiv preprint: 0bs/2206.048117 2022.
[60] Kushal Tirumala, Aram Markosyan, Luke Zettlemoyer, and Armen Aghajanyan. Memorization without overfitting: Analyzing the training dynamics of large language models. Advances in Neural Information Processing Systems, 35:38274-38290, 2022.
[61] Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. In arXiv preprint: abx/2309.02390) 2023.
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[63] Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal mediation analysis. Advances in neural information processing systems, 33:12388-12401, 2020.
[64] Boshi Wang, Xiang Deng, and Huan Sun. Iteratively prompt pre-trained language models for chain of thought. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2714-2730, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[65] Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023.
[66] Xinyi Wang, Alfonso Amayuelas, Kexun Zhang, Liangming Pan, Wenhu Chen, and William Yang Wang. Understanding the reasoning ability of language models from the perspec- tive of reasoning paths aggregation. In arXiv preprint: abs/2402.03326,, 2024.
[67] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[68] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, October 2020. Association for Computational Linguistics.
[69] Wilson Wu, John X. Morris, and Lionel Levine. Do language models plan ahead for future tokens? In arXiv preprint: abs/2404. 00859, 2024.
[70] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation models, 2023.
[71] Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel. Do large language models latently perform multi-hop reasoning? arXiv preprint: abs/2402.168377 2024.
15

---

[72] Eric Zelikman, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, and Noah D. Goodman. Quiet-star: Language models can teach themselves think before speaking, arXiv preprint: abx/2403.09629, 2024.
[73] Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. STar: Bootstrapping reasoning with reasoning. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[74] Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metrics and methods. In arXiv preprint: abs/2309. 16042, 2024.
[75] Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmen- tation. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 5657-5673, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[76] Zexuan Zhong, Zhengxuan Wu, Christopher Manning, Christopher Potts, and Danqi Chen. MQuAKE: Assessing knowledge editing in language models via multi-hop questions. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 15686-15702, Singapore, December 2023.
Association for Computational Linguistics.
[77] Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, LILI YU, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. LIMA: Less is more for alignment. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[78] Xuckai Zhu, Yao Fu, Bowen Zhou, and Zhouhan Lin. Critical data size of language models from a grokking perspective. In arXiv preprint: 0bs/2401.10463 2024.
16

---

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 hours at maximum.
B Effect of Model Scale
We run the experiments on composition with larger model scales with [8] 2000 and ø e (5.4, 9.0, 18.03. The results are shown in Figure 7 Overall, could seen that scaling up the model won't '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
0.8
Accuracy
Layers, 760 Dim Layers, 36 Layers, 1280 Dim
0.6
Layers, 24 Layers, 1024 Dim Layers, 1280 Tran HI)
0.4
Test (ID) (OOD)
0.2
Test (OOD)
0.0
0.0
5000
10000
15000
20000
25000
30000
2500
5000
7500 10000 12500 15000 17500 20000 2 2250
Optimization step
Optimization step
(a) Ø = 18.0.
(b) Φ = 5.4.
1.0
0.8
8 Layers, 768 Dim Layers, 1024 Dim Layers, Train (ID)
Accuracy
0.6
0.4
Test (ID)
Test (OOD)
0.2
0.0
5000
10000
15000
20000
25000
30000
Stti
(c) Φb = 9.0.
Figure 7: Results for different model scales across different Φ. Larger models converge in fewer optimization steps, but have no qualitative changes on the learned behaviors.
C Effect of Tokenizations
In the experiments in main content, tokenization is done by having unique token for each entity 12 This different from how real-world entities are typically tokenized practice, entities are usually multi-token, and different entities could share tokens different positions. We investigate
Prelimiiary experiments show that tokenization of the relations does not exhibit notable impacts, which is expected since relations always explicitly given.
17

---

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 and Φ = 12.6. We generate 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, the token multiplicity would be 2000/50 = 40.
Figure 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 seen that 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ı] throughout training to predict the second token of the bridge entity b in the setting with token multiplicity 40 [3] where the results are shown in Figure ().. It can be seen that the second token of b can be perfectly decoded from 5(5, r1] 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 into S|5, ri] 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 4
encoded in the hidden states (69
In summary, different tokenizations affect the results in rather expected ways, and do not influence our main findings and conclusions.
1.0
1.0
0.8
0.8
Accuracy
Accuracy
Multiplicity:
Multiplicity 4
Multiplicity:
Multiplicity 16
Test (ID)
Multiplicity:
Probing accuraty
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
Optimization on step
1e6
Optimization step
1e6
(a)
(b)
Figure 8: (a) ID generalization across different token multiplicity. (b) Probing accuracy on the second token of b at S(5, r.).
D
More Details on Circuit Analysis
D.1
Composition
We run causal tracing on hidden states in layer 1-7 and every position, where the target is the final prediction state r2), The changes in the strengths are monotone and smooth, and show in Figure 9 the strengths for the model checkpoint at the start and end of grokking, and also their difference (same as Figure (b)) We also find that after grokking, the state S(5, (which encodes (2) is not affected by perturbing the input / or T1-
Deriving the generalizing circuit. Starting from the 9 states in layers 0, 5, S(s, fj and 5(8, rı] since they have no computational paths connecting to S(8, r2). S|5, ht] can be eliminated could be seen by Figure (c... The connections from S h] and S ril to S|5, ral could be eliminated as mentioned earlier.
directly eliminate
13 Recall that S(5, rı] is the state that encodes the bridge entity in the generalizing circuit (Figure a))
18

---

h
h
Layer
Layer
Layer
Layer 2-
Layer 2-
Layer 2
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
-0.50
Layer
Layer
Layer
Layer
Layer
Layer
-1.00
-1.00
-1.00
(a)
(b)
(c)
Figure 9: Causal strengths on composition task, with the final prediction r(8 as the target. (a) Start of grokking. (b) Change during grokking. (c) End of grokking.
D.2 Comparison
Figure 10 includes the causal tracing results where the target is the prediction state S[8, e2]. and Figure π includes the results with target state S|5, e2]. can be seen that after grokking, S e2] does not depend on e1, which gives the generalization circuit in Figure 5]a).
Figure shows the rank (via logit lens) of the three relations {a<. a= a> in the label space at state S(7, a). where we use Recall @ as the measure. It can be seen that S(7, al encodes the label space throughout grokking.
e2
e1
e2
e2
1.00
Layer
Layer
Layer
Layer
Layer 2
Layer 2
Layer
Layer
Layer
Layer
Layer
Layer
0.00
Layer 5-
-0.25
Layer
Layer 5
0.50
0.50
0.50
Layer
Layer
Layer
0.75
0.75
0.75
Layer?
Layer
Layer
11.00
-1.00
11.00
(a)
(b)
(c)
Figure 10: Causal strengths on comparison task, with the final prediction Si8, ez] as the target. (a) Start of grokking. (b) Change during grokking. (c) End of grokking.
E
Additional Results
E.1
Weight decay
Figure 13 shows the ID generalization performance when varying the degree of weight decay (( = 2000 and á 9.0). It can be seen that larger weight decay can improve the speed of grokking, and vice versa.
19

---

1.00
1.00
1.00
0.75
0.75
0.75
0.50
Layer
Layer 1-
Layer
Layer
0.00
Layer
0.00
Layer
0.00
-0.25
-025
-025
Layer
Layer
-0.50
-0.0
-0.5
Layer
(a)
(b)
(c)
Figure 11: Causal strengths on comparison task, with S|5, ez] as the target. (a) Start of grokking. (b) Change during grokking. (c) End of grokking.
1.0
0.9
0.8
0.7
0.6
3
Test (ID)
Test (OOD)
0.5
0
1
4
Optimization step (1e5)
Figure 12: For the comparison task, S[7, a] encodes the label space throughout grokking,
E.2 Transformer with parameter sharing
We share the parameters the first layers and the last layers, similar as Universal Trans- former [7]. This would allow the model to share the knowledge in the upper and lower layers. The results on the setting with |8| = 2000 and ø = 12.6 are shown in Figure 14 could be seen that the parameter-sharing scheme can unlock OOD generalization, even though is gained much more slowly than ID generalization during grokking.
E.3 Comparison task across different inferred/atomic ratio
Figure 15 includes the result for & E {3.6, 9.0, 12.6) for the comparison task. It can be seen that a higher ratio & would give a higher generalization speed, consistent with the results in the composition task.
Complex Reasoning Task
For the complex reasoning task, for each attribute, we include 3% random facts from the comparisons between ID and OOD entities and the comparisons between ID and ID entities. In total, the training set contains 18K ID atomie facts, 437K (ID, ID) comparisons, and 108K (ID, OOD) comparisons, altogether 563K facts (28K average for each attribute). For translating the facts into natural
20

---

1.0
0.8
Accuracy
0.6
0.4
weight decay: 0.03
weight decay: 0.1
0.2
weight decay: 0.3
0.0
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Optimization step (1e5)
Figure 13: Effect of weight decay. larger weight decay can improve the speed of grokking, and Vice versa.
1.0
0.8
Accuracy
0.6
0.4
Test (ID)
0.2
Test (OOD)
0.0
0.00
0.25
0.75
1.00
1.25
1.50
Optimization step
1e6
Figure 14: OOD accuracy of the shared-parameter transformer model.
language for testing LLMs with non con-paraettric memory, we always use the attribute age (we find that the choice of attribute does not exhibit notable impact) and the templates "The age of (entity) is (attribute value)." and "(entity 1) [younger than/older than/in the same age as) (entity 2)." for atomic facts and comparisons. The entities are mapped to distinct random names generated by random generator We also try different mappings (e.g., unique IDs) and templates, and find the results to be consistent. All (retrieved) facts randomly permuted and concatenated before being loaded into the LLM context.
The train/test accuracy, and also the accuracy of inferring the attribute values of the query entities (which we test using the same format as the atomic facts in training) are included Figure could be seen that, during grokking, the model gradually locates the ground truth attribute values of the query entities (note that the model not explicitly encouraged or trained do this), allowing the model to solve the problem efficiently with near-perfect accuracy.
14 MRecall that we test each attribute separately by grouping the facts/queries. "anpss/pppp.orr/ppojecaaaam
21

---

1.00
1.00
0.95
0.90
Accuracy
Accuracy
0.85
0.85
0.80
0.80
#: 12.6
12.6
0.75
9.0
0.75
9.0
0.70
0: 3.6
4:
0.65
0.65
2
Optimization step (1e5)
Optimization step (1e5)
(a) accuracy.
(b) OOD accuracy.
Figure 15: Results for the comparison task across different ratio @.
1.0
0.8
Accuracy
0.6
0.4
Train
0.2
Test
Atomic (OOD)
3
Optimization step (1e5)
Figure 16: Accuracy on the train and test split, and also the accuracy of inferring the attribute values of the query entities (Atomic (OOD) for the complex reasoning task in $5
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All the findings and claims that make are based on evidence from set of controlled experiments, as presented in the different main sections 4 We also strive to make sure that the results robust alternative setups through additional experiments included in Appendix C
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
22

---

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have discussed the potential limitations of our work in Appendix We do not include it in the main content for the submission due to space limits, and will add it to the main content in the final version (if the paper is accepted).
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
We do not include theoretical results in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
23

---

Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide full details our data, model and training configurations, and evaluation setup in 4 and the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open tions to faithfully reproduce the main experimental results, as described in supplemental material?
the data and code, with sufficient instruc-
Answer: [Yes]
Justification: We have released the data and code with sufficient documentations alongside the release.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (bttps://nips... cc/
/ ss sswssssooosssoosoooo olicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new spun-s benchmark).
24

---

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https} des Cttee/CoeeSSbnnnn bmm ss i onPol icyy for more details.
//nips. cc/
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide full details our data, model and training configurations, and evaluation setup in $2 and the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: All of our experimental results are obtained with specified random seeds for reproducibility. We do not perform extensive additional re-runs of the experiments, however, we find preliminarily that the results are very robust randomness in the setup. In addition, our results exhibit strong patterns from controlled variations in the parameter settings, which we believe side-proof that the results are statistically significant.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
25

---

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com-- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide details on computing resources in Appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We believe our research conforms with the Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: Our research is at the foundational level which is not tied to particular applica- tions deployments, and hence believe there negative societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
26

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We believe our work poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer:
[Yes]
Justification: We have credited the original creators of assets such as the code framework in Appendix A
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, saperswithhon com/ dat aset s has curated licenses for some datasets. Their licensing guide help determine the license of a dataset.
For existing datasets that are re- ce-packaed, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
27

---

Justification:
We have documented our assets alongside the release.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work not involved with such subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our work is involved with such subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "pDdgUN7sl8", "reviewer": "Reviewer_jEH8", "summary": "This paper explores the ability of transformers to learn two synthetic tasks when trained from scratch. The tasks consist of rigidly structured data and effectively represent: 1. following a path of length 2 in a graph (the \"composition\" task), and 2. looking up + comparing two items in a dataset (\"comparison\"). \n\nThe investigations find that the transformers only manage to generalise the tasks after \"grokking\", a phenomenon in which test performance improves long after train performance has saturated, and moreover that the composition task does not generalise as strongly as the comparison task. \n\nCausal tracing is done on the transformers to gain an intuition of how they solve the tasks, i.e. revealing their internal \"circuits\". For the composition task, it is found that the network looks up individual edges of the graph at two different stages, whereas for the comparison task all information can be stored in a single subsection for the solution to work, providing some intuition on their different generalisation potential.\n\nSome manipulations of the datasets are done to see how this influences training, with conclusions drawn on dataset size and on the ratio of types of data. The ability of two commercially available LLMs to perform the task in a zero shot setting, when given the same data in-context, is also evaluated.", "questions": "1. The (recurring) statement that GPT4-Turbe and Gemini-1.5.Pro are \"based on non parametric memory\" (lines 16-17, and later 64-65, 280-281, 307) is confusing and potentially misleading - the two obviously also have significant parametric memory. The intention here is more about their ability to reason over *additional* knowledge provided in a non parametric way, i.e., this is about comparing reasoning over in context knowledge vs trained (parametric) knowledge. I would rephrase to be clear on that: this has been a comparison of zero shot in context learning vs task- and dataset- specific pretraining. I imagine there are works already talking about the fact that fine tuning is stronger than in context learning.\n2. line 25, English: an implication is not an impact or vice versa, rephrase to align\n3. line 28 define systematic generalisation \n4. intro style comment: try to reserve italics only for when a new term is being introduced (in which case it should be accompanied by a useful definition). can use very sparingly for emphasis but this is a lot. generally, a lot of terms do fly around here that would benefit from description, but are not accompanied by one - implicit reasoning, systematic generalization, mechanistic analysis, parametric memory, etc.\n5. general comment: can we really say that a model has learned to reason (as claimed in this paper) when it does not generalise out of distribution (composition task)?\n6. fig 1 last sentence shouldn't be here - caption is for explaining figure, not promoting other parts of paper\n7. lines 55-59 too vague by themselves, i get nothing from this part of the intro. subsequently, lines 59-61 not meaningful until after reading entire paper - which is not the goal of an intro! be more concrete.\n8. missing: clear numbers on dataset sizes and in particular repetitions while training. when you made phi larger did that make the dataset larger or did it mean more repetition or..? do i understand correctly that the maximum number of possible inferred facts is |E| x 20 x 20 (subject, outgoing relation 1, outgoing relation 2)? if so for |E|=2000 as in line 113, that means 800k inferred facts (and slightly less for the actual trainable inferred facts ie those not set aside as either OOD or even ID test) and 20 x |E|=40k atomic facts. given batch size of 512, thats only enough for <2k optimisation steps .... was there a huge amount of epochs here? (line 120: training goes to 2 million training steps)\n9. from things like figure 3 and equation 1, i infer that the input sequences are presented as short structured 3- or 4- (for atomic vs inferred facts) token sequences, but this should have been explicitly stated and presented.\n10. regarding the facts graph: in a real set of facts, some relations do not make sense: for example, a number entity should not be the target of a person relation (e.g. 'barack's wife is 1964'). it seems the graphs considered in this work do not reflect such constraints. that's fine for an exploration of graphs, but not great for drawing conclusions on general reasoning skills. again, it is my opinion that this paper should be reframed to reflect its more straightforward graph processing nature. it can of course be accompanied by some discussion on how this may relate to reasoning in more 'general' models, and what insights may be taken from this research for work on such models.\n11. lines 122,3: \"ID accuracy\", \"other splits\" - be more explicit (\"test ID accuracy\", \"train ID and test OOD\") else hard to read\n12. figure 2: what happens at \"ratio\" (caveat: see next comment) greater than 18, i.e., some of the atomic facts not given? what happens at ratio infinity? could be interesting. especially given the fact that it doesn't generalise to the OOD facts, i think its building at least one of the layers of atomic facts from the composed samples.\n13. in the experiments framed around ratio between atomic and inferred facts, it is something else that is being directly varied: the percentage of the inferred facts given in the dataset. this is probably also why the investigation was capped at the highly specific ratio 18, instead of continuing to increase the ratio and find a sweet spot: ratio 18 is the maximum for this dataset when doing this variation (effectively it means, take \"90%\" of the data, i.e., include all facts that have not been reserved for the test/validation sets). Reporting it in terms of percentage of inferred facts included would be a much clearer and more generalisable framing - the magic number 18 is obviously specific to this dataset. To actually investigate ratio, one should hold fixed the percentage of inferred and of atomic facts being included in the data, and then repeat them different amounts of times in the dataset to change the ratio between them (recall there is a huge amount of repetition happening in training here anyway, see comment 8). the conclusion of line 125-126 now becomes the much clearer \"providing maximum samples of inferred data speeds up generalisation\". \n14. the statement on line 127 also becomes less convincing once understanding the above. in one investigation (fig2a) the dataset size is manipulated by adding/removing inferred samples while keeping a fixed size of atomic samples. in the other (fig2b), it is manipulated by changing the number of atomic samples. a more likely conclusion, which also can only really be drawn if drawing a figure 2a for multiple sizes of |E| and verifying that it stays the same, is that *percentage of [all inferrable data] that gets into training* is what matters for generalisation. for conclusions on distribution, one would have to hold that percentage (and similarly the percentage of atomic data presented in training) constant, and manipulate only their relative repetitions (which is possible given that both are repeated a lot anyway in this setup).\n15. fig 2a: explicitly state |E| in caption/image/somewhere. did you check this for multiple |E| and get the same results?\n16. fig 2 caption: at this point in the paper,  test_inferred_id has not been introduced, define things on time\n17. line 137 - scaling the data has no effect on relative speed - not very convincing once youve scaled for epochs and not steps (i.e. normalised for dataset size). Unless you mean relative speed between generalisations. Some numbers would be more convincing - define when things have happened and see if ratios hold not just visually - e.g. seems |E|=10k takes twice as long to saturate train than smaller |E|, but not so for test.\n18. line 161 S[i,a] - poor phrasing specifically in description of 'a', it is not the input token, it is a marker of input position. I.e. it does not take the value of the actual input tokens (and should not, as then we may have ambiguity because r1 and r2 could even be equal)\n19. lines 167-172 do you take care that the perturbed run should be valid, i.e. that the path (h,r1',r2) exists in the graph?\n20. 175-176 \"the ratio of such alterations (between 0 and 1)\" - be more explicit, dont understand\n21. line 198 claim that there exists a memorizing circuit: this was not shown, or i missed it. it's a reasonable hypothesis but i dont think it has (or can?) actually been explicitly concluded from anything here.\n22. line 208: i havent read ref. [43] in detail, but it seems to me that they say what they find doesnt apply to ADAM, which is what is used here. in fact they mention a 2017 paper suggesting there will not be generalisation with adam. clarify discussion and frame accordingly (is this in agreement or contradiction with previous results)\n23. footnote pg 6: too vague, cant get anything from this\n24. line 216-217: reasonable hypothesis, unfortunately presented as fact and not verified at all. what if the network is actually using the upper layers to complete the atomic facts? following discussion also not sufficiently hedged as a result.\n25. generally: a simple circuit was found for this inference, and using only 2 computation 'steps' in practice (layer 0 to 5, and 5 to 8). if correct, such a finding suggests that most of the layers are redundant. why no experiments on transformers with less layers?\n26. in general, consider doing multiple experiments for each task, to see if the same circuits arise every time for these tasks, rather than these just being the solutions two specific models learned. \n27. generally: in abstract and in conclusion, the results are said to guide potential improvements to data, training, and architecture. the actual concrete suggestions are spread out through the paper. the paper would be easier to act on if these suggestions were also put together in a clear list at some point in the paper, each with references to its supporting evidence in the paper.\n28. section 4: presentation of comparison task general - there is some kind of conflation of facts (things that are true about numbers or about the \"underlying data\" (for lack of a given term)), samples (tuples that go into the dataset for this underlying data), and rules (how facts and samples relate) here that starts to get fuzzy as the section progresses, tidy up. In particular, for equation 3 to work, there also has to be an understanding of what a sample implies about the underlying data, which is currently missing. (Specifically, the missing part is that there are no 'incorrect' samples. This can be corrected by phrasing eq. 2 with an if-and-only-if instead of only an if (i.e., a double headed arrow), the quantifiers need to be carefully updated for this too.)\n29. lines 288-289 took a while to parse, be more explicit\n30. lines 294-296 be more explicit about whether talking about atomic or inferred facts\n31. line 296 \"no surface form clues\" - explain/be more explicit\n32. 297-298 \"unlike most conventienal QA benchmarks where..\" - do you have any support for this?\n33. 310 \"model tends to give up ...\" - would appreciate examples and numbers. how did you conclude it was doing this? i want to be convinced too\n34. the suggestion on parameter sharing following the results on composition is reasonable, but not explored beyond that (i.e. no parameter sharing transformer was trained to see if it would generalise better), so this is more like a general hypothesis/thought for future work.\n35. the comparison between parametric and non parametric memory has a confounding factor that is not addressed: the transformers trained in this paper are trained both on the specific data *and on the task* at once. Meanwhile the 'non parametric models' they are compared to are trained on general NLP tasks, and given the data and task in context. Effectively, the non parametric memory is being evaluated for how well transformers can utilise it *in a zero shot setting*. It is conceivable that a transformer fine tuned specifically for performing the task on data given in parametric memory would have much better performance, and this would be a fairer evaluation of the potential of parametric memory.", "limitations": "concluded guidelines on data/training/architecture settings not sufficiently supported, comparison to non parametric memory unfair", "rating": 5, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "Through causal tracing, the manner in which a transformer may solve one of two types of tasks - amounting to chained retrieval (\"composition\") or parallel retrieval (\"comparison\") - is found. As may be expected, in chained retrieval the transformer must access data at two different points in computation, while in parallel retrieval it may retrieve all data from a single 'area' (part of its weights) - making parallel retrieval more easy to generalise. While this is not surprising, it is nice to see explicitly, and shows an interesting new case where weight sharing or universal transformers may (i note may, as this is not investigated) have an advantage over vanilla transformers (i previously thought the benefit of universal transformers were restricted to their adjustable computation depth, so this is nice). The conclusion on the more complicated composition mechanism is supported also by the investigations of generalisation ability (in and out of distribution test performances) of transformers on the two tasks.", "weaknesses": "1. Overclaiming: Framing should be adjusted to more clearly reflect the work done.  \n- Title overblown for an investigation of two rigidly structured synthetic tasks. I highlight: the structure is so rigid that all input samples have one of two short fixed lengths: 3 tokens for single datapoints, and 4 for demonstrations of composition/comparison.\n- In general I feel people (not just in this paper) are too quick to make claims on 'reasoning' from experiments on straightforward formal tasks (which are valid and interesting in themselves without such aggrandizations, not to mention neater to discuss without them). In the case of this paper, the framing is around reasoning, but in practice it investigates the ability to learn instances of two highly structured synthetic tasks (instances of: the tasks are not learned in a general sense, but rather, with respect to specific datasets).\n- The suggestions for how transformers/training should be improved are either poorly supported by the experiments or not investigated beyond speculation: see comments 12, 13, 14, and 34. Correcting these analyses and doing the investigations to support the various suggestions would significantly improve the paper.\n- Unfair comparison to parametric memory, see comment 35\n2. Poor analysis of results: in particular, see comments 12, 13, 14 regarding manipulating the ratio of inferred data in the dataset, or comment 24 regarding conclusions on internal mechanics of the model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the very informative and constructive comments! We will incorporate them to improve our work in the revised version. Overall, we believe that all the major technical concerns raised here are due to some misunderstandings of our paper, and we will group your comments into different topics and respond in detail below. \n\n**[Comparing non-parametric and parametric memory - Question 1,35]** Our setting on comparing non-parametric memory and parametric memory is very different from comparing in-context learning and fine-tuning, and we believe our experimental setting is far from being “unfair”. To explain in more detail:\n- First, both the grokked model and the LLMs understand the rules in the task and the task objective. This is clear for the grokked model and also true for the strong LLMs we tested given the simplicity of the rules and task objectives involved (we also intentionally made the natural language templates very simple s.t. LLMs have no difficulty in understanding them; see lines 299-300 and Appendix F). We also confirmed in our preliminary experiments by testing the LLMs on easy small-scale instances of the task, where they succeed in no doubt. Based on this, what we are comparing here is solely the capabilities of the different models to make difficult deductions (as those in our test set) based on the given facts.\n- Now, it is important to note that we are **not including any examples constructed in the same way as the test examples in the given facts** (see lines 287-293). In other words, **both** the grokked model and the LLM are doing zero-shot prediction here (which is perfectly fine given they both understand the rules and task objective). We believe the confusion here mainly comes from the given facts (specifically, the ID-ID comparisons as “train_inferred_id” in earlier sections) also serving the role of teaching the rules and task objective to the transformer. A clearer setting would be to take a pretrained (grokked) transformer (e.g., one from Section 4) and then continue training it on the given facts here, however, this won’t make qualitative differences since the given facts can teach the model these things by themselves anyways.\n- Overall, we admit that it is very hard to be perfectly fair, especially given that we don’t have control over the LLM’s training data, but still, we believe our experimental settings are far from being “unfair”.\n\nIn terms of terminology: thank you for the suggestion and we will improve the terminologies used in the revised version. It is important to note that prior literature (e.g., [1,2] cited below at the end of the response) has used terms including “non-parametric” and “memory” for the “in-context knowledge” you mentioned, and we are not inventing completely new terms here in the paper.\n\n**[Experiments on models with parameter sharing - Question 34, also in “Strengths”]** We **did** perform experiments on the parameter-shared transformer model (akin to Universal Transformers), and showed that it can achieve highly non-trivial systematic generalization for the composition task. The relevant details are included in Appendix E.2 and mentioned in lines 226-228 in the main content. We believe this is a firm initial step showcasing the potential of the suggested directions implied by our findings and analysis.\n\n**[Ratio between inferred & atomic facts, and repetitions of facts during training - Question 8,12,13,14,17]**\n- We are in the setting where each example in the training set is unique, and we train the model for a very large number of steps/epochs (basically training the model forever). Yes, increasing the ratio (while holding the atomic facts fixed) means that the dataset will be larger (we also mentioned this in line 128), and your calculations on the number of facts are correct. We are not adding these explicitly in the paper due to space limits and will add them in the revised version.\n- Thank you for the suggestion on using percentage instead of ratio - indeed, this could help avoid quantities specific to the dataset construction, and we will consider switching to it in the revised version. Nevertheless, in our setting, the percentage of included inferred facts is always the ratio divided by the KG’s outgoing degree (20), and hence they are only different in units. In other words, it is not that there is “something else that is being directly varied” (which may suggest the existence of unconsidered confounding factors to some degree); we are directly varying it in the first place.\n- Regarding repetitions of facts - having different amounts of repetitions of the facts is another interesting factor to study, but also goes beyond our scope here. Regarding your comment on the necessity of changing the relative repetitions of the facts while holding the percentage fixed (Q13, 14), here we are **not restricting ourselves to comparing different distributions over the same support** - rather, we are talking about **variations in characteristics of the distributions (which could have different supports)**. We believe this should resolve your confusion here. Also, it should be clear at this point that your “more likely conclusion” mentioned in Q14 is really equivalent to our conclusion in the first place.\n- By “relative speed” (Q17), we mean the relative speed of improvements in generalization and training (line 137) for a dataset. So it is valid to rescale the steps (including ones based on dataset size as we did). We believe the visual evidence should already be clear enough, but will consider adding quantitative scores in the revised version.\n- Overall, we believe that your major technical concerns here are mostly caused by misunderstandings potentially due to the writing being a bit abstract at certain places (mostly due to space limit), and we don’t see real pitfalls in our findings or analysis. We will improve the writing and incorporate your suggestions in the revised version.\n\n(continued in comment \"Part 2\")", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LcTlK9een2", "reviewer": "Reviewer_f5kN", "summary": "This paper explains the reasoning ability in LMs is acquired through grokking,  which requires extended training beyond overfitting. Through analytical experiments, the authors explore the mechanisms behind grokking, the formation of generalizing circuits, and the impact of systematicity in the configuration of these circuits.", "questions": "1. Can you provide the detail of your data sizes ? i can find the statics and samples of your data.\n2. any other reasoning task can confirm your assumptions ?", "limitations": "See in Questions.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Novel Insight. critical data distribution decides the characteristics of grokking\n2. Sufficient Mechanistic Analysis. The detailed mechanistic analysis of the internal workings of transformers during grokking offers valuable insights into the formation of generalizing circuits and their efficiency.\n3. Clear Presentation. The study is well-structured, with clear explanations of the experimental setup", "weaknesses": "1. Limited Scope. The focus on only two types of reasoning (composition and comparison) may limit the generalizability of the findings to other reasoning tasks.\n2. Dependence on Synthetic Data. The experiments primarily use synthetic data, which may not fully capture the complexities and nuances of real-world data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive comments! We will incorporate them to improve our work in the revised version.\n\n**[Limited scope and dependence on synthetic data]** We do admit these limitations in the paper (Appendix H), however, we believe that despite these, we are taking the initial steps toward studying the problem by laying out a clear and well-defined formulation, and also conducting rather deep investigations on the tasks we study. This could serve as a solid foundation for future work to build on and move to more realistic settings.\n\n**[Details of data sizes]** For the task of composition, the number of total atomic facts is 20*|E|, where 95% are utilized as ID facts and the remaining ones are OOD facts. The specific quantities of the facts would then depend on 1) the number of entities |E| and 2) the inferred/atomic ratio, which we vary in our experiments to see their effects on the model. For example, when |E|=2K and ratio=9.0, there will be 38K ID atomic facts, 2K OOD atomic facts, and 342K ID inferred facts that go into the training data. The case of comparison is similar. We will make some of these quantities explicit in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cg4xdiksfG", "reviewer": "Reviewer_Xiqz", "summary": "This paper investigates whether Transformer models can learn implicit reasoning through the phenomenon of \"grokking\" focusing on two types of reasoning: composition and comparison. Also, the paper reveals the mechanisms behind grokking and the differences in systematic generalization across different reasoning tasks through internal model analysis.", "questions": "1.\tIn line 203, what does 'stores... in the weights' mean? Does it imply that during the memory phase, the logit lens is not 'b'?\n2.\tIs the grokked transformer a newly trained model or is it based on the model described in the previous sections? If it is a newly trained model, how was it trained?\n3.\tWhat will happen when we input out-of-distribution data in Figure 4a? Will the logit lens at S[5, r1] not show 'b'?\n4.\tIn my opinion, the failure of the OOD in the composition is due to the unseen embedding of OOD data during the training.", "limitations": "See weaknesses.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1.\tThe paper explores the relatively unexplored phenomenon of \"grokking\" in the context of Transformer models, offering new insights into how extended training can lead to implicit reasoning capabilities.\n2.\tBy focusing on both composition and comparison reasoning tasks, the study provides a thorough examination of the different ways Transformers can generalize based on the circuit mechanism.\n3.\tThe paper analyzes in-distribution (ID) and out-of-distribution (OOD) generalization for composition and comparison reasoning tasks. It explains why the composition task cannot generalize systematically while the comparison task can, by examining circuit mechanisms.\n4.\tThe paper trained a grokked transformer to finish the more complex reasoning task which is far beyond the SOTA like GPT-4-Turbo and Gemini-1.5-Pro.", "weaknesses": "1.\tThe paper does not provide how to tokenize the whole sentence in (1)(2) and how to evaluate the loss function. \n2.\tThe beginning epoch of the abscissa in Fig.1 needs clarification. Why the first point in Fig1b Train accuracy is about 0.9?\n3.\tSince the accuracy is a discontinuous metric, it may sometimes mislead for indicating grokking. The loss for each training needs to show.\n4.\tThe interval over which the change in causal strengths in Figure 4b is calculated is not specified. Additionally, it is unclear how the change is defined.\n5.\tThe term 'same type' in line 172 needs clarification. Does it refer to data within the same entity set or add noise?\n6.\tHow to calculate the ratio in line 175 needs clarification.\n7.\tFig.4b,4c only show the correlation between the causal strengths and the grokking. There is no evidence indicating that the formation of the structure causes the grokking. In other words, this correlation is shown in (Reddy, 2023).\n8.\tThe paragraph titled 'Why does grokking happen?'  in line 197 is inappropriate. It only describes the correlation between grokking and the formation of the circuit, explaining what happens during grokking but not why it occurs. The main question that needs to be addressed is why transformers tend to form this circuit when the training accuracy is high (the training loss is low). \n9.\tIn the circuit of Fig5a, the analysis through the logit lens only shows the output at this layer, but how can the transformer maintain the output to layer 5 is not shown.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the informative and constructive comments. We will incorporate them to improve our work in the revised version.\n\n**[Tokenization and loss - W1]** If the “(1)(2)” here means the rules in Equations (1) and (2), these rules are latent and the model only sees the atomic and inferred facts (deduced from the atomic facts via the rules). This is explicitly mentioned in Section 2. We also stated how we tokenize the facts in the main experiments (lines 102-104, 246-248) and have some further investigations on tokenization (Appendix C). For the loss, we use the standard cross-entropy loss as in normal language modeling. We will add these details in the revised version.\n\n**[Initial point in Figure 1 - W2]** The initial point here is for 2K optimization steps. For reasons why the comparison task has a high performance in the beginning: the comparison task has a significantly smaller effective label space (3 comparative relations for the specific attribute) compared to the composition task (all entities), and also the model could more easily guess the answer correctly, e.g., if an entity’s attribute usually appears to be larger than others’, when the model is asked to compare it with someone else, it could just guess that this entity has the larger attribute without really comparing them.\n\n**[Metric used - W3]** Thank you for the suggestion. We plot the loss curves in Figure 1 of the added PDF under the same setting as Figure 1 in our paper. It could be seen that the loss curves have the same trends (upside down since small means better for loss) as the accuracy curves. This clears out the concern here.\n\n**[Details in causal tracing - W4,5,6]**\n- The interval is taken to start from the beginning of grokking (when training performance saturates) to the end of grokking (when test performance converges). The change is defined in the natural way, which is the causal strengths at the end minus the strengths at the beginning of grokking.\n- ‘Type’ here means entity/relation. For example, we are not perturbing an entity into a relation.\n- The ratio is calculated as the percentage of the 300 random examples where the causal intervention changes the prediction.\n\nWe will make these details more explicit in the revised version. \n\n**[Explanation of grokking - W7,8]** We are suggesting a plausible explanation of what we observe here through the lens of circuit efficiency, with concrete evidence (e.g., comparison of the amount of facts different circuits store) and further experimental consolidations (Appendix E.1). A more fine-grained and rigorous answer to why grokking happens would require further efforts such as detailed reverse-engineering of the weights and analyzing the training dynamics. These are interesting follow-up directions based on our work, but we believe they are also beyond the scope and amount of effort of this paper. We are not sure which paper you are referring to by \"(Reddy, 2023)\"; it would be great if you could further clarify this.\n\n**[Detailed mechanisms within the circuit - W9]** Related to the last point, while a more detailed and lower-level understanding of how the model does the computations within the circuits is interesting to have, these won’t really affect our conclusions as we are not doing fine-grained (theoretical) analysis in this work. We believe that our current work already presents significant efforts and provides highly insightful findings and analyses, which could serve as a solid foundation for subsequent study at more detailed levels.\n\n**[Q1]** The logits lens at S[5, r1] is ‘b’ from the beginning, however, the causal strength between S[5, r1] and the target state is very low (lines 189-190, Appendix D) which suggests the model is not traversing the bridge and directly memorizing the inferred facts.\n\n**[Q2]** It is a model trained (from scratch) on the given facts, not directly taken from previous experiments. We will make this clearer in the revised version.\n\n**[Q3]** We did compute the logit lens results for the OOD setting, and found that S[5, r1] encodes the bridge entity (with MRR 0.98), which is a strong indication for our subsequent explanation in lines 216-217. We will include these results and expand the discussion in the revised version.\n\n**[Q4]** The OOD facts actually share the same set of entities and relations as the ID facts (see Section 3.1). Moreover, we did experiments on an alternative architecture with cross-layer parameter-sharing and found that it could achieve impressive OOD generalization in composition (lines 226-228, Appendix E.2), which implies that embeddings are not the major issue and also consolidates our explanations in the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i2G8Kesyun", "reviewer": "Reviewer_smMo", "summary": "The work mainly focuses on the systematic generalization (specificially, in the paper, two implicit reasoning types: composition and comparison of facts) of the grokked (i.e., training far beyond overfitting so that the model can finally learn some specific generalization skills and achieve high testing performance) transformer-based language models. \n\nThe paper presents empirical findings that the trained transformers: (1) for composition reasoning, grokking is observed in in-distribution generalization testing but not in out-of-distribution generalization testing; and (2) for comparison reasoning, have grokking is observed in both in-distribution and out-of-distribution generalization testing. \n\nThe paper also interprets the phenomenon of grokking by investigating the inner working patterns by correlating generalization performance with different reasoning circuits implicitly performed in the models. \n\nBesides, the paper designs a synthetic complex reasoning task, where state-of-the-art LLMs (GPT-4-Turbo and Gemini-1.5-Pro) achieve low performance even with advanced prompting techniques and retrieval augmentation while grokked small-size transformer achieve good performance, stressing the power and the potential of parametric memory inside language models.", "questions": "(1) For my major questions and concerns, please refer to the \"Weaknesses\" part.\n\n(2) minor question: In figure 4 (a), how do you demonstrate it is the bridge entity (**b**) encoded in the hidden state of (Layer 5, r1) position that causally affect the final prediction (**t**). In the causal tracing procedure described in the paper, I can only infer that the whole hidden state has causal effect on the final prediction (what if something else in the hidden state plays the important role?)\n\n(3) minor question: I am wondering what if we generalize the composition reasoning reasoning from two-hop to three-hop and even more hops? Do the grokking phenomenon still emerges? It seems that the complexity of circuits implemented by transformers is bounded by their layer.\n\n(4) minor question: I find detailed description, settings and results for Table 4 (e.g., CoT prompt examples, specific error cases) are missed in the paper. I am wondering why the performances with chain-of-thought prompting are even worse. Could the authors please provide more detailed results and insights to this point?\n\n(5) A minor suggestion: this work [1] also discusses the mechanism of compositional factual reasoning in the LLMs, suggesting LLMs can implicitly perform mutli-hop reasoning and the reasoning errors might stem from improperly leveraging the bridge reasoning results (which are also in align with the observations presented in this work). Hence it is good (not necessary) to also discuss this work in the Related Work section.\n\nTo conclude, though I have concerns on the generality (to widely-used LLMs) and novelty (from the perspective of grokking) of the paper, I appreciate its organic combination of grokking, implicit reasoning and mechanistic circuit analysis. I think the work will clearly have its impact on the LLMs' implicit reasoning (with the parametric knowledge) field and hence lean towards accepting it.\n\n[1]: Understanding and Patching Compositional Reasoning in LLMs, https://arxiv.org/abs/2402.14328 .", "limitations": "Yes, the authors do include a section (Appendix H) to discuss the main limitations of the work. I would also suggest the authors discuss the limitations of Logit Lens [1,2] (used for analyzing the internal reasoning circuits) in this section.\n\n[1] Logit Lens: https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens .\n\n[2] Eliciting Latent Predictions from Transformers with the Tuned Lens, https://arxiv.org/abs/2303.08112 .", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The topics of this paper, the grokking and systematic generalization of transformer-based language models, are very important, timely and of interest to the NeurIPS community. Besides, the paper is overall well-written with fluent presentation and clear logic structure, thus being easy-to-follow.\n\n2. The study presented in the paper is comprehensive: from observing generalization performance to investigating the internal reasoning circuits. The derived reasoning circuits in the grokked transformers can explain the different generalization behaviours in different settings (for both in-distribution testing and out-of-distribution testing; for both the composition reasoning task and the comparison reasoning task) very well.\n\n3. *From the perspective of implicit reasoning in LMs*, I appreciate that the authors correlate the implicit reasoning (for both composition reasoning and comparison reasoning) with the grokking, demonstrating that LLMs have potential to perform implicit step-by-step reasoning rather than simple memorization, which is a good complement to existing works [1, 2, 3]\n\n4. The results for complex reasoning are quite impressive to me, demonstrating that well-trained (\"grokked\") transformers with parametric knowledge have potential to largely surpass the state-of-the-art large models with techniques like chain-of-thought prompting and retrieval augmentation.\n\n[1]: Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models. https://arxiv.org/abs/2310.14491 .\n\n[2]: Do Large Language Models Latently Perform Multi-hop Reasoning? https://arxiv.org/abs/2402.16837 .\n\n[3]: Understanding and Patching Compositional Reasoning in LLMs, https://arxiv.org/abs/2402.14328 .", "weaknesses": "My major concerns are two-fold:\n\n1. Though the mechansitic interpretation results to the grokking of implicit reasoning patterns in this paper are beatiful and persuasive to me, I think its inspiration to analyzing real LLMs pre-trained on a large corpus can be quite limited, given the settings of this work (i.e., structured synthetic data, single task, random initalized GPT-2 transformer training from scratch.), as the authors discuss in the Appendix H.\n\n2. *From the perspective of grokking*, there have existing works (to name a few, [1,2]) discussing the mechanism of the grokking phenomenon: memorization versus generalization, the transition and competition of the implemented circuits inside models (most of which focused on different tasks, e.g., algorithmic tasks). I think though this work specifically focuses on knowledge reasoning, its conclusions and insights partly share with previous works so that they may not that novel to me.\n\n[1] PROGRESS MEASURES FOR GROKKING VIA MECHANISTIC INTERPRETABILITY, https://arxiv.org/abs/2301.05217 .\n\n[2] Do Machine Learning Models Memorize or Generalize? https://pair.withgoogle.com/explorables/grokking/ .", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the informative and constructive comments! We will incorporate them to improve our work in the revised version.\n\n**[Abstract nature and connection to practice - W1]** Yes, we do admit that one limitation of our work is its synthetic and abstract nature (Appendix H). Still, we believe that our work is a firm initial step that future work could build on and gradually move to more realistic settings. \n\n**[Insights on grokking - W2]** While we indeed borrow ideas and existing developments in grokking to understand and explain our findings, we note that we also bring novel aspects and potential corrections of prior explanations of grokking through our controlled study. Specifically, through controlled experiments (lines 122-143), we find that existing explanations based on “critical data size” are potentially flawed or not general enough, and it could instead be “critical data distribution” that is really the crucial factor.\n\n**[Limitation of logit lens - Q2, Limitations]** This is one limitation of logit lens, or more generally, all interpretability methods that are based on “lossy” projections (including further improvements of it such as tuned lens as cited, etc.). This is still an open problem for interpretability research to the best of our knowledge, and we currently don’t have a good way to resolve this. Some potential directions may be to do further interventions on the “remaining components” (e.g., projection of the state to a subspace orthogonal to the bridge entity’s embedding) but these could again require further justifications. Nevertheless, despite the possibility of alternatives, we believe that it is reasonable to conceptualize the state as the bridge entity in our case here given the strong correlation, and will discuss these limitations in the revised version.\n\n**[Composition with more hops - Q3]** This is one interesting direction for follow-up studies; here based on our findings, we believe that grokking should still happen generally (depending on the inferred/atomic ratio) and the vanilla transformer would certainly still suffer in systematic generalization. Indeed it would be difficult for a vanilla transformer to resolve a large number of hops given their bounded computations; one related interesting direction would be to investigate whether cross-layer memory-sharing models such as universal transformers (we also briefly investigated this in Appendix E.2) could solve this and maybe even generalize to more steps than those seen in training (when the layers are allowed to be executed with unbounded steps).\n\n**[Details about the results and settings in Table 4 -  Q4]** Indeed we omitted a lot of details here (we also put some details in Appendix F) in the paper due to space limit, and we will add them in the revised version. For the CoT prompt, we are using the zero-shot instruction which encourages the model to “think step by step”. Regarding why adding CoT makes the performance worse, first of all, the model’s performance without CoT is still pretty close to random guesses, so a summary would be both settings fail badly. Now specific to the performance drop, we found that with CoT, 70.7% of Gemini’s responses ended up saying that the answer cannot be decided (which we treat as wrong since the answer can be decided). The ratio drops a bit to 58.7% when augmenting with retrieval. One typical example model response is included in Table 1 in the added PDF. Intuitively this does make a lot of sense because, when the model is instructed to verbalize the reasoning, it is harder for it to “guess” the answer because the generated rationales are present in the context and the model would have more explicit clues that the logic doesn’t really work out.\n\n**[Related work - Q5]** Thank you for referring to the related work here. We will discuss it in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OLtDriYhL7", "reviewer": "Reviewer_EZN5", "summary": "Recent work has shown that LLMs are bad at implicit reasoning over parametric knowledge, and this work asks whether that's a fundamental limitation of the transformer architecture or not. Through thoroughly investigating the performance of an 8-layer transformer trained from scratch on 2-hop reasoning tasks, this paper shows that the models can only learn to reason implicitly over parametric knowledge when they learn a \"generalising circuit\" (as opposed to a \"memorising circuit\"). They find that in most settings, the model first learns a memorising circuit, and only after grokking (training for many steps after the training set has been fitted), the model learns the generalising circuit. By varying aspects of the training data distribution the authors find that, contrary to what prior work claims, it's not dataset size contributing to grokking happening, but the data distribution (ratio of inferred facts/atomic facts, higher means grokking happens faster). Reasoning about the generalising circuit and when it would be preferred over the memorising circuit, the authors come up with testable predictions for when grokking should happen, why the model does/does not generalise OOD, and what could help it generalise OOD. They empirically show (in the appendix) that these predictions are true in their domain. One of the main insights here is that the generalising circuit is a more efficient representation of data generating distribution, and if you bias a model towards favouring efficient representations with things like weight decay, the model will find it faster.  Another insight is that parameter sharing can help OOD generalisation. The authors finally construct a more difficult, three-hop implicit reasoning task, that requires searching over a larger space to find the \"bridges\" that allow \"hopping\" from one fact to another and answer a comparison question. They show that SotA LLMs like GPT-4 cannot do this, but a grokked transformer can do it almost perfectly. The authors take their results to show that it's important for future work on robustness / systematicity of LLM reasoning to consider the power of reasoning over parametric memory with a generalising circuit that more efficiently encodes the task than a memorising circuit.", "questions": "Questions\n- Do you think the recommendations to make finding the generalising circuit easier will transfer to LLM training? \n- Do you have an explanation for why the model learns a parallel pathway for the comparison task and not the composition task?\n- What does it mean for the model to \"prepare the label space\"? Is it that, according to the unembedding at those hidden states, the top tokens are always those in the label space?\n- Why does the memorising circuit need to also store the atomic facts and not just the inferred facts for the composition task?\n- Line 85-96 about \"what happens during grokking\". Why do you say that these results indicate the model is mainly memorising without going through the hop before grokking? From the observations you mention it could still be that b is encoded/retrieved elsewhere in the model no? Or did you find that $b$ doesn't get strongly encoded anywhere before grokking?\n\nWriting suggestions:\n- Can you give an example of an actual input and output pair you train the models on in the main text?\n- Some details about how you discover the generalising circuit are still missing. How many examples do you use to determine the ratio of alterations that is called the \"causal strength\"? You say you eliminate / prune certain connections, what does that mean, you leave them out and the accuracy of the remaining model is still high? What does it mean that the lower layers retrieve the first-hop fact? Is it that they all unembed to top-1 (or top-3) tokens of the first-hop fact (h, r1, b)?", "limitations": "Addressed in the appendix.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper explains the grokking phenomenon very effectively. The biggest strength of this paper is that they use their mechanistic analyses to make testable predictions about when and why grokking should happen, and then manage to use these predictions to manipulate the phenomenon to show up earlier in training. Additionally, they use their findings to propose improvements to LLM training / architecture that might help it reason more robustly in the future. Of course, it remains to be seen whether this will hold at scale, but the authors convincingly show that these improvements help a smaller single-task transformer to learn to generalise OOD and robustly reason, and that reasoning robustly (in this particular rule-based way) doesn't need to be a general limitation of transformers.\n\n- Each figure in this paper is exceptional and supplements understanding of the work very well, making the presentation overall great.", "weaknesses": "- The paper does a lot, which is a strength, but also means that some details that are required to understand the paper better seem left out. I will ask questions below about these details.\n\n- Not so much a limitation of this work itself, but of one of the methods used (causal tracing). I'm wondering how well causal tracing would work for problems that require using multiple different hidden states in a distributed way, and whether it's mainly for small transformers on relatively simple / rule-based tasks. It seems unlikely LLMs rely that much on a single hidden layer for predictions.\n\n- It's not entirely clear how/if the findings from this work would transfer to very general models like LLMs. Is it reasonable to expect the same recommendations (on weight decay, parameter sharing, using parametric memory to efficiently store atomic facts and learn to apply rules over them) to hold for models like LLMs, that  currently do one-epoch training and do not overfit, let alone far exceed overfitting, on a specific task.  Intuitively, it seems unlikely that we can get both very general models and \"grok\" each reasoning problem it needs to be able to do, even if we make the generalising circuit more favorable by parameter sharing or weight decay. I do think the authors somewhat address this point in the limitations, and also do not claim their recommendations should hold for every reasoning problem, just those that are well described by rule-based deduction and inference. Nonetheless, I'd be interested to see some more discussion on this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the informative and constructive comments! We will incorporate them to improve our work in the revised version.\n\n**[Causal tracing - W2]** First of all, as you mentioned, we believe this is not a weakness of our work and more of a general question about causal tracing. To explain here, causal tracing (also called activation patching) is a general method and there are a lot of different instantiations based on different views of the causal graph of transformers in different granularity levels. For example, one can do casual tracing at finer granularity (e.g., attention heads, attention QKV matrices, and MLP neurons) or coarser granularity (e.g., groups of hidden states across different layers and token positions). The case you mentioned could definitely be discovered through the method in principle, though this may be quite challenging (due to a large search space). Causal tracing also does work in larger-scale pretrained models and less rigid tasks (e.g., [1] and also many others). \n\n**[Connecting to general models - W3, Q1]** Yes we do admit that our study overall has some distance in connecting to large general models trained on broad corpora and tasks (Appendix H). Nevertheless, currently we have rather little understanding of even the simplest tasks/settings/models, and we believe our work lays a good foundation for subsequent work to study more general and realistic settings. A good concrete direction related to “generality” here would be to study how things behave when we do multi-task training , where the memorizing/generalizing circuits from different tasks are mixed together. For the problem you mentioned (hard to make a general model “grok” on every task in practice), there are also many interesting directions such as how to accelerate grokking (e.g., [2]) and how to better prepare the training data. For example, if the model architecture is good for a rule where it could systematically generalize, then we don’t really need to gather a lot of data for the rule - instead, we just need to curate a small but high-quality subset of data (presumably with high inferred/atomic ratios), and train the model extensively on this data till it gets the rule well (which shouldn’t cost a lot of compute). Then the model should be able to apply the rules when it is learning about other domains. Overall, there are many interesting follow-ups of our work on moving to practical settings, and we will add more discussions in the revised version. \n\n**[Parallel circuit - Q2]** This is determined by the nature of the tasks and the transformer architecture (which processes tokens in parallel). For the comparison task, the two underlying atomic facts have no order dependence on each other and both the query entities are present in the context tokens, while for the composition task, the second-hop fact depends on the (object entity of) first-hop fact, which is why the model can only learn a “serial” circuit.\n\n**[Label space - Q3]** Yes. We included these in Appendix D.2 (lines 575-577). We will add them to the main content in the revised version.\n\n**[Memorizing circuit for composition - Q4,5]** The atomic facts always need to be stored in the weights, and here since the model is very likely directly memorizing the inferred facts without utilizing the stored atomic facts, it separately stores all the inferred facts. To explain further our reasoning around the memorizing circuit: first, the state S[5,r1] encodes the bridge entity throughout grokking, however, its causal connection with the prediction state is very low (Figure 4c, Appendix D.1) in the beginning (which grew significantly during grokking). Indeed we cannot rule out the possibility that the model is somehow using the bridge entity information in some other hidden way elsewhere before grokking, but this is unlikely since then the model has to “force itself” not to use the state S[5,r1] which we know already encodes the bridge entity very well. One future direction to check this further is to understand in finer granularity the circuits and their evolvement throughout grokking, which may also help understand deeper the findings here and inspire theoretical investigations.\n\n\n**[Writing suggestions]** Thank you for the suggestions;  we will incorporate them to improve the writing in the revised version. For the ratio, we mean the percentage of the 300 random examples where performing causal intervention changes the target prediction. More details regarding the circuit discovery are included in Appendix D. Yes, by pruning we mean intervening (via the perturbed run), and we prune an edge/node if pruning it negligibly affects the accuracy. By “lower layers retrieve the first-hop fact”, we mean the “left-right” component of the circuit (Figure 4a) where the input states of h and r connect to S[5, r1] which encodes the bridge (via logit lens). \n\n====\n\nReferences\n- [1] Lieberum et al. Does Circuit Analysis Interpretability Scale? Evidence from Multiple Choice Capabilities in Chinchilla. arXiv-23.\n- [2] Lee et al. Grokfast: Accelerated Grokking by Amplifying Slow Gradients. arXiv-24.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for the constructive comments. We have provided detailed responses individually, with additional figures/tables (referred to as the "added PDF") attached here. There are two recurring topics across the reviews which we would like to briefly reiterate here.

**[Synthetic nature of the study and connection to practice]** One clear limitation of our work is its synthetic and abstract nature (highly structured inputs, single-task setting, etc.), which we also admit and discuss in Limitations (Appendix H). Despite our further efforts on ensuring that the results are robust to different setups (Appendix B,C,E), there are still substantial distances from our settings to those in practice. However, we would like to emphasize that our current understandings of the models and their behaviors are very limited even for the idealistic settings we consider here, and we believe our findings and analyses in this work are highly non-trivial and lay a good foundation for future research to build upon, expand, and transition toward more realistic settings.

**[Abstractness of writing]** We admit that our current draft is overall a bit abstract and lacks some concrete illustrations, where it may take multiple passes or some time to think between the words at different places. This is primarily due to the limited space where we need to make the language compact and also move lots of experiments and details to the Appendix, which may have also caused certain confusions and misunderstandings. We believe that these could be easily resolved (and we will) in the revised version.

We look forward to further discussions!

Meta-review: The paper investigates the reasoning capabilities of transformers in two synthetic scenarios, measuring compositionality and the ability to compare items. The main takeaway is that reasoning capabilities on unseen cases (generalization) are more evident when the transformers enter the "grokking" phenomenon. Furthermore, they highlight how compositionality is harder than comparing.

Reviewers appreciated the paper's direction (one reviewer was enthusiastic and gave 8, but it's an outlier), but highlighted also a number of shortcomings such as overclaiming results (starting from the title and continuing with the tone of writing thorough the paper) and the limited experimental setting (no ablations). As it is, despite the positive average score, the paper is borderline. I still believe that it can add an interesting perspective to the conference, as it allows to quantify reasoning in some (limited) reasoning scenarios where one can check for latent confounders.  

I recommend acceptance *conditioned* on the fact that authors include all reviewers' feedback as promised, especially the one from reviewer jEH8, starting by amending the title.
Decision: Accept (poster)