Language Models Use Trigonometry to Do Addition — part 2/4
Cleaned-text code points [15679, 31564). Bibliography omitted here; source.md includes it.
To leverage this technique, we follow Engels et al. (2024) and input our fit for h a , clean l h_{a,\text{clean}}^{l} when patching. This allows us to causally determine if our fit preserves the information the model uses for the computation. We compare our k k Fourier feature helical fit with four baselines: using the actual h a , c l e a n l h_{a,clean}^{l} (layer patch), the first 2 k + 1 2k+1 PCA components of h a , c l e a n l h_{a,clean}^{l} (PCA), a circular fit with k k Fourier components (circle), and a polynomial fit with basis terms B ( a ) = [ a , a 2 , … a 2 k + 1 ] B(a)=[a,a^{2},...a^{2k+1}] (polynomial). For each value of k k , we choose the combination of Fourier features that maximizes 1 L ∑ l L D a l \frac{1}{L}\sum_{l}LD_{a}^{l} as the best set of Fourier features.
In Fig. 4 , we see that the helical fit is most performant against baselines, closely followed by the circular fit. This implies that Fourier features are predominantly used to compute addition. Surprisingly, the k = 4 k=4 full helical and circular fits dominate the strong PCA baseline and approach the effect of layer patching, which suggests that we have identified the correct “variables” of computation for addition. Additionally, we note a sharp jump between the fit for layer 0’s input (the output of the embedding) and layer 1’s input, aligning with evidence from Nikankin et al. (2024) that layer 0 is necessary for numerical processing.
In Appendix C.2 , we provide evidence that Llama3.1-8B and Pythia-6.9B also use helical numerical representations. Additionally, we provide evidence that our fits are not overfitting by using a train-test split with no meaningful effect on our results. The helix functional form is not overly expressive, as a helix trained on a randomized order of a a is not causally relevant. We also observe continuity when values of a a that the helix was not trained on are projected into the helical subspace. This satisfies the definition of a nonlinear feature manifold proposed by Olah & Jermyn (2024) , and provides additional evidence for the argument of Engels et al. (2024) against the strongest form of the Linear Representation Hypothesis.
### 4.5 Is the Helix the Full Picture?
To identify if the helix sufficiently explains the structure of numbers in LLMs, we test on five additional tasks. 1. a − 23 a-23 for a ∈ [ 23 , 99 ] a\in[23,99]
2. a / / 5 a//5 (integer division) for a ∈ [ 0 , 99 ] a\in[0,99]
3. a ∗ 1.5 a*1.5 for even a ∈ [ 0 , 98 ] a\in[0,98]
4. a mod 2 a\mod 2 for a ∈ [ 0 , 99 ] a\in[0,99]
5. If x − a = 0 x-a=0 , what is x = x= for a ∈ [ 0 , 99 ] a\in[0,99]
For each task, we fit full helices with T = [ 2 , 5 , 10 , 100 ] T=[2,5,10,100] and compare against baselines. In Table 1 , we describe our results on these tasks by listing max l L D a l \max_{l}LD_{a}^{l} , which is the maximal causal power of each fit (full plot and additional task details in Appendix C.2 ). Notably, while the helix is causally relevant for all tasks, we see that it underperforms the PCA baseline on tasks 2, 3, and 5. This implies that there is potentially additional structure in numerical representations that helical fits do not capture. However, we are confident that the helix is used for addition. When ablating the helix dimensions from the residual stream (i.e. ablating C † C^{\dagger} from h a l h_{a}^{l} ), performance is affected roughly as much as ablating h a l h_{a}^{l} entirely (Fig. 21 , Appendix C.1 ).
Thus, we conclude that LLMs use a helical representation of numbers to compute addition, although it is possible that additional structure is used for other tasks.
## 5 LLMs Use the Clock Algorithm to Compute Addition
### 5.1 Introducing the Clock Algorithm
Taking inspiration from Nanda et al. (2023a) , we propose that LLMs manipulate helices to compute addition using the “Clock” algorithm.
Since we have already shown that models represent a a and b b as helices (Appendix C.2 ), we provide evidence for the last three steps in this section. In Fig. 5 we observe that last token hidden states are well modeled by h = l = helix ( a , b , a + b ) h_{=}^{l}=\mathrm{helix}(a,b,a+b) , where helix ( x , y ) \mathrm{helix}(x,y) is shorthand to denote helix ( x ) + helix ( y ) \mathrm{helix}(x)+\mathrm{helix}(y) . Remarkably, despite only using 9 parameters, at some layers helix ( a + b ) \mathrm{helix}(a+b) fits last token hidden states better than a 27 dimensional PCA. The a + b a+b helix having such causal power implies it is at the heart of the computation.
In Appendix D , we show that other LLMs also use helix ( a + b ) \mathrm{helix}(a+b) . Since the crux of the Clock algorithm is computing the answer helix for a + b a+b , we take this as compelling evidence that all three models use the Clock algorithm. However, we would like to understand how specific LLM components implement the algorithm. To do so, we focus on GPT-J.
In Fig. 6 , we use activation patching to determine which last token MLP and attention layers are most influential for the final result. We also present path patching results, which isolates how much components directly contribute to logits. For example, MLP18’s total effect (TE, activation patching) includes both its indirect effect (IE), or how MLP18’s output is used by downstream components like MLP19, and its direct effect (DE, path patching), or how much MLP18 directly boosts the answer logit. 1 1 1 For more on path patching, we refer readers to Goldowsky-Dill et al. (2023) ; Wang et al. (2023) In Fig. 6 , we see that MLPs dominate direct effect.
We now investigate specific attention heads, MLPs, and individual neurons.
### 5.2 Investigating Attention Heads
In GPT-J, every attention layer is the sum of 16 attention heads whose outputs are concatenated. We activation and path patch each attention head on the last token and rank them by total effect (TE). To determine the minimal set of attention heads required, we activation patch k k attention heads at once, and find the minimum k k such that their combined total effect approximates patching in all attention heads. In Appendix D.1 , we see that patching k = 17 k=17 heads achieves 80% of the effect of patching in all 448 attention heads, and we choose to round up to k = 20 k=20 heads (83.9% of effect).
Since attention heads are not as influential as MLPs in Fig. 6 , we hypothesize that they primarily serve two roles: 1) moving the a , b a,b helices to the last token to be processed by downstream components ( a , b a,b heads) and 2) outputting the a + b a+b helix directly to logits ( a + b a+b heads). Some mixed heads output all three a , b , and a + b a,b,\text{ and }a+b helices. We aim to categorize as few attention heads as mixed as possible.
To categorize attention heads, we turn to two metrics. c a , b c_{a,b} is the confidence that a certain head is an a , b a,b head, which we quantify with c a , b = ( 1 − DE TE ) helix ( a , b ) helix ( a , b , a + b ) c_{a,b}=(1-\frac{\mathrm{DE}}{\mathrm{TE}})\frac{\mathrm{helix}(a,b)}{\mathrm{helix}(a,b,a+b)} . The first term represents the fractional indirect effect of the attention head, and the second term represents the head’s total effect recoverable by just using the a , b a,b helices instead of helix ( a , b , a + b ) \mathrm{helix}(a,b,a+b) . Similarly, we calculate c a + b c_{a+b} as the confidence the head is an a + b a+b head, using c a + b = DE TE helix ( a + b ) helix ( a , b , a + b ) c_{a+b}=\frac{\mathrm{DE}}{\mathrm{TE}}\frac{\mathrm{helix}(a+b)}{\mathrm{helix}(a,b,a+b)} .
We sort the k = 20 k=20 heads by c = max ( c a , b , c a + b ) c=\max({c_{a,b},c_{a+b}}) . If a head is an a + b a+b head, we model its output using the a + b a+b helix and allow it only to output to logits (no impact on downstream components). If a head is an a , b a,b head, we restrict it to outputting helix ( a , b ) \mathrm{helix}(a,b) . For m = [ 1 , 20 ] m=[1,20] , we allow m m heads with the lowest c c to be mixed heads, and categorize the rest as a , b a,b or a + b a+b heads. We find that categorizing m = 4 m=4 heads as mixed is sufficient to achieve almost 80% of the effect of using the actual outputs of all k = 20 k=20 heads. Thus, most important attention heads obey our categorization. We list some properties of each head type below.
• a , b a,b heads (11/20): In layers 9-14 (but two heads in l = 16 , 18 l=16,18 ), attend to the a , b a,b tokens, and output a , b a,b helices which are used mostly by downstream MLPs.
• a + b a+b heads (5/20): In layers 24-26 (but one head in layer 19), attend to the last token, take their input from preceding MLPs, and output the a + b a+b helix to logits.
• Mixed heads (4/20): In layers 15-18, attend to the a , b , and a + b a,b,\text{ and }a+b tokens, receive input from a , b a,b attention heads and previous MLPs, and output the a , b , and a + b a,b,\text{ and }a+b helices to downstream MLPs.
For evidence of these properties refer to Appendix D.1 . Notably, only mixed heads are potentially involved in creating the a + b a+b helix, which is the crux of the computation, justifying our conclusion from Fig. 6 that MLPs drive addition.
### 5.3 Looking at MLPs
GPT-J seems to predominantly rely on last token MLPs to compute a + b a+b . To identify which MLPs are most important, we first sort MLPs by total effect, and patch in k = [ 1 , L = 28 ] k=[1,L=28] MLPs to find the smallest k k such that we achieve 95% of the effect of patching in all L L MLPs. We use a sharper 95% threshold because MLPs dominate computation and because there are so few of them. Thus, we use k = 11 k=11 MLPs in our circuit, specifically MLPs 14-27, with the exception of MLPs 15, 24, and 25 (see Appendix D.2 for details).
We hypothesize that MLPs serve two functions: 1) reading from the a , b a,b helices to create the a + b a+b helix and 2) reading from the a + b a+b helix to output the answer in model logits. We make this distinction using two metrics: helix ( a + b ) \mathrm{helix}(a+b) /TE, or the total effect of the MLP recoverable from modeling its output with helix ( a + b ) \mathrm{helix}(a+b) , and DE/TE ratio. In Fig. 7 , we see that the outputs of MLPs 14-18 are progressively better modeled using helix ( a + b ) \mathrm{helix}(a+b) . Most of their effect is indirect and thus their output is predominantly used by downstream components. At layer 19, helix ( a + b ) \mathrm{helix}(a+b) becomes a worse fit and more MLP output affects answer logits directly. We interpret this as MLPs 14-18 “building” the a + b a+b helix, which MLPs 19-27 translate to the answer token a + b a+b .
However, our MLP analysis has focused solely on MLP outputs. To demonstrate the Clock algorithm conclusively, we must look at MLP inputs. Recall that GPT-J uses a simple MLP: MLP ( x ) = σ ( x W up ) W down \mathrm{MLP}(x)=\sigma\left({xW_{\mathrm{up}}}\right)W_{\mathrm{down}} . x x is a vector of size ( 4096 , ) (4096,) representing the residual stream, and W up W_{\mathrm{up}} is a ( 4096 , 16384 ) (4096,16384) projection matrix. The input to the MLP is thus the 16384 dimensional x W up xW_{\mathrm{up}} . We denote the n n th dimension of the MLP input as the n n th neuron preactivation, and move to analyze these preactivations.
### 5.4 Zooming in on Neurons
Activation patching the 27 ∗ 16384 27*16384 neurons in GPT-J is prohibitively expensive, so we instead use the technique of attribution patching to approximate the total effect of each neuron using its gradient (see Kramár et al. (2024) ). We find that using just 1% of the neurons in GPT-J and mean ablating the rest allows for the successful completion of 80% of prompts (see Appendix D.2 ). Thus, we focus our analysis on this sparse set of k = 4587 k=4587 neurons.
#### 5.4.1 Modeling Neuron Preactivations
For a prompt a + b a+b , we denote the preactivation of the n n th neuron in layer l l as N n l ( a , b ) N_{n}^{l}(a,b) . When we plot a heatmap of N n l ( a , b ) N_{n}^{l}(a,b) for top neurons in Fig. 8 , we see that their preactivations are periodic in a , b a,b , and a + b a+b . When we Fourier decompose the preactivations as a function of a + b a+b , we find that the most common periods are T = [ 2 , 5 , 10 , 100 ] T=[2,5,10,100] , matching those used in our helix parameterization (Appendix D.2 ). This is sensible, as the n n th neuron in a layer applies W u p n W_{up}^{n} of shape ( 4096 , ) (4096,) to the residual stream, which we have effectively modeled as a helix ( a , b , a + b ) \mathrm{helix}(a,b,a+b) . Subsequently, we model the preactivation of each top neuron as
N n l ( a , b ) = ∑ t = a , b , a + b c t t + ∑ T = [ 2 , 5 , 10 , 100 ] c T t cos ( 2 π T ( t − d T t ) ) N_{n}^{l}(a,b)=\sum_{t=a,b,a+b}c_{t}t+\sum_{T=[2,5,10,100]}c_{Tt}\cos\left(\frac{2\pi}{T}(t-d_{Tt})\right) (3)
For each neuron preactivation, we fit the parameters c c and d d in Eq. 3 using gradient descent (see Appendix D.2 for details). In Fig. 8 , we show the highest magnitude fit component for a selection of top neurons.
We evaluate our fit of the top k k neurons by patching them into the model, mean ablating all other neurons, and measuring the resulting accuracy of the model. In Fig. 9 , we see that our neuron fits provide roughly 75% of the performance of using the actual neuron preactivations. Thus, these neurons are well modeled as reading from the helix.
#### 5.4.2 Understanding MLP Inputs
We use our understanding of neuron preactivations to draw conclusions about MLP inputs. To do so, we first path patch each of the top k k neurons to find their direct effect and calculate their DE/TE ratio. For each neuron, we calculate the fraction of their fit that helix ( a + b ) \mathrm{helix}(a+b) explains, which we approximate by dividing the magnitude of c T , a + b c_{T,a+b} terms by the total magnitude of c T t c_{Tt} terms in Eq. 3 . For each circuit MLP, we calculate the mean of both of these quantities across top neurons, and visualize them in Fig. 10 .
Once again, we see a split at layer 19, where earlier neurons’ preactivation fits rely on a , b a,b terms, while later neurons use a + b a+b terms and write to logits. Since the neuron preactivations represent what each MLP is “reading” from, we combine this result with our evidence from Section 5.3 to summarize the role of MLPs in addition. • MLPs 14-18 primarily read from the a , b a,b helices to create the a + b a+b helix for downstream processing.
• MLPs 19-27 primarily read from the a + b a+b helix to write the answer to model logits.
Thus, we conclude our case that LLMs use the Clock algorithm to do addition, with a deep investigation into how GPT-J implements this algorithm.
### 5.5 Limitations of Our Understanding
There are several aspects of LLM addition we still do not understand. Most notably, while we provide compelling evidence that key components create helix ( a + b ) \mathrm{helix}(a+b) from helix ( a , b ) \mathrm{helix}(a,b) , we do not know the exact mechanism they use to do so. We hypothesize that LLMs use trigonometric identities like cos ( a + b ) = cos ( a ) cos ( b ) − sin ( a ) sin ( b ) \cos(a+b)=\cos(a)\cos(b)-\sin(a)\sin(b) to create helix ( a + b ) \mathrm{helix}(a+b) . However, like the originator of the Clock algorithm Nanda et al. (2023a) , we are unable to isolate this computation in the model. This is unsurprising, as there is a large solution space for how models choose to implement low-level details of algorithms. For example, Yip et al. (2024) finds that in Zhong et al. (2023) ’s “Pizza” algorithm for modular addition, MLPs in one layer transformers implement numerical integration techniques to transform cos ( k 2 ( a + b ) ) \cos(\frac{k}{2}(a+b)) to cos ( k ( a + b ) ) \cos(k(a+b)) .