A Mechanistic Interpretation of Arithmetic Reasoning in Language Models using Causal Mediation Analysis (2023)

Abstract: Mathematical reasoning in large language models (LMs) has garnered significant attention in recent work, but there is a limited understanding of how these models process and store information related to arithmetic tasks within their architecture. In order to improve our understanding of this aspect of language models, we present a mechanistic interpretation of Transformer-based LMs on arithmetic questions using a causal mediation analysis framework. By intervening on the activations of specific model components and measuring the resulting changes in predicted probabilities, we identify the subset of parameters responsible for specific predictions. This provides insights into how information related to arithmetic is processed by LMs. Our experimental results indicate that LMs process the input by transmitting the information relevant to the query from mid-sequence early layers to the final token using the attention mechanism. Then, this information is processed by a set of MLP modules, which generate result-related information that is incorporated into the residual stream. To assess the specificity of the observed activation dynamics, we compare the effects of different model components on arithmetic queries with other tasks, including number retrieval from prompts and factual knowledge questions.

## introduction mathematical reasoning with transformer-based models @xcite is challenging as it requires an understanding of the quantities and the mathematical concepts involved. while large language models (lms) have recently achieved impressive performance on a set of math-based tasks @xcite @xcite , their behavior has been shown to be inconsistent and context-dependent @xcite . recent literature shows a multitude 1 our code and data is available at https://github.com/ alestolfo/lm-arithmetic . we trace the flow of numerical information within transformerbased lms: given an input query, the model processes the representations of numbers and operators with early layers @xcite . then, the relevant information is conveyed by the attention mechanism to the end of the input sequence (b). here, it is processed by late mlp modules, which output result-related information into the residual stream (c). however, there is a limited understanding of the inner workings of these models and how they store and process information to correctly perform mathbased tasks. insights into the mechanics behind lms' reasoning are key to improvements such as inference-time correction of the model's behavior @xcite and safer deployment. therefore, research in this direction is critical for the development of more faithful and accurate next-generation lm-based reasoning systems. in this paper, we present a set of analyses aimed at mechanistically interpreting lms on the task of answering simple arithmetic questions (e.g., "what is the product of 11 and 17?"). in particular, we hypothesize that the computations involved in reasoning about such arithmetic problems are carried out by a specific subset of the network. then, we test this hypothesis by adopting a causal mediation analysis framework @xcite , where the model is seen as a causal graph going from inputs to outputs, and the model components (e.g., neurons or layers) are seen as mediators @xcite . within this framework, we assess the impact of a mediator on the observed output behavior by conducting controlled interventions on the activations of specific subsets of the model and examining the resulting changes in the probabilities assigned to different numerical predictions. through this experimental procedure, we track the flow of information within the model and identify the model components that encode information about the result of arithmetic queries. our findings show that the model processes the input by conveying information about the operator and the operands from mid-sequence early layers to the final token using attention. at this location, the information is processed by a set of mlp modules, which output result-related information into the residual stream (shown in figure 1 ). we verify this finding for biand tri-variate arithmetic queries across four pretrained language models with different sizes: 2.8b, 6b, and 7b parameters. finally, we compare the effect of different model components on answering arithmetic questions to two additional tasks: a synthetic task that involves retrieving a number from the prompt and answering questions related to factual knowledge. this comparison validates the specificity of the activation dynamics observed on arithmetic queries. causality-based interpretability. causal mediation analysis is an important tool that is used to effectively attribute the causal effect of mediators on an outcome variable @xcite . this paradigm was applied to investigate lms by @xcite , who proposed a framework based on causal mediation analysis to investigate gender bias. variants of this approach were later applied to mechanistically interpret the inner workings of pre-trained lms on other tasks such as subject-verb agreement @xcite , natural language inference @xcite , indirect object identification @xcite , and to study their retention of factual knowledge @xcite . ## causal effects on arithmetic queries our analyses address the following question: q1 what are the components of the model that mediate predictions involving arithmetic computations? we address this question by first studying the flow of information throughout the model by measuring the effect of each component (mlp and attention block) at each point of the input sequence for two-operand queries ( §4.1). then, we distinguish between model components that carry information about the result and about the operands of the arithmetic computations ( §4.2 and §4.3). finally, we consider queries involving three operands ( §4.4) and present a measure to quantify the changes in information flow ( §4.5). ## causal effects on different tasks in order to understand whether the patterns in the effect of the model components that we observed so far are specific to arithmetic queries, we compare our observations on arithmetic queries to two different tasks: the retrieval of a number from the prompt ( §5.1), and the prediction of factual knowledge ( §5.2). with this additional set of experiments, we aim to answer the question: q2 are the activation patterns observed so far specific to the arithmetic setting? ## conclusion we proposed the use of causal mediation analysis to mechanistically investigate how lms process information related to arithmetic. through controlled interventions on specific subsets of the model, we assessed the impact of these mediators on the model's predictions. we posited that models produce predictions to arithmetic queries by conveying the math-relevant information from the mid-sequence early layers to the last token, where this information is then processed by late mlp modules. we carried out a causality-grounded experimental procedure on four different transformer-based lms, and we provided empirical evidence supporting our hypothesis. furthermore, we showed that the information flow we observed in our experiments is specific to arithmetic queries, compared to two other tasks that do not involve arithmetic computation. our findings suggest potential avenues for research into model pruning and more targeted training/fine-tuning by concentrating on specific model components associated with certain queries or computations. moreover, our results offer insights that may guide further studies into using lms' hidden representations to correct the model's behavior on math-based tasks at inference time @xcite and to estimate the probability of the model's predictions to be true @xcite .