SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Hanna*  
ILLC
University of Amsterdam m.w.hanna@uva.nl
Ollie Liu*  
University of Southern California  
zliu2898@usc.edu
Alexandre Variengien†  
Redwood Research  
alexandre.variengien@gmail.com
Pre-trained language models can be surprisingly adept at tasks they were not explicitly trained on, but how they implement these capabilities is poorly understood. In this paper, we investigate the basic mathematical abilities often acquired by pre-trained language models. Concretely, we use mechanistic interpretability techniques to explain the (limited) mathematical abilities of GPT-2 small. As a case study, we examine its ability to take in sentences such as "The war lasted from the year 1732 to the year 17", and predict valid two-digit end years (years  $>32$ ). We first identify a circuit, a small subset of GPT-2 small's computational graph that computes this task's output. Then, we explain the role of each circuit component, showing that GPT-2 small's final multi-layer perceptrons boost the probability of end years greater than the start year. Finally, we find related tasks that activate our circuit. Our results suggest that GPT-2 small computes greater-than using a complex mechanism that activates across diverse contexts.
As pre-trained language models (LMs) have grown in both size and effectiveness, their abilities have expanded to include a wide range of tasks, even without fine-tuning [6]. Such abilities can range from translation to text classification and multi-step reasoning [52]. Yet despite heavy study of these models [41, 42, 26], how LMs implement these abilities is still poorly understood.
In this paper, we study one such LM ability, performing mathematics. Mathematical ability has long been of interest in natural language processing: models have been trained to perform tasks such as simple arithmetic and word problems [51, 46]. Researchers have also fine-tuned pre-trained LMs on these tasks, instead of training from scratch [18, 25]. Recently, however, LMs seem to have acquired significant mathematical abilities without explicit training on such tasks [6, 34, 8, 16].
How these mathematical abilities arise in LMs is largely unknown. While studies have investigated pre-trained LMs' mathematical abilities [32, 38], existing work is behavioral: it explains what models can do, rather than how they do it. Most work that delves into model internals does so using models trained directly on such tasks: Hupkes et al. [23] probe such models for hierarchical structure, while Liu et al. [27] and Nanda et al. [31] study toy models trained on modular addition. Some studies do examine the structure of number representations in pre-trained models [30, 49]; however, they do not provide a causal explanation about how these models leverage these representations to perform math. The mechanisms underlying pre-trained LMs' mathematical abilities thus remain unclear.
*Work performed as part of Redwood Research's REMIX program
†Work performed during an internship. Now at Conjecture
Figure 1: Year-span prediction example (XX=17 and YY=32) with sample (in)valid output years.
To understand the roots of these mathematical abilities, we study them in GPT-2 small,[1] [39] which we show still possesses such abilities, despite its small size. This small size enables us to investigate its mathematical abilities at a very low level. Concretely, we adopt a circuits perspective [35, 14], searching for a minimal subset of nodes in GPT-2's computational graph responsible for this ability. To do so, we use fine-grained, causal methods from mechanistic interpretability, that allow us to identify nodes in GPT-2 that belong in our circuit, and then prove our circuit's correctness, through carefully designed causal ablations [21]. We also use mechanistic methods to pinpoint how each circuit component contributes to the mathematical task at hand. The end result of this case study is a detailed description of GPT-2's ability to perform one simple mathematical operation: greater-than.
Our investigation is structured as follows. We first define year-span prediction, a task that elicits mathematical behavior in GPT-2 (Section 2). We give the model input like "The war lasted from the year 1732 to the year 17"; it assigns higher probability to the set of years greater than 32. We next search for the circuit responsible for computing this task, and explain each circuit component's role (Section 3). We find a set of multi-layer perceptrons (MLPs) that computes greater-than, our operation of interest. We then investigate how these MLPs compute greater-than (Section 4). Finally, we find other tasks requiring greater-than to which this circuit generalizes (Section 5).
Via these experiments, we accomplish two main goals. First, we find a circuit for greater-than in GPT-2. This brings new mechanistic insights into math in pre-trained LMs, and builds on the limited existing work on circuits in pre-trained LMs [50] by examining a new task with a wide output space and rich structure. Second, we show that GPT-2's greater-than relies on a complex circuit that activates across contexts. This mechanism surpasses simple memorization, but does not reflect full mathematical competence; it lies between memorization and generalization. We thus add nuance to the memorization-generalization dichotomy, and take the first step towards a rich characterization of the states in between them.
2 Year-Span Prediction in GPT-2
GPT-2's size is ideal for low-level study, especially with potentially resource-intensive techniques like those in Section 3.1. However, this small size poses a challenge: GPT-2 is less capable than larger LMs, which still often struggle with mathematical tasks [29]. With this in mind, we craft a simple task to elicit a mathematical behavior in GPT-2, and verify that GPT-2 produces said behavior.
Task and Dataset We focus on a simple mathematical operation, greater-than, framed as it might naturally appear in text: an incomplete sentence following the template "The <noun> lasted from the year XXYY to the year XX" (Figure 1). The model should assign higher probability to years  $> \mathrm{YY}$ . We automatically generate sentences using this template. We draw the nouns from a pool of 120 nouns that could have a duration, found using FrameNet [1]; Appendix G lists the full pool of nouns. We sample the century XX of the sentence from  $\{11, \ldots, 17\}$ , and the start year YY from  $\{02, \ldots, 98\}$ .
We impose the latter constraints because we want GPT-2 to be able to predict a target as it would naturally be tokenized. However, GPT-2 uses byte-pair encoding, in which frequent strings more often appear as single tokens [44]. Thus, more frequent years—multiples of 100 or those in the 20th century—are tokenized as single tokens; less frequent years are broken into two. This causes a problem: GPT-2 could predict "[00]" after "[17]", but "1700" is always tokenized as "[1700]" in normal data and never as "[17][00]". So, we exclude all single-token years from our year pool. Finally, we want each example to have at least one correct and one incorrect validly tokenized answer, so we exclude each century's highest and lowest validly tokenized year from the pool of start years.
Input: "The war lasted from the year 1732 to the year 17"
GPT-2 small: 00X 12X 32X 33√ 63√ 99√
Further references to GPT-2 refer to GPT-2 small
Figure 2: Left: Probability heatmap of GPT-2 for year-span prediction. Y-axis: the sentence's start year (YY). X-axis: the two-digit output year candidate. (X,Y): Mean probability assigned by GPT-2 to output year X given input year Y. Right: GPT-2's average output distribution when YY=41
Qualitative Evaluation We first qualitatively analyze GPT-2's baseline behavior on this task by running it on a dataset of 10,000 examples. Each example has a noun randomly drawn from our 120 nouns, and a year drawn randomly from the 768 valid years from 1000 to 1899. For each YY in  $\{2,\ldots ,98\}$ , we take the average of GPT-2's probability distribution over predicted years, for all examples with start year YY; we visualize these average distributions in Figure 2.
GPT-2 appears to perform greater-than on our year-span prediction task: it creates a sharp cutoff between invalid end years ( $\leq$  YY), and valid end years ( $>$  YY). It assigns higher probability to the latter years, though not all of them: 15-20 years after YY, probabilities drop. The exact length of the year-span receiving higher probability likely reflects patterns in GPT-2's training data. In the real-world, and likely also in GPT-2's training data, our prompts' nouns, such as a "war," "dynasty", or "pilgrimage", have average durations that GPT-2 may have learned, influencing its output.
Quantitative Evaluation We design two numerical measures of model performance for the purpose of quantitative assessment. Let  $\mathrm{YY} \in \{02, \dots, 98\}$  be the start year of our sentence, and  $p_y$  be the probability of a two-digit output year  $y$ . We define the following two metrics:
- Probability difference:  $\sum_{y > \mathrm{YY}} p_y - \sum_{y \leq \mathrm{YY}} p_y$
- Cutoff sharpness:  $p_{\mathrm{YY} + 1} - p_{\mathrm{YY} - 1}$
Probability difference verifies that model output reflects a greater-than operation by measuring the extent to which GPT-2 assigns higher probability to years  $> \mathrm{YY}$ . It ranges from -1 to 1; higher is better. In contrast, cutoff sharpness is not intrinsically connected to greater-than. However, it quantifies an interesting behavior of GPT-2: the sharp cutoff between valid and invalid years. In doing so, it checks that the model depends on YY, and does not produce constant (but valid) output, e.g. by always outputting  $p(99) = 1$ . Cutoff sharpness ranges from -1 to 1; larger values indicate a sharper cutoff.
We perform this evaluation with our same 10,000-element dataset; on this dataset, GPT-2 achieves  $81.7\%$  probability difference (SD:  $19.3\%$ ) and a cutoff sharpness of  $6.0\%$  (SD:  $7.2\%$ ). Overall, both qualitative and quantitative results indicate that GPT-2 performs the greater-than operation on the year-span prediction task. For more study of GPT-2's behavior on this task, see Appendix A.
3 A Circuit for Year-Span Prediction
Having defined our task, we now aim to understand how our model performs it internally. Since the advent of pre-trained models, many methods, such as attention analysis [9] and probing [3], have sought to answer this question. Probing, which trains auxiliary models (probes) to extract information from model representations, has been particularly popular; it has been used to localize syntactic and semantic processing in many pre-trained LMs [37, 45, 15]. However, it has significant pitfalls: most crucially, probes can sometimes extract information from model representations that is irrelevant to model behavior [13, 40, 22]. To avoid this issue, other work has used causal interventions [36] to intervene on model internals, and observe changes in model behavior [20, 48, 17]. This ensures that our insights about model internals are actually functionally relevant to model behavior.
In light of this, we examine GPT-2's task performance by using causal techniques to identify a circuit: a minimal computational subgraph of our model that suffices to compute the task [35, 14]. While many causal methods aim to identify important components (nodes) of models [48, 11, 17], the circuits methodology instead considers important edges. Circuits thus examine not only important nodes, but also their interactions, and how they work together to support model behavior.
Below, we explain the path patching technique, and how to use it to find circuits (Section 3.1). We then find a circuit for greater-than, and prove its correctness (Section 3.2). Finally, we assign semantics to the nodes and edges of our circuit (Section 3.3). All of our experiments use the rust-circuit library, and our code is available at https://github.com/hannamw/gpt2-greater-than. For information on how to apply this methodology to other problems, see Appendix H.
3.1 Path Patching
To find a circuit, we use path patching, introduced by Wang et al. [50] and further described by Goldowsky-Dill et al. [21]. This technique determines how important a model component (e.g. an attention head or MLP) is to a task, by altering that component's inputs and observing model behavior post-alteration. It is much like causal mediation analysis or interchange interventions [48, 17]; however, unlike these, it allows us to constrain our intervention's effects to a specific path.
To illustrate this, consider a model's forward pass on its inputs as a directed acyclical graph. Its nodes are components such as attention heads or MLPs. The input of a node  $v$  is the sum of the outputs of all nodes with a direct edge to  $v$ . GPT-2 can be thought of as such a graph flowing from its input tokens to its logits (and thereafter, its probabilities), as depicted in Figure 3.
In path patching, we specify new input tokens, and a path of components through which they will reach the logits. For example, if we want to ascertain the effects of MLP 10 on the logits, we can patch the direct path (MLP 10, logits) with new input, which we call the 01-input: "The war lasted from the year 1701 to the year 17". We thus alter MLP 10's direct effects on the logits without changing its output to the attention and MLP of layer 11 (Figure 3). If the model's behavior (as indicated by its logits) changes, we can be sure that this is because MLP 10 is important to that behavior; it is not due to downstream components. Earlier methods like interchange interventions lack this distinction—when they alter a component, they affect all components downstream from it.
The specificity of path patching allows us to test detailed hypotheses. For example, imagine that we know that MLP 10 affects the logits both directly and via its effects on MLP 11. We want to know how important layer 10's attention is to the circuit via MLP 10. We can test this by patching two paths at once: (Attn 10, MLP 10, logits) and (Attn 10, MLP 10, MLP 11, logits), as in Figure 3. This allows us to pinpoint the relationship between precisely these two components, Attn 10 and MLP 10. This technique underpins our circuits approach: we search for a path starting in the inputs and ending in the logits that explains how our model performs the greater-than task.
To perform path patching, we need a new dataset that replaces a node's original inputs. To this end, we create the "01-dataset": we take each example in the original dataset and replace the last two digits YY of the start year with "01". If a component normally boosts logits of years  $>YY$ , patching it with the 01-dataset will cause it to boost the logits of years  $>01$ , inducing a larger error in the model.
Figure 3: A. The computational graph of GPT-2, run on our normal dataset. B: GPT-2, where the (MLP 10, logits) path is patched to receive 01-input. C. GPT-2, where the (Attn 10, MLP 10, logits) path receives 01-input. Nodes receiving normal input have blue output; nodes receiving 01-input have red output; nodes receiving both have purple output. Note that in B and C, there is no longer an edge connecting MLP 10 and the logits.
Figure 4: Iterative path-patching (IPP) heatmaps. Y-axis: layer of the component. X-axis: attention head number, or MLP. (X,Y): Change in probability difference induced by patching the corresponding component. A: Heatmap for the path ((X,Y), logits). B: Heatmaps for MLPs 8-11.
3.2 Circuit Components
MLPs We search for a circuit by identifying components that perform year-span prediction via their direct connection to the logits. We consider as potential components GPT-2's 144 attention heads (12 heads/layer  $\times$  12 layers), and 12 MLPs (1 per layer). We do so because the residual stream [14] that serves as input to the logits is simply the sum of these components' direct contributions (along with the token embeddings; we ignore these as they contain no YY information). If we consider each of these, we will not miss any components that contribute to this task. For details, see Appendix C.
We iteratively path patch each component's direct contributions to the logits, replacing its inputs with the 01-dataset. In our earlier notation, for a component of interest  $C$ , we patch the path  $(C, \logits)$  as in Figure 3 B, where  $C = \text{MLP} \, 10$ . We patch only one component at a time, and only at the end of the sentence; at other positions, these components cannot affect the logits directly.
After we patch a component, we run the model and record the probability difference, comparing it to that of the unpatched model. If patching a component caused model performance to change significantly, that component contributed to the model's computation of year-span prediction.
Figure 4 shows the results of this experiment; for computational reasons we run it using a smaller dataset (490 datapoints, 5 per year YY). The heatmap indicates that MLPs 8-11 are the most important direct contributors to the logits, along with a9.h1: attention layer 9's head 1. However, the MLPs cannot act alone: to compute years  $>$  YY, these MLPs at the end of the sentence must know the value of YY. But unlike attention heads, MLPs cannot attend to earlier tokens such as the YY token. Thus, we search for nodes that contribute to the circuit via these MLPs.
Attention Heads We find components that contribute to the circuit via the MLPs using more path patching. We start by patching components through MLP 11, since it is the furthest downstream; for a component of interest  $C$ , we patch  $(C, \text{MLP} 11, \text{logits})$ . We find that MLP 11 relies mostly on the 3 MLPs upstream of it (Figure 4), so we search for components that act via those MLPs.
We next find components that contribute to the circuit through MLP 10. For a given  $C$ , we patch  $(C, \text{MLP} 10, \text{logits})$  and  $(C, \text{MLP} 10, \text{MLP} 11, \text{logits})$ , as in Figure 3 C. We do so because MLP 10 contributes directly to two nodes in our circuit, the logits and MLP 11, and we want to know which nodes contribute via MLP 10 to the entire circuit. We repeat this procedure for MLPs 9 and 8.
The results in Figure 4 indicate that MLPs rely heavily on other MLPs upstream of them. MLPs 8 and 9, the furthest upstream of our MLPs, also rely on attention heads. MLP 9 relies on a9.h1, while MLP 8 relies on a8.h11, a8.h8, a7.h10, a6.h9, a5.h5, and a5.h1; we add these to our circuit. Many of these attention heads can be seen to contribute to the logits directly, though more weakly than the MLPs do. For this reason, we also add these heads' direct connections to the logits to our circuit.
Figure 5 visualizes the circuit we have found. We could further develop this by specifying a circuit from the token inputs to the logits; indeed, we do so in Appendix B. However, the present circuit
Figure 5: Left: Diagram of the year-span prediction circuit. Center: Diagram showing which GPT-2 components receive our standard dataset vs. our 01-dataset in the circuit evaluation experiment. Right: The probability heatmap (as in Figure 2) for the patched model.
already captures the most interesting portion of the model: the MLPs that compute greater-than. So, we instead provide evidence that our circuit is correct, and then analyze its constituent parts.
Evaluation Having defined our circuit, we perform another path-patching experiment to ensure it is correct. In this experiment, we give most of the model inputs from the 01-dataset. The model only receives our standard dataset via the paths specified in our circuit. So, our attention heads' contributions to the logits are backed by the standard dataset, as are their contributions to the MLPs, and the MLPs' contributions to one another. But, some components of the MLPs' inputs (those that come from model components not in the circuit) receive input from the 01-dataset as well. We stress that this is a difficult task, where the large majority of the model receives input that should push it to poor performance. For a diagram of the circuit and our evaluation, see Figure 5.
We perform this evaluation using the larger dataset, and almost entirely recover model performance. The probability difference is  $72.7\%$  (89.5% of the original) and the cutoff sharpness is  $8\%$  — sharper than pre-patching. This indicates that our circuit is mostly sufficient to compute this task. The circuit is also necessary: performing the opposite of the prior experiment, giving nodes in our circuit the 01-dataset, and those outside it the normal dataset, leaves GPT-2 unable to perform the task: it achieves a probability difference of  $-36.6\%$ .
If we target other circuits of a size and location similar to our circuit's, performance is related to the preservation of the paths from the input, to our attention heads, our MLPs (especially MLPs 9 and 10), to the logits. If the paths are interrupted (i.e. no attention head or no MLP overlap with the original circuit), performance is very low. If at least one path is preserved, performance improves with each additional component in common with the original circuit; MLPs have the biggest impact.
3.3 Circuit Semantics
Now, we interpret each circuit component, starting with the attention heads. We first perform a simple attention-pattern analysis of the heads in our circuit. Figure 6 shows which tokens our attention heads attend to at which positions. At the relevant (end) position, in-circuit attention heads attend to the YY position, suggesting that they detect the year which the output year must be greater than.
Figure 6: Attention patterns of a7.h11 and a8.h10. <bos> denotes GPT-2's start of sentence token.
Figure 7: Logit lens of a7.h11 and a8.h10. Axes as in Figure 2; blue indicates that the module upweights the output year, and red, that it downweights the year.
Next, we examine the contributions of attention heads using the logit lens approach [33]: we multiply each head's output by GPT-2's unembedding matrix, translating this output into unembedding (vocabulary) space. Note that here, we do not only view logit lens as a tool for obtaining intermediate estimates of model predictions [4]. Rather, we also use it to understand components' outputs more generally: the logit lens can capture how such outputs shape model predictions, but it can also capture how these outputs add information to the residual stream in unembedding space.
We visualize the heads' outputs for each sentence in our small dataset in Figure 7. Attention head outputs for a sentence with start year YY have a high dot product with the embedding vector for YY, as shown by the blue diagonal in the plots; this indicates that the head upweights YY, making it a more likely output. Note that YY was not just upweighted highly compared other 2-digit years; YY was the most highly upweighted token across all tokens. Given our earlier analysis, we therefore hypothesize that these attention heads identify the start year (at the YY position), and indicate it via a spike in unembedding space of the residual stream at the end position; they thus communicate YY to downstream components.
We similarly apply the logit lens to the outputs of MLPs 8-11 (Figure 8). The results indicate that MLPs of 9 and 10 directly specify which years are greater than YY: the logit lens of each layer's output has an upper triangular pattern, indicating that they upweight precisely those years greater than YY. MLP 11 plays a similar role, but seems to upweight roughly the first 50 years after YY, enforcing a maximum duration for the event in the sentence. However, MLP 8 is unusual: its logit lens shows a diagonal pattern, but no upper triangular pattern that would indicate that it computes greater-than.