A Mechanistic Analysis of a Transformer Trained on a Symbolic Multi-Step Reasoning Task (2024)

Abstract: Transformers demonstrate impressive performance on a range of reasoning benchmarks. To evaluate the degree to which these abilities are a result of actual reasoning, existing work has focused on developing sophisticated benchmarks for behavioral studies. However, these studies do not provide insights into the internal mechanisms driving the observed capabilities. To improve our understanding of the internal mechanisms of transformers, we present a comprehensive mechanistic analysis of a transformer trained on a synthetic reasoning task. We identify a set of interpretable mechanisms the model uses to solve the task, and validate our findings using correlational and causal evidence. Our results suggest that it implements a depth-bounded recurrent mechanisms that operates in parallel and stores intermediate results in selected token positions. We anticipate that the motifs we identified in our synthetic setting can provide valuable insights into the broader operating principles of transformers and thus provide a basis for understanding more complex models.

## introduction transformer-based language models @xcite demonstrate impressive performance on reasoning 2 tasks @xcite , mathematical problem-solving @xcite , and planning @xcite . however, despite strong performance on certain reasoning benchmarks, it remains unclear to what extent these abilities are a result of actual reasoning or simple pattern memorization @xcite ). ## related work expressiveness of transformers to understand the capabilities of transformers, one line of work characterizes their theoretical properties @xcite @xcite @xcite . these studies answer questions about the expressiveness of the transformer architecture by treating them as approximators of different classes of functions @xcite , or by characterizing the class of formal languages that a transformer can recognize @xcite . ## symbolic reasoning using backward chaining in this section, we present a mechanistic analysis of the internal mechanisms of the model and provide correlational and causal evidence. our findings suggest that the model uses an interpretable and meaningful backward chaining algorithm to perform pathfinding in a tree. to help guide the reader, we present an intuitive explanation before breaking down the individual algorithmic steps in the following sections. the backward chaining algorithm first, the model aggregates the source and target nodes of each edge in the edge list into the target node position (see section 5.1). then, the model starts at the goal node and moves up the tree one level with each layer of the model. this mechanism is implemented using attention heads, which we term "deduction heads" (see section 5.2). by the composition of multiple attention heads in consecutive layers, the model can traverse the tree upwards for l -1 edges, where l is the number of layers in the model. we refer to this mechanism as backward chaining, inspired by the use of the term in the symbolic artificial intelligence literature @xcite . in more complex scenarios, where the required path exceeds the depth of the model, it relies on backward chaining from multiple nodes in the tree in parallel. this creates multiple subpaths that can be used to find the correct next step in cases where the goal node is more than l -1 edges away (see section 5.3). in addition, the model uses a simple heuristic as a fallback mechanism, where it identifies child nodes of the current position and evaluates whether these are leaf nodes of the tree. this enables the model to make informed guesses when backward chaining and parallelization are insufficient to solve the problem (see section 5.4). ## discussion register tokens our model uses some token positions as a form of working memory to store intermediate results. this observation aligns with @xcite which found that image models use some image patches to accumulate global information while discarding spatial information. similarly, @xcite show that adding uninformative tokens at the end of each prompt can enhance language model performance on downstream tasks without introducing additional parameters. our findings suggest that these techniques enable the model to store intermediate results and perform more computations in parallel. this is consistent with theoretical insights from @xcite which highlights how the effective state of a transformer depends on the number of tokens in the sequence. ## conclusion in this paper, we conducted a mechanistic analysis of a transformer trained on pathfinding in trees. our results suggest that the model implements a backward chaining algorithm that starts at the goal node and climbs the tree one level per layer. to solve more complex problems, where the required reasoning depth exceeds the number of layers, it executes the same mechanism in parallel across multiple register tokens and combines the resulting subpaths on the final token position. in addition, it performs a simple one-step lookahead in which it finds the child nodes of the current position and evaluates whether they are leaf nodes. our findings in this synthetic setting demonstrate the ability of a transformer to perform deductive reasoning up to a certain reasoning depth, after which it resorts to simple heuristics. by using parallelized computations to store intermediate results in register tokens and then merging these results in the final token position, the model demonstrates a form of deductive reasoning that, while effective within a given setting, is constrained by its architecture. these observations suggests that transformers may exhibit a inductive bias towards adopting highly parallelized strategies for tasks involving search, planning, or reasoning.