How Transformers Reason: A Case Study on a Synthetic Propositional Logic Problem

Authors: guan zhe hong
nishanth  dikkala
enming  luo
cyrus  rashtchian
xin  wang
rina  panigrahy
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-26 11:46:44

Abstract: Large language models (LLMs) have demonstrated remarkable performance in tasks that require reasoning abilities. Motivated by recent works showing evidence of LLMs being able to plan and reason on abstract reasoning problems in context, we conduct a set of controlled experiments on a synthetic propositional logic problem to provide a mechanistic understanding of how such abilities arise. In particular, for a decoder-only Transformer trained solely on our synthetic dataset, we identify the specific mechanisms by which a three-layer Transformer solves the reasoning task. In particular, we identify certain "planning" and "reasoning" circuits which require cooperation between the attention blocks to in totality implement the desired reasoning algorithm. To expand our findings, we then study a larger model, Mistral 7B. Using activation patching, we characterize internal components that are critical in solving our logic problem. Overall, our work systemically uncovers novel aspects of small and large transformers, and continues the study of how they plan and reason.

Main paper: 
 introduction: 
 Language models using the transformer architecture [Vaswani et al., 2017] have shown remarkable capabilities on many natural language tasks [Brown et al., 2020 , Radford et al., 2019b] . Trained with causal language modeling wherein the goal is next-token prediction on huge amounts of text, these models exhibit deep language understanding and generation skills. An essential milestone in the pursuit of models which can achieve a human-like artificial intelligence, is the ability to perform human-like reasoning and planning in complex unseen scenarios. While some recent works using probing analyses have shown that the activations of the deeper layers of a transformer contain rich information about certain mathematical reasoning problems [Ye et al., 2024] , the question of what mechanisms inside the model enables such abilities remains unclear.While the study of how transformers reason in general remains a daunting task, in this work, we aim to improve our mechanistic understanding of how a Transformer reason through simple propositional logic problems. For concreteness' sake, consider the following problem:Rules: A or B implies C. D implies E. Facts: A is true. B is false. D is true.Question: what is the truth value of C?An answer with minimal proof is "A is true. A or B implies C; C is true."The reasoning problem, while simple-looking on the surface, requires the model to perform several actions that are essential to more complex reasoning problems, all without chain of thought (CoT).Before writing down any token, the model has to first discern the rule which is being queried: in this case, it is "A or B implies C". Then, it needs to rely on the premise variables A and B to the locate the relevant facts, and find "A is true" and "B is false". Finally, it needs to decide that "A is true" is the correct one to invoke in its answer due to the nature of disjunction. It follows that, to write down the first token "A", the model already has to form a "mental map" of the variable relations, value assignments and query! Therefore, we believe that this is close to the minimal problem to examine how a model internalizes and plans for solving a nontrivial mathematical reasoning problem where apparent ambiguities in the problem specification cannot be resolved trivially.To understand the internal mechanisms of how a transformer solves problems resembling the minimal form above, we perform two flavors of experiments. The first is on shallow transformers trained purely on the synthetic propositional logic problems. This enables a fine-grained analysis in a controlled setting. The other set of experiments are on a pre-trained LLM (Mistral-7B), where we primarily rely on activation patching to uncover necessary circuits for solving the reasoning problem, including specialized roles of certain components. At a high level, we make the following discoveries based on our two fronts of analysis:1. We discover that small transformers, trained purely on the synthetic problem, utilize certain "routing embeddings" to significantly alter the information flow of the deeper layers when solving different sub-categories of the reasoning problem. We also characterize the different reasoning pathways: we find that problems querying for reasoning chains involving logical operators typically require greater involvement of all the layers in the model.2. We uncover properties of the circuit which the pretrained LLM Mistral-7B-v0.1 employs to solve the minimal version of the reasoning problem. We find four families of attention heads, which have surprisingly specialized roles in processing different sections of the context: queried-rule locating heads, queried-rule mover heads, fact-processing heads, and decision heads. We find evidence suggesting that the model follows the natural reasoning path of "QUERY→Relevant Rule→Relevant Fact(s)→Decision".We discuss related works and scope of this work in detail in Appendix A. 
problem setting: 
 In this section, we present the core properties of the synthetic propositional logic problem which shall be the data model of this paper. We delay finer details and more examples of the problem to Appendix B.Our problem follows an implicit causal structure, as illustrated in Figure 1 . The structure consists of two distinct chains: One containing a logical operator at the end of the chain, and the other forming a purely linear chain.K D V E Logical operator A P T SFigure 1 : Synthetic data model. The causal structure has two chains: one with a logical operator (LogOp) at the end and the other being purely a linear causal chain. This example is the length-3 case.We require the model to generate a minimal reasoning chain, consisting of "relevant facts", proper rule invocations, and intermediate truth values, to answer the truth-value query. Consider an example constructed from the causal graph in Figure 1 , written in English:• Rules: K implies D. D or E implies A. V implies E. T implies S. P implies T. • Facts: K is true. P is true. V is false.• Query: A.• Answer: K is true. K implies D; D is true.D or E implies A; A is true.In this example, the QUERY token A is the terminating node of the OR chain. Since any true input to an OR gate (either D or E) results in A being true, the minimal solution chooses only one of the starting nodes from the OR chain to construct its argument: in this case, node K is chosen.3 Mechanisms of planning and reasoning: a case study of the length-3 problemIn this section, we discuss the internal mechanisms of a small transformer trained purely on the synthetic problem. While there are many parts of the answer of the transformer which can lead to interesting observations, due to space limitations, we primarily focus on the model's "mental process" for producing the most important part of the answer, namely the first token. To further justify this choice, we find that on our problem, a model's full-answer accuracy strongly correlates with its accuracy of the first answer token, as detailed in Figure 4 in the Appendix.Architecture choice. We study a decoder-only attention-only transformer closely resembling the form of GPT-2 [Radford et al., 2019a] . We discuss training and architecture details in the Appendix.We select the smallest transformer that can achieve 100% accuracy (or sufficiently close to it) to initiate our analysis, a 3-layer 3-head variant. 
empirical observations: 
 We begin our analysis with mechanisms that are universal to how the model plans and reasons for predicting the first token. Then we describe the mechanisms that only arise when the model needs to deal with specific situations. We discuss the main observations here, and leave the quantitative details to Appendix D.Mental notes at the QUERY position. The QUERY token is likely the most important token in the context: it determines which chain is being queried. The transformer makes use of this token in its answer in an intriguing way.Observation 1: chain-type disentanglement at QUERY. We observe that, the second layer's attention block exhibit disentanglement in its output direction dependent on whether it is the linear chain that is being queried. Intriguingly, the third layer's attention heads place greater than 90% of their attention weights on the QUERY position on average when the linear chain is queried.Based on Observation 1, we hypothesize that given a chain type (linear or LogOp), there exists certain directions at the second attention block which somehow change the behavior of the third attention block: attracting its attention to QUERY when it is the linear chain, and pushing its attention away from QUERY when it is the LogOp chain. We confirm the existence and role of this "routing" signal.Observation 2: existence of an abstract "routing signal". We compute the average of the second attention block's output on 1k samples whose QUERY is for the linear chain, which we denote as h route . There are two interesting properties of this embedding direction:1. (Linear→LogOp intervention) We generate 500 test samples where QUERY is for the linear chain. Subtracting the embedding h route from the second attention block's output results in the model outputting the correct first token for the LogOp chain of the problem 100% of the time on the test samples. In other words, the "mode" in which the model reasons is flipped from "linear" to "LogOp". 2. (LogOp→linear intervention) We generate 500 test samples where QUERY is for the LogOp chain. Adding h route to the second attention block's output causes the three attention heads in layer 3 to focus on the QUERY position: greater than 99% of the attention weights are on this position averaged over the test samples. In this case, however, the model does not output the correct starting node for the linear chain on more than 90% of the test samples.It follows that there indeed exists an abstract embedding direction inside the transformer which significantly changes the information flow depending on the chain type being queried.Linear chain. At this point, it is clear to us that, when QUERY is for the linear chain, the third layer mainly serves a simple "message passing" role at the QUERY position. A natural question arises: does the input to the third layer truly contain the information to determine the first token of the answer, namely the starting node of the linear chain? The answer is yes. Attention head type (Layer idx, head idx)Figure 2 : High-level properties of Mistral-7B's reasoning circuit. The (chunks of) input tokens are on the left, which are passed into the residual stream and processed by the attention heads. We illustrate the information flow manipulated by the different types of attention heads we identified to be vital to the reasoning task.Observation 3: linearly-decodable linear-chain answer at layer 2. We train an affine classifier with the same input as the third attention block, with the target being the start of the linear chain; the training samples only query for the linear chain, and we generate 5k of them. We obtain a test accuracy above 97% for this classifier (on 5k test samples), confirming that layer 2 already has the answer at the QUERY position.LogOp chain: partial answer in layers 1 & 2 + refinement in layer 3. To predict the correct starting node of the LogOp chain, the model employs the following strategy:1. The first two layers encode the LogOp and only a "partial answer". More specifically, we find evidence that (1) when the LogOp is an AND gate, layers 1 and 2 tend to pass the node(s) with FALSE assignment to layer 3, (2) when the LogOp is an OR gate, layers 1 and 2 tend to pass node(s) with TRUE assignment to layer 3. 2. The third layer, combining information of the two starting nodes of the LogOp chain, and the information in the layer-2 residual stream at the ANSWER position, output the correct answer.We delay the full set of evidence for the above claims to Appendix D.2.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] Justification: This is not a theory work. 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 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. 
the reasoning circuit in mistral-7b: 
 We now turn to examine how a pretrained LLM, namely Mistral-7B solves this reasoning problem. We choose this LLM as it is amongst the smallest accessible model which achieves above 70% accuracy on (a minimal version of) our problem. We present a hypothesis for the reasoning circuit inside the model for predicting the crucial first token of the length-2 problem in Figure 2 , and provide evidence relying on a popular technique in mechanistic interpretability, activation patching.We describe the main properties of the reasoning circuit inside the model for this prediction task in Figure 2 . At a high level, there are several intriguing properties of the reasoning circuit of the LLM:2 1. Compared to the attention blocks, the MLPs are relatively unimportant to correct prediction. 2. There is a sparse set of attention heads that are found to be central to the reasoning circuit:• (Queried-rule locating head) Attention heads (9,25;26), (12,9), (14,24;26) locate the queried rule using the QUERY token, and stores this information at the QUERY position. • (Queried-rule mover head) Attention heads (13,11), (15,8) move QUERY and the queried-rule information from the QUERY position to the ":" position. • (Fact processing heads) Attention heads (16,12;14), (17,25) locate the relevant facts, and move information to the ":" position.• (Decision head) Attention head (19, 8) , relying on the aggregated information, makes a decision on which token to output.We only discuss high-level intuitions and results here due to space limitations, and delay the full set of experiments and their interpretations to Appendix E.Intuitively speaking, to support our hypothesis for the reasoning circuit employed by Mistral-7B to solve the reasoning problem, we rely on activation patching to discover the attention heads which have the greatest influence on the model's output distribution (recall that the MLPs are not as important in this problem). We combine such "causal-mediation" evidence with inspections on these heads' attention patterns. This leads to the set of evidence that is (partially) visualized in Figure 3 below. 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 detail our experimental procedures, including architecture choice, training procedure, hyperparameters and metrics used in 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 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. (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. 
conclusion: 
 We studied the reasoning mechanisms of both small transformers and LLMs on a synthetic propositional logic problem. We analyzed a shallow decoder-only attention-only transformer trained purely on this problem as well as a pretrained Mistral-7B LLM. We uncovered interesting mechanisms the small and large transformers adopt to solve the problem. For the small models, we found the existence of "routing" signals that significantly alter the model's reasoning pathway depending on the sub-category of the problem instance. For Mistral-7B, we found four families of attention heads that implement the reasoning pathway of "QUERY→Relevant Rule→Relevant Facts→Decision". These findings provide valuable insights into the inner workings of LLMs on mathematical reasoning problems.NeurIPS Paper Checklist Justification: We explicitly discuss all the limitations of this work in the conclusion section, from architecture choice, to data model limitations, to limitations in our analysis. 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.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: [No] Justification: Our experiments rely on internal code, which can be difficult to release.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) 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 a new open-source benchmark). • 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 discuss such details in the Appendix.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?Answer: [Yes] Justification: The research conforms to the Code of Ethics of Neurips.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). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss potential broader impact in section F.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.• 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 or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer:[NA] Justification: The paper poses no such risk.• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: 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?Answer: [NA] Justification: The paper does not use existing assets. Guidelines:• The answer NA means that the paper does not use 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 is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: The paper does not release new assets Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of 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 is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.Appendix / supplemental material A Related works, and scope of this work Mechanistic interpretability. Our work falls in the area of mechanistic interpretability, which aims to understand the mechanisms that enable capabilities of the LLM; such studies involve uncovering certain "circuits" in the network [Elhage et al., 2021 , Olsson et al., 2022 , Meng et al., 2022 , Vig et al., 2020 , Feng and Steinhardt, 2024 , Wu et al., 2023 , Wang et al., 2023 , Hanna et al., 2024 , Merullo et al., 2024 , McGrath et al., 2023 , Singh et al., 2024 , Feng et al., 2024] . While the definition of a "circuit" varies across different works, in this paper, our definition is similar to the one in Wang et al.[2023]: it is a collection of model components (attention heads, neurons, etc.) with the "edges" in the circuit indicating the information flow between the components in the forward pass; the "excitation" of the circuit is the input tokens.Evaluation of reasoning abilities of LLMs. Our work is also related to the line of work which focus on empirically evaluating the reasoning abilities of LLMs across different types of tasks [Xue et al., 2024 , Chen et al., 2024 , Patel et al., 2024 , Morishita et al., 2023 , Seals and Shalin, 2024 , Zhang et al., 2023 , Saparov and He, 2023 , Saparov et al., 2024 , Luo et al., 2024 , Han et al., 2024 , Tafjord et al., 2021 , Hendrycks et al., 2021 , Dziri et al., 2024 , Yang et al., 2024] . While these studies primarily benchmark their performance on sophisticated tasks, our work focuses on understanding "how" transformers reason on logic problems accessible to fine-grained analysis. ] type problem in conjunction with analysis on both small transformers trained purely on the synthetic problem, and large language models trained on a large corpus of internet data. 
 
 Appendix: Activation patching. At its core, activation patching, a.k.a. causal mediation analysis [Vig et al., 2020 , Meng et al., 2022 , Hase et al., 2024 , Heimersheim and Nanda, 2024 , Zhang and Nanda, 2024] , uses causal interventions for uncovering the internal mechanisms or "circuits" of LLMs that enable them to perform certain tasks. Typically, the LLM is run on pairs of "source" and "destination" prompts, and we search for components inside the model that "recover" the model's behavior on the source prompts by replacing parts of the model's activation with "source activations" when running on the destination prompt. The opposite "destination→source" intervention can also be adopted.Scope of this work. We define the scope of our analysis as follows. First, in the shallow transformer experiments, we focus on the variant which only has self-attention layers in addition to layer normalization, positional encoding, embedding and softmax parameters. While we could have also included MLP layers, we choose not to because the no-MLP models already achieve 100% accuracy on the problem, and adding MLPs would unnecessarily complicate the analysis. As a second way to focus the scope of paper, in the Mistral-7B experiments, we do not seek to uncover every model component that participates in solving the reasoning problem. We focus more on finding and analyzing the components that are necessary to the model's reasoning circuit, and necessary towards implementing the reasoning pathway as described before. By doing so, we can fully justify the necessity of these key components, without guessing about the roles of less-impactful sub-circuits.In this section, we provide a more detailed description of the propositional logic problem we study in this paper, and list representative examples of the problem.At its core, the propositional logic problem requires the reasoner to (1) distinguish which chain type is being queried (LogOp or linear), and ( 2 Remark. In this case, the answer "A true. A or B implies C; C True" is also correct. Remark. In this case, the answer "B false. A and B implies C; C undetermined" is also correct.The length-3 case is a simple generalization of this set of examples, so we do not cover those examples here.The architecture definition follows that of GPT-2 closely.Define input x = (x 1 , x 2 , ..., x T ) ∈ N T , a sequence of tokens with length T . It is converted into a sequence of (trainable) token embeddings X token = (e(x 1 ), e(x 2 ), ..., e(x T )) ∈ R din×T . Adding to it the (trainable) positional embeddings P = (p 1 , p 2 , ..., p T ) ∈ R din×T , we form the zero-th layer embedding of the transformer X 0 = (e(x 1 ) + p 1 , ..., e(x T ) + p T ). The input is processed by the attention blocks as follows.Let the model have L layers and H heads. For layer index ℓ ∈ [L] and head index j ∈ [H], atten-tion head A ℓ,j is computed by A ℓ,j (X ℓ-1 ) = S causal Xℓ-1 Q T ℓ,j K ℓ,j XT ℓ-1 / √ d H Xℓ-1 V T ℓ , with Xℓ-1 = LayerNorm(X ℓ-1 ), S(•) being the softmax operator, causal[•] the causal mask opera- tor. The output of the attention block is A ℓ = X ℓ-1 + Concat[A ℓ,1 (X ℓ-1 ), ..., A ℓ,H (X ℓ-1 )]W TO,ℓ , with W O,ℓ the square projection matrix (with bias). Finally, we apply an affine classifier (with softmax) f (x) = S( XL,T W T class + b class ) to predict the next word. In this paper, we set the hidden space embedding to 768.In all of our experiments, we set the learning rate to 10 -4 , and weight decay to 10 -4 . For models with depth less than 6, we use a batch size of 512, and train the model for 60k iterations; for models with depth greater than or equal to 6, we use a batch size of 256, and train for 80k iterations. We use the AdamW optimizer in PyTorch, with 5k iterations of linear warmup, following by cosine annealing to a learning rate of 0. Each model is trained on a single V100 GPU; the full set of models take around 2 -3 days to finish training.Problem specification. In each logic problem instance, the proposition variables are randomly sampled from a pool of 80 variables (tokens). The truth values in the fact section are also randomly chosen. The linear chain is queried 20% of the time; the LogOp chain is queried 80% of the time.Accuracy of different variants of the model for the length-3 problem. We show in Figure 4 below that the 3-layer 3-head variant is the smallest model which achieves 100% accuracy on the problem. Observation 1: chain-type disentanglement at QUERY. We construct 200 samples, with the first half querying the linear chain, the second half querying the logical-operator chain. We record the second-layer's self-attention block output on these samples, and compute the cosine similarity between each pair. We show the result in Figure 5 . Observation 2a: existence of an abstract "routing signal". We make the following experimental observations. 1. We generate 1,000 samples whose QUERY is for the linear chain, and compute the average output embedding of layer-2 self-attention block (post projection matrix). We denote this embedding as h route .2. (Linear→LogOp intervention) Sample a set of 500 validation samples, all of which query for the linear chain. In the forward pass of the model on every validation sample, we subtract the h route to the output of the second attention block -note that this "corrupted" signal from the second layer is also received by the third layer. We observe that the model's first token prediction is 100% of the time the correct first token for the LogOp chain. 3. (LogOp→linear intervention) We repeat the above experiment the other way around. We sample another 500 validation samples, but in this case they all query the LogOp chain. On every sample, during the forward pass we add h route to the output of the second attention block. We record the attention weight of the third layer's attention blocks at the ANSWER token position. We then average these attention weights for each head. We find that all three attention heads place greater than 99% attention weight on the QUERY position on average, behaving exactly like when the sample naturally queries for the linear chain.Observation 2b: linearly-decodable linear-chain answer at layer 2. We simply frame the learning problem as a linear classification problem. The input vector of the classifier is the same as the input to the layer-3 self-attention block, equivalently the layer-2 residual-stream embedding. The output space is the set of proposition variables (80-dimensional vector). We train the classifier on 5k training samples (all whose QUERY is for the linear chain) using the AdamW optimizer, with learning rate set to 5 × 10 -3 and weight decay of 10 -2 . We verify that the trained classifier obtains an accuracy greater than 97% on an independently sampled test set of size 5k (all whose QUERY is for the linear chain too).Evidence 3a: Distinct behaviors of affine predictors at different layers. We train two affine classifiers at two positions inside the model (each with 10k samples): W resid,ℓ=2 at layer-2 residual stream, and W attn,ℓ=3 at layer-3 attention-block output, both at the position of ANSWER, with the target being the correct first token. In training, if there are two correct answers possible (e.g. OR gate, starting nodes are both TRUE or both FALSE), we randomly choose one as the target; in testing, we deem the top-1 prediction "correct" if it coincides with one of the answers. We observe the following predictor behavior on the test samples:1. W attn,ℓ=3 predicts the correct answer 100% of the time.2. W resid,ℓ=2 always predicts one of the variables assigned FALSE (in the fact section) if LogOp is the AND gate, and predicts one assigned TRUE if LogOp is the OR gate.Evidence 3b: linearly decodable LogOp information from first two layers. We train an affine classifier at the layer-2 residual stream to predict the LogOp of the problem instance, over 5k samples (and tested on another 5k samples). The classifier achieves greater than 98% accuracy. We note that training this classifier at the layer-1 residual stream also yields above 95% accuracy.Evidence 3c: identification of LogOp-chain starting nodes at layer 3. Attention heads (3,1) and (3,3), when concatenated, produce embeddings which we can linearly decode the two starting nodes of the LogOp chain with test accuracy greater than 98%. We also find that they focus their attention in the rule section of the context (as shown in Figure 6 ). Due to causal attention, this means that they determine the two starting nodes from the LogOp-relevant rules.Remark. The above pieces of observations suggest the "partial information→refinement" process. 3 To further validate that the embedding from the first two layers are indeed causally linked to the correct answer at the third layer, we perform an activation patching experiment.Evidence 3d: layer-2 residual stream at ANSWER is important to correct prediction. We verify that layer-3 attention does rely on information in the layer-2 residual stream (at the ANSWER position):• Construct two sets of samples D 1 and D 2 , each of size 10k: for every sample X 1,n ∈ D 1 and X 2,n ∈ D 2 , the context of the two samples are exactly the same, except the LogOp is flipped, i.e. if X 1,n has disjunction, then X 2,n has the conjunction operator. If layer 3 of the model has no reliance on the Resid ℓ=2 (layer-2 residual stream) for LogOp information at the ANSWER position, then when we run the model on any X 2,n , patching Resid ℓ=2 (X n,2 ) with Resid ℓ=2 (X n,1 ) at ANSWER should not cause significant change to the model's accuracy of prediction. However, we observe the contrary: the accuracy of prediction degrades from 100% to 70.87%, with standard deviation 3.91% (repeated over 3 sets of experiments).Observation: LogOp-relevant reasoning at the third layer. We show that the output from attention heads (3,1) and (3,3) (before the output/projection matrix of the layer-3 attention block), namely A 3,1 (X 2 ) and A 3,3 (X 2 ), when concatenated, contain linearly decodable information about the two starting nodes of the LogOp chain. We frame this as a multi-label classification problem, detailed as follows:1. We generate 5k training samples and 5k test samples, each of whose QUERY is for the LogOp chain. For every sample, we record the target as a 80-dimension vector, with every entry set to 0 except for the two indices corresponding to the two proposition variables which are the starting nodes of the LogOp chain.2. Instead of placing softmax on the final classifier of the transformer, we use the Sigmoid function. Moreover, instead of the Cross-Entropy loss, we use the Binary Cross-Entropy loss (namely the torch.nn.functional.binary_cross_entropy_with_logits in PyTorch, which directly includes the Sigmoid for numerical stability). 3. We train an affine classifier, with its input being the concatenated Concat[A 3,1 (X 2 ), A 3,3 (X 2 )] (a 512-dimensional vector) on every training sample, and with the targets and training loss defined above. We use a constant learning rate of 0.5 × 10 -3 , and weight decay of 10 -2 . The optimizer is AdamW in PyTorch. 4. We assign a "correct" evaluation of the model on a test sample only if it correctly outputs the two target proposition variable as the top-2 entries in its logits. We observe that the classifier achieves greater than 98% once it converges.Figure 6 : Attention statistics, averaged over 500 samples, all of which query for the LogOp chain. The x-axis is simply an example prompt that helps illustrate where the attention is really placed at.Observe that only attention head (3,2) pays significant attention to the fact section. The other two heads focus on the rule section. Reminder: due the the design of the problem, the rule, fact and query sections all have consistent length for every sample!Observation 3 supplement: linearly-decodable linear-chain answer at layer 2. We simply frame the learning problem as a linear classification problem. The input vector of the classifier is the same as the input to the layer-3 self-attention block, equivalently the layer-2 residual-stream embedding. The output space is the set of proposition variables (80-dimensional vector). We train the classifier on 5k training samples (all whose QUERY is for the linear chain) using the AdamW optimizer, with learning rate set to 5 × 10 -3 and weight decay of 10 -2 . We verify that the trained classifier obtains an accuracy greater than 97% on an independently sampled test set of size 5k (all whose QUERY is for the linear chain too).Remarks on truth value determination. Evidence suggests that determining the truth value of the simple propositional logic problem is easy for the model, as the truth value of the final answer is linearly decodable from layer-2 residual stream (with 100% test accuracy, trained on 10k samples) when we give the model the context+chain of thought right before the final truth value token. This is expected, as the main challenge of this logic problem is not about determining the query's truthWe provide evidence in this part of the paper primarily relying on a popular technique in mechanistic interpretability: causal mediation analysis. Our methodology is roughly as follows:1. Suppose we are interested in the role of the activations of certain components of the LLM in a certain (sub-)task. For a running example, say we want to understand what role the attention heads play in processing and passing QUERY information to the ":" position for inference. Let us denote the activations as A ℓ,h;t (X), representing the activation of head h in layer ℓ, at token position t. 2. Typically, the analysis begins by constructing two sets of prompts which differ in subtle ways. A natural construction in our example is as follows: define sets of samples D orig and D alt , where X orig,n and X alt,n have exactly the same context, except in X orig,n , QUERY is for the LogOp chain, while in X alt,n , QUERY is for the linear chain. Moreover, denote the correct targets y orig,n and y alt,n respectively. 3. We run the LLM on D orig and D alt , caching the attention-head activations. We also obtain the logits of the model. We can compute the model's logit differences∆ orig,n = logit(X orig,n )[y orig,n ] -logit(X orig,n )[y alt,n ].For a high-accuracy model, ∆ orig,n should be large for most n's, since it must be able to clearly tell that on an X orig,n , it is the LogOp chain which is being queried, not the linear chain. 4. We now perform intervention for all n, ℓ, h and t:(a) Run the model on X orig,n , however, replacing the original activation A ℓ,h;t (X orig,n ) by the altered A ℓ,h;t (X alt,n ). Now let the rest of the run continue. 4 Let us denote the logits obtained in this intervened run as logit →alt;(ℓ,h,t) (X orig,n ). (b) Now compute the intervened logit difference∆ orig→alt,n;(ℓ,h,t) = logit →alt;(ℓ,h,t) (X orig,n )[y alt,n ] -logit →alt;(ℓ,h,t) (X orig,n )[y orig,n ].5. Now average the ∆ orig→alt,n;(ℓ,t) 's over n for every ℓ, h and t (recall that n is the sample index). 6. This procedure helps us identify components that are significant in processing and passing the QUERY information for inference. Intuitively, an activation that result in a positive and large ∆ orig→alt,n;(ℓ,t) play a significant role in this subtask, because this activation helps "altering" the model's "belief" from "QUERY is for the LogOp chain" to "QUERY is for the linear chain". 7. Remark: due to the symmetry of this running example, it is perfectly sensible to perform alt → orig interventions too, by mirroring the above procedures.Each of our experiments are done on 60 samples unless otherwise specified -we repeat some experiments (especially the attention-head patching experiments) to ensure statistical significance when necessary.Calibrated metric. Please note that in this work, we adopt a calibrated/normalized version of the intervened logit difference (aimed at keeping the score's magnitude between 0 and 1). In particular, we compute the following metric for head (ℓ, h) at token position t:EQUATION)where∆ orig = 1 N n∈[N ] logit(X n )[y alt,n ] -logit(X n )[y orig,n ], and ∆ alt = 1 N n∈[N ] logit(X ′ n )[y alt,n ] -logit(X ′ n )[y orig,n ].The closer to 1 this score is, the stronger the model's belief is altered.Our analysis relies on QUERY-based patching, following the same procedure as detailed in subsection E.2. In this set of experiments, we discover the main attention heads responsible for processing the context and performing inference as introduced in the beginning of this Section. Why is QUERY-based patching important to reasoning circuit discovery? To answer this question, there are two points to emphasize first. (1) We know that to solve the reasoning problem, the QUERY token is critical to initiating the reasoning chain: without it, the rules and facts are completely useless; with it, the reasoner can then proceed to identify the relevant rules and facts to predict the answer.(2) The prompt pairs differ only by the QUERY token. Based on (1) and(2), we know that if performing the aforementioned QUERY-based causal intervention on a model component leads to a large intervened logit difference (i.e. it alters the model's "belief"), then this component must be integral to the reasoning circuit, because the component is now identified to be QUERY-sensitive and has causal influence on (parts of) the model's reasoning actions.High-level interventions. We begin by presenting higher-level patching results in Figure 7 , where we intervene at the level of residual streams, attention blocks, and MLPs. We can draw a few insights from these results:1. A sharp transition of "QUERY processing" occurs from layer 12 to layer 13 (indexed from 0) in Figure 7 (a) and (b). 2. Figure 7 (b) shows that a small set of attention blocks are observed to be significant for the "belief altering" action, namely those in layers 9, 12, 13, 14, 16 and 19. 3. The MLPs, shown in Figure 7 (c), play little role in this circuit, except for MLP-0. However, MLP-0 had been observed to act more as a "nonlinear token embedding" than a complex high-level processing unit [Wang et al., 2023] . In the rest of this section, we primarily devote our analysis to the attention heads, and leave the exact role of the MLPs to future work.Attention-head interventions. Figure 7 helps us locate a small set of attention blocks which are important to the task. However, these results alone still leave us with far too many components to examine in detail. Therefore, we run a set of finer-grained experiments, intervening on the attention heads (over the relevant context). The results are shown in Figure 8 . We find that, interestingly, only a very small set of attention heads are central to the "belief altering" of the LLM. More specifically, in Figure 8 (a), only attention heads (12,9), (13,11), (14,24;26), (16,0;12;14), (17,25) , (19,8) are observed with relatively high intervened logit differences.We note that Grouped-Query Attention used by Mistral-7B adds subtlety to the analysis5 : patching a single head might not yield a high logit difference since other heads in the same group (which possibly perform a similar function) could overwhelm the patched head and maintain the model's previous "belief". Therefore, we also run a coarser-grained experiment which simultaneously patches the attention heads sharing the same key and value activations, shown in Figure 8 (b). This experiment reveals that heads belonging to the group (9, 24 -27) also have high intervened logit difference.We aim to understand why the attention heads identified in the last sub-section are important. For now, we continue with QUERY altering in the prompt pairs. Through intervening on the sub-components of each attention head, namely their value, key, and query, and through examining details of their attention weights, we find that there are roughly four types of attention heads. We show the results in Figure 9 (repeated here in the Appendix for convenience):1. Queried-rule locating head. Attention head (12,9)'s query activation has a large intervened logit difference according to Figure 3 (a), therefore, its query and attention patterns are QUERYdependent and contribute to altering the model's "belief". Furthermore, at the QUERY position, we find that on average, its attention weight is above 90% at the "conclusion" variable of the rule being queried. In other words, it is responsible for locating the queried rule, and storing that rule's information at the QUERY position.6 2. Queried-rule mover head. Attention head (13,11)'s value activations have large intervened logit difference, and intriguingly, its query and key activations do not share that tendency. This already suggests that its attention pattern performs a fixed action on both the original and altered prompts, and only the value information is sensitive to QUERY. Furthermore, within the relevant context (excluding the 6 in-context examples given), head (13,11) assigns above 50% attention weight to the QUERY position, and its attention weight at QUERY is about 10 times larger than the second largest one on average. Recalling the role of layer 12, we find evidence that head (13,11) moves the QUERY and queried-rule information to the ":" position.7 3. Fact processing heads. Attention heads (16,12), (16,14) and (17,25)'s query activations have large intervened logit differences. Within the relevant context, they place greater than 56%, and 70% of their attention respectively in the fact section of the context (starting from "Fact" and ending on "." before "Question").4. Decision head. Attention head (19,8)'s query activations have large intervened logit differences. Its attention pattern suggests that it is a "decision" head: within the relevant context, when the model is correct, the head's top-2 attention weights are always on the correct starting node of the queried rule and the correct variable in the fact section, and the two token positions occupy more than 60% of its total attention in the relevant context on average. In other words, it already has the answer.In this subsection, we provide finer details on the attention patterns of the attention heads we discovered in Section E.3. Note that the attention weights percentage we present in this section are calculated by dividing the observed attention weight at a token position by the total amount of attention the head places in the relevant context, i.e. the portion of the prompt which excludes the 6 in-context examples.Queried-rule locating heads. Figure 10 presents the average attention weight the queried-rule locating heads place on the "conclusion" variable and the period "." immediately after the queried rule at the QUERY token position (i.e. the query activation of the heads come from the residual stream at the QUERY token position) -(12,9) is an exception to this recording method, where we only record its weight on the conclusion variables alone, and already observe very high weight on average. The heads (12,9), (14,24), (14,26), (9,25), (9,26) indeed place the majority of their attention on the correct position consistently across the test samples. The reason for counting the period after the correct conclusion variable as "correctly" locating the rule is that, it is known that LLMs tend to use certain "register tokens" to record information in the preceding sentence. We can observe that head (12,9) has the "cleanest" attention pattern out of the ones identified, placing on average 93.0 ± 9.4% of it attention on the correct conclusion variable alone. The more diluted attention patterns of the other heads likely contribute to their weaker intervened logit difference score shown in Section E.3 in the main text.Queried-rule mover heads. Figure 11 shows the attention weight of the queried-rule mover heads. While they do not place close to 100% attention on the QUERY location consistently (when the query activation comes from the residual stream from token ":", right before the first answer token), the top-1 attention weight consistently falls on the QUERY position, and the second largest attention weight is much smaller. In particular, head (13,11) places 54.2 ± 12.5% attention on the QUERY position on average, while the second largest attention weight in the relevant context is 5.2 ± 1.1% on average (around 10 times smaller; this ratio is computed per sample and then averaged).An extra note about head (16,0): it does not primarily act like a "mover" head, as its attention statistics suggest that it processes the mixture of information from the QUERY position and the ":" position. Therefore, while we present its statistics along with the other queried-rule mover heads here since it does allocate significant attention weight on the QUERY position on average, we do not list it as such in the circuit diagram of Figure 2 .Figure 11 : Average attention weights of the queried-rule mover heads, along with the standard deviations. The raw attention patterns are obtained at token position ":" (i.e. the query activation comes from the residual stream at the ":" position), right before the first answer token, and the exact attention weight (indicated by the blue bars) is taken at the QUERY position; for head (16,0), we also obtain its attention weight at the ":" position, as we found that it also allocates a large amount of attention weight to this position in addition to the QUERY position. Note: for (15,8), we found that it only acts as a "mover" head when the linear chain is being queried, so we are only reporting its attention weight statistics in this specific scenario; the other heads do not exhibit this interesting behavior, so we report those heads' statistics in all query scenarios.Fact processing heads. Figure 12 below shows the attention weights of the fact processing heads; the attention patterns are obtained at the ":" position, right before the first answer token, and we sum the attention weights in the Fact section (starting at the first fact assignment, ending on the last "." in this section of the prompt). It is clear that they place significant attention on the Fact section of the relevant context. Remark. There is only one "decision head" which we identified, i.e. head (19, 8) . Since there is no further subtleties with how we recorded its attention weights or peculiar behaviors of the attention patterns observed, we do not elaborate further on it in the Appendix.This paper presents work whose goal is to advance the field of Machine Learning, particularly the area of mechanistic interpretability. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here We use (ℓ, h) to denote an attention head. When referencing multiple heads in the same layer, we write (ℓ, h1; h2; ...; hn) for brevity.In fact, the observations suggest that layer 3 performs a certain "matching" operation. Take the OR gate as an example. Knowing which of the three starting nodes (for LogOp and linear chain) are TRUE, and which two nodes are the starting nodes for the LogOp chain are sufficient to determine the first token! This exact algorithm, however, is not fully validated by our evidence; we leave this as part of our future work.Please note that layers ℓ + 1 to L are influenced at and after token position t, and technically speaking, now operate "out of distribution".In Mistral-7B-v0.1, each attention layer has 8 key and value activations, and 32 query activations. Therefore, heads (ℓ, h × 4) to (ℓ, h × 4 + 3) share the same key and value activation.Heads (9,25;26), (14,24;26) exhibit similar tendencies, albeit with smaller intervened logit differences.Heads (15,8), (16,0) also appear to belong to this type, albeit with smaller intervened logit difference. 


Decision: The paper provides an interesting mechanistic analysis of how transformers reason. The authors are suggested to improve the clarity of the paper and try out some other domains to enhance the soundness of their work.

Reviews: review: title: title: Review of How Transformers Reason: A Case Study on a Synthetic Propositional Logic Problem
main_review: main_review: Summary:  
This paper explores how transformers reason through a synthetic propositional logic task, focusing on specific mechanisms in a small 3-layer model and the mistral-7b model. It identifies key patterns, such as query-type disentanglement and reasoning circuits, using methods like activation patching. 

Strengths:  
- mechanistic insight: the detailed analysis of attention heads and routing signals provides valuable insights into how transformers handle reasoning tasks.
- empirical results: the observations around ''planning'' and ''reasoning'' circuits, especially the disentanglement at the query level, are interesting and well supported by experiments.
- good use of interpretability methods: the use of activation patching in larger models like mistral-7b is a nice point.

Suggestions:
- connect small and large models: it would help to clarify how findings from small transformers extend to larger models. 
   
- expand literature review: the current review feels limited. i would strongly suggest the authors to expand it, providing readers with a more comprehensive understanding of existing works and contexts in mechanistic interpretability, and implications on future works.
   
- explore robustness of routing signals: the routing signals in the second attention block are interesting. additional experiments testing how these signals behave with different logic tasks or architectures could offer more insight into their general applicability.

 
Please note that i am not a domain expert in this specific topic and its related literature, but i found the paper’s explanations clear and the results convincing in general, thus i lean towards acceptance.
title: title: Review
main_review: main_review: **Summary**:  
This paper investigates how transformers solve a synthetic propositional logic task, providing a detailed analysis of the mechanisms by which a 3-layer transformer and the larger Mistral-7B model process reasoning tasks. The authors identify key attention circuits responsible for reasoning, showing how transformers use specific layers and attention heads to handle different types of queries (linear vs. logical operator chains). The study uses activation patching to analyze Mistral-7B's reasoning processes.

**Strengths**:
- **Problem Selection**: The synthetic logic problem is a well-chosen, controlled task for understanding reasoning in transformers.
- **In-Depth Analysis**: The paper offers rigorous experimentation, particularly with activation patching, providing deep insights into transformer architectures.
- **Interesting Findings**: The discovery of "routing signals" and the analysis of Mistral-7B's layer-specific reasoning are valuable contributions.

**Weaknesses**:
- **Limited Task Scope**: The study focuses on only one synthetic propositional logic problem, limiting the generalizability of the findings to other reasoning domains or tasks.
title: title: Review
main_review: main_review: This paper presents a study on a synthetic propositional logic problem to show whether Transformers can reason and they adopt a three-layer Transformer to decipher this mechanism.

The main result of this manuscript is: 
*They discover that small transformers trained purely on the synthetic problem repurposes certain tokens of the context to filter out information that are less relevant to the answer.* 


After reading this paper, I find that this paper is hard to follow for the following reasons: 
1. The organization of this paper is bad, the results of this paper are not easy to understand.
2. I tried to catch the main idea but could not understand what this paper did. 
For example, in section 2.1: 

> For example, as shown in Figure 1, only the fact "D False" ... 

However, no `D` is shown in Figure 1. 

> if QUERY is for the LO chain ... 

So what is LO chain? Is it different from the linear chain? Why are they different? 

**Pros**
- The problem studied is very interesting, and the use of propositional logic is fundamental and potentially a good choice.

**Cons**
- Poor writing and organization hinder the reader's understanding of this paper.
- The experiments seem interesting but are not convincing since the authors only conducted them on a very specific case.

Overall, I find this paper interesting but it needs improvement in organization and detail. It would be helpful if the authors clarified their main claims and improved the writing.