Arithmetic Without Algorithms: Language Models Solve Math With a Bag of Heuristics — part 2/5
Cleaned-text code points [14979, 29819). Bibliography omitted here; source.md includes it.
where 𝐡 i n l , 𝐡 o u t l ∈ ℝ d \mathbf{h}_{in}^{l},\mathbf{h}_{out}^{l}\in\mathbb{R}^{d} are the input and output representations of the MLP at layer l l , respectively. 𝐡 p o s t l ∈ ℝ d m l p \mathbf{h}_{post}^{l}\in\mathbb{R}^{d_{mlp}} is the output of the up-projection of the MLP, 2 2 2 The up-projection MLP i n \operatorname{MLP}_{in} is implemented differently in each LLM we analyze. See Appendix C . where we define the n th n^{\text{th}} value 𝐡 p o s t l , n ∈ ℝ \mathbf{h}_{post}^{l,n}\in\mathbb{R} as a neuron. 𝐖 o u t l ∈ ℝ d m l p × d \mathbf{W}_{out}^{l}\in\mathbb{R}^{d_{mlp}\times d} is the output projection matrix, and 𝐯 o u t l , n \mathbf{v}_{out}^{l,n} is its n th n^{\text{th}} row vector. Biases are omitted. By expressing the output representation 𝐡 o u t l \mathbf{h}_{out}^{l} as a linear combination of row vectors 𝐯 o u t l , n \mathbf{v}_{out}^{l,n} and their corresponding neuron activations 𝐡 p o s t l , n \mathbf{h}_{post}^{l,n} , we can identify the neurons that most affect the completion of arithmetic prompts.
To measure the effect of each neuron 𝐡 p o s t l , n \mathbf{h}_{post}^{l,n} , we perform activation patching experiments on individual neurons (as described in Section 2.1 ), and measure the average effect across prompts. We observe that very few neurons have a high effect; an example for layer l = 17 l=17 is shown in Figure 4(a) . Additionally, we notice the neurons with the highest effect are different between operators. In fact, roughly 45% of the important neurons for each operator are unique ( Appendix D ). Thus, when analyzing the circuit at the neuron level, we analyze it as 4 separate circuits—one for each arithmetic operator. We hypothesize that for each operator, the highest-effect neurons are sufficient to explain most of the model’s arithmetic behavior. To verify this, we measure the faithfulness of the arithmetic circuit when mean ablating non-circuit components and lower effect MLP neurons in middle- and late-layer MLPs. The results ( Figure 4(b) ) confirm that only 200 neurons (roughly 1.5%) per layer are needed to achieve high faithfulness and correctly compute arithmetic prompts .
### 3.2 MLP neurons act as memorized heuristics
To understand how the top- 200 200 middle- and late-layer important MLP neurons contribute to the generation of correct answers, we view them as key-value memories ( Geva et al., 2021 ) . 3 3 3 Not to be confused with the attention heads’ keys and values. In this view, the input to each MLP layer 𝐡 i n l \mathbf{h}_{in}^{l} is multiplied by a key (a row vector in the MLP input weight matrix) to generate a neuron activation 𝐡 p o s t l , n \mathbf{h}_{post}^{l,n} that determines how strongly does a value (a row vector 𝐯 o u t l , n \mathbf{v}_{out}^{l,n} in the MLP output weight matrix) gets written to the MLP output ( Equation 3 ). Geva et al. (2021) demonstrated that keys correspond to specific topics or n-grams, triggering high neuron activations when these are given as input, and their corresponding values represent tokens that can serve as appropriate completions for these topics or n-grams. Building on this insight, we hypothesize that (i) in arithmetic contexts, keys correspond to numerical patterns, e.g., a neuron might activate strongly when both operands in an arithmetic operation are odd numbers; and (ii) the associated value vectors encode numerical tokens that represent plausible answers to the key patterns.
To test the first hypothesis, we investigate the activation pattern of the top- 200 200 neurons in each layer. For each neuron l , n l,n , we plot the activations 𝐡 p o s t l , n \mathbf{h}_{post}^{l,n} (at position p = 4 p=4 ) as a function of operand values, separately for each operator. We find that many neurons in the arithmetic circuit exhibit distinct, human-identifiable patterns. For instance, in “ 226 − 68 = 226-68= ”, neuron 𝐡 p o s t 24 , 12439 \mathbf{h}_{post}^{24,12439} shows high activation values for subtraction prompts with results between 150 150 and 180 180 ( Figure 1 ). Additional examples are provided in Appendix J .
To verify the second hypothesis, we check whether the tokens embedded in the value vectors of the top neurons relate to their activation patterns. Using the Logit Lens ( nostalgebraist, 2020 ) , a method of projecting a vector 𝐯 ∈ ℝ d \mathbf{v}\in\mathbb{R}^{d} onto a probability distribution over the vocabulary space ℙ d v o c a b \mathbb{P}^{d_{vocab}} , we project each value vector 𝐯 o u t l , n \mathbf{v}_{out}^{l,n} to find the numerical tokens whose logits are highest. This reveals two distinct patterns: First, in some neurons, the activation pattern depends on both operands and the value vector encodes the expected result of the arithmetic calculation ( Figure 1 b,c). We term such neurons direct heuristics . Second, in neurons where the activation pattern depends on a single operand, the value vectors often encode features for downstream processing, rather than the result tokens directly ( Figure 1 a). We term such neurons indirect heuristics . Next, we demonstrate how these heuristic neurons combine to produce correct arithmetic answers.
## 4 Arithmetic prompts are answered with a bag of heuristics
Observing the example prompt, “ 226 − 68 = 226-68= ”, we have shown that it satisfies the pattern of several heuristic neurons, where each such neuron slightly increases the logit of the result token, r = 158 r=158 ( Figure 1 ). These small increases combine to promote the correct token as the final answer. We hypothesize that a combination of independent heuristics—termed a bag of heuristics —emerges across arithmetic prompts, comprising the mechanism used by the model to produce correct answers.
### 4.1 Classifying neurons to heuristic types
To present evidence for the causal effect of the bag of heuristics on generating correct answers, we first systematically classify neurons into heuristic types. Through manual observation of key activation patterns, we identify several categories of human-identifiable heuristics, exemplified in Figure 5 , and further detailed in Appendix E . To determine if a neuron n n at layer l l implements a specific heuristic, we examine the intersection between the prompts that activate the neuron and the prompts expected to be activated for this heuristic. A visual example of this procedure is shown in Figure 6 . An automated algorithm of this approach is described in Appendix F .
We apply this algorithm to each pair of important MLP neuron n n in layer l l and heuristic H H . Through this method, we classify as arithmetic heuristics 91% of the 3,200 top neurons for each operator (200 per layer across 16 layers). Manual inspection of neurons that fail to classify into one of the defined heuristics reveals patterns that are not clearly identifiable ( Appendix J ).
### 4.2 Heuristic types are combined to answer arithmetic prompts
Following the classification of the important neurons in each middle- and late-layer MLP into one or several heuristic types, we provide evidence that the bag of heuristics is the primary mechanism the model uses to correctly answer arithmetic prompts. We show this through two ablation experiments.
##### Knocking out neurons by heuristic type.
We now verify that the neurons in each heuristic type contribute to the accuracy of associated prompts by knocking out entire heuristic types and observe the resulting changes in model accuracy. We define a prompt as associated with a heuristic if and only if its components meet the conditions specified by that heuristic. (For instance, the example prompt “ 226 − 68 = 226-68= ” is associated with the heuristic “ o p 1 ≡ 0 ( mod 2 ) op1\equiv 0\pmod{2} ”.) For each heuristic, we sample two sets of 100 100 correctly completed prompts each, one containing prompts associated with the heuristic and the other containing prompts not associated with that heuristic. For each heuristic, we knock out all neurons classified into it (by setting each h p o s t l , n h_{post}^{l,n} activation to zero) and then remeasure the accuracy on both sets of prompts. We expect a higher decrease in accuracy on the associated prompts, since we claim each heuristic is causally linked only to its associated prompts.
The results ( Figure 7 ) show that ablating neurons of a specific heuristic causes a significant accuracy drop on associated prompts, more than on not associated prompts, on average. This confirms the causal importance of heuristic neurons in promoting correct answer logits specifically in prompts that are associated with their heuristic type, verifying the targeted functional role of these heuristics.
However, the ablation does not result in a complete accuracy drop; it causes an average drop of 29% out of 95% average pre-ablation accuracy. We find two reasons for this. First, some heuristics have low recall: they do not apply to all associated prompts as they should (see Appendix J ). Second, each prompt relies on several unrelated heuristic types, so even when one is ablated, others still contribute to increasing the correct answer’s logit. In the following ablation experiment, we verify that this interplay of heuristics provides a fuller image than focusing on one heuristic at a time.
##### Knocking out neurons by prompt.
To provide further evidence that the bag of heuristics is causally linked to correct arithmetic completion, we conduct a second ablation experiment. For each correct prompt, we identify the heuristic types that should affect it based on its operands and ground truth result. We then ablate the neurons with the highest classification scores ( Section 4.1 ) in these heuristics, up to a certain neuron count, and check if the model’s completion changes.
The results ( Figure 8 ) show that ablating neurons from associated heuristics significantly drops the model’s accuracy, much more than the accuracy drop caused by ablating the same number of randomly chosen neurons from unassociated heuristics. This demonstrates that we can identify the neurons important to a given prompt solely based on its associated heuristics. This also indicates a causal link between the neurons belonging to several heuristics and the prompt’s correct completion. This supports our bag of heuristics claim: each heuristic only slightly boosts the correct answer logit, but combined, they cause the model to produce the correct answer with high probability.
### 4.3 Failure modes of the bag of heuristics
The bag of heuristics mechanism employed in Llama3-8B does not generalize perfectly: it fails to achieve perfect accuracy across all arithmetic prompts ( Appendix H ). This limitation contrasts with the theoretical robustness of a genuine algorithmic approach. Here, we aim to elucidate the specific failures of this mechanism, focusing on why it falters for some prompts.
We hypothesize that the bag of heuristics mechanism completes prompts incorrectly in two ways. (1) The “bag” might not be big enough; i.e., a prompt might lack sufficient associated neurons. (2) the heuristics might have imperfect recall (e.g., a neuron that fires for most prompts where the first operand is even, but does not fire for the prompt “ 226 − 68 = 226-68= ”) or have low logits for the correct answer token in the value vectors.
To test these hypotheses, we randomly sample 50 correctly completed and 50 incorrectly completed prompts. To test hypothesis (i), we count the number of heuristic neurons associated with each prompt. We find that on average, incorrect prompts have more heuristic neurons associated with them than correct prompts. Therefore, we find no support for this hypothesis. To check hypothesis (ii), we calculate the total contribution of all heuristic neurons to the logit of the correct answer for each prompt. This measurement considers both the specific activation of each neuron for the prompt, as well as the logit of the correct answer token embedded in each neuron’s value vector. On average, there is indeed a slight advantage in the total logit contribution for correct prompts over incorrect prompts ( Figure 9 ). This suggests that the primary reason for the bag of heuristics failure on certain prompts is poor promotion of correct answer logits, rather than a lack of heuristics.
## 5 Tracking heuristics development across training steps
Does the bag of heuristics emerge as the primary arithmetic mechanism from the onset of training, or does it override an earlier, different mechanism that initially drives arithmetic performance? We conduct an analysis of heuristic development across the training trajectory of the Pythia-6.9B model ( Biderman et al., 2023 ) , due to the public availability of its training checkpoints. Specifically, we analyze the model at its final checkpoint (143K steps) and at 10K-step intervals down to 23K steps. The 23K checkpoint is the earliest checkpoint showing good arithmetic performance; Thus, we begin our analysis at this checkpoint. To guide this analysis, we aim to answer three sub-questions:
When do the final heuristic neurons first appear? We examine when each heuristic neuron from the final checkpoint first appears during training. For each neuron classified into a particular heuristic type, we check if the same neuron gets classified into the same heuristic in earlier checkpoints. Averaging this measure across all heuristic types and operators provides insight into when the final heuristics initially appear during training. We observe ( Figure 10(a) ) that the model develops its final heuristic mechanism gradually across training, starting from an early checkpoint.
Do additional heuristic neurons exist mid-training? Next, for each mid-training checkpoint, we investigate whether its heuristic neurons that are mutual with the final checkpoint make up the entire heuristic mechanism in that checkpoint, or whether other heuristics exist that later become vestigial. We examine the faithfulness ( Section 2.1 ) of the arithmetic circuit at each checkpoint—once when including only the neurons mutual with the final checkpoint, and once when including all heuristic neurons in that checkpoint. The difference between these two measurements gives us an estimate of the importance of the mutual heuristics. Using this metric, we observe that these final heuristics explain most of the circuit performance for each intermediate checkpoint: they account for an average of 79% of the total heuristics’ contribution to accuracy at each checkpoint. This indicates that, while other non-mutual heuristics exist in each checkpoint, these are less important to the circuit’s accuracy and slowly become vestigial as the circuit converges to its final form.