Pre-trained Large Language Models Use Fourier Features to Compute Addition — part 3/4
Cleaned-text code points [31346, 47089). Bibliography omitted here; source.md includes it.
In this paper, we only consider the output tokens to be numbers. Hence, we have the unembedding matrix W U ∈ ℝ p × D W^{U}\in\mathbb{R}^{p\times D} , where p p is the size of the number space. As we are given the length- N N input sequences and predict the ( N + 1 ) (N+1) -th, we only consider h N ( ℓ ) = h N ( ℓ − 1 ) + Attn N ( ℓ ) + MLP N ( ℓ ) h_{N}^{(\ell)}=h_{N}^{(\ell-1)}+\mathrm{Attn}_{N}^{(\ell)}+\mathrm{MLP}_{N}^{(\ell)} . For simplicity, we ignore the subscript N N in the following paper, so we get Eq. ( 1 ).
###### Definition A.2 (Intermediate Logits) .
Let ℒ Attn ( ℓ ) := W U Attn ( ℓ ) \mathcal{L}_{\mathrm{Attn}}^{(\ell)}:=W^{U}\mathrm{Attn}^{(\ell)} denote the intermediate logits of the attention module at the ℓ \ell -th layer. Let ℒ MLP ( ℓ ) := W U MLP ( ℓ ) \mathcal{L}_{\mathrm{MLP}}^{(\ell)}:=W^{U}\mathrm{MLP}^{(\ell)} denote the intermediate logits of the MLP module at the ℓ \ell -th layer. Let ℒ ( ℓ ) := W U h ( ℓ ) \mathcal{L}^{(\ell)}:=W^{U}h^{(\ell)} denote the logits on intermediate state h ( ℓ ) h^{(\ell)} .
Throughout the model, h h undergoes only additive updates (Eq. ( 1 )), creating a continuous residual stream [ 9 ] , meaning that the token representation h h accumulates all additive updates within the residual stream up to layer t t .
To analyze the logits in Fourier space, we give the formal definition of the Fourier basis as follows:
###### Definition A.3 (Fourier Basis) .
Let p p denote the size of the number space. Let 𝐱 → := ( 0 , 1 , … , ( p − 1 ) ) \overrightarrow{\mathbf{x}}:=(0,1,\ldots,(p-1)) . Let ω k := 2 π k p − 1 \omega_{k}:=\frac{2\pi k}{p-1} . We denote the normalized Fourier basis F F as the p × p p\times p matrix: F := [ 1 p − 1 ⋅ 𝟏 → 2 p − 1 ⋅ sin ( ω 1 𝐱 → ) 2 p − 1 ⋅ cos ( ω 1 𝐱 → ) 2 p − 1 ⋅ sin ( ω 2 𝐱 → ) ⋮ 2 p − 1 ⋅ cos ( ω ( p − 1 ) / 2 𝐱 → ) ] ∈ ℝ p × p F:=\left[\begin{array}[]{c}\sqrt{\frac{1}{p-1}}\cdot\overrightarrow{\mathbf{1}}\\ \sqrt{\frac{2}{p-1}}\cdot\sin\left(\omega_{1}\overrightarrow{\mathbf{x}}\right)\\ \sqrt{\frac{2}{p-1}}\cdot\cos\left(\omega_{1}\overrightarrow{\mathbf{x}}\right)\\ \sqrt{\frac{2}{p-1}}\cdot\sin\left(\omega_{2}\overrightarrow{\mathbf{x}}\right)\\ \vdots\\ \sqrt{\frac{2}{p-1}}\cdot\cos\left(\omega_{(p-1)/2}\overrightarrow{\mathbf{x}}\right)\end{array}\right]\in\mathbb{R}^{p\times p} The first component F [ 0 ] F[0] is defined as a constant component. For i ∈ [ 0 , p − 1 ] i\in[0,p-1] , F [ i ] F[i] is defined as the k k -th component in Fourier space, where k = ⌊ i + 1 2 ⌋ k=\lfloor\frac{i+1}{2}\rfloor . The frequency of the k k -th component is f k := k p − 1 f_{k}:=\frac{k}{p-1} . The period of the k k -th component is T k := p − 1 k T_{k}:=\frac{p-1}{k}
We can compute the discrete Fourier transform under that Fourier basis as follows:
###### Remark A.4 (Discrete Fourier transformer (DFT) and inverse DFT) .
We can transform any logits u ∈ ℝ p u\in\mathbb{R}^{p} to Fourier space by computing u ^ = F ⋅ u \widehat{u}=F\cdot u . We can transform u ^ \widehat{u} back to u u by u = F ⊤ ⋅ u ^ u=F^{\top}\cdot\widehat{u}
Next, we define the logits in Fourier space.
###### Definition A.5 (Logits in Fourier Space) .
Let ℒ ( L ) \mathcal{L}^{(L)} , ℒ Attn ( ℓ ) \mathcal{L}_{\mathrm{Attn}}^{(\ell)} and ℒ MLP ( ℓ ) \mathcal{L}_{\mathrm{MLP}}^{(\ell)} denote the logits (Definition A.2 ). The output logits before softmax in Fourier space is defined as: ℒ ^ ( L ) = F ⋅ ℒ ( L ) \widehat{\mathcal{L}}^{(L)}=F\cdot\mathcal{L}^{(L)} . The logits of the MLP and attention modules in Fourier space are defined as: ℒ ^ Attn ( ℓ ) = F ⋅ ℒ Attn ( ℓ ) and ℒ ^ MLP ( ℓ ) = F ⋅ ℒ MLP ( ℓ ) . \displaystyle\widehat{\mathcal{L}}_{\mathrm{Attn}}^{(\ell)}=F\cdot\mathcal{L}_{\mathrm{Attn}}^{(\ell)}\quad\text{and}\quad\widehat{\mathcal{L}}_{\mathrm{MLP}}^{(\ell)}=F\cdot\mathcal{L}_{\mathrm{MLP}}^{(\ell)}.
We ignore the first elements in ℒ ^ ( L ) , ℒ ^ Attn ( ℓ ) \widehat{\mathcal{L}}^{(L)},\widehat{\mathcal{L}}_{\mathrm{Attn}}^{(\ell)} and ℒ ^ MLP ( ℓ ) \widehat{\mathcal{L}}_{\mathrm{MLP}}^{(\ell)} for the Fourier analysis in this paper as they are the constant terms. Adding a constant to the logits will not change the prediction.
Let τ ∈ ℝ \tau\in\mathbb{R} denote a constant threshold. The low-frequency components for the logits in Fourier space are defined as ℒ ^ ( ℓ ) [ 1 : 2 τ ] \widehat{\mathcal{L}}^{(\ell)}[1:2\tau] . The high-frequency components for the logits in Fourier space are defined as ℒ ^ ( ℓ ) [ 2 τ : ] \widehat{\mathcal{L}}^{(\ell)}[2\tau:] . For the following analysis, we choose τ = 50 \tau=50 (the specific choice of τ = 50 \tau=50 is explained in Appendix B ).
Next, we propose the formal definition of low-pass/high-pass filter that is used in the following ablation study.
###### Definition A.6 (Loss-pass / High-pass Filter) .
Let x ∈ ℝ D x\in\mathbb{R}^{D} denote the output of MLP or attention modules. Let F F denote the Fourier Basis (Definition A.3 ). Let τ ∈ R \tau\in R denote the frequency threshold. Let W U ∈ R p × D W^{U}\in R^{p\times D} denote the output embedding. For low-pass filter, we define a diagonal binary matrix B ∈ { 0 , 1 } p × p B\in\{0,1\}^{p\times p} as b i i = { 1 if i ≥ τ 0 otherwise . b_{ii}=\begin{cases}1&\text{if }i\geq\tau\\ 0&\text{otherwise}\end{cases}. For high-pass filter, we define a diagonal binary matrix B ∈ { 0 , 1 } p × p B\in\{0,1\}^{p\times p} as b i i = { 1 if 1 ≤ i < τ 0 otherwise . b_{ii}=\begin{cases}1&\text{if }1\leq i<\tau\\ 0&\text{otherwise}\end{cases}. Note that we retain the constant component, so b i , i = 0 b_{i,i}=0 . The output of the filter ℱ ( x ) : ℝ D → ℝ D \mathcal{F}(x):\mathbb{R}^{D}\rightarrow\mathbb{R}^{D} is defined by the following objective function: min y \displaystyle\min_{y}\quad ‖ x − y ‖ 2 2 \displaystyle\|x-y\|_{2}^{2} subject to \displaystyle\mathrm{subject~to}\quad B F W U y = 0 \displaystyle BFW^{U}y=0
The solution to the above optimization problem is given by a linear projection.
###### Remark A.7 .
The result of the optimization problem defined in Definition A.6 is the projection of x x to the null space of B F W U BFW^{U} . Let 𝒩 ( B F W U ) \mathcal{N}(BFW^{U}) denote the null space of B F W U BFW^{U} . We have ℱ ( x ) = 𝒩 ( B F W U ) ⋅ 𝒩 ( B F W U ) ⊤ ⋅ x ⊤ \displaystyle\mathcal{F}(x)=\mathcal{N}(BFW^{U})\cdot\mathcal{N}(BFW^{U})^{\top}\cdot x^{\top}
## Appendix B Fourier Components Separation and Selection of τ \tau
Following Definition A.6 , we define single-pass filter as follows:
###### Definition B.1 (Single-Pass Filter) .
Let x ∈ ℝ D x\in\mathbb{R}^{D} denote the output of MLP or attention modules. Let F F denote the Fourier Basis (Definition A.3 ). Let γ ∈ R \gamma\in R denote the γ \gamma -th Fourier component (Definition A.3 ) that we want to retain. Let W U ∈ R V × D W^{U}\in R^{V\times D} denote the output embedding. We define a diagonal binary matrix B ∈ { 0 , 1 } V × V B\in\{0,1\}^{V\times V} as b i i = { 0 if ⌊ i + 1 2 ⌋ = γ or i = 0 , 1 otherwise . b_{ii}=\begin{cases}0&\text{if }\lfloor\frac{i+1}{2}\rfloor=\gamma\text{ or }i=0,\\ 1&\text{otherwise}.\end{cases}
The output of the filter ℱ γ ( x ) : ℝ D → ℝ D \mathcal{F}_{\gamma}(x):\mathbb{R}^{D}\rightarrow\mathbb{R}^{D} is defined as the following objective function: min y \displaystyle\min_{y}\quad ‖ x − y ‖ 2 2 \displaystyle\|x-y\|_{2}^{2} subject to \displaystyle\mathrm{subject~to}\quad B F W U y = 0 \displaystyle BFW^{U}y=0
###### Remark B.2 .
The result of the optimization problem defined in Definition B.1 is the projection of x x to the null space of B F W U BFW^{U} . Let 𝒩 ( B F W U ) \mathcal{N}(BFW^{U}) denote the null space of B F W U BFW^{U} . We have ℱ γ ( x ) = 𝒩 ( B F W U ) ⋅ 𝒩 ( B F W U ) ⊤ ⋅ x ⊤ \displaystyle\mathcal{F}_{\gamma}(x)=\mathcal{N}(BFW^{U})\cdot\mathcal{N}(BFW^{U})^{\top}\cdot x^{\top}
For the single-pass filter, we only retrain one Fourier component and analyze how this component affects the model’s prediction. The residual stream is then updated as follows: h ( ℓ ) = h ( ℓ − 1 ) + ℱ γ ( Attn ( ℓ − 1 ) ) + ℱ γ ( MLP ( ℓ − 1 ) ) \displaystyle h^{(\ell)}=h^{(\ell-1)}+\mathcal{F}_{\gamma}(\mathrm{Attn}^{(\ell-1)})+\mathcal{F}_{\gamma}(\mathrm{MLP}^{(\ell-1)})
We evaluated the fine-tuned GPT-2-XL model on the addition dataset with the Fourier components period 520 520 and 2 2 . Given that T k := V − 1 k T_{k}:=\frac{V-1}{k} (Definition A.3 ), we retained only the Fourier components with γ = 1 \gamma=1 and 260 260 , respectively.
As shown in Figure 10 a, with only one frequency component, whose period is 2 2 , the model accurately predicts the parity with 99.59 % 99.59\% accuracy. As depicted in Figure 10 b, with a single frequency component of period 520 520 , the model fails to accurately predict with 96.51 % 96.51\% accuracy. We consider the frequency component with a period of 2 2 as the model’s prediction for the mod 2 task, and the frequency component with a period of 520 520 as its prediction for the mod 520 task. Figures 10 and 11 suggest that the model effectively learns the mod 2 task, as it involves a two-class classification, but struggles with the mod 520 task, which requires classifying among 520 520 classes. As the model does not need to be trained to converge to the optimal for these low-frequency components as explained at the end of Section 3.2 , predicting with the period- 520 520 component leads to predictions that normally distributed around the correct answers.
The Fourier components with larger periods present greater difficulty in solving the corresponding modular addition task compared to those with smaller periods. As demonstrated in Figure 11 , components with large periods serve primarily as approximations of the correct answer. Consequently, we categorize the Fourier components into low-frequency and high-frequency groups. The low-frequency components approximate the magnitude of the answer, whereas the high-frequency components are employed to enhance the precision of the predictions.
In reference to Figure 4 , to elucidate the contribution of these distinct Fourier components to our final prediction and the rationale behind their separation, consider the example: “Put together 15 15 and 93 93 . Answer: 108 108 ”. We selected the top-10 Fourier components of ℒ ^ ( L ) \widehat{\mathcal{L}}^{(L)} based on their magnitudes and converted them back to logits in the numerical space by multiplying with F ⊤ F^{\top} . We plotted the components with components index less than 50 in Figure 12 a and those with components index greater than 50 in Figure 12 b. Leveraging the constructive and destructive inference for different waves, the components with low periods assign more weight to the correct answer, 108 108 , and less weight to numbers close to 108 108 . These high-frequency (low-period) components ensure the prediction’s accuracy at the unit place. For the low-frequency (large-period) components, the model fails to precisely learn the magnitude of the factor between the cos \cos and sin \sin components, which results in failing to peak at the correct answer. Thus, the low-frequency (large-period) components are used to approximate the magnitude of the addition results.
## Appendix C Does Fourier Features Generalize?
### C.1 Token Embedding for Other LMs
We first show that other pre-trained LMs also have Fourier features in their token embedding for the numbers [ 0,520 ] [0,520] .
### C.2 Multiplication Task
A key question is whether pre-trained models utilize Fourier Features solely for solving addition tasks or if they generalize to other arithmetic tasks. We hypothesize the latter, knowing that numbers are represented by their Fourier features in the token embeddings after pre-training. Consequently, this Fourier representation should be leveraged in a variety of number-related tasks. To validate this hypothesis, we perform a Fourier analysis on the GPT-2-XL model fine-tuned for the multiplication task.
Considering a maximum number of 520 520 for multiplication would result in an insufficient dataset size. Therefore, we set the maximum allowable product to 10000 10000 . For each pair of numbers where the product does not exceed this limit, we generate various phrasings of multiplication questions and their corresponding answers in base 10 10 . The different phrasings used include: “What is the product of num1 and num2?”, “Find the product of num1 multiplied by num2.”, “Calculate num1 times num2.”, “num1 multiplied by num2 equals what?”, and “Multiplication of num1 with num2.” The dataset is then shuffled to ensure randomness and split into training ( 80 % 80\% ), validation ( 10 % 10\% ), and test ( 10 % 10\% ) sets. We finetune the model for 25 25 epochs with a learning rate of 1 e − 4 1e-4 . Upon convergence, the validation accuracy reaches 74.58 % 74.58\% .
As the primary objective is to determine whether the Fourier features are utilized in tasks other than addition, Figure 15 displays the logits in Fourier space for each layer, as in Figure 3 . It is evident that the logits are sparse in Fourier space.
### C.3 Same Results for other format
To demonstrate that our observations are not confined to a specific description of the mathematical problem, we conducted experiments on another format of addition problem and obtained consistent results. From Figure 16 , we can see that there are also periodic structures in the intermediate logits.
From Figure 17 , we can also see the Fourier features for the MLP and attention output. These two experiments validate that our observations are not confined to a specific format of the addition problems.
### C.4 Fourier Features in Other Pre-trained LM
Using the Fourier analysis framework proposed in Section 3.2 , we demonstrate that for GPT-J, the outputs of MLP and attention modules exhibit approximate sparsity in Fourier space across the last 15 15 layers (Figure 18 )
## Appendix D Supporting Evidence For the Fourier Features
We selected the layers that clearly show the periodic pattern in Figure 1 b and Figure 1 c and plot their logits in Figure 19 .
Figure 20 illustrates that the errors resulting from the ablation study (Section 3.3 ) correspond with our theoretical insights. Removing low-frequency parts from the MLP results in errors such as off-by 10 10 , 50 50 , and 100 100 . Without these low-frequency components, the MLP is unable to accurately approximate, although it still correctly predicts the unit digit. In contrast, removing high-frequency components from the attention modules results in smaller errors, all less than 6 6 in magnitude. These findings support our statement that low-frequency components are essential for accurate approximation, whereas high-frequency components are key for precise classification tasks. Consequently, the primary function of MLP modules is to approximate numerical magnitudes using low-frequency components, and the essential function of attention modules is to facilitate precise classification by identifying the correct unit digit.
## Appendix E More Experiments on GPT-2-XL Trained from Scratch
Following the methodology proposed in Section 3 , we plotted the logits of the MLP and attention modules for each layer, as shown in Figure 21 . The prediction is solely determined by the 40 40 -th layer MLP. Unlike Figure 3 , there is no observable periodic structure across all layers.
For the model trained from scratch on the created addition dataset, all of the predictions on the test dataset deviate from the correct answer within 2 2 as shown in Figure 22 .
## Appendix F Details of Experimental Settings
#### Fine-tuned GPT-2-XL