Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models (2023)

Abstract: Recent work has shown that language models (LMs) have strong multi-step (i.e., procedural) reasoning capabilities. However, it is unclear whether LMs perform these tasks by cheating with answers memorized from pretraining corpus, or, via a multi-step reasoning mechanism. In this paper, we try to answer this question by exploring a mechanistic interpretation of LMs for multi-step reasoning tasks. Concretely, we hypothesize that the LM implicitly embeds a reasoning tree resembling the correct reasoning process within it. We test this hypothesis by introducing a new probing approach (called MechanisticProbe) that recovers the reasoning tree from the model’s attention patterns. We use our probe to analyze two LMs: GPT-2 on a synthetic task (k-th smallest element), and LLaMA on two simple language-based reasoning tasks (ProofWriter & AI2 Reasoning Challenge). We show that MechanisticProbe is able to detect the information of the reasoning tree from the model’s attentions for most examples, suggesting that the LM indeed is going through a process of multi-step reasoning within its architecture in many cases.

## introduction large language models (lms) have shown impressive capabilities of solving complex reasoning problems @xcite . yet, what is the underlying "thought process" of these models is still unclear (figure 1 ). do they cheat with shortcuts memorized from pretraining corpus @xcite @xcite ? or, do they follow a rigorous reasoning process and solve the problem procedurally @xcite ? answering this question is not only critical to our understanding of these models but is also critical for the development of next-generation faithful language-based reasoners @xcite @xcite . a recent line of work tests the behavior of lms by designing input-output reasoning examples @xcite . however, it is expensive and challenging to construct such high-quality examples, making it hard to generalize these analyses to other tasks/models. another line of work, mechanistic interpretability @xcite @xcite @xcite , directly analyzes the parameters of lms, which can be easily extended to different tasks. inspired by recent work @xcite @xcite ) that uses attention patterns for linguistic phenomena prediction, we propose an attention-based mechanistic interpretation to expose how lms perform multi-step reasoning tasks @xcite . we assume that the reasoning process for answering a multi-step reasoning question can be represented as a reasoning tree (figure 2 ). then, we investigate if the lm implicitly infers such a tree when answering the question. to achieve this, we designed a probe model, mechanisticprobe, that recovers the reasoning tree from the lm's attention patterns. to simplify the probing problem and gain a more fine-grained understanding, we decompose the problem of discovering reasoning trees into two figure 2 : illustration of our mechanisticprobe with one example from each of the three reasoning tasks considered in this work. we are given a number of input statements: @xmath0, s 2 , ..., } and a question: q. mechanisticprobe recovers the reasoning tree describing the ground-truth reasoning process to answer the question. mechanisticprobe works in two stages: in the first stage, mechanisticprobe detects if the lm can select the set of useful statements required in reasoning, and then in the second stage, mechanisticprobe detects if the lm can predict the reasoning tree given the useful statements. we conduct experiments with gpt-2 @xcite on a synthetic task (finding the k-th smallest number in a sequence of numbers) and with llama @xcite on two natural language reasoning tasks: proofwriter @xcite and ai2 reasoning challenge (i.e., arc @xcite . for most examples, we successfully detect reasoning trees from attentions of (finetuned) gpt-2 and (few-shot & finetuned) llama using mechanisticprobe. we also observe that lms find useful statements immediately at the bottom layers and then do the subsequent reasoning step by step ( §4). to validate the influence on the lm's predictions of the attention mechanisms that we identify, we conduct additional analyses. first, we prune the attention heads identified by mechanisticprobe and observe a significant accuracy degradation ( §5). then, we investigate the correlation between our probing scores and the lm's performance and robustness ( §6). our findings suggest that lms exhibit better prediction accuracy and tolerance to noise on examples with higher probing scores. such observations highlight the significance of accurately capturing the reasoning process for the efficacy and robustness of lms. ## mechanisticprobe in this section, we introduce mechanisticprobe: a probing approach that analyzes how lms solve multi-step reasoning tasks by recovering reasoning trees from the attentions of lms. ## mechanistic probing of lms we use the probe to analyze how lms perform the reasoning tasks. we first verify the usefulness of attentions in understanding the reasoning process of lms by visualizing a simp ( §4.1). then, we use the probe to quantify the information of g contained in a simp ( §4.2). finally, we report layer-wise probing results to understand if lms are reasoning procedurally across their architecture ( §4.3). 7 ## correlating probe scores with model accuracy and robustness our results show that lms indeed reason mechanistically. but, is mechanistic reasoning necessary for lm performance or robustness? we attempt to answer this question by associating the probing scores with the performance and robustness of lms. given that finetuned gpt-2 has a very high test accuracy on the synthetic task and llama does not perform as well on arc, we conduct our analysis mainly with llama on the proofwriter task. accuracy. we randomly sample 64 to 128 examples from the dataset and test 4-shot llama on these examples. we calculate their test accuracy and the two probing scores s p1 and s p2 . we repeat this experiment 2048 times. then, we calculate the correlation between the probing scores and test accuracy. from robustness. following the same setting, we also associate the probing scores with lm robustness. in order to quantify model robustness, we randomly corrupt one useless input statement for each example, such that the prediction would remain unchanged. 11 we measure robustness by the decrease in test accuracy after the corruption. figure 10 shows that if s p2 is small (less than 0.7), the prediction of llama could be easily influenced by the noise (test accuracy decreases around 10%). however, if the probing s p2 is high, llama is more robust, i.e., more confident in its correct prediction (test accuracy increases around 4%). this provides evidence that if the lm encodes the gold reasoning trees, its predictions are more reliable (i.e., robust to noise in the input). ## related work attention-based analysis of lms. attention has been popularly used to interpret lms (vig and 11 we do it by adding negation: corrupting si as ["that"• si • " is false"], where • means string concatenation. @xcite @xcite . a direct way for interpreting an attentionbased model is to visualize attentions @xcite . but the irrelevant, redundant, and noisy information captured by attentions makes it hard to find meaningful patterns. alternatively, accumulation attentions that quantify how information flows across tokens can be used for interpretation @xcite . however, for casual lms, information flows in one direction, and it causes an over-smoothing problem when the model is deep. 12 to tackle this, other works propose new metrics and analyze attentions using them @xcite . however, these metrics are proposed under specific scenarios, and these are not useful for detecting the reasoning process in lms. we address this challenge by designing a more structured probe that predicts the reasoning tree in lms. ## conclusion in this work, we raised the question of whether lms solve procedural reasoning tasks step-by-step within their architecture. in order to answer this question, we designed a new probe that detects the oracle reasoning tree encoded in the lm architecture. we used the probe to analyze gpt-2 on a synthetic reasoning task and the llama model on two natural language reasoning tasks. our empirical results show that we can often detect the information in the reasoning tree from the lm's attention patterns, lending support to the claim that lms may indeed be reasoning "mechanistically".