SlowGuess/NeurIPS_2023_OCR huggingface dataset:
G Mathematical Description of Constant-Attention transformer
In this section, we examine the structure of constant-attention transformers loosely following the notation of [10].
Denote the weight of embedding layer as  $W_{E}$ , the weight of positional embedding as  $W_{\mathrm{pos}}$ , the weight of the value and output matrix of the  $j$ -th head of the  $t$ -th layer as  $W_{V}^{t,j}$  and  $W_{O}^{t,j}$ , the weights and biases of the input linear map of MLP in the  $t$ -th layer as  $W_{\mathrm{in}}^{t}$  and  $b_{\mathrm{in}}^{t}$ , the corresponding weights and biases of the output linear map as  $W_{\mathrm{out}}^{t}$  and  $b_{\mathrm{out}}^{t}$ , and the weight of the unembedding layer as  $W_{U}$ . Notice that the query and the key matrices are irrelevant as the attention matrix is replaced with an all-one matrix. Denote  $x^{j}$  as the value of residue stream vector after the first  $j$  layers and denote  $c_{i}$  as the character in the  $i$ -th position. We use subscripts like  $x_{t}$  to denote taking a specific element of vector.
We can formalize the logit calculation as the following:
- Embedding:  $x_{i}^{0} = W_{E,c_{i}} + W_{\mathrm{pos},i}$ .
- For each layer  $t$  from 1 to  $n_{\mathrm{layer}}$ :
- Constant Attention:  $w_{i}^{t} = x_{i}^{t - 1} + \sum_{j}W_{O}^{t,j}W_{V}^{t,j}\sum_{k}x_{k}^{t - 1}$ .
- MLP:  $x^{t} = w^{t} + b_{\mathrm{out}}^{t} + W_{\mathrm{out}}^{t}\mathrm{ReLU}(b_{\mathrm{in}}^{t} + W_{\mathrm{in}}^{t}w^{t})$ .
Output:  $O = W_{U}x^{n_{\mathrm{layer}}}$
In the particular case where the input length is 2, the number of layer is 1, and we focus on the logit of the last position, we may restate as the following (denote  $z$  as  $x^{1}$  and  $y$  as  $w^{1}$ ):
- Embedding:  $x_{1} = W_{E,c_{1}} + W_{\mathrm{pos},1}$ ,  $x_{2} = W_{E,c_{2}} + W_{\mathrm{pos},2}$ .
- Constant Attention:  $y = x_{2} + \sum_{j} W_{O}^{j} W_{V}^{j}(x_{1} + x_{2})$ .
- MLP:  $z = y + b_{\mathrm{out}}^{t} + W_{\mathrm{out}}^{t}\mathrm{ReLU}(b_{\mathrm{in}}^{t} + W_{\mathrm{in}}^{t}y)$ .
Output:  $o = W_Uz$
If we remove the skip connections, the network after embedding could be seen as
$$
o = L _ {U} \left(L _ {\text {o u t}} \left(\operatorname {R e L U} \left(L _ {\text {i n}} \left(\sum_ {j} L _ {O} ^ {j} \left(L _ {V} ^ {j} \left(x _ {1} + x _ {2}\right)\right)\right)\right)\right)\right)
$$
where  $L_V^j, L_O^j, L_{\mathrm{in}}, L_{\mathrm{out}}, L_U$  are a series of linear layers corresponding to the matrices.
H Pizza with Attention
Extrapolating from Figure 7, we trained transformers with width 1024 and attention rate 1 (normal attention). After several tries, we are able to observe a trained circular model with distance irrelevance 0.156 and gradient symmetricity 0.995, which fits our definition of Pizza (Figure 17).
Figure 17: Correct logits of the trained model in Section H after circle isolation (Section 3.3).
I Results on Slightly Different Setups
We considered the following variations of our setups (Appendix F.1, Section 4), for which the existence of pizzas and clocks as well as the phase changes are still observed.
GeLU instead of ReLU We conducted the same 1-layer transformer experiment with activation function GeLU instead of ReLU. Very similar results are observed (Figure 18).
Encode Two Tokens Differently We conducted the 1-layer transformer experiments but with different embedding for the two tokens. Again very similar results are observed (Figure 19). We also discovered that the two tokens' embeddings are often aligned to implement the Pizza and Clock algorithm (Figure 20).
Adding Equal Sign We conducted the 1-layer transformer experiment with an equal sign added. Very similar results are observed (Figure 21).
J Pizza Occurs Early in the Clock Training
We plotted intermediate states during the training of a model with attention (attention rate 1). Pizza-like pattern was observed early in the training, but the pattern gradually disappeared during the run (Figure 22).
Figure 18: Training results from 1-layer transformers with GeLU instead of ReLU as the activation function. Each point in the plots represents a training run that reached  $100\%$  validation accuracy.
Figure 19: Training results from 1-layer transformers where the two tokens use different embeddings (feed  $[a, b + p]$  to the model on input  $(a, b)$ ;  $2p$  tokens are handled in the embedding layer). Each point in the plots represents a training run that reached  $100\%$  validation accuracy. We did not use circularity to filter the result because it is no longer well-defined.
Figure 20: Correct logits after circle isolation from a trained model where two tokens use different embeddings. The blue points represent the embeddings for the first token and the green points represent the embeddings for the second token. The model is implementing the Pizza algorithm. The correct logit pattern is shifted comparing to the previous patterns because the embeddings of two tokens do not line up exactly. For example, the third circle has near-maximum correct logit for  $a = 6$ ,  $b = 3$  (the two points lining up on the top) and  $(a - b) / 18 \equiv 10 \pmod{59}$ . This is the reason that the correct logit pattern appears to be shifted 10 units down.
Figure 21: Training results from 1-layer transformers where an equal sign is added (feed  $[a, b, =]$  to the model on input  $(a, b)$  where  $=$  is a special token;  $p + 1$  tokens are handled in the embedding layer; context length of the model becomes 3). Each point in the plots represents a training run that reached  $100\%$  validation accuracy. We did not use circularity to filter the result because it is no longer well-defined.
K Accompanying Pizza Occurs Early in the Pizza Training
We plotted intermediate states during the training of a model without attention (attention rate 0). We observed the early emergence of a pattern similar to accompanying pizza in training runs (Figure
90 Steps
210 Steps
300 Steps
510 Steps
840 Steps
Figure 22: For a 1-layer transformer with attention, correct logits after principal component (possibly non-circle) isolations at various states during the training. The pizza-like pattern gradually resolved.
23) and removing that circle brings accuracy down from  $99.7\%$  to  $97.9\%$ . They are less helpful later in the network (removing accompanying pizzas in trained Model A only brings accuracy down to  $99.7\%$ ).
Figure 23: Immediate state after 600 epochs of training for a 1-layer transformer with constant attention.
L A Closer Look at a Linear Pizza Model
In this section, we provide a full picture of the linear model shown in Figure 15 by investigating the actual weights in the model.
L.1 Model Structure
As described in Appendix E, on input  $(a,b)$ , the output logits of the model is computed as
$$
L _ {3} \left(\operatorname {R e L U} \left(L _ {2} \left(\operatorname {R e L U} \left(L _ {1} \left(\operatorname {E m b e d} [ a ] + \operatorname {E m b e d} [ b ]\right)\right)\right)\right)\right).
$$
Denote the weight of embedding layer as  $W_{E}$ , the weight of the unembedding layer ( $L_{3}$ ) as  $W_{U}$ , and the weights and biases of  $L_{1}$  and  $L_{2}$  as  $W_{1}, b_{1}$  and  $W_{2}, b_{2}$ , respectively, then the output logits on input  $(a, b)$  can be written as
$$
W _ {U} \operatorname {R e L U} \left(b _ {2} + W _ {2} \operatorname {R e L U} \left(b _ {1} + W _ {1} \left(W _ {E} [ a ] + W _ {E} [ b ]\right)\right)\right).
$$
L.2 General Picture
We first perform principal component visualizations on the embedding and unembedding matrices. From Figure 24, we can see that the embedding and unembedding matrices formed matching circles (circles with the same gap  $\delta$  between adjacent entries).
We now give the general overview of the circuit. Each pair of matching circles forms an instance of Pizza and they operate independently (with rather limited interference). Specifically for each pair,
- The embedding matrix first places the inputs  $a, b$  on the circumference:  $W_E'[a] \approx (\cos(w_k a), \sin(w_k a))$  and  $W_E'[b] \approx (\cos(w_k b), \sin(w_k b))$  ( $w_k = 2\pi k / p$  for some integer  $k \in [1, p - 1]$  as in Section 2.1;  $W_E'$  stands for the two currently considered principal components of  $W_E$ ; rotation and scaling omitted for brevity).
- The embeddings are added to get
$$
\begin{array}{l} \left(\cos \left(w _ {k} a\right) + \cos \left(w _ {k} b\right), \sin \left(w _ {k} a\right) + \sin \left(w _ {k} b\right)\right) \\ = \cos \left(w _ {k} (a - b) / 2\right) \cdot \left(\cos \left(w _ {k} (a + b) / 2\right), \sin \left(w _ {k} (a + b) / 2\right)\right) \\ \end{array}
$$
- It is then passed through the first linear layer  $L_{1}$ . Each result entry pre-ReLU will thus be a linear combination of the two dimensions of the aforementioned vectors, i.e.  $\cos(w_{k}(a - b)/2) \cdot (\alpha \cos(w_{k}(a + b)/2) + \beta \sin(w_{k}(a + b)/2))$  for some  $\alpha, \beta$ , which will become  $|\cos(w_{k}(a - b)/2)||\alpha \cos(w_{k}(a + b)/2) + \beta \sin(w_{k}(a + b)/2)|$  after ReLU.
- These values are then passed through the second linear layer  $L_{2}$ . Empirically the ReLU is not observed to be effective as the majority of values is positive. The output entries are then simply linear combinations of aforementioned outputs of  $L_{1}$ .
- The unembedding matrix is finally applied. In the principal components we are considering,  $W_U'[c] \approx (\cos(w_k c), \sin(w_k c))$ . ( $W_U'$  stands for the two currently considered principal components of  $W_U$ ; rotation and scaling omitted for brevity) and these two principal components correspond to a linear combination of the output entries of  $L_2$ , which then correspond to a linear combination of the outputs of  $L_1$  (thanks to the non-functional ReLU).
- Similar to the formula  $|\sin(t)| - |\cos(t)| \approx \cos(2t)$  discussed in Appendix A, these linear combinations provide good approximations for  $|\cos(w_k(a - b)/2)|\cos(w_k(a + b))$  and  $|\cos(w_k(a - b)/2)|\sin(w_k(a + b))$ . Finally we arrive at
$$
\begin{array}{l} | \cos (w _ {k} (a - b) / 2) | (\cos (w _ {k} c) \cos (w _ {k} (a + b)) + \sin (w _ {k} c) \sin (w _ {k} (a + b))) \\ = | \cos (w _ {k} (a - b) / 2) | \cos (w _ {k} (a + b - c)) \\ \end{array}
$$
Figure 24: Visualization of the principal components of the embeddings and unembedding matrices.
L.3 Aligning Weight Matrices
We first verify that the ReLU from the second layer is not functional. After removing it, the accuracy of the model remains  $100\%$  and the cross-entropy loss actually decreased from  $6.20 \times 10^{-7}$  to  $5.89 \times 10^{-7}$ .
Therefore, the model output can be approximately written as
$$
W _ {U} \left(b _ {2} + W _ {2} \operatorname {R e L U} \left(b _ {1} + W _ {1} \left(W _ {E} [ a ] + W _ {E} [ b ]\right)\right)\right) = W _ {U} b _ {2} + W _ {U} W _ {2} \operatorname {R e L U} \left(b _ {1} + W _ {1} \left(W _ {E} [ a ] + W _ {E} [ b ]\right)\right).
$$
We now "align" the weight matrices  $W_{1}$  and  $W_{2}$  by mapping through the directions of the principal components of the embeddings and unembeddings. That is, we calculate how these matrices act on and onto the principal directions (consider  $W_{1}v$  for every principal direction  $v$  in  $W_{E}$  and  $v^{T}W_{2}$  for every principal direction  $v$  in  $W_{U}$ ). We call the other dimension of aligned  $W_{1}$  and  $W_{2}$  output and source dimensions, respectively (Figure 25).
In the aligned weight matrices, we can see a clear domino-like pattern: in most output or source dimensions, only two principal components have significant non-zero values, and they correspond to a pair of matching circle, or a pizza. In this way, every immediate dimension serves for exactly one pizza, so the pizzas do not interfere with each other.
Figure 25: Visualization of the aligned  $W_{1}$  and  $W_{2}$ .
L.4 Approximation
Everything becomes much clearer after realigning the matrices. For a pizza and its two corresponding principal embedding / unembedding dimensions,  $W_E'[a] + W_E'[b] \approx \cos(w_k(a - b)/2) \cdot (\cos(w_k(a + b)/2), \sin(w_k(a + b)/2))$  will be mapped by realigned  $W_1$  into its corresponding columns (which are different for every pizza), added with  $b_1$  and apply ReLU. The result will then be mapped by the realigned  $W_2$ , added with realigned  $b_2$ , and finally multiplied by  $(\cos(w_kc), \sin(w_kc))$ .
For the first two principal dimensions, realigned  $W_{1}$  has 44 corresponding columns (with coefficients of absolute value  $>0.1$ ). Let the embedded input be  $(x,y) = W_E'[a] + W_E'[b] \approx \cos (w_k(a - b) / 2) \cdot (\cos (w_k(a + b) / 2),\sin (w_k(a + b) / 2))$ , the intermediate columns are
ReLU([0.530x - 1.135y + 0.253, -0.164x - 1.100y + 0.205, 1.210x - 0.370y + 0.198, -0.478x - 1.072y + 0.215, -1.017x + 0.799y + 0.249, 0.342x - 0.048y + 0.085, 1.149x - 0.598y + 0.212, -0.443x + 1.336y + 0.159, -1.580x - 0.000y + 0.131, -1.463x + 0.410y + 0.178, 1.038x + 0.905y + 0.190, 0.568x + 1.188y + 0.128, 0.235x - 1.337y + 0.164, -1.180x + 1.052y + 0.139, -0.173x + 0.918y + 0.148, -0.200x + 1.060y + 0.173, -1.342x + 0.390y + 0.256, 0.105x - 1.246y + 0.209, 0.115x + 1.293y + 0.197, 0.252x + 1.247y + 0.140, -0.493x + 1.252y + 0.213, 1.120x + 0.262y + 0.239, 0.668x + 1.096y + 0.205, -0.487x - 1.302y + 0.145, 1.134x - 0.862y + 0.273, 1.143x + 0.435y + 0.171, -1.285x - 0.644y + 0.142, -1.454x - 0.285y + 0.218, -0.924x + 1.068y + 0.145, -0.401x + 0.167y + 0.106, -0.411x - 1.389y + 0.249, 1.422x - 0.117y + 0.227, -0.859x - 0.778y + 0.121, -0.528x - 0.216y + 0.097, -0.884x - 0.724y + 0.171, 1.193x + 0.724y + 0.131, 1.086x + 0.667y + 0.218, 0.402x + 1.240y + 0.213, 1.069x - 0.903y + 0.120, 0.506x - 1.042y + 0.153, 1.404x - 0.064y + 0.152, 0.696x - 1.249y + 0.199, -0.752x - 0.880y + 0.106, -0.956x - 0.581y + 0.223]).
For the first principal unembedding dimension, it will be taken dot product with
[1.326, 0.179, 0.142, -0.458, 1.101, -0.083, 0.621, 1.255, -0.709, 0.123, -1.346, -0.571, 1.016, 1.337, 0.732, 0.839, 0.129, 0.804, 0.377, 0.078, 1.322, -1.021, -0.799, -0.339, 1.117, -1.162, -1.423, -1.157, 1.363, 0.156, -0.165, -0.451, -1.101, -0.572, -1.180, -1.386, -1.346, -0.226, 1.091, 1.159, -0.524, 1.441, -0.949, -1.248].
Call this function  $f(x, y)$ . When we plug in  $x = \cos(t)$ ,  $y = \sin(t)$ , we get a function that well-approximated  $8\cos(2t + 2)$  (Figure 26). Therefore, let  $t = w_k(a + b)/2$ , the dot product will be approximately  $8|\cos(w_k(a - b)/2)|\cos(w_k(a + b) + 2)$ , or  $|\cos(w_k(a - b)/2)|\cos(w_k(a + b))$  if we ignore the phase and scaling. This completes the picture we described above.
Figure 26:  $f(\cos(t), \sin(t))$  well-approximates  $8\cos(2t + 2)$ .