Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization

Authors: boshi  wang
♠ xiang yue
yu  su
♠  huan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:32:00

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.

Main paper: 
 introduction: 
 Large language models (LLMs) have been shown deficient in implicit reasoning with their parametric memory of knowledge and rules. For example, a range of LLMs are found to be incapable of robustly composing internalized facts [48, 71] , and even GPT-4 [42] cannot adequately compare entities' attributes despite knowing them [1] .Deficiency in implicit reasoning has profound implications. It implies the models' limitations in inducing structured and compressed representations of facts and rules, which lead to redundant knowledge storage and difficulty in propagating knowledge updates [76] , and importantly, fundamentally impede the model from systematic generalization over knowledge [25] . While explicit verbalizations of reasoning steps (e.g., chain-of-thought rationales) can improve task performance [67, 64, 73, 55, 31] , they are not available during large-scale (pre-)training where the model's core capabilities are acquired [77, 29] .Is implicit reasoning doomed given that even the most capable models struggle? Can it be resolved by further scaling data and compute, or are there fundamental limitations of the transformer [62] that prohibit robust acquisition of this skill? Figure 1 : We find that transformers can learn to reason implicitly, but this skill is only robustly acquired through grokking, i.e., an extended period of training far beyond overfitting. Moreover, the transformer fails to systematically generalize for composition, yet succeeds for comparison. We conduct a mechanistic study into the model internals throughout grokking, which reveals distinct generalizing circuits across the two 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 a mixture of "atomic facts" and "inferred facts" (which are deduced from the atomic facts via a set of latent rules), resembling "axioms" and "theorems" in a formal system. To evaluate how well the model learns the rules, we test its ability to make novel deductions (i.e., completing unseen inferred facts) in both in-distribution (ID) and out-of-distribution (OOD) scenarios. 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 1 ), a phenomenon known as grokking [47] . We find that the speed of improvement in generalization correlates with the ratio between inferred and atomic facts in training, and depends little on the absolute size of the training data (Figure 2 ). This suggests a correction of prior explanations of grokking based on critical data size [33, 61, 78, 21] , in that it should instead be the critical data distribution that decides the characteristics of grokking. Our findings extend prior observations of the grokking phenomenon primarily in algorithmic and linguistic tasks [47, 38] to the domain of knowledge-based reasoning, and deepen our understanding of the grokking phenomenon.Moreover, we find that the transformer exhibits different levels of systematicity across reasoning types. While ID generalization is consistently observed, in the OOD setting, the model fails to systematically generalize for composition but succeeds in comparison (Figure 1 ). To understand why this happens, we conduct mechanistic analysis of the internal mechanisms of the model. The analysis uncovers the gradual formation of the generalizing circuit throughout grokking and establishes the connection between systematicity and its configuration, specifically, the way atomic knowledge and rules are stored and applied within the circuit. Our findings imply that proper cross-layer memory-sharing mechanisms for transformers such as memory-augmentation [54, 17] and explicit recurrence [7, 22, 57] are needed to further unlock transformer's generalization.Finally, to demonstrate the power and potential of parametric memory for complex reasoning, we show that for a reasoning task with a large search space, a fully grokked transformer can achieve near-perfect accuracy, while state-of-the-art LLMs like GPT-4-Turbo [43] and Gemini-1.5-Pro [16] based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation. 
general setup: 
 Training data & ID/OOD evaluation. As stated in §1, we are interested in whether transformers can induce and apply latent rules over knowledge implicitly in a generalizable way. We create a datagenerating process consisting of 1) sampling a set of basic atomic facts, and 2) using the atomic facts and latent rules to deduce inferred facts. To better characterize the level of generalization acquired by the model, we evaluate the model's in-distribution (ID) and out-of-distribution (OOD) performance. We prepare two separate sets of atomic facts: atomic ID and atomic OOD . Our training set includes all the atomic facts and a uniformly random portion of the inferred facts deduced from atomic ID , which we call train_inferred ID . For evaluation, (1) ID generalization aims to evaluate whether the model learns the latent rules correctly, by testing its ability to complete unseen inferred facts also deduced from atomic ID , which we denote by test_inferred ID . (2) OOD generalization aims to evaluate the systematicity [25] acquired by the model, namely, the ability to apply rules over knowledge regardless of its distribution. To do this, we test the model on the facts deduced from atomic OOD , denoted by test_inferred OOD .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 AdamW [34] with learning rate 10 -4 , 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.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 7. 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 in 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-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 on 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
composition-delayed generalization without systematicity: 
 We begin our investigation with composition, where a model needs to "chain" different pieces of facts, e.g., "Barack's wife is Michelle" and "Michelle is born in 1964", to successfully complete a compositional sentence, e.g., "Barack's wife is born in [1964] ". Prior work extensively studied whether transformer-based language models can perform implicit composition, and negative results are consistently reported [48, 1, 71] . Specifically, there exists a "compositionality gap" [48] , i.e., the frequency at which the model knows all the underlying basic facts but fails to compose them, which is considerable across different LLMs and does not decrease as models scale. Are transformers doomed to fail on such kind of reasoning, and if so, why?We focus on two-hop composition in this work. For atomic facts, we generate a random knowledge graph G consisting of |E| entities and |R| = 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 ID and atomic OOD (95%: 5%). The rule of (two-hopEQUATIONwhich is used to deduce the ID and OOD inferred facts from atomic ID and atomic OOD , 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 a unique token to each relation/entity by default, and also find that the results are robust to different tokenizations (details in Appendix C).We study the influence of the following two aspects on the model's learned behaviors:• Ratio between inferred and atomic facts. Grokking observed in ID generalization but not in OOD generalization. We find that the model can generalize to ID test examples, but high performance is only achieved through extended training far beyond overfitting, a phenomenon called grokking [47] . Specifically, the training performance saturates (over 99% 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 2 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 ϕ. We omit the other splits since for all settings, the training performance saturates quickly and the OOD test accuracy remains at zero as earlier. 4 It could be seen that the ratio ϕ strongly correlates with the speed of generalization. A very large ratio can push generalization to improve at a similar pace as the model fits the training data, reducing the need for extended training. Training data distribution, instead of training data size, qualitatively influences generalization behavior. When ϕ increases and |E| holds constant, the size of training data also gets larger. Prior studies hypothesize that training data size plays a central role in order for grokking to happen. In particular, previous work connects grokking with the notion of critical data size (CDS) [33, 61, 78, 21] , where it is hypothesized that CDS marks the shift from memorization to generalization (via grokking), and the speed of generalization improves as the training data further scales. However, results from our controlled experiments seem to contradict such a hypothesis. Figure 2 (b) shows the results of varying |E| with a fixed ϕ = 9.0, where we change the horizontal axis from optimization step to epoch for better visualization. 6 When fixing the ratio ϕ, the training data size does not qualitatively affect the model's generalization. Specifically, scaling the data affects neither the relative speed of ID generalization and training improvement (as seen by the rather constant "gap" between train_inferred ID and test_inferred ID 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] .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 still remain: what happens during grokking, why does it happen, and why do transformers struggle with OOD examples? Answering these questions requires a deeper understanding of (the changes in) the model's inner workings, which we investigate next.We analyze the internal mechanisms within the model via a combination of two prevalent approaches: logit lens and causal tracing. We apply our analysis to the setting with |E| = 2000, ϕ = 9.0 on 300 random examples from train_inferred ID . Logit lens. We interpret individual hidden states via logit lens [40, 15, 71] , where the activation is converted into a set of logits for each vocabulary token by multiplying with the output embedding matrix. We follow the recent practice [71] where the activation first goes through the transformer's final normalization layer before multiplying with the output embedding (Figure 3 , top right).EQUATIONCausal tracing. The transformer could be viewed as a causal graph [46] that propagates information from the input to the output through a grid of intermediate states, which allows for a variety of causal analyses on its internal computations [63, 35, 19, 65, 12] . For convenience, we will refer to a hidden state by S[i, a], where i is the layer index and a marks the role of the input token at the same position as the state (one of {h, r 1 , r 2 }). We illustrate our method in Figure 3 , where the hidden state of interest is S[4, r 1 ] and the target is the model's prediction state S[8, r 2 ]. There are in total three steps:1. The normal run records the model's hidden state activations on a regular input (h, r 1 , r 2 ).Note that since the model maintains perfect training performance throughout grokking, the final prediction is always the ground truth tail entity t.7 2. In the perturbed run, a slightly perturbed input is fed to the model which changes the prediction, where again the hidden state activations are recorded. For the perturbation, prior work has explored adding noise to the input [35] and replacing key tokens with semantically close ones [63, 12] . We adopt token replacement which avoids unnecessary distribution shifts [74] . Specifically, for the hidden state of interest, we replace the input token at the same position as the state to be a random alternative of the same type (e.g., r 1 → r ′ 1 ) that leads to a different target prediction (e.g., t → t ′ ). 3. Intervention. During the normal run, we intervene the state of interest by replacing its activation with its activation in the perturbed run. We then run the remaining computations and measure if the target state (top-1 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. What happens during grokking? To understand the underlying mechanism behind grokking, we track the strengths of causal connections and results from logit lens across different model checkpoints during grokking (the "start" of grokking is the point when training performance saturates). We observe two notable amplifications (within the identified graph) that happen during grokking. The first is the causal connection between S[5, r 1 ] and the final prediction t, which is very weak before grokking (Appendix D) and grows significantly during grokking (Figure 4 (b)). The second is the r 2 component of S[5, r 2 ] via logit lens, for which we plot its mean reciprocal rank (MRR) (Figure 4 (c)). Additionally, we find that the state S[5, r 1 ] has a large component of the bridge entity b throughout grokking (Figure 4 (c)). These observations strongly suggest that the model is gradually forming the second hop in the upper layers (5) (6) (7) (8) during grokking. This also indicates that, before grokking, the model is very likely mostly memorizing the examples in train_inferred ID by directly associating (h, r 1 , r 2 ) with t, without going through the first hop.Why does grokking happen? These observations suggest a natural explanation of why grokking happens through the lens of circuit efficiency [61] . Specifically, as illustrated above, there exist both a memorizing circuit C mem and a generalizing circuit C gen that can fit the training data. While C mem is learned first (which causes training performance to saturate quickly), C gen is relatively more efficient, in the sense that it could fit the data with a lower complexity. To see this, we can compare the amount of facts C mem and C gen need to store (denoted as N mem and N gen ) as a proxy for their complexity. 8 C mem stores both atomic facts and inferred facts in the weights. C gen (Figure 4(a) ) stores the atomic facts in the lower layers, and another copy of the atomic facts that appear as the second hop in the inferred facts in the upper layers. As the inferred/atomic ratio ϕ increases, N mem would increase rapidly while N gen increases slowly and is always bounded by two times the total amount of atomic facts, and hence, the relative efficiency of C gen increases. In the long run, the model will be incentivized to transition from C mem to C gen due to implicit bias of the optimization [53] and explicit regularization such as weight decay which prefers more efficient circuits, and the transition would happen faster as ϕ increases. This also explains why the training data size does not affect the speed of grokking, since solely increasing the size does not change the relative efficiency of C mem and C gen . The explanation also implies that a larger regularization factor should accelerate grokking (and vice versa), which we confirm by varying the degree of weight decay (Appendix E.1).Explaining and mitigating the deficiency in OOD generalization. The configuration of C gen also has another important implication: while the model does acquire compositionality through grokking, it does not have any incentive to store atomic facts in the upper layers that do not appear as the second hop during training. This explains why the model fails in the OOD setting where facts are only observed in the atomic form, not in the compositional form-the OOD atomic facts are simply not stored in the upper layers when queried during the second hop. 9 Such issue originates from the non-recurrent design of the transformer architecture which forbids memory sharing across different layers. Our study provides a mechanistic understanding of existing findings that transformers seem to reduce compositional reasoning to linearized pattern matching [10] , and also provides a potential explanation for the observations in recent findings that LLMs only show substantial positive evidence in performing the first hop reasoning but not the second [71] . Our findings imply that proper cross-layer memory-sharing mechanisms for transformers such as memory-augmentation [54, 17] and explicit recurrence [7, 22, 57] are needed to improve their generalization. We also show that a variant of the parameter-sharing scheme in Univeral Transformer [7] can improve OOD generalization in composition (Appendix E.2).Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced. • 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 to 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. 
comparison-systematic generalization via parallel circuit: 
 We have just shown that the vanilla transformer fails to achieve OOD generalization for composition, but is the vanilla transformer generally incapable of acquiring systematic implicit reasoning skills?We show that for comparison, a task where SoTA LLMs such as GPT-4 also struggle [1] , the vanilla transformer does have the capability to acquire systematic generalization, again through grokking. On the surface, it seems that the comparison task is no different than the composition task-both require retrieving and reasoning over two pieces of facts. However, as it turns out through our analysis, the comparison task emits a "parallel circuit" that is learned by the transformer during grokking, which allows atomic facts to be stored and retrieved in the same region and enables systematicity to happen.Setup. The comparison task involves comparing the attribute values of entities. We assume there are |E| = 1000 entities, |A| = 20 attributes and |V| = 20 ordinal values for the attributes. Each attribute a ∈ A has a label space {a < , a = , a > }, a set of relations specifying its comparative form. For example, an attribute age would have a < , a = , a > to be younger, contemporary, older, respectively.The atomic facts are (entity, attribute, value) triplets, where we assign a random value v ∈ V for each (e, a) ∈ E × A. Again, we randomly partition the atomic facts into atomic ID and atomic OOD (90%: 10%). The rules of comparison are:EQUATIONTake the attribute age as an example, the first rule means if the age of e 1 is smaller than the age of e 2 , then we can infer "In terms of age, the relation between e 1 and e 2 is younger". Each entity/attribute/value/label is assigned a unique token, and training/testing is done by having the model predict the last token (attribute value for atomic facts; comparative relation for inferred facts).Results & analysis. The findings here showcase transformer's ability to learn parallel solutions to seemingly sequential problems, akin to the findings in Liu et al. [30] where it is shown that transformers can learn "shortcuts" to automata. The difference in the acquired generalization across the two tasks that we study also emphasizes the need for controlled and mechanistic study on understanding the transformer's reasoning before making general claims on its limitations.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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 on our data, model and training configurations, and evaluation setup in §2,3,4,5 and the Appendix. 
the power of parametric memory for complex reasoning: 
 At the high level, our study so far paves the way towards better understanding and improving transformer's reasoning with parametric representation of knowledge and rules. But why is parametric memory practically important? Can we not simply enhance LLMs with non-parametric memory, e.g., by using their long-context modes and/or doing explicit retrieval, to solve the tasks at hand?We believe parametric memory has its unique capability to perform deep compression and integration of information for complex reasoning. To showcase the potential of parametric memory for complex reasoning, we create a difficult reasoning task with a large search space, and show that 1) it is far out of reach even for current SoTA models (e.g., GPT-4-Turbo [43] and Gemini-Pro-1.5 [16] ) based on non-parametric memory; 2) a fully grokked transformer can solve the task with near-perfect accuracy.Our task is a variation of the comparison task above where we use a simple way to massively expand the search space, based on an additional set of rules that are already contained within the task itself, namely, the (anti-)symmetry and transitivity of comparison: ∀e 1 , e 2 ∈ E, ∀a ∈ A, (a, e 1 , e 2 , a < /a = /a > ) =⇒ (a, e 2 , e 1 , a > /a = /a < ), ∀e 1 , e 2 , e 3 ∈ E, ∀a ∈ A, ∀y ∈ {a < , a = , a > }, (a, e 1 , e 2 , y) ∧ (a, e 2 , e 3 , y) =⇒ (a, e 1 , e 3 , y).(EQUATIONIn the original setting ( §4) of the task, for the OOD test set, one can simply retrieve the two OOD facts and compare the attribute values, which requires no further search. We change the task setting via the following. For each attribute, 1) we do not add the OOD atomic facts into training, and 2) we add a random portion of the comparisons between ID entities and OOD entities into training. We test the models on queries consisting of derivable (from all training facts) comparisons between OOD entities where any possible proof would involve rules from both Eqs.(2) and Eqs.(3). Consequently, answering a test query would require the model to successfully locate two ID bridge entities which can connect the two query entities into a proof (Figure 6 ). We select a balanced (by a < , a = , a > ) subset from these queries for evaluation. More details are included in Appendix F.The difficulty of such a task is two-fold. First, the search space is large. For example, on average, each query entity connects with more than 50 facts, and each bridge entity in the ground truth proof 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 retrieval-augmented setup (5.4K facts retrieved on average) where the two-hop neighborhoods of the two query entities are retrieved, which includes enough facts to deduce the answer. We also try both standard prompting where the model answers directly, and chain-of-thought (CoT) prompting where the model is prompted to verbalize the reasoning. We test GPT-4-Turbo and Gemini-Pro-1.5, where for GPT-4-Turbo we only test the retrieval-augmented setup due to context length limit.Table 1 : Results on the complex reasoning task. Direct/CoT: predict the answer directly/verbalize the reasoning steps. "+R": retrieval augmentation.Gemini-Pro-1. 1 , all models based on non-parametric memory fail badly, where the only setting that surpasses random guess (33.3%) is the retrieval-augmented setting with Gemini-Pro-1.5 and direct answer prediction. Intriguingly, LLMs perform worse (especially Gemini) when prompted to reason verbally. Through closer examinations, we find that with CoT, 70.7% of Gemini's responses conclude that the answer cannot be decided (which we treat as wrong since the answer can be decided) after a series of search steps. 11 More shockingly, most of the CoT rationales that achieve the correct final answer are actually wrong due to either hallucinating underivable facts or logical errors. This illustrates the current models' inability to reason deeply with non-parametric memory. On the other hand, a grokked transformer, which is trained extensively on the given facts to compress and integrate the information to the extreme, could achieve near-perfect accuracy. By examining the model, we find that it acquires the same generalizing circuit as in Figure 5 (a), and remarkably, even though not explicitly encouraged/trained to do this, the model successfully infers most of the OOD entities' attribute values by integrating the observed training facts (Appendix F).Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?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. • The experimental setting should be presented in the core of the paper to a 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, or as supplemental material. 
related work: 
 Knowledge and reasoning in language models. Numerous work finds that transformer language models, even SoTA ones such as GPT-4, struggle in implicit reasoning over their parametric knowledge [56, 23, 52, 51, 48, 1, 71] , suggesting their limitations in inducing structured and compressed representations of facts and rules during training. A series of efforts try to understand transformer's knowledge and reasoning through controlled experiments [24, 49, 10, 66] , which is also our focus. We find that transformers can learn implicit reasoning over knowledge through grokking, and characterize the connection between the acquired systematicity level and the inductive bias of transformer."Chain-of-Thought" and verbalized reasoning. A series of studies prompt/fine-tune language models to verbalize (i.e., generate) the intermediate knowledge and reasoning steps [67, 64, 73, 55, 31, 72] during inference, which has been shown to improve performance especially for large models with strong generation capabilities. There are also theoretical results showing the benefits of such verbalizations [11, 27] . Our focus here on implicit reasoning is orthogonal, and it is an interesting open problem to have principled understandings of the role of such verbalizations in reasoning problems, and also develop methods that can decide the appropriate balance between implicit and explicit reasoning to handle challenging problems with large intrinsic complexity. Relatedly, recent work also finds that explicit verbalizations could be a useful medium for teaching models to reason implicitly via distillation or curriculum [9, 8] .Grokking is first discovered by Power et al. [47] on a set of small algorithmic reasoning tasks. [13] . We apply causal tracing [63, 35, 19, 65, 12] and logit lens [40, 15, 71] to discover interpretable circuits in the model to understand the grokking process and how/why generalization happens.Parametric and non-parametric memory. Our focus in this work is on parametric memory in language models, and an orthogonal direction is to enhance models with non-parametric memory, such as extending the effective context length [70, 43, 14, 16] and augmenting with retrieval [18, 26, 3, 58, 75, 37] . The two types of memory are largely complementary to each other-parametric memory has its unique ability to compress and integrate information but is also inevitably lossy and subject to hallucination, while non-parametric memory is lossless and could also provide attribution. Similarly for humans-a human acquires expertise in a domain by acquiring and structuring knowledge in the brain (parametric), but he/she also wouldn't memorize all pieces of details and could refer to the source when necessary (non-parametric). How to decide the tradeoff between parametric and non-parametric memory (or, how to define the objective for such a tradeoff) is another interesting open problem for future work. 
conclusion: 
 We find that transformers are capable of learning to implicitly reason over parametric knowledge, however, such a skill is only robustly acquired through extended training far beyond the point of overfitting, or grokking. Mechanistic analysis into the model's internals reveals the configuration and gradual formation of the generalizing circuit, and also explains the different levels of systematicity the model acquires across tasks. These findings guide data and training setup to better induce implicit reasoning, and suggest potential improvements to the transformer architecture to further unlock its generalization. We conclude by showcasing the unique power of parametric memory on a challenging reasoning task with a large search space.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? 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 to randomness in the setup. In addition, our results exhibit strong patterns from controlled variations in the parameter settings, which we believe is a 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 error bars, confidence 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 are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 of range (e.g. negative error rates).• 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. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?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 as 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 make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
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 applications or deployments, and hence we believe there are no negative societal impacts.Guidelines:• The answer NA means that there is 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 or 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 a 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 disinformation. 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. 
 
 Appendix: Scope of our task formulation. As mentioned in §1, we formulate the implicit reasoning problem as induction and application of inference rules from a mixture of atomic and inferred facts. This may not apply to the full spectrum of reasoning which has a range of different types and meanings [20] . Still, our formulation could capture the nature of a wide range of reasoning problems, and crucially, we believe that it is a good conceptualization of certain aspects of language model (pre-)training, where the model needs to both memorize the "atomic" world knowledge and also induce generalizable rules from the massive amount of records of human activities, which allow the model to connect and reason over knowledge, and ultimately help with humans.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.We run the experiments on composition with larger model scales with |E| = 2000 and ϕ ∈ {5.4, 9.0, 18.0}. The results are shown in Figure 7 . Overall, it could be seen that scaling up the model won't qualitatively change the model's generalization behaviors, and the main pattern is that larger models converge in fewer optimization steps, which shares with prior findings [60, 28] . 5000In the experiments in our main content, tokenization is done by having a unique token for each entity. 12 This is different from how real-world entities are typically tokenized-in practice, entities are usually multi-token, and different entities could share tokens at different positions. We investigate the effect of tokenizations on the composition task by having two tokens for each entity (resembling the first name and last name of a person) in the setting with |E| = 2000 and ϕ = 12.6. We generate a set of unique first names and a set of unique last names with equal size from which the two tokens for each entity are randomly chosen (we make sure each entity gets a unique ordered pair of tokens). We define token multiplicity to be the number of entities that share the same first/last name. For example, when the size of the set of first/last names is 50, the token multiplicity would be 2000/50 = 40. It can be seen that a larger token multiplicity would delay the generalization, which is expected to a certain degree since the scale of the model is effectively smaller due to having fewer tokens in the vocabulary. Nevertheless, ID generalization always happens. We also run linear probing on S[5, r 1 ] throughout training to predict the second token of the bridge entity b in the setting with token multiplicity 40, 13 the results are shown in Figure 8 (b). It can be seen that the second token of b can be perfectly decoded from S[5, r 1 ] after grokking, and the decodability improves throughout grokking. This suggests that for the multi-token case, the model is additionally storing the second token of b into S[5, r 1 ] throughout grokking, which may be another factor that further delays the speed of grokking. These results also share with recent findings that in many cases, tokens beyond the immediate next token are linearly encoded in the hidden states [69, 44, 2, 4] .In summary, different tokenizations affect the results in rather expected ways, and do not influence our main findings and conclusions. We run causal tracing on hidden states in layer 1-7 and every position, where the target is the final prediction state S[8, r 2 ]. The changes in the strengths are monotone and smooth, and we show in Figure 9 the strengths for the model checkpoint at the start and end of grokking, and also their difference (same as Figure 4 (b)). We also find that after grokking, the state S[5, r 2 ] (which encodes r 2 ) is not affected by perturbing the input h or r 1 .Deriving the generalizing circuit. Starting from the 9 states in layers 0, 5, 8 we can directly eliminate S[8, h] and S[8, r 1 ] since they have no computational paths connecting to S[8, r 2 ]. S[5, h] can be eliminated as could be seen by Figure 9 (c). The connections from S[0, h] and S[0, r 1 ] to S[5, r 2 ] could be eliminated as mentioned earlier. Figure 13 shows the ID generalization performance when varying the degree of weight decay (|E| = 2000 and ϕ = 9.0). It can be seen that a larger weight decay can improve the speed of grokking, and vice versa. We share the parameters of the first 4 layers and the last 4 layers, similar as in Universal Transformer [7] . This would allow the model to share the knowledge in the upper and lower layers. The results on the setting with |E| = 2000 and ϕ = 12.6 are shown in Figure 14 . It could be seen that the parameter-sharing scheme can unlock OOD generalization, even though it is gained much more slowly than ID generalization during grokking.Figure 15 includes the result for ϕ ∈ {3.6, 7.2, 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.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 atomic facts, 437K (ID, ID) comparisons, and 108K (ID, OOD) comparisons, altogether 563K facts (28K on average for each attribute). For translating the facts into natural language for testing LLMs with non-parametric memory, we always use the attribute age 14 (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} is {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 a random generator. 15 We also try different mappings (e.g., unique IDs) and templates, and find the results to be consistent. All (retrieved) facts are 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 in Figure 16 . It could be seen that, during grokking, the model gradually locates the ground truth attribute values of the query entities (note that the model is not explicitly encouraged or trained to do this), allowing the model to solve the problem efficiently with near-perfect accuracy.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.• 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 to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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.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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Definitions of ID/OOD are introduced in §2.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 use a fixed batch size.For convenience, when we refer to a state as a token, we mean the top token of the state via logit lens.While the circuits also consist of other components, they pale in comparison as the number of facts scales.We verified that in the OOD setting, S[5, r1] and S[5, r2] encode b and r2 respectively as in the ID case.This could also be thought of as performing a retrieval step into the memory based on the attribute.The ratio drops a bit to 58.7% when augmenting with retrieval.Preliminary experiments show that tokenization of the relations does not exhibit notable impacts, which is expected since relations are always explicitly given.Recall that S[5, r1] is the state that encodes the bridge entity in the generalizing circuit (Figure4(a))Recall that we test each attribute separately by grouping the facts/queries.https://pypi.org/project/names/ 


Decision: 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.

Reviews: review: paper_summary: paper_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"). 

The 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. 

Causal 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.

Some 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.
strength_weakness: strength_weakness: 1. Overclaiming: Framing should be adjusted to more clearly reflect the work done.  
- 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.
- 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).
- 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.
- Unfair comparison to parametric memory, see comment 35
2. 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.
questions: 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.
2. line 25, English: an implication is not an impact or vice versa, rephrase to align
3. line 28 define systematic generalisation 
4. 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.
5. 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)?
6. fig 1 last sentence shouldn't be here - caption is for explaining figure, not promoting other parts of paper
7. 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.
8. 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)
9. 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.
10. 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.
11. lines 122,3: "ID accuracy", "other splits" - be more explicit ("test ID accuracy", "train ID and test OOD") else hard to read
12. 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.
13. 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". 
14. 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).
15. fig 2a: explicitly state |E| in caption/image/somewhere. did you check this for multiple |E| and get the same results?
16. fig 2 caption: at this point in the paper,  test_inferred_id has not been introduced, define things on time
17. 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.
18. 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)
19. 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?
20. 175-176 "the ratio of such alterations (between 0 and 1)" - be more explicit, dont understand
21. 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.
22. 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)
23. footnote pg 6: too vague, cant get anything from this
24. 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.
25. 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?
26. 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. 
27. 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.
28. 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.)
29. lines 288-289 took a while to parse, be more explicit
30. lines 294-296 be more explicit about whether talking about atomic or inferred facts
31. line 296 "no surface form clues" - explain/be more explicit
32. 297-298 "unlike most conventienal QA benchmarks where.." - do you have any support for this?
33. 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
34. 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.
35. 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: limitations: concluded guidelines on data/training/architecture settings not sufficiently supported, comparison to non parametric memory unfair
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. Dependence on Synthetic Data. The experiments primarily use synthetic data, which may not fully capture the complexities and nuances of real-world data.
questions: questions: 1. Can you provide the detail of your data sizes ? i can find the statics and samples of your data.
2. any other reasoning task can confirm your assumptions ?
limitations: limitations: See in Questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	The paper does not provide how to tokenize the whole sentence in (1)(2) and how to evaluate the loss function. 
2.	The beginning epoch of the abscissa in Fig.1 needs clarification. Why the first point in Fig1b Train accuracy is about 0.9?
3.	Since the accuracy is a discontinuous metric, it may sometimes mislead for indicating grokking. The loss for each training needs to show.
4.	The 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.
5.	The term 'same type' in line 172 needs clarification. Does it refer to data within the same entity set or add noise?
6.	How to calculate the ratio in line 175 needs clarification.
7.	Fig.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).
8.	The 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). 
9.	In 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.
questions: questions: 1.	In line 203, what does 'stores... in the weights' mean? Does it imply that during the memory phase, the logit lens is not 'b'?
2.	Is 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?
3.	What will happen when we input out-of-distribution data in Figure 4a? Will the logit lens at S[5, r1] not show 'b'?
4.	In my opinion, the failure of the OOD in the composition is due to the unseen embedding of OOD data during the training.
limitations: limitations: See weaknesses.
paper_summary: paper_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. 

The 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. 

The 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. 

Besides, 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.
strength_weakness: strength_weakness: My major concerns are two-fold:

1. 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.

2. *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.

[1] PROGRESS MEASURES FOR GROKKING VIA MECHANISTIC INTERPRETABILITY, https://arxiv.org/abs/2301.05217 .

[2] Do Machine Learning Models Memorize or Generalize? https://pair.withgoogle.com/explorables/grokking/ .
questions: questions: (1) For my major questions and concerns, please refer to the "Weaknesses" part.

(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?)

(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.

(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?

(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.

To 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.

[1]: Understanding and Patching Compositional Reasoning in LLMs, https://arxiv.org/abs/2402.14328 .
limitations: 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.

[1] Logit Lens: https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens .

[2] Eliciting Latent Predictions from Transformers with the Tuned Lens, https://arxiv.org/abs/2303.08112 .
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

- 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.

- 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.
questions: questions: Questions
- Do you think the recommendations to make finding the generalising circuit easier will transfer to LLM training? 
- Do you have an explanation for why the model learns a parallel pathway for the comparison task and not the composition task?
- 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?
- Why does the memorising circuit need to also store the atomic facts and not just the inferred facts for the composition task?
- 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?

Writing suggestions:
- Can you give an example of an actual input and output pair you train the models on in the main text?
- 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: limitations: Addressed in the appendix.

Comments: title: Thanks for the authors' response
Comment on the Synthetic nature of the study and connection to practice
Rebuttal by Authors (Part 3)
Rebuttal by Authors (Part 2)
comment: Thank you very much for your support of our efforts. Also many thanks for clarifying the reference. It is a very insightful analysis for in-context learning, which is very different from our focus (i.e., implicit reasoning), though. The progress measures defined in the paper are inspiring. It'd be very interesting follow-up work to investigate similar quantitative measures and visualize them during the grokking process of our studied tasks. We will discuss this work and point out related future directions in the revised version. Thanks agains for your efforts reviewing our work and for all the great comments!
Thank you very much for your appreciation of our work and for other constructive comments earlier!  In our revised version, we will make things clearer based on feedback from all reviews. We believe our work will inspire more future work in this space, and thank all reviewers again for their time and effort reviewing our work!
Dear authors,

Thanks for the detailed response, which resolve part of my conerns. I maintain my initial score to this paper.

Reviewer smMo.
I would like to add that I strongly agree that even though the work is synthetic in nature, this is a fair limitation (not weakness) for a paper to have as long as it's discussed, which the authors do. In a common pursuit of trying to understand how complex models work, controlling the setup to investigate phenomena is a very important part of research that is complementary to approaches that look at the often more complicated and fuzzy realistic non-synthetic tasks and models. The authors do that in a way that yields testable predictions, which they confirm in supplementary experiments in the appendix (e.g. higher weight decay makes grokking happen faster, and this is because the generalising circuit is more efficient). All in all, I believe this paper gives important insight into how transformers learn to do a type of reasoning.
Thanks for answering my questions. I am going to raise to an 8 and will strongly recommend accepting this paper. It very neatly shows how the grokking phenomenon works.
I refer to the paper 'The mechanistic basis of data dependence and abrupt learning in an in-context classification task' (Reddy 2023). Several hidden progress measures are implemented to show the grokking of ICL (In context learning) in that paper. If there is any relationship with your results?

And I have no further inquiries. Based on the responses provided, I will be revising my evaluation favorably.
- Q25: It’s not accurate to conclude “many layers are redundant” from our findings, since it is not the general case that one layer can perform one step, e.g., many times it may take multiple layers to perform one step. Investigating whether we could actually compress the computations into fewer layers is interesting future work.
- Q26: We did confirm that the same circuits arise across the different settings we experimented with. We will add relevant details in the revised version.
- Q28: We believe it is implied from our writing (e.g., Section 2) that all the given examples/facts are correct ones. We will make this explicit in the revised version.
- Q31,32: Here we mean that the query itself does not leak information about the ground-truth proof structure. For conventional QA benchmarks such as NaturalQuestions, TriviaQA, HotpotQA, MuSiQue, etc., the ground truth proof structure can mostly be obtained already by directly parsing/decomposing the query, which is not the case for our constructed task.
- Q33: 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). It drops a bit to 58.7% when augmenting with retrieval. One typical example of such cases is included in Table 1 in the added PDF. We will add more discussions and examples in the revised version.

**[Suggestions for writing - Question 2,3,4,6,7,11,16,23,27,29,30]** Thank you for pointing out these issues and suggestions for our writing. These are greatly helpful for improving the paper. We do admit that our current draft is overall a bit abstract and lacks some concrete illustrations, but this is mainly due to the limited space where we need to make the language compact. We still had to put a lot of experiments and details and discussion of limitations/related work in the Appendix, which may have also caused certain confusions and misunderstandings above. These could be rather easily resolved (and we will) in the revised version.

====

Citations

[1] Zhong et al. Training Language Models with Memory Augmentation. EMNLP-22.

[2] Min et al. Nonparametric Masked Language Modeling. Findings of ACL-23.
**[Tokenizations & rigid input structures - Weakness 1, Question 9]** We actually explicitly mentioned the tokenization scheme in lines 103-104, which immediately implies the input representations. Additionally, we investigated alternative tokenizations in Appendix C to ensure the robustness of our findings to tokenizations. We also performed other preliminary experiments (e.g., different templates) not included in the paper where the number of tokens is larger, and find the results to be consistent. We will add relevant content in the revised version. On the other hand, since we are training from scratch for full control and rigor of the results, the model does not have any pretrained language understanding and hence the inputs we use are indeed (and need to be) quite rigid without natural language variations (like those in real corpora). We do admit that this is a limitation (Appendix H), but note that since we have rather poor understandings even for the simple settings, we believe our work lays a solid foundation for future work to build on and gradually move to more realistic (non-rigid) settings.

**[Overlaiming]** Indeed we are only studying two reasoning types in synthetic settings, but on the other hand, our evaluations are very clean and rigorous, and we believe our experiments with the complex reasoning task (Section 5) where the grokked transformer is shown capable of performing very difficult reasoning are arguably strong and exciting findings. We will consider adjusting the title to a less strong one.

**[Other technical questions]** We respond to other relatively minor technical questions here.
- Q5: In some sense this is more of a problem of definition; but generally we feel OOD generalization (or, systematicity to be specific) is a very desirable property for models to have but not a requirement for saying that the models can reason.
- Q10: Indeed the KG here is an abstract one and doesn’t reflect such constraints. One future direction is to make the setting more realistic by adding more natural variations and constraints into the setting.
- Q15: Thank you for the suggestion. Here |E| is 2000 as in the earlier setting. We did perform experiments across different |E| and found the results to be consistent (Figure 2 in added PDF). We will include these in the revised version.
- Q18: The second coordinate here could be regarded as the marker for position but it is also the abstract variable name for the corresponding token, which could be more informative and easier to conceptualize. Indeed it obviously should not be the actual token as you mentioned, and we will fix the writing in the revised version.
- Q19: Yes, we do take care of this, and also made sure to choose from the perturbed examples that end with a different tail entity.
- Q20: Here we mean the ratio of examples (from the 300 random examples we studied) where the target prediction is altered. We will make this explicit in the revised version.
- Q21: Indeed we are not able to explicitly show the evidence for the memorizing circuits (which is rather difficult and may require new analysis techniques), but we have convincing evidence for such a hypothesis (lines 188-196 and Appendix D). To reiterate here, the causal strength between S[5, r1] (which is the state that encodes the bridge entity throughout grokking, Figure 4c) and the prediction is very weak at the beginning of grokking (Figure 9a), and grew significantly during grokking (Figure 4b). This suggests that in the beginning, the model is very likely not traversing through the bridge entity when predicting the tail, and hence directly associating h, r1, r2 with t, which is our definition of “memorization”.
- Q22: Here our purpose is to connect with prior work on implicit regularizations, which could be a potential direction for explaining our observations. We are not trying to be very technically solid here, since even though there is a wide range of prior results along this line, they are all in somewhat restricted settings and no general mathematically proven theory of implicit regularization for deep networks exists to our best knowledge. The specifically cited work in our paper also mostly focuses on linear and restricted ReLU network settings, which are very different from multi-layer transformers. Again, we don’t think deep technical discussion is necessary here for our purpose, and our work could potentially inspire further theoretical investigations along this line.
- Q24: We did compute the logit lens results for the OOD setting, and found that S[5, r1] encodes the bridge entity (MRR 0.98) and S[5, r2] encodes the second relation (MRR 1.0) as in the ID case, which are strong indications that the lower layers in the OOD setting are performing similar roles as in the ID setting, including storing the OOD atomic facts. We will include these results and expand the discussion in the revised version.

(continued in comment "Part 3")