SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vlademyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858, 2023.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating Turing machines with transformers. Advances in Neural Information Processing Systems, 35:12071-12083, 2022a.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682, 2022b.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022c.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03846, 2023.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumont, Clement Delangue, Anthony Moi, Pierrick Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771, 2019.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
Zheng Yuan, Hongyi Yuan, Chuanqi Tan, Wei Wang, and Songfang Huang. How well do large language models perform in arithmetic tasks? arXiv preprint arXiv:2304.02015, 2023.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077, 2019.
Xiao Zhang, Yaodong Yu, Lingxiao Wang, and Quanquan Gu. Learning one-hidden-layer relu networks via gradient descent. In The 22nd international conference on artificial intelligence and statistics, pages 1524-1534. PMLR, 2019.
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. arXiv preprint arXiv:2211.09066, 2022.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv:2310.16028, 2023.
A Construction 14
A.1 The Transformer Architecture 14
A.2 Positional Encodings 14
A.3 Constructing Some Useful "Black-box" Functions 15
A.4 Results on Filtering 16
B Experimental Details 22
B.1 Model Evaluation 22
B.2 Implementation 22
C Additional Experimental Results 22
C.1 Out-of-distribution Experiments 22
C.2 Further Evidence of Model Expressivity 23
C.3 Filtering Evidence in 2-layer MLPs 24
C.4 Comparison of Filtered CoT with ICL 25
C.5 CoT across Different Sizes of GPT-2 25
C.6 Comparison of Transformer Prediction and Linear Regression 26
A Construction
A.1 The Transformer Architecture
For the purpose of this proof, we consider encoder-based transformer architectures and assume that the positional encodings are appended to the input<sup>4</sup>. We also consider that the heads are added and each one has each own key, query and value weight matrices. Formally, we have
$$
\operatorname {a t t n} (\boldsymbol {X}) = \boldsymbol {X} + \sum_ {h = 1} ^ {H} \boldsymbol {W} _ {\mathbf {V}} ^ {h} \boldsymbol {X} \operatorname {s o f t m a x} \left(\left(\boldsymbol {W} _ {\mathbf {K}} ^ {h} \boldsymbol {X}\right) ^ {\top} \boldsymbol {W} _ {\mathbf {Q}} ^ {h} \boldsymbol {X}\right) \tag {A.1}
$$
$$
\operatorname {T F} (\boldsymbol {X}) = \operatorname {a t t n} (\boldsymbol {X}) + \boldsymbol {W} _ {2} \left(\boldsymbol {W} _ {1} \operatorname {a t t n} (\boldsymbol {X}) + \boldsymbol {b} _ {1} \mathbf {1} _ {n}\right) _ {+} + \boldsymbol {b} _ {2} \mathbf {1} _ {n} \tag {A.2}
$$
where  $\mathbf{X} \in \mathbb{R}^{d \times n}$ ,  $H$  is the number of heads used and  $f(x) = (x)_+$  is the ReLU activation. We also make use of the temperature  $\lambda$ , which is a parameter of the softmax. Specifically,  $\text{softmax}(\mathbf{x}) = \left\{\mathrm{e}^{\lambda x_i} / \sum_j \mathrm{e}^{\lambda x_j}\right\}_i$ . Notice that as  $\lambda \to \infty$  we have  $\text{softmax}(\mathbf{x}) \to \max_i x_i$ . We also assume that the inputs are bounded; we denote with  $N_{\text{max}}$  the maximum sequence length of the model.
Assumption 2 Each entry is bounded by some constant  $c$ , which implies that  $\| X \| \leq c'$ , for some large  $c'$  that depends on the maximum sequence length and the width of the transformer.
A.2 Positional Encodings
In the constructions below we use a combination of the binary representation of each position, as well as some additional information (e.g. 0-1 bits) as described in the following sections. The binary representations and in general encodings we construct, require only logarithmic space with respect
<sup>4</sup>We note here that in terms of our construction adding the encodings or appending them to the input can be viewed in a similar manner. Specifically, we can consider that the up-projection step projects the input to zero padded vectors, while the encodings are orthogonal to that projection in the sense that the have zero in the non-zero coordinates of the input. In that case adding the positional encodings corresponds to appending them to the input.
to the sequence length. Notice that binary representation of the positions satisfy the following two conditions:
1. Let  $\mathbf{r}_i$  be the binary representation of the  $i$ -th position, then there exists  $\varepsilon > 0$  such that  $\mathbf{r}_i^\top \mathbf{r}_i > \mathbf{r}_i^\top \mathbf{r}_j + \varepsilon$  for all  $j \neq i$ .
2. There exists a one layer transformer that can implement the addition of two pointers (see Lemma 2).
A.3 Constructing Some Useful "Black-box" Functions
We follow the construction of previous work [Giannou et al., 2023] and use the following individual implementations as they do. We repeat the statements here for convenience. The first lemma is also similar to Akyurek et al. [2022], we however follow a slightly different proof.
Lemma 1 A transformer with one layer, two heads, and embedding dimension of  $\mathcal{O}(\log n + d)$ , where  $d$  is the input dimension and  $n$  is the sequence length, can copy any block of the input to any other block of the input with error  $\varepsilon$  arbitrarily small.
Proof. Consider an input of the following form
$$
\boldsymbol {X} = \left[ \begin{array}{l l l l} \boldsymbol {x} _ {1} & \boldsymbol {x} _ {2} & \dots & \boldsymbol {x} _ {n} \\ \boldsymbol {y} _ {1} & \boldsymbol {y} _ {2} & \dots & \boldsymbol {y} _ {n} \\ \mathbf {r} _ {1} & \mathbf {r} _ {2} & \dots & \mathbf {r} _ {n} \\ \mathbf {r} _ {k} & \mathbf {r} _ {2} & \dots & \mathbf {r} _ {n} \end{array} \right] \tag {A.3}
$$
where  $\pmb{x}_i, \pmb{y}_i \in \mathbb{R}^d$  are column vectors for all  $i = 1, \dots, n$ . We present how we can move a block of data from the  $(d + 1 : 2d, k)$  position block to the  $(1 : d, 1)$  position block, meaning to move the point  $\pmb{y}_k$  to the position of the point  $\pmb{x}_1$ . It is straightforward to generalize the proof that follows to move blocks from  $(i : i + k, j)$  to  $(i' : i' + k, j')$  for arbitrary  $i, i', j, j', k$ .
Let in Eq. A.1  $W_{\mathbf{K}}^{1} = \left[ \begin{array}{llll}0 & 0 & I & 0 \end{array} \right]$  and  $W_{\mathbf{Q}}^{1} = \left[ \begin{array}{llll}0 & 0 & 0 & I \end{array} \right]$  and  $W_{\mathbf{K}}^{2} = \left[ \begin{array}{llll}0 & 0 & 0 & I \end{array} \right]$  and  $W_{\mathbf{Q}}^{2} = \left[ \begin{array}{llll}0 & 0 & 0 & I \end{array} \right]$  then we have
$$
\left(\boldsymbol {W} _ {\mathbf {K}} ^ {1} \boldsymbol {X}\right) ^ {\top} \left(\boldsymbol {W} _ {\mathbf {Q}} ^ {1} \boldsymbol {X}\right) = \left[ \begin{array}{c c c c c c} \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {k} & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {2} & \dots & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {k} & \dots & \mathbf {r} _ {1} ^ {\top} \mathbf {r} _ {n} \\ \mathbf {r} _ {2} ^ {\top} \mathbf {r} _ {k} & \mathbf {r} _ {2} ^ {\top} \mathbf {r} _ {2} & \dots & \mathbf {r} _ {2} ^ {\top} \mathbf {r} _ {k} & \dots & \mathbf {r} _ {2} ^ {\top} \mathbf {r} _ {n} \\ \vdots & \vdots & \ddots & \vdots & \ddots & \vdots \\ \mathbf {r} _ {k} ^ {\top} \mathbf {r} _ {k} & \mathbf {r} _ {k} ^ {\top} \mathbf {r} _ {2} & \dots & \mathbf {r} _ {k} ^ {\top} \mathbf {r} _ {k} & \dots & \mathbf {r} _ {k} ^ {\top} \mathbf {r} _ {n} \\ \vdots & \vdots & \ddots & \vdots & \ddots & \vdots \\ \mathbf {r} _ {n} ^ {\top} \mathbf {r} _ {k} & \mathbf {r} _ {n} ^ {\top} \mathbf {r} _ {2} & \dots & \mathbf {r} _ {n} ^ {\top} \mathbf {r} _ {k} & \dots & \mathbf {r} _ {n} ^ {\top} \mathbf {r} _ {n} \end{array} \right] \tag {A.4}
$$
As  $\lambda \to \infty$  and after the application of the softmax operator the above matrix becomes equal to  $\left[\mathrm{e}_k\quad \mathrm{e}_2\quad \ldots \quad \mathrm{e}_k\quad \ldots \mathrm{e}_n\right] + \epsilon \mathbf{M}$ , where  $\mathbf{e}_i$  is the one-hot vector with 1 in the  $k$ -th position,  $\| M\| \leq 1$  and  $\epsilon$  is controllable by the temperature parameter and can be arbitrary small. Let finally  $\boldsymbol{W}_{\mathbf{V}}^{1} = \begin{bmatrix} \mathbf{0} & \boldsymbol{I} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} \end{bmatrix}$  we get that
$$
\boldsymbol {W} _ {\mathbf {V}} ^ {1} \boldsymbol {X} \operatorname {s o f t m a x} \left(\left(\boldsymbol {W} _ {\mathbf {K}} ^ {1} \boldsymbol {X}\right) ^ {\top} \left(\boldsymbol {W} _ {\mathbf {Q}} ^ {1} \boldsymbol {X}\right)\right) = \left[ \begin{array}{c c c c} \boldsymbol {y} _ {k} & \boldsymbol {x} _ {2} & \dots & \boldsymbol {x} _ {n} \\ \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \end{array} \right] + \epsilon \mathbf {M} \tag {A.5}
$$
By repeating the exact same steps for the second head and letting  $W_{\mathbf{V}}^{2} = \left[ \begin{array}{cccc} - I & 0 & 0 & 0\\ 0 & 0 & 0 & 0 \end{array} \right]$  and adding back the residual we get the desired result.
A slightly different implementation can be found in Giannou et al. [2023], in which the ReLU layers are also used, together with indicator vectors that define whether a position should be updated or not.
Lemma 2 There exists a 1-hidden layer feedforward, ReLU network, with 8d activations in the hidden layer and  $d$  neurons in the output layer that when given two  $d$ -dimensional binary vectors representing two non-negative integers, can output the binary vector representation of their sum, as long as the sum is less than  $2^{d + 1}$ .
Lemma 3 Let  $\mathbf{A} \in \mathbb{R}^{d \times m}$  and  $\mathbf{B} \in \mathbb{R}^{d \times n}$ . Then for any  $\epsilon > 0$  there exists a transformer-based function block with 2 layers, 1 head and width  $r = O(d)$  that outputs the multiplication  $\mathbf{A}^{\top} \mathbf{B} + \epsilon \mathbf{M}$ , for some  $\| \mathbf{M} \| \leq 1$ .
Remark 2 Notice that based on the proof of this lemma, the matrices/scalars/vectors need to be in the same rows, i.e.,  $\mathbf{Q} = [\mathbf{A}\quad \mathbf{B}]$ . By also appending the appropriate binary encodings we can move the output at any specific place we choose as in Lemma 1. Also, following the proof of the paper the input matrix is
$$
\boldsymbol {X} = \left[ \begin{array}{c c c} \mathbf {Q} & \mathbf {0} & \mathbf {0} \\ \mathbf {0} & \mathbf {1 1} ^ {\top} & \mathbf {0} \\ \boldsymbol {I} & \mathbf {0} & \mathbf {0} \\ & \mathbf {r} ^ {(1)} \\ & \mathbf {r} ^ {(2)} \end{array} \right]
$$
where  $\mathbf{r}^{(1)},\mathbf{r}^{(2)}$  are chosen as to specify the position of the result.
A.4 Results on Filtering
Lemma 4 Assume that the input to a transformer layer is of the following form
$$
\boldsymbol {X} = \left[ \begin{array}{c c c c c} \boldsymbol {x} _ {1} & \boldsymbol {x} _ {2} & \dots & \boldsymbol {x} _ {n - 1} & \boldsymbol {x} _ {n} \\ \mathbf {0} & \mathbf {0} & \dots & \mathbf {0} & \mathbf {0} \\ b _ {1} & b _ {2} & \dots & b _ {n - 1} & b _ {n} \\ b _ {1} ^ {\prime} & b _ {2} ^ {\prime} & \dots & b _ {n - 1} ^ {\prime} & b _ {n} \end{array} \right] \tag {A.6}
$$
where  $b_{i}, b_{i}^{\prime} \in \{0,1\}$ , with zero indicating that the corresponding point should be ignored. Then there exists a transformer TF consisting only of a ReLU layer that performs this filtering, i.e.,
$$
T F (\boldsymbol {X}) = \left[ \begin{array}{c c c c} \mathbf {1} \left\{b _ {1} \neq 0 \right\} \boldsymbol {x} _ {1} & \mathbf {1} \left\{b _ {2} \neq 0 \right\} \boldsymbol {x} _ {2} & \dots & \mathbf {1} \left\{b _ {n - 1} \neq 0 \right\} \boldsymbol {x} _ {n - 1} & \mathbf {1} \left\{b _ {n} \neq 0 \right\} \boldsymbol {x} _ {n} \\ \mathbf {1} \left\{b _ {1} ^ {\prime} \neq 0 \right\} \boldsymbol {x} _ {1} & \mathbf {1} \left\{b _ {2} ^ {\prime} \neq 0 \right\} \boldsymbol {x} _ {2} & \dots & \mathbf {1} \left\{b _ {n - 1} ^ {\prime} \neq 0 \right\} \boldsymbol {x} _ {n - 1} & \mathbf {1} \left\{b _ {n} ^ {\prime} \neq 0 \right\} \boldsymbol {x} _ {n} \\ b _ {1} & b _ {2} & \dots & b _ {n - 1} & b _ {n} \\ b _ {1} ^ {\prime} & b _ {2} ^ {\prime} & \dots & b _ {n - 1} ^ {\prime} & b _ {n} \end{array} \right] \tag {A.7}
$$
Proof. The layer is the following:
$$
\operatorname {T F} \left(\boldsymbol {x} _ {i}\right) = x _ {i} + (- C b _ {i} - \boldsymbol {x} _ {i}) _ {+} - (- C b _ {i} + \boldsymbol {x} _ {i}) _ {+} \tag {A.8}
$$
for some large constant  $C$ . Notice that if  $b_{i} = 1$  the output is just  $\pmb{x}_i$ . But if  $b_{i} = 0$  then the output is zero. For the second set instead of using the bits  $b_{i}$ , we use the  $b_{i}'$ .
Remark 3 Notice that if some  $b_{i}$  is instead of  $1, 1 \pm \varepsilon, \varepsilon < c / C$ . Then the output of the above layer would be
$$
\begin{array}{l} T F \left(\boldsymbol {x} _ {i}\right) = \boldsymbol {x} _ {i} + (- C \pm C \varepsilon - \boldsymbol {x} _ {i}) - (- C \pm C \varepsilon + \boldsymbol {x} _ {i}) (A.9) \\ = \boldsymbol {x} _ {i} (A.10) \\ \end{array}
$$
while if some  $b_{i} = \pm \varepsilon$  instead of zero, and assuming that  $\pmb{x}_i > c > 0$  or  $\pmb{x}_i < -c < 0$  the output would be
$$
\begin{array}{l} T F \left(\boldsymbol {x} _ {i}\right) = \boldsymbol {x} _ {i} + (- C \varepsilon - \boldsymbol {x} _ {i}) _ {+} - (- C b _ {i} + \boldsymbol {x} _ {i}) _ {+} (A.11) \\ = \boldsymbol {x} _ {i} \pm C \varepsilon - \boldsymbol {x} _ {i} (A.12) \\ \leq c (A.13) \\ \end{array}
$$
If  $|\pmb{x}_i| \leq c$ , again the output would be less than or equal to  $c$ , where  $c$  can be arbitrarily small.
Our target is to create these binary tokens, as to perform the filtering. We now describe for clarity how next word prediction is performed. An  $d \times N$ -dimensional input is given to the transformer architecture, which is up-projected using the embedding layer; the positional encodings are also added/ appended to the input. At the last layer, the last token predicted is appended to the initial input to the transformer architecture. The only difference of the new input including the positional encodings (the input for the next iteration) is the  $n + 1$ -th token. This is a property that our construction maintains, i.e., the positional encodings used are oblivious to the prediction step performed and are
always the same for each individual token. We consider encoder-based architectures in all of our lemmas below.
In the subsequent lemma we construct an automated process that works along those guidelines. To do so we assume that the input to the transformer contains the following information:
1. An enumeration of the tokens from 1 to  $N$ .
2. The ln of the above enumeration.
3. Zeros for the tokens that correspond to the data points, 1 for each token that is the  $x_{\mathrm{test}}$  or it is a prediction based on it.
4. An enumeration 1 to  $L$  for each one of the data points provided. For example, if we are given three sets of data we would have:  $1 \dots L$  1 ...  $L$  1 ...  $L$ .
5. Some extra information that is needed to implement a multiplication step as described in Lemma 3 and to move things to the correct place.
The above information can be viewed as part of the encodings that are appended to the input of the transformer. Formally, we have
Lemma 5 Consider that a prompt with  $n$  in-context samples is given and the  $\ell - 1$ -th prediction has been made, and the transformer is to predict the  $\ell$ -th one. Assume the input to the transformer is:
$$
\boldsymbol {X} = \left[ \begin{array}{c c c c c c c c c c c} \boldsymbol {x} _ {1} & \dots & \boldsymbol {s} _ {1} ^ {\ell - 1} & \boldsymbol {s} _ {1} ^ {\ell} & \dots & \boldsymbol {s} _ {2} ^ {\ell - 1} & \boldsymbol {s} _ {2} ^ {\ell} & \dots & \boldsymbol {s} _ {n} ^ {L} & \boldsymbol {x} _ {t e s t} & \dots & \hat {\boldsymbol {s}} ^ {\ell - 1} \\ \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & L + \ell + 1 & L + \ell + 2 & \dots & n (L + 1) & n (L + 1) + 1 & \dots & N \\ \ln (1) & \dots & \ln (\ell) & \ln (\ell + 1) & \dots & \ln (L + \ell + 1) & \ln (L + \ell + 2) & \dots & \ln (n (L + 1)) & \ln (n (L + 1) + 1) & \dots & \ln (N) \\ 0 & \dots & 0 & 0 & \dots & 0 & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ \boldsymbol {m} _ {1} & \dots & \boldsymbol {m} _ {\ell} & \boldsymbol {m} _ {\ell + 1} & \dots & \boldsymbol {m} _ {L + \ell + 1} & \boldsymbol {m} _ {L + \ell + 2} & \dots & \boldsymbol {m} _ {n (L + 1)} & \boldsymbol {m} _ {n (L + 1) + 1} & \dots & \boldsymbol {m} _ {N} \end{array} \right]
$$
where  $(\hat{\pmb{s}}^i)_{i = 1}^{\ell -1}$  denote the first  $\ell -1$  recurrent outputs of TF and for simplicity, let  $N\coloneqq n(L + 1) + \ell$  denote the total number of tokens. Then there exists a transformer TF consisting of 7 layers that has as output
$$
T F (\boldsymbol {X}) = \left[ \begin{array}{c c c c c c c c c c} \mathbf {0} & \dots s _ {1} ^ {\ell - 1} & \mathbf {0} & \dots & s _ {2} ^ {\ell - 1} & \mathbf {0} & \dots & \mathbf {0} & \mathbf {0} & \dots \hat {\boldsymbol {s}} ^ {\ell - 1} \\ \mathbf {0} & \dots & \mathbf {0} & s _ {1} ^ {\ell} & \dots & \mathbf {0} & s _ {2} ^ {\ell} & \dots & \mathbf {0} & \dots & \mathbf {0} \\ \hline 0 & \dots & 1 & 0 & \dots & 1 & 0 & \dots & 0 & 0 & \dots & 1 \\ 0 & \dots & 0 & 1 & \dots & 0 & 1 & \dots & 0 & 0 & \dots & 0 \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & L + \ell + 1 & L + \ell + 2 & \dots & n (L + 1) & n (L + 1) + 1 & \dots & N \\ \ln (1) & \dots \ln (\ell) & \ln (\ell + 1) & \dots \ln (L + \ell + 1) & \ln (L + \ell + 2) & \dots \ln (n (L + 1)) & \ln (n (L + 1) + 1) & \dots \ln (N) \\ 0 & \dots & 0 & 0 & \dots & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ m _ {1} & \dots & m _ {\ell} & m _ {\ell + 1} & \dots & m _ {L + \ell + 1} & m _ {L + \ell + 2} & \dots & m _ {n (L + 1)} & m _ {n (L + 1) + 1} & \dots & m _ {N} \end{array} \right]
$$
with error up to  $\delta \mathbf{M}$ , where  $\| \mathbf{M} \| \leq 1$  and  $\delta > 0$  is a constant that is controlled and can be arbitrarily small.
Step 1: Extract the sequence length (1 layer). Let
$$
\boldsymbol {W} _ {\mathbf {K}} = \left[ \begin{array}{l l l l l l l l l} \mathbf {0} & \mathbf {0} & 0 & 0 & 0 & 1 & 0 & 0 & \mathbf {0} \end{array} \right] \quad \boldsymbol {W} _ {\mathbf {Q}} = \left[ \begin{array}{l l l l l l l l l l} \mathbf {0} & \mathbf {0} & 0 & 0 & 0 & 0 & 0 & 0 & 1 & \mathbf {0} \end{array} \right] \tag {A.14}
$$
$$
\left(\boldsymbol {W} _ {\mathbf {K}} \boldsymbol {X}\right) ^ {\top} \left(\boldsymbol {W} _ {\mathbf {Q}} \boldsymbol {X}\right) = \left[ \begin{array}{c c c c} \ln (1) & \ln (1) & \dots & \ln (1) \\ \ln (2) & \ln (2) & \dots & \ln (2) \\ \vdots & \vdots & \ddots & \vdots \\ \ln (\mathrm {N}) & \ln (\mathrm {N}) & \dots & \ln (\mathrm {N}) \end{array} \right]. \tag {A.15}
$$
So after the softmax is applied we have
$$
\operatorname {s o f t m a x} \left(\left(\boldsymbol {W} _ {\mathbf {K}} \boldsymbol {X}\right) ^ {\top} \left(\boldsymbol {W} _ {\mathbf {Q}} \boldsymbol {X}\right)\right) = \left[ \begin{array}{c c c c} \frac {1}{\sum_ {i = 1} ^ {\mathrm {N}} i} & \frac {1}{\sum_ {i = 1} ^ {\mathrm {N}} i} & \dots & \frac {1}{\sum_ {i = 1} ^ {\mathrm {N}} i} \\ \overline {{\sum_ {i = 1} ^ {\mathrm {N}} i}} & \overline {{\sum_ {i = 1} ^ {\mathrm {N}} i}} & \dots & \overline {{\sum_ {i = 1} ^ {\mathrm {N}} i}} \\ \vdots & \vdots & \ddots & \vdots \\ \overline {{\sum_ {i = 1} ^ {\mathrm {N}} i}} & \overline {{\sum_ {i = 1} ^ {\mathrm {N}} i}} & \dots & \overline {{\sum_ {i = 1} ^ {\mathrm {N}} i}} \end{array} \right]. \tag {A.16}
$$
We then set the weight value matrix as to zero-out all lines except for one line as follows
$$
\boldsymbol {W} _ {\mathbf {V}} \boldsymbol {X} = \left[ \begin{array}{l l l l} \mathbf {0} & \mathbf {0} & \dots & \mathbf {0} \\ 1 & 2 & \dots & \mathrm {N} \\ \mathbf {0} & \mathbf {0} & \dots & \mathbf {0} \end{array} \right]. \tag {A.17}
$$
After adding the residual and using an extra head where the softmax returns identity matrix and the value weight matrix is minus the identity, we get attention output
$$
\left[ \begin{array}{c c c c c c c c c c c c} \boldsymbol {x} _ {1} & \dots & \boldsymbol {s} _ {1} ^ {\ell - 1} & \boldsymbol {s} _ {1} ^ {\ell} & \dots & \boldsymbol {s} _ {2} ^ {\ell - 1} & \boldsymbol {s} _ {2} ^ {\ell} & \dots & \boldsymbol {s} _ {n} ^ {L} & \boldsymbol {x} _ {\text {t e s t}} & \dots & \hat {\boldsymbol {s}} ^ {\ell - 1} \\ \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} & \boldsymbol {0} & \dots & \boldsymbol {0} \\ \hline 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ 1 & \dots & \ell & \ell + 1 & \dots & \ell & \ell + 1 & \dots & L + 1 & 1 & \dots & \ell \\ \hline \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {i = 1} ^ {N} i} & \dots & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {i = 1} ^ {N} i} & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {i = 1} ^ {N} i} & \dots & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {i = 1} ^ {N} i} & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {1 = 1} ^ {N} i} & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {i = 1} ^ {N} i} & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {i = 1} ^ {N} i} & \dots & \frac {\sum_ {i = 1} ^ {N} i ^ {2}}{\sum_ {\mathrm {i n}} ^ {N} i} \\ \ln (1) & \dots & \ln (\ell) & \ln (\ell + 1) & \dots & \ln (L + \ell + 1) & \ln (L + \ell + 2) & \dots & \ln (n (L + 1)) & \ln (n (L + 1) + 1) & \dots & \ln (N) \\ 0 & \dots & 0 & 0 & \dots & 0 & 0 & \dots & 0 & 1 & \dots & 1 \\ 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 & 1 & \dots & 1 \\ \boldsymbol {m} _ {1} & \dots & \boldsymbol {m} _ {\ell} & \boldsymbol {m} _ {\ell + 1} & \dots & \boldsymbol {m} _ {L + \ell + 1} & \boldsymbol {m} _ {L + \ell + 2} & \dots & \boldsymbol {m} _ {n (L + 1)} & \boldsymbol {m} _ {n (L + 1) + 1} & \dots & \boldsymbol {m} _ {\mathrm {N}} \\ \end{array} \right]
$$
Notice that  $\frac{\sum_{i=1}^{N} i^2}{\sum_{i=1}^{N} i} = \frac{2N(N + 1)(2N + 1)}{6N(N + 1)} = \frac{2N + 1}{3}$ . We then use the ReLU layer as to multiply with  $3/2$  and subtract 1 from this column. This results in attention output