Pre-trained Large Language Models Use Fourier Features to Compute Addition — part 2/4

Cleaned-text code points [15633, 31346). Bibliography omitted here; source.md includes it.

To understand the role various frequency components play for the addition task, we introduce low-pass and high-pass filters ℱ \mathcal{F} . For an intermediate state h h , and a set of frequencies Γ = { γ 1 , … , γ k } \Gamma=\{\gamma_{1},\dotsc,\gamma_{k}\} , the filter ℱ ⁡ ( h , Γ ) \mathcal{F}(h;\Gamma) returns the vector h ~ \widetilde{h} that is closest in L 2 L_{2} distance to h h subject to the constraint that the Fourier decomposition of W U ​ h ~ W^{U}\widetilde{h} at every frequency γ i \gamma_{i} is 0 0 . We show in Appendix A that this has a simple closed-form solution involving a linear projection. We then apply either a low-pass filter by taking Γ \Gamma to be all the components whose frequencies are greater than the frequency of the τ \tau -th component for some threshold τ \tau (i.e., removing high-frequency components), and a high-pass filter by taking Γ \Gamma to be all the components whose frequencies are less than the frequency of the τ \tau -th component (i.e., removing low-frequency components). As in the previous subsection, we take the high-frequency threshold τ = 50 \tau=50 for the following experiments (see Appendix B for more details).

#### Different roles of frequency components in approximation and classification tasks.

We evaluated the fine-tuned GPT-2-XL model on the test dataset with different frequency filters applied to all of the output of MLP and attention modules. The results, presented in Table 1 , indicate that removing low-frequency components from attention modules or high-frequency components from MLP modules does not impact performance. This observation suggests that attention modules are not crucial for approximation tasks, and MLP modules are less significant for classification tasks.

Eliminating high-frequency components from attention results in a noticeable decrease in accuracy. Furthermore, removing high-frequency components from both the attention and MLP modules simultaneously leads to an even greater reduction in accuracy. This finding corresponds with observations from Figure 1 b,c and Figure 3 , which indicate that both MLP and attention modules are involved in classification tasks due to the presence of high-frequency components in the logits. However, the approximation tasks are primarily performed by the MLP modules alone.

The errors induced by these ablations align with our mechanistic understanding. Ablating low-frequency parts of MLPs leads to off-by 10 10 , 50 50 , and 100 100 errors: the model fails to perform the approximation subtask, though it still accurately predicts the unit digit. Conversely, ablating high-frequency parts of attention leads to small errors less than 6 6 in magnitude: the model struggles to accurately predict the units digit, but it can still estimate the overall magnitude of the answer. See Figure 20 in the Appendix for more details. These observations validate our hypothesis that low-frequency components are crucial for approximation, while high-frequency components are vital for classification. The primary function of MLP modules is to approximate the magnitude of outcomes using low-frequency components, while the primary role of attention modules is to ensure accurate classification by determining the correct unit digit.

## 4 Effects of Pre-training

The previous section shows that pre-trained LLMs leverage Fourier features to solve the addition problem. Now, we study where the models’ reliance on Fourier features comes from. In this section, we demonstrate that LLMs learn Fourier features in the token embeddings for numbers during pre-training. These token embeddings are important for achieving high accuracy on the addition task: models trained from scratch achieve lower accuracy, but adding just the pre-trained token embeddings fixes this problem. We also show that pre-trained models leverage Fourier features not only when fine-tuned, but also when prompted.

### 4.1 Fourier features in Token Embedding

#### Number embedding exhibits approximate sparsity in the Fourier space.

Let W E ∈ ℝ p × D W^{E}\in\mathbb{R}^{p\times D} , where p = 521 p=521 and D D is the size of the token embeddings, denote the token embedding for numbers. We apply the discrete Fourier transform to each column of W E W^{E} to obtain a matrix V ∈ ℝ p × D V\in\mathbb{R}^{p\times D} , where each row represents a different Fourier component. Then we take the L 2 L_{2} norm of each row to yield a p p -dimensional vector. Each component j j in this vector measures the overall magnitude of the j j -th Fourier component across all the token embedding dimensions. Figure 5 a shows the magnitude of different Fourier components in the token embedding of GPT-2-XL. We see that the token embedding has outlier components whose periods are 2 , 2.5 , 5 2,2.5,5 , and 10 10 . Therefore, similar to how the model uses different Fourier components to represent its prediction (as shown in Section 3.2 ), the token embeddings represent numbers with different Fourier components. Figure 14 in the Appendix shows that the token embeddings of other pre-trained models have similar patterns the Fourier space. This suggests that Fourier features are a common attribute in the token embedding of pre-trained LLMs. In Figure 5 b, we use t-SNE and k k -means to visualize the token embedding clustering. We can see that numbers cluster not only by magnitude but also by their multiples of 10 10 .

### 4.2 Contrasting Pre-trained Models with Models Trained from Scratch

To understand the necessity of Fourier features for the addition problem, we trained the GPT-2-XL model from scratch on the addition task with random initialization. After convergence, it achieved only 94.44 % 94.44\% test accuracy (recall that the fine-tuned GPT-2-XL model achieved 99.74 % 99.74\% accuracy).

#### Fourier features are learned during pre-training.

Figure 6 shows that there are no Fourier features in the intermediate logits of the GPT-2-XL model trained from scratch on the addition task. Furthermore, Figure 7 a shows that the token embeddings also have no Fourier features. Without leveraging Fourier features, the model merely approximates the correct answer without performing modular addition, resulting in frequent off-by-one errors between the prediction and the correct answer (see details in Figure 22 ).

#### Pre-trained token embeddings improve model training.

We also trained GPT-2-small, with 124 124 million parameters and 12 12 layers, from scratch on the addition task. GPT-2-small often struggles with mathematical tasks [ 27 ] . This model achieved a test accuracy of only 53.95 % 53.95\% after convergence. However, when we freeze the token embedding layer and randomly initialize the weights for all other layers before training on the addition task, the test accuracy increases to 100 % 100\% , with a significantly faster convergence rate. This outcome was consistently observed across five different random seeds, as illustrated in Figure 7 b. This demonstrates that given the number embeddings with Fourier features, the model can effectively learn to leverage these features to solve the addition task.

### 4.3 Fourier Features in Prompted Pre-Trained Models

Finally, we ask whether larger language models use similar Fourier features during prompting.

#### Pre-trained LLMs use Fourier features to compute addition during in-context learning.

We first test on the open-source models GPT-J [ 43 ] with 6B parameters, and Phi-2 [ 21 ] with 2.7B parameters on the test dataset. Without in-context learning, the model cannot perform addition tasks. Therefore, we use 4-shot in-context learning to test its performance. Their absolute errors are predominantly multiples of 10 10 : 93% of the time for GPT-J, and 73% for Phi-2 . Using the Fourier analysis framework proposed in Section 3.2 , we demonstrate that for Phi-2 and GPT-J, the outputs of MLP and attention modules exhibit approximate sparsity in Fourier space across the last 15 15 layers (Figure 8 and Figure 18 ). This evidence strongly suggests that these models leverage Fourier features to compute additions.

#### Closed-source models exhibit similar behavior.

We study the closed-source models GPT-3.5 [ 33 ] , GPT-4 [ 34 ] , and PaLM-2 [ 16 ] . While we cannot analyze their internal representations, we can study whether their behavior on addition problems is consistent with reliance on Fourier features. Since closed-source LLMs are instruction tuned and perform well without in-context learning, we conduct error analysis with 0-shot. Most absolute errors by these models are also multiples of 10 10 : 100% of the time for GPT-3.5 and GPT-4, and 87% for PaLM-2. The similarity in error distribution to that of open-source models leads us to hypothesize that Fourier features play a critical role in their computational mechanism.

## 5 Related Work

#### Learning mathematical tasks.

Previous studies primarily explore what pre-trained LMs can achieve on arithmetic tasks, with less emphasis on the underlying mechanisms [ 30 , 37 ] . For instance, [ 23 ] demonstrates that small Transformer models can effectively learn arithmetic by altering the question format and utilizing a scratchpad method [ 28 ] . [ 20 ] identifies activation patterns for the “greater-than” operation in GPT-2, and [ 6 ] focuses on the enumeration and selection processes in GCD computation. In this paper, we dive into the specific roles of MLP and attention layers in solving mathematical tasks. Our research analyzes these components’ distinct contributions to integer addition tasks.

#### Mechanisms of pre-trained LMs.

Recent studies have significantly advanced our understanding of the underlying mechanisms of pre-trained Transformer models. For instance, research on “skill neurons” by [ 45 ] and “knowledge neurons” by [ 8 ] underscores the development of specialized neural components that encode task-specific capabilities or hold explicit factual information in the pre-trained LMs, enhancing model performance on related tasks. [ 25 ] and [ 14 ] discuss how MLPs and FFNs transform and update token representations for general language tasks. In contrast, we show that the pre-trained LMs use multiple layers to compute addition by combining the results of approximation and classification. Additionally, [ 46 ] demonstrated the capacity of GPT-2 to consolidate similar information through pre-training in the model weights, which aligns with our observations on the importance of pre-training in developing effective number embedding and arithmetic computation strategies in LMs.

#### Fourier features in Neural Networks.

Fourier features are commonly observed in image models, particularly in the early layers of vision models [ 32 , 31 , 13 ] . These features enable the model to detect edges, textures, and other spatial patterns effectively. Recently, Fourier features have been noted in networks trained for tasks that allow cyclic wraparound, such as modular addition [ 29 , 24 ] , general group compositions [ 5 ] , or invariance to cyclic translations [ 38 ] . [ 29 ] demonstrates that learning Fourier features can induce ‘grokking’ [ 35 ] . Furthermore, [ 26 ] provides a mathematical framework explaining the emergence of Fourier features when the network exhibits invariance to a finite group. We extend these insights by observing Fourier features in tasks that do not involve cyclic wraparound. [ 40 ] found that by selecting problem-specific Fourier features, the performance of MLPs can be improved on a computer vision-related task.

## 6 Conclusion

In this paper, we provide a comprehensive analysis of how pre-trained LLMs compute numerical sums, revealing a nuanced interplay of Fourier features within their architecture. Our findings demonstrate that LLMs do not simply memorize answers from training data but actively compute solutions through a combination of approximation and classification processes encoded in the frequency domain of their hidden states. Specifically, MLP layers contribute to approximating the magnitude of sums, while attention layers contribute to modular operations.

Our work also shows that pre-training plays a critical role in equipping LLMs with the Fourier features necessary for executing arithmetic operations. Models trained from scratch lack these crucial features and achieve lower accuracy; introducing pre-trained token embeddings greatly improves their convergence rate and accuracy. This insight into the arithmetic problem-solving capabilities of LLMs through Fourier features sets the stage for potential modifications to training approaches. By imposing specific constraints on model training, we could further enhance the ability of LLMs to learn and leverage these Fourier features, thereby improving their performance in mathematical tasks.

## Acknowledgments

DF and RJ were supported by a Google Research Scholar Award. RJ was also supported by an Open Philanthropy research grant. VS was supported by NSF CAREER Award CCF-2239265 and an Amazon Research Award. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of the funding agencies.

## Appendix

#### Roadmap.

In Appendix A , we introduce some formal definitions that used in our main content. In Appendix B , we show why we separate the Fourier components into the high-frequency part and the low-frequency part and why we choose τ \tau to be 50 50 . In Appendix C , we show our observation generalizes to another format of dataset, another arithmetic task and other models. In Appendix D , we provide more evidence that shows the Fourier features in the model when computing addition. In Appendix E , we provide more evidence that shows the GPT-2-XL trained from scratch does not use Fourier feature to solve the addition task. In Appendix F , we give the details of our experimental settings.

## Appendix A Formal Definition of Transformer and Logits in Fourier Space

We first introduce the formal definition of the Transformer structure that we used in this paper.

###### Definition A.1 (Transformer) .

An autoregressive Transformer language model G : 𝒳 → 𝒴 G:\mathcal{X}\rightarrow\mathcal{Y} over vocabulary Vocab \mathrm{Vocab} maps a token sequence x = [ x 1 , … , x N ] ∈ 𝒳 , x t ∈ Vocab x=\left[x_{1},\ldots,x_{N}\right]\in\mathcal{X},x_{t}\in\mathrm{Vocab} to a probability distribution y ∈ 𝒴 ⊂ ℝ | Vocab | y\in\mathcal{Y}\subset\mathbb{R}^{|\mathrm{Vocab}|} that predicts next-token continuations of x x . Within the Transformer, the i i -th token is embedded as a series of hidden state vectors h t ( ℓ ) h_{t}^{(\ell)} , beginning with h t ( 0 ) = emb ⁡ ( x t ) + pos ⁡ ( i ) ∈ ℝ D h_{t}^{(0)}=\operatorname{emb}\left(x_{t}\right)+\operatorname{pos}(i)\in\mathbb{R}^{D} . Let W U ∈ ℝ | Vocab | × D W^{U}\in\mathbb{R}^{|\mathrm{Vocab}|\times D} denote the output embedding. The final output y = softmax ⁡ ( W U ​ ( h N ( L ) ) ) y=\operatorname{softmax}(W^{U}\left(h_{N}^{(L)}\right)) is read from the last hidden state. In the autoregressive case, tokens only draw information from past tokens: h t ( ℓ ) = h t ( ℓ − 1 ) + Attn t ( ℓ ) + MLP t ( ℓ ) \displaystyle h_{t}^{(\ell)}=h_{t}^{(\ell-1)}+\mathrm{Attn}_{t}^{(\ell)}+\mathrm{MLP}_{t}^{(\ell)} where Attn t ( ℓ ) := Attn ( ℓ ) ​ ( h 1 ( ℓ − 1 ) , h 2 ( ℓ − 1 ) , … , h t ( ℓ − 1 ) ) and MLP t ( ℓ ) := MLP t ( ℓ ) ​ ( Attn t ( ℓ ) , h t ( ℓ − 1 ) ) . \displaystyle\mathrm{Attn}_{t}^{(\ell)}:=\mathrm{Attn}^{(\ell)}\left(h_{1}^{(\ell-1)},h_{2}^{(\ell-1)},\ldots,h_{t}^{(\ell-1)}\right)\quad\text{and}\quad\mathrm{MLP}_{t}^{(\ell)}:=\mathrm{MLP}_{t}^{(\ell)}(\mathrm{Attn}_{t}^{(\ell)},h_{t}^{(\ell-1)}).