SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathrm {b u c k e t} (\left[ \begin{array}{c c c c c c c c c c} 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 2 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 3 & 2 & 1 & 0 & 0 & 0 & 0 & 0 & 0 \\ 4 & 3 & 2 & 1 & 0 & 0 & 0 & 0 & 0 \\ 5 & 4 & 3 & 2 & 1 & 0 & 0 & 0 & 0 \\ 6 & 5 & 4 & 3 & 2 & 1 & 0 & 0 & 0 \\ 7 & 6 & 5 & 4 & 3 & 2 & 1 & 0 & 0 \\ 8 & 7 & 6 & 5 & 4 & 3 & 2 & 1 & 0 & 0 \\ 9 & 8 & 7 & 6 & 5 & 4 & 3 & 2 & 1 & 0 \end{array} \right]) = \left[ \begin{array}{c c c c c c c c c c} 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 2 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline \hline
$$
ALiBi Similar to T5's Relative PE, ALiBi subtracts a scalar bias from the attention score. As the distance between the query and key tokens increases, the bias grows linearly. Specifically, the dot product in every layer can be written as:
$$
\left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {i} \right\rangle = \boldsymbol {q} _ {t} ^ {\top} \boldsymbol {k} _ {i} - (t - i). C ^ {(m + 1)} \tag {13}
$$
where  $m$  is head index and  $C$  is a constant defined as:
$$
C = 2 ^ {- 2 ^ {- \log_ {2} (\# \text {h e a d s} + 3)}
$$
For example, if the number of heads is 8, then we have  $\frac{1}{2}, \frac{1}{2^2}, \ldots, \frac{1}{2^8}$  (Press et al., 2022).
Rotary The Rotary is a relative PE that applies a rotation to the query and key representations based on their absolute positions before dot product attention. Due to this rotation, the attention dot product relies solely on the relative distance between tokens.
First, we formulate Rotary for model dimension  $d = 2$ . Rotary positional encoding defines the dot product as:
$$
\begin{array}{l} \left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {i} \right\rangle = \left\langle \operatorname {R o t} \left(\boldsymbol {q} _ {t}, t\right), \operatorname {R o t} \left(\boldsymbol {k} _ {i}, i\right) \right\rangle \\ = \left\langle \boldsymbol {R} ^ {t \theta} \boldsymbol {q} _ {t}, \boldsymbol {R} ^ {i \theta} \boldsymbol {k} _ {i} \right\rangle \\ = \left(\boldsymbol {R} ^ {t \theta} \boldsymbol {q} _ {t}\right) ^ {\top} \left(\boldsymbol {R} ^ {i \theta} \boldsymbol {k} _ {i}\right) \\ = \boldsymbol {q} _ {t} ^ {\intercal} (\boldsymbol {R} ^ {t \theta}) ^ {\intercal} \boldsymbol {R} ^ {i \theta} \boldsymbol {k} _ {i} \\ = q _ {t} ^ {\intercal} R ^ {(i - t) \theta} k _ {i} \tag {14} \\ \end{array}
$$
where  $\pmb{R}^{t\theta}$  is a rotation matrix that rotates  $\pmb{x}$  by  $t\theta$  radians:
$$
\boldsymbol {R} ^ {n \theta} = \left[ \begin{array}{c c} \cos (n \theta) & - \sin (n \theta) \\ \sin (n \theta) & \cos (n \theta) \end{array} \right] \tag {15}
$$
for  $d > 2$ , Rotary applies the same approach on every two consecutive dimensions of  $\mathbf{q}_t$  and  $\pmb{k}_i$ , but with different  $\theta$  angles. Refer to Su et al. (2021) for the exact formulation.
NoPE NoPE does not explicitly encode positional encodings. So, the dot product in every layer can be written as:
$$
\left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {i} \right\rangle = \boldsymbol {q} _ {t} ^ {\intercal} \boldsymbol {k} _ {i} \tag {16}
$$
In this section, we provide proof of why NoPE can implicitly learn both absolute and relative positions. We refer the readers to Appendix B.1 for the notation and definitions used in this section.
C.1 Absolute Positional Encoding in NoPE
This section discusses how NoPE can recover absolute positions in the hidden state. Our proof is inspired by Weiss et al. (2021); Lindner et al. (2023) and relies on the causal attention mask in the decoder-only Transformer and the softmax function to recover absolute positions.
Theorem 1 (Absolute Encoding). Let  $\mathbf{x} = [\langle bos\rangle, x_1, \ldots, x_T]$  be an input sequence of length  $T + 1$  to the model. Then, the first layer of  $f_{\theta}$  can recover absolute positions  $[1, \ldots, T + 1]$  in the hidden state  $H^{(1)}$ . That is, there exist  $W_Q, W_K, W_V, W_O, W_1,$  and  $W_2$  such that the self-attention and feedforward operations in the first layer compute absolute positions and write it to the next hidden state.
Our proof only specifies the weights of a single attention head in the first layer (and additionally the parameterization of feedforward sub-layer). In this parameterization, we only require the first three dimensions of the hidden states. The rest of the heads, as long as they do not override the first three dimensions, can be arbitrary. This does not impose any challenges as Transformers used in practice usually have a very large model dimension  $d$ . In the rest, we provide the construction of the weights and then verify that they can recover absolute positions.
First, we construct the word embedding matrix  $W_{E} \in \mathbb{R}^{d \times V}$ , where each column is the embedding of a token in the vocabulary. We construct  $W_{E}$  such that it always sets the first dimension of every embedding vector to be 1. Additionally, it sets the second dimension to 1 if and only if the token is  $<\text{bos}>$ . Otherwise, it sets it to zero. The third dimension of all embedding vectors is set to zero. Other dimensions can take any arbitrary values. Without loss of generality, assume  $<\text{bos}>$  is the first token in the vocabulary, i.e. The first column. Then, we have:
$$
\boldsymbol {W} _ {E} = \left[ \begin{array}{c c c c c} 1 & 1 & 1 & \dots & 1 \\ 1 & 0 & 0 & \dots & 0 \\ 0 & 0 & 0 & \dots & 0 \\ e _ {4, 1} & e _ {4, 2} & e _ {4, 3} & \dots & e _ {4, V} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ e _ {d, 1} & e _ {d, 2} & e _ {d, 2} & \dots & e _ {d, V} \end{array} \right] _ {d \times V} \tag {17}
$$
where  $e_{d,i}\in \mathbb{R}$
Secondly, for head dimensions  $h \geq 1$ , we construct the weights  $W_{Q}, W_{K}, W_{V}, W_{O}$  of the first attention head in the first layer. Specifically,
$$
\boldsymbol {W} _ {K} = \left[ \begin{array}{c c c c} 1 & 0 & \dots & 0 \\ 1 & 0 & \dots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 1 & 0 & \dots & 0 \end{array} \right] _ {h \times d} \quad \boldsymbol {W} _ {V} = \left[ \begin{array}{c c c c c} 0 & 1 & 0 & \dots & 0 \\ 0 & 0 & 0 & \dots & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \dots & 0 \end{array} \right] _ {h \times d} \tag {18}
$$
$W_{K}$  reads from the first dimension of the hidden state, which is initialized with 1 using the embedding matrix. Since all word embeddings have one in their first dimension, this parameterization will result all key vectors to be the same. Moreover,  $W_{V}$  reads from the second dimension of the hidden state, which is initialized with 1 if the token is  $\langle \text{bos} \rangle$ . So, the value vector will have 1 in its first dimension only if the corresponding token is  $\langle \text{bos} \rangle$ .
$W_{Q}$  can be any arbitrary matrix.  $W_{O}$  will write the result of the attention to the third dimension of the hidden state and can be constructed as:
$$
\boldsymbol {W} _ {O} = \left[ \begin{array}{c c c c c c} 0 & 0 & 0 & 0 & \dots & 0 \\ 0 & 0 & 0 & 0 & \dots & 0 \\ 1 & 0 & 0 & 0 & \dots & 0 \\ 0 & 0 & 0 & 0 & \dots & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & 0 & \dots & 0 \end{array} \right] _ {d \times h} \tag {19}
$$
Now, we verify that for any input sequence  $\pmb{x} = [<\mathrm{bos}], x_1, \ldots, x_T]$ , the first layer can recover absolute positions  $[1, \ldots, T + 1]$  in the hidden state  $H^{(1)}$ . We verify this for column  $t$  of  $H^{(1)}$ . That is, we show that absolute position information is available in the third dimension of  $h_t^{(1)}$ .
First, we use the word embedding matrix  $\mathbf{W}_E$  to compute the embedding  $\mathbf{H}^{(0)}$ :
$$
\boldsymbol {H} ^ {(0)} = \boldsymbol {W} _ {E} \boldsymbol {X} = \left[ \begin{array}{c c c c c} 1 & 1 & 1 & \dots & 1 \\ 1 & 0 & 0 & \dots & 0 \\ 0 & 0 & 0 & \dots & 0 \\ e _ {4, 1} & e _ {4, 2} & e _ {4, 3} & \dots & e _ {4, V} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ e _ {d, 1} & e _ {d, 2} & e _ {d, 2} & \dots & e _ {d, V} \end{array} \right] _ {d \times (T + 1)} \tag {20}
$$
We now provide the attention computation at position  $1 \leq t \leq T + 1$ . First, we use  $\mathbf{W}_Q$  to compute the query vector  $\mathbf{q}_t$  by applying  $\mathbf{q}_t = \mathbf{W}_Q \mathbf{h}_t^{(0)}$ :
$$
\boldsymbol {q} _ {t} = \left[ q _ {1}, q _ {2}, q _ {3}, \dots , q _ {h} \right] ^ {\intercal} \tag {21}
$$
Recall that  $W_{Q}$  can be any arbitrary matrix. So,  $q_{j} \in \mathbb{R}$  can take any arbitrary value. Next, we compute the key vectors by applying  $\pmb{k}_{i} = \pmb{W}_{K}\pmb{h}_{i}^{(0)}$ :
$$
\boldsymbol {k} _ {1} = \left( \begin{array}{c} 1 \\ 1 \\ \vdots \\ 1 \end{array} \right) \quad \boldsymbol {k} _ {2} = \left( \begin{array}{c} 1 \\ 1 \\ \vdots \\ 1 \end{array} \right) \quad \dots \quad \boldsymbol {k} _ {t} = \left( \begin{array}{c} 1 \\ 1 \\ \vdots \\ 1 \end{array} \right) \tag {22}
$$
Note that all key vectors are the same and we only need to compute them up to position  $t$  as the attention mask is causal, i.e. query can only look at positions  $\leq t$ . Next, we compute the attention weight vectors  $\alpha$ :
$$
\begin{array}{l} \boldsymbol {\alpha} = \left[ \left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {1} \right\rangle , \left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {2} \right\rangle , \dots , \left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {t} \right\rangle \right] ^ {\intercal} (23) \\ = \left[ \alpha^ {*}, \alpha^ {*}, \dots , \alpha^ {*} \right] ^ {\mathsf {T}} (24) \\ \end{array}
$$
where  $\alpha^{*} = q_{1} + q_{2} + \ldots + q_{h}$ . Next, we apply softmax to compute the attention probabilities. Since all  $\alpha^{i}$ 's are the same, we have:
$$
\hat {\boldsymbol {\alpha}} = \operatorname {s o f t m a x} (\boldsymbol {\alpha}) = \left[ \frac {1}{t}, \frac {1}{t}, \dots , \frac {1}{t} \right] ^ {\intercal} \tag {25}
$$
Now, we compute the value vectors by applying  $\boldsymbol{v}_i = \boldsymbol{W}_V \boldsymbol{h}_i^{(0)}$ :
$$
\boldsymbol {v} _ {1} = \left( \begin{array}{c} 1 \\ 0 \\ \vdots \\ 0 \end{array} \right) \quad \boldsymbol {v} _ {2} = \left( \begin{array}{c} 0 \\ 0 \\ \vdots \\ 0 \end{array} \right) \quad \dots \quad \boldsymbol {v} _ {t} = \left( \begin{array}{c} 0 \\ 0 \\ \vdots \\ 0 \end{array} \right) \tag {26}
$$
Finally, we compute the output of the attention head by applying  $\mathbf{W}_O$ :
$$
\boldsymbol {o} _ {t} = \boldsymbol {W} _ {O} \left(\sum_ {i \leq t} \hat {\alpha} _ {i} \boldsymbol {v} _ {i}\right) = \boldsymbol {W} _ {O} \left(\frac {1}{t} \sum_ {i \leq t} \boldsymbol {v} _ {i}\right) = \boldsymbol {W} _ {O} \binom {1 / t} {0} _ {h} = \binom {0} {0} _ {d} \tag {27}
$$
Thus, the output of our constructed attention head recovers the absolute position information and writes it to the third dimension of output.
We used the decoder-only property of Transformer implicitly in Equation (23), which helped us to only attend to position  $\leq t$ . So, the lengths of the attended sequence are always  $t$ . Moreover, the presence of <bos> token in the input sequence helped us to anchor the absolute position information. This is not a problem as in practice models are often prompted with some instructions which can act as <bos> token.
With this information available to the rest of the network, the feedforward sub-layer, with sufficient hidden width, can recover the absolute positions  $[1,2,\dots,T + 1]$  from the third dimension of attention output. This is because the feedforward sub-layer is MLP with ReLU activation. So, it can learn any arbitrary function (Park et al., 2020). Note that the layer-norm operation can be bypassed as explained by Akyurek et al. (2023).
C.2 Relative Positional Encoding in NoPE
In this section, we show if the hidden state contains absolute positional information as explained in the previous section, then the attention mechanism in all subsequent layers can implement a relative positional encoding. We refer the readers to Appendices B.1 and C.1 for the notation and definitions used in this section.
Theorem 2 (Relative Encoding). Suppose that the hidden state  $H^{(1)}$  contains absolute positional information, as stated in Theorem 1, and assume that it is not overwritten by any subsequent layers. Then, the self-attention in all subsequent layers can implement a relative positional encoding: there exists a parameterization of  $f_{\theta}$  such that, for  $l \geq 2$ , the attention dot product between query  $q_{t}$  and key  $k_{i}$  at positions  $t$  and  $i$  ( $t \geq i$ ) can be expressed as:
$$
\left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {i} \right\rangle = f _ {\text {c n t}} \left(\boldsymbol {q} _ {t}, \boldsymbol {k} _ {i}\right) + f _ {\text {r e l}} (t - i) \tag {1}
$$
Our proof only specifies a few entries of weight matrices for attention heads in layers  $l \geq 2$ , which does not impose any challenges for Transformers used in practice as they usually have a very large model dimension  $d$ . Moreover, we require to have absolute positions in the third dimension of the hidden state as explained in Theorem 1. To show NoPE can implement relative encoding, we only need to prove that its attention dot product depends on the relative distance between tokens (See
Appendix B.1 for an overview of relative encoding methods). In the rest, we provide the construction of the weights and then verify that they can implement relative position encoding.
For head dimension  $h \geq 2$ , we construct the weights  $W_{Q}, W_{K}$  of the attention heads in the second layers and above. Specifically,
$$
\boldsymbol {W} _ {Q} = \left[ \begin{array}{c c c c c c} 1 & 0 & 0 & 0 & \dots & 0 \\ 0 & 0 & - 1 & 0 & \dots & 0 \\ w _ {3, 1} & w _ {3, 2} & w _ {3, 3} & w _ {3, 4} & \dots & w _ {3, d} \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \\ w _ {h, 1} & w _ {h, 2} & w _ {h, 3} & w _ {h, 4} & \dots & w _ {h, d} \end{array} \right] _ {h \times d} \tag {28}
$$
$$
\boldsymbol {W} _ {V} = \left[ \begin{array}{c c c c c c} 0 & 0 & 1 & 0 & \dots & 0 \\ 1 & 0 & 0 & 0 & \dots & 0 \\ w _ {3, 1} ^ {\prime} & w _ {3, 2} ^ {\prime} & w _ {3, 3} ^ {\prime} & w _ {3, 4} ^ {\prime} & \dots & w _ {3, d} ^ {\prime} \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \\ w _ {h, 1} ^ {\prime} & w _ {h, 2} ^ {\prime} & w _ {h, 3} ^ {\prime} & w _ {h, 4} ^ {\prime} & \dots & w _ {h, d} ^ {\prime} \end{array} \right] _ {h \times d} \tag {29}
$$
where  $w_{i,j}, w_{i,j}^{\prime} \in \mathbb{R}$  can take any arbitrary value. Their corresponding  $\pmb{W}_{V}$  and  $\pmb{W}_{O}$  can take any arbitrary values as long as they do not override the first three dimensions of the residual stream.
Now we verify that for any input sequence  $\pmb{x} = [\langle \mathrm{bos} \rangle, x_1, \dots, x_T]$ , the attention dot product between query  $\pmb{q}_t$  and key  $\pmb{k}_i$  at positions  $t$  and  $i$  ( $t \geq i$ ) will depend the relative distance between tokens.
First, assume that absolute positions are computed in the hidden state  $\pmb{H}^{(l)}$  for  $l \geq 1$ , as stated in Theorem 1. Specifically,
$$
\boldsymbol {H} ^ {(l)} = \left[ \begin{array}{c c c c c c} 1 & 1 & 1 & 1 & \dots & 1 \\ 1 & 0 & 0 & 0 & \dots & 0 \\ 1 & 2 & 3 & 4 & \dots & T + 1 \\ h _ {4, 1} & h _ {4, 2} & h _ {4, 3} & h _ {4, 4} & \dots & h _ {4, T + 1} \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \\ h _ {d, 1} & h _ {d, 2} & h _ {d, 3} & h _ {d, 4} & \dots & h _ {d, T + 1} \end{array} \right] _ {d \times (T + 1)} \tag {30}
$$
where  $h_{i,j} \in \mathbb{R}$  can be any arbitrary value as the first three dimensions of the hidden state are reserved for PE computation. The rest of the dimensions can take any arbitrary values as in regular computation of Transformers.
We now present the attention computations at position  $1 \leq t \leq T + 1$ . We use  $\mathbf{W}_Q$  to compute the query vector  $\mathbf{q}_t$  by applying  $\mathbf{q}_t = \mathbf{W}_Q \mathbf{h}_t^{(l)}$ :
$$
\boldsymbol {q} _ {t} = [ 1, - t, q _ {3}, \dots , q _ {h} ] ^ {\intercal} \tag {31}
$$
where  $q_{j} \in \mathbb{R}$  can take any arbitrary value. Next, we compute the key vectors by applying  $\pmb{k}_{i} = \pmb{W}_{K}\pmb{h}_{i}^{(l)}$ :
$$
\boldsymbol {k} _ {1} = \left( \begin{array}{c} 1 \\ 1 \\ k _ {3, 1} \\ \vdots \\ k _ {h, 1} \end{array} \right) \quad \boldsymbol {k} _ {2} = \left( \begin{array}{c} 2 \\ 1 \\ k _ {3, 2} \\ \vdots \\ k _ {h, 2} \end{array} \right) \quad \boldsymbol {k} _ {3} = \left( \begin{array}{c} 3 \\ 1 \\ k _ {3, 3} \\ \vdots \\ k _ {h, 3} \end{array} \right) \quad \dots \quad \boldsymbol {k} _ {t} = \left( \begin{array}{c} t \\ 1 \\ k _ {3, t} \\ \vdots \\ k _ {h, t} \end{array} \right) \tag {32}
$$
where  $k_{(\cdot ,\cdot)}\in \mathbb{R}$  can have any arbitrary value. So, for  $\pmb{k}_i$  we have:
$$
\boldsymbol {k} _ {i} = [ i, 1, k _ {3, i}, \dots , k _ {h, i} ] ^ {\intercal} \tag {33}
$$
Next, we let us present the attention dot product between  $\pmb{q}_t$  and  $\pmb{k}_i$ :
$$
\begin{array}{l} \left\langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {i} \right\rangle = 1. i + (- t). 1 + q _ {3}. k _ {3, i} + \dots + q _ {h}. k _ {h, i} (34) \\ = i - t + \sum_ {j = 3} ^ {h} q _ {j}. k _ {j, i} (35) \\ = \left(\sum_ {j = 3} ^ {h} q _ {j}. k _ {j, i}\right) - (t - i) (36) \\ = f _ {\text {c n t}} \left(\boldsymbol {q} _ {t}, \boldsymbol {k} _ {i}\right) + f _ {\text {r e l}} (t - i) (37) \\ \end{array}
$$
Thus, the dot product between  $\mathbf{q}_t$  and  $\mathbf{k}_i$  depends on the relative distance between tokens (assuming the rest of the terms do not cancel out which can be easily avoided by setting the respective weights in Equations (28) and (29)). Note that our proof uses the linear spacing between tokens, but the MLP the first layer can write any arbitrary function of absolute positions to the third dimension of the hidden state, which enables more complex relative encoding schemes.
D.1 Tasks
Here we provide the details and more examples of the tasks and datasets we used in our evaluation. For each task, we sample 100K examples for the training set and 10K for the test. Also, we use  $15\%$  of the train as the validation set.
Addition The addition task (Nye et al., 2021) asks the model to compute the sum of two numbers. Each number is represented as a sequence of digits that are separated by space. So, the model has access to the exact digits.
Input   
Compute: 53726+1917?   
Output   
The answer is 55643.
we create each length bucket based on the number of digits in each number, e.g. 6-by-3, 6-by-4, etc. For the training set, we use buckets where one of the numbers has at most  $L$  digits. For the test set, we use buckets where any of the numbers have at most  $L$  digits. The model is evaluated on the correctness of its predicted result.
Polynomial Evaluation The polynomial evaluation task (Nye et al., 2021) asks the model to evaluate a polynomial expression at a given value of  $x$ . The polynomial terms and digits are separated to make just the tokenizer does not glue symbols together.
Input   
Evaluate  $\mathbf{x} = 3$  in  $(3\texttt{x} *\texttt{0} +1\texttt{x} *\texttt{1} +1\texttt{x} *\texttt{2})\% 10?$  Output   
The answer is 5.
The length bucket is created based on the number of terms in the polynomial expression. We sample  $x$  from  $\mathcal{U}(-2, 2)$ , the degree of each term from  $\mathcal{U}(0, 3)$ , and the coefficient of each term from  $\mathcal{U}(-3, 3)$ . We take the modulo of the result by 10 to make the task easier for the model and make sure we only measure the generalization of the length of the problem instance not the value of the polynomial. The model is evaluated on the correctness of its predicted result.
Sorting The sorting task (Saxton et al., 2019) asks the model to sort a sequence of input numbers. We use this task in two variants: Single Token and Multi Digit. In the Single Token variant, we create an alphabet of 50 tokens from the model's vocabulary and fix some canonical ordering among them
through task. Each instance is a sequence of tokens from the alphabet in a random order, and the model is asked to sort them in the canonical order.
Input   
Sort the following numbers: 31415? Output   
The answer is 11345.
In the Multi Digit variant, we simply present a sequence of multi digit/tokens numbers to the model, and ask it to sort them in ascending order. Each number is represented by its digits and they are separated by a space.
Input   
Sort the following numbers: 31,41,59,126,533? Output   
The answer is 31,41,59,126,533.
In this case, we sample each number from  $\mathcal{U}(0,10000)$ . In both cases, the length bucket is created based on the length of the input sequence. The model is evaluated on the correctness of its predicted result.
Summation In this task (Saxton et al., 2019), we ask the model to compute the sum of a sequence of input numbers modulo 10 as we want to specifically measure how the model generalizes to longer sequences not the value of summation result:
Input   
Compute:  $(1 + 2 + 3 + 4 + 7)\% 10?$  Output   
The answer is 7.
Each digit is randomly sampled from  $\mathcal{U}(1,9)$ . The length bucket is created based on the length of the input sequence. The model is evaluated on the correctness of its predicted result.
Parity In the parity task (Anil et al., 2022), we ask the model to compute the parity of a binary sequence.
Input   
Is the number of 1's even in [10011] ? Output   
The answer is No.
LEGO In the LEGO task (Zhang et al., 2023), the model is provided with a simple computation graph (DAG), where each node represents a variable, and variables are connected by simple operations which created the edges in the computation graph. We refer to Zhang et al. (2023) for a detailed description.
Input   
If  $a = -1$  .  $\texttt{b} = \texttt{-a};\texttt{c} = +\texttt{b};\texttt{d} = +\texttt{c}$  . Then what is c? Output   
The answer is  $+1$
To sample each example, we first sample the list of variables based on the length of the example, and then we uniformly sample the value of each variable to make sure all variables are represented with both -1 and +1. Finally, given the value of variables, we deterministically compute the operation on each edge. For each example, we query all variables from the middle of the computation graph to the end. The model is evaluated on the correctness of its predicted result.
Copy The copy task is straightforward. The model has to repeat the input sequence in the output.
Input   
Copy the following words: <w1> <w2> <w3> <w4> <w5> . Output   
<w1> <w2> <w3> <w4> <w5>
We create multiple variants of this task to better understand the models' generalization behavior. In the first variant, the input tokens are the same, so the model has to basically count the number of input tokens. In the second variant, the model has to replace the input tokens with another token sampled from the vocabulary. In the third variant, we sample the input tokens from the model's vocabulary, and the model has to predict them in the same order. We also create 2x versions of variants 1 and 3 to make the tasks more challenging.
Reverse In this task the model, the model has to reverse the order of input tokens in its output.
Input   
Reverse the following words: <w1> <w2> <w3> <w4> <w5> . Output   
<w5> <w4> <w3> <w2> <w1> .
As in the copy task, we create multiple variants of this task. In the first variant, the model has to reverse the order of input tokens, where the tokens are randomly sampled from the model's vocabulary. In the second variant, the model has to reverse the order of input tokens, as in the first variant, but also it has to reverse it one more time, recreating the original input.