SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To account for those, we consider five components in each step of scratchpad: <input>, <computation>, <output>, <variable_update>, and <remaining_input> (Figure 5). In our experiments, we create different variations of scratchpad format by enabling or disabling each component, which allows us to systematically study their impact. Figure 6 summarizes our results. Similar to the remarks made by (Nye et al., 2021; Anil et al., 2022) we observe that across all PEs and regardless of the format, scratchpad is beneficial solely for the
addition task. Additionally, our findings indicate that having a positional encoding with robust length generalization is crucial since scratchpad/CoT alone may not enhance the generalization.
Figure 5: Example of an addition task depicted with its first scratchpad step. Each step consists of five components: Step Input  $\mathcal{I}$ , Step Computation  $\mathcal{C}$ , Step Output  $\mathcal{O}$ , Intermediate Variable Updates  $\mathcal{V}$ , and Remaining Input  $\mathcal{R}$ .
6.1 Which part of the sequence is attended to?
The scratchpad format that is often used (Nye et al., 2021), similar to Figure 5, contains redundant information. One such example is the repetition of the remaining portion of an input  $(\mathcal{R})$  in each step of the scratchpad. But, the attention can attend to this information directly from the main input. So, it remains unclear which specific part of the scratchpad different PEs rely on to solve the task.
To address this question, we take the models trained with full Format on addition, the case in which scratchpad is helpful across all PEs, and examine their attentions. Specifically, for tokens in
<sup>1</sup>Since using scratchpad creates very long sequences, we follow Nye et al. (2021) and set the length threshold  $L = 8$  for tasks that use it to avoid out-of-memory errors.
Figure 6: Mean ranks of scratchpad format aggregated across all models per each dataset. The effectiveness of scratchpad is task dependent.
the output sequence, we calculate the distance  $d$  between the current query  $\pmb{q}_t$  and the attended key  $\pmb{k}_n$  as  $(t - n + 1)$  and subsequently normalize it based on the length of the sequence at the present step. The normalized value is denoted as  $\bar{d}$ .
Figure 7 depicts the distribution of  $\bar{d}$ . Values of  $\bar{d}$  close to 0 indicate attention to tokens near the current position (e.g. current scratchpad step), while values close to 1 signify attention to distant tokens (e.g. the input). NoPE and T5's Relative PE resemble each other and exhibit a bimodal distribution, reflecting both short-range and long-range attention. Conversely, ALiBi (due to its recency bias) strongly favors short-range attention. Rotary, on the other hand, produces a distribution resembling APE, which is more uniformly distributed. Notably, NoPE and T5's RPE are the top-performing PEs in this setup, which suggest the bimodal distribution to be more optimal.
Practitioners have to make important choices about the nuances of the Transformer architecture like positional encoding before undertaking the costly pretraining process. In the I.I.D evaluation of PEs, we demonstrate similar performance across different PEs, in line with observations of Haviv et al. (2022) and Scao et al. (2022b), which makes the choice of optimal positional encoding challenging.
In our paper, we utilize length generalization in downstream tasks as a mean to assess the expressivity of positional encodings. Our setup, in contrast to the I.I.D.
evaluation, reveals a clear distinction among approaches of encoding positions. We find that NoPE outperforms explicit PEs, and within explicit PEs, commonly used methods lag behind T5's Relative PE. In fact, the recent release of LLMs (Touvron et al., 2023; Chowdhery et al., 2022) suggests a shift towards adopting Rotary as a replacement for APE in the Transformer architecture. However, our result in Section 4 clearly demonstrates that Rotary marginally outperforms APE at length generalization. Furthermore, it exhibits similar behavior to APE, as shown in Section 6.1, indicating potential susceptibility to the same limitations.
The disadvantages of explicit PEs over NoPE in length extrapolation contribute to the growing evidence that positional encodings pose challenges for Transformers (Sinha et al., 2022; Luo et al., 2021). Our empirical results and theoretical analysis suggest that removing positional encoding holds promise as a modification to the widely used decoder-only Transformer architecture.
Figure 7: Distribution of the normalized distance between the query and the key of the self-attention (addition task + full scratchpad), averaged across all layers and heads.
Scaling up to 1B models In order to study the behavior of position embeddings at scale, we trained three 1B variants post-submission - ALiBi, Rotary and NoPE - with context length of 1024 tokens on a subset of StarCoder training set (Li et al., 2023). For more details, refer to Appendix F. Our results on language modeling show that at I.I.D all variants have similar perplexity, but at length generalization, Rotary fails to generalize as its perplexity explodes. NoPE and ALiBi generalize similarly to larger context sizes up to almost twice their training context size, and for larger contexts ALiBi is relatively more stable than NoPE (see the discussion on perplexity vs. downstream performance). Preliminary exploration of fine-tuning the pretrained models, on datasets in Section 3, yielded to identical performance among PE variants as the training context size of the 1.3B models is much larger than instance lengths in our datasets. A comprehensive downstream evaluation of these models remains a subject for future research.
Perplexity vs. downstream Performance Due to human cognitive constraints (Gibson, 1998; Kiyono et al., 2021), language modeling data might encompasses short-range dependencies. The combination of this naturally occurring structure (which can be abundant in internet-based corpora) with the Recency Bias inherent in positional encodings like ALiBi could portray an unrealistic representation of models' length generalization performance. In fact, Chi et al. (2023) recently demonstrated that ALiBi's length generalization performance could be replicated using a window attention mask, where tokens beyond a window size  $w$  are masked out. Interestingly, we also observe that T5's Relative PE, which can be regarded as trainable version of ALiBi, learns to attend both large and short range dependencies (Figure F.3). This is in line with Tay et al. (2022) observation and underscores the importance of evaluation setups on downstream tasks as compared to solely relying on perplexity.
Length Generalization Failure In Transformers The length generalization problem has been a topic of interest in the study of neural sequence models for a long time (Graves et al., 2016; Kaiser and Sutskever, 2016; Lake and Baroni, 2018; Hupkes et al., 2020; Yehudai et al., 2021). Transformers, being state-of-the-art sequence models, have been no exception. A group of studies showed the generalization failure of conventional Transformers with APE on specific datasets such as PCFG (Hupkes et al., 2020), LEGO (Zhang et al., 2023), or CLUTRR (Sinha et al., 2019; Gontier et al., 2020). The length generalization problem has been reported even in pretrained Transformers such as T5 (Furrer et al., 2020) and LaMDA (Anil et al., 2022). Csordás et al. (2021) and Ontanon et al. (2022) study the effect of positional encoding on length generalization but mainly focus on showing relative PE outperforms APEs. Press et al. (2022), on the other hand, propose a new encoding method, ALiBi, and demonstrate that it outperforms popular PEs on extrapolation but only in the context of human language modeling. Most relevant is Deletang et al. (2023)'s recent study on length generalization in various neural sequence models (including RNNs and Stacked-RNNs) for tasks from Chomsky hierarchy. However, they do not focus on the difference among positional encoding or on autoregressive models. Unlike these studies, our work extensively compares length generalization in popular PEs for a wide range of tasks, specifically focusing on autoregressive models, which represent many contemporary LLMs.
Positional Encoding A core component of Transformers is the positional encoding mechanism, which helps the model represent the order of the input sequence. Self-attention mechanism in the encoder of Transformers is order-invariant and requires PE to avoid becoming a bag-of-word model. Many methods have been proposed for this purpose. Originally, Vaswani et al. (2017) introduced absolute positional encoding sinusoidal functions (a learned variant popularized by Devlin et al. (2019)). Relative approach for encoding positional information was further introduced by Shaw et al. (2018), which gave rise to a number of pre-trained LM with relative encodings such as TransformerXL (Dai et al., 2019) and T5 (Raffel et al., 2020) that perform well in length generalization. More recently, Su et al. (2021) takes the concept of sinusoidal functions and suggests a new way of encoding positional information by rotating the hidden representations before applying self-attention. This method, referred to as Rotary, has become a popular choice in the recent LLMs. Press et al. (2022) simplify the T5's Relative encoding and introduced a more efficient variant called ALiBi, while keeping the same or improving extrapolation performance. Decoder-only Transformers, due to their causal attention mask, are not order-agnostic and can operate without explicit positional information.
This was observed early on by Shen et al. (2018) and later explained by Tsai et al. (2019). The observation that Transformers without positional encoding can perform on par with explicit PE has been made in various domains such as machine translation (Yang et al., 2019), language modelling (Irie et al., 2019; Haviv et al., 2022), and even other domains like vision or speech (Likhomanenko et al., 2021). In our work, not only we demonstrate that Transformers can operate without explicit position information, but also we present an important setup where they outperform explicit PEs. Furthermore, we theoretically show how they are capable of learning both absolute and relative encodings.
We studied the robustness of different positional encodings, in decoder-only Transformers, at length generalization on various downstream mathematical and reasoning tasks. Our extensive empirical study shows the effectiveness of NoPE, and further demonstrates that widely used explicit PEs are not suited for length generalization. We also prove that NoPE can implicitly learn both absolute and relative positions, but uses the latter in practice. Finally, we find the effectiveness of scratchpad is task-dependent, and is not a reliable solution for length generalization.
Our work primarily focuses on positional encodings as a design choice in the Transformers decoder architecture. We could not study how large-scale pretraining affects different PEs because there are no publicly available large language models trained with various PEs under similar conditions. We leave this for future work due to our limited compute budget.
We are grateful to our anonymous reviewers, Nicolas Chapados, and Omer Levy for their invaluable suggestions and discussions The Mila-IBM grant program provided the funding for this project. SR acknowledges the support provided by the NSERC Discovery Grant program and the Facebook CIFAR AI Chair program. This research was enabled in part by compute resources provided by Mila, the Digital Research Alliance of Canada and ServiceNow.
Ekin Akyurek and Afra Feyza Akyurek. 2022. Notes on teaching gpt-3 adding numbers.
Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. 2023. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations.
Loubna Ben Allal, Raymond Li, Denis Kocetkov, Chenghao Mou, Christopher Akiki, Carlos Munoz Ferrandis, Niklas Muennighoff, Mayank Mishra, Alex Gu, Manan Dey, Logesh Kumar Umapathi, et al. 2023. Santacoder: don't reach for the stars! ArXiv, abs/2301.03988.
Cem Anil, Yuhuai Wu, Anders Johan Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Venkatesh Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. 2022. Exploring length generalization in large language models. In Advances in Neural Information Processing Systems.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer normalization. *ArXiv*, abs/1607.06450.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. 2023. Unveiling transformers with lego: a synthetic reasoning task. ArXiv, abs/2206.04301.
Figure A.1: Histogram of instruction lengths in two instruction finetuning datasets: FLAN (CoT subset) (Longpre et al., 2023) and Super Natural Instructions (Wang et al., 2022). The dotted line indicates the median length of the instructions in each dataset.
A Number of instances decreases rapidly as sequence length grows
The recent trend of SFT-RLHF pipeline (Ouyang et al., 2022) relies on finetuning LLMs on the instruction following tasks. However, the training data of these datasets is often skewed towards shorter sequences. Figure A.1 shows the distribution of instruction lengths in two instruction finetuning datasets: FLAN (CoT subset) (Longpre et al., 2023) and Super Natural Instructions (Wang et al., 2022). The median length of instructions in these datasets is quite short compared to the maximum length that exists. Such distribution shape highlights the importance of length generalization in these tasks. In fact, the models are supposed to learn from short instructions and generalize to ones during inference that might be much longer.
B Background
In this section, we lay the groundwork and introduce the notation we use throughout the paper. We will refer to this in Appendices C.1 and C.2.
Let  $f_{\theta}$  be a decoder-only Transformer model, where  $\theta$  denotes the full set of model parameters.  $f_{\theta}$  processes the input sequence  $\mathbf{x} = [x_0, x_1, \dots, x_T]$  and maps it to the output sequence  $\mathbf{y} = [y_0, y_1, \dots, y_T]$  by applying a sequence of Transformer layers. Note that being decoder-only means the attention mechanism in each layer is causal, i.e. the attention weights are computed based on the previous positions only.
The layer  $\mathrm{TLayer}^{(l)}(\pmb{H}^{(l-1)};\theta_l)$ , consisting of self-attention heads and a feed-forward sub-layer, reads the previous hidden state  $\pmb{H}^{(l-1)}$  and produces the hidden state at layer  $l$ :  $\pmb{H}^l$ , where  $l$  is the layer index, and  $\theta_l$  is the set of parameters of the  $l$ -th layer. Each hidden state  $\pmb{H}^{(l)} \in \mathbb{R}^{d \times (T+1)}$  is matrix where column  $t$ , denoted as  $\pmb{h}_t^{(l)}$ , is the hidden state at position  $t$ .
A layer  $l$  is parameterized by a set of parameters  $\theta_{l} = \{(\pmb{W}_{Q}^{m},\pmb{W}_{K}^{m},\pmb{W}_{V}^{m},\pmb{W}_{O}^{m})_{m},\pmb{W}_{1},\pmb{W}_{2}\}$ , where  $\pmb{W}_{Q}^{m},\pmb{W}_{K}^{m},\pmb{W}_{V}^{m}\in \mathbb{R}^{h\times d}$  and  $\pmb{W}_{O}^{m}\in \mathbb{R}^{d\times h}$  are the query, key, value, and output matrices of the  $m$ -th head, respectively.  $\pmb{W}_{1},\pmb{W}_{2}\in \mathbb{R}^{d\times k.d}$  are the weight matrices of the feed-forward sub-layer.  $d$  denotes the model's hidden state size,  $h$  is the attention dimension (where  $h = \frac{d}{\# \text{heads}}$ ), and  $k$  is a multiplier of the hidden state size in the feed-forward sub-layer (it is usually set to 4 in common implementations of the Transformer). Note that we drop the layer index  $l$  and the attention head index  $m$  where it is clear from the context.
The Transformer layer  $\mathrm{TLayer}^{(l)}$  processes each column of  $H^{(l-1)}$  independently and in parallel to produce the output. The computation of the  $t$ -th column of  $H^{(l)}$  is as follows:
$$
\boldsymbol {h} _ {t} ^ {(l)} = \operatorname {F F} \left(\lambda \left(\boldsymbol {a} _ {t} + \boldsymbol {h} _ {t} ^ {(l - 1)}\right)\right) + \boldsymbol {a} _ {t} + \boldsymbol {h} _ {t} ^ {(l - 1)} \tag {3}
$$
where FF is the feed-forward sub-layer,  $\lambda$  is layer normalization, and  $a_{t} \in \mathbb{R}^{d}$  is the output of the multi-head self-attention sub-layer at position  $t$ . Specifically,  $a_{t}$  is computed as:
$$
\boldsymbol {a} _ {t} = \sum_ {m} \operatorname {A t t n} ^ {(m)} \left(\boldsymbol {h} _ {t} ^ {(l - 1)}, \boldsymbol {H} ^ {(l - 1)}\right) \tag {4}
$$
where  $\mathrm{Attn}^{(m)}$  is the  $m$ -th attention head. Let  $\pmb{o}_t \in \mathbb{R}^d$  denote the output of an attention head at position  $t$ . Then,  $\pmb{o}_t$  is computed as:
$$
\boldsymbol {o} _ {t} = \boldsymbol {W} _ {O} \left(\sum_ {i \leq t} \hat {\alpha} _ {i} \boldsymbol {v} _ {i}\right) \tag {5}
$$
where  $\hat{\alpha} = \mathrm{softmax}(\alpha) \in \mathbb{R}^{(t + 1)}$ , and  $\alpha$  is the attention weight vector such that:
$$
\boldsymbol {\alpha} = \left[ \langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {0} \rangle , \langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {1} \rangle , \dots , \langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {t} \rangle \right] ^ {\intercal} \tag {6}
$$
where  $\pmb{q}_t = \pmb{W}_Q\pmb{h}_t^{(l-1)} \in \mathbb{R}^h$ ,  $\pmb{k}_i = \pmb{W}_K\pmb{h}_i^{(l-1)} \in \mathbb{R}^h$ , and  $\pmb{v}_i = \pmb{W}_V\pmb{h}_i^{(l-1)} \in \mathbb{R}^h$ .  $\langle \cdot, \cdot \rangle$  denotes the dot product operation.
The feed-forward sub-layer  $\mathrm{FF}(\cdot)\in \mathbb{R}^d$  is a two-layer MLP:
$$
\operatorname {F F} (\boldsymbol {x}) = \boldsymbol {W} _ {2} \sigma \left(\boldsymbol {W} _ {1} ^ {\mathsf {T}} \boldsymbol {x}\right) \tag {7}
$$
where  $\sigma$  is a non-linear activation function (usually ReLU or GeLU (Hendrycks and Gimpel, 2020)). Additionally,  $\lambda (\cdot)\in \mathbb{R}^d$  is layer normalization (Ba et al., 2016). Note that we take the additive (Elhage et al., 2021) view of attention heads in Equation (4) instead of concatenate and multiple view (Vaswani et al., 2017) as it is easier to understand and analyze. But, they are mathematically equivalent (Elhage et al., 2021).
The hidden state is initialized with a learned embedding of the input sequence  $\pmb{H}^{(0)} = \pmb{W}_E\pmb{X}$ , where  $\pmb{W}_E \in \mathbb{R}^{d \times V}$  is the embedding matrix and  $\pmb{X} \in \mathbb{R}^{V \times (T + 1)}$  is the one-hot encoded input sequence.  $V$  is the vocabulary size.
B.2 Positional Encoding
Almost all positional encoding methods can be explained and formulated as how they implement the dot product operation in Equation (6). So, in this section, we explain how the dot product  $\langle \pmb{q}_t,\pmb{k}_i\rangle$  is implemented in different positional encoding schemes.
Absolute Positional Encoding (APE) The process of Absolute Positional Encoding (APE) involves assigning a position vector  $\pmb{p}_i$  to each absolute position  $i$  and combining them with word embeddings before inputting them into the model. So, APE first modifies how the hidden state is initialized:
$$
\boldsymbol {H} ^ {(0)} = \boldsymbol {W} _ {E} \boldsymbol {X} + \boldsymbol {W} _ {P} \boldsymbol {P} \tag {8}
$$
where  $\mathbf{W}_P \in \mathbb{R}^{d \times T}$  is the positional embedding matrix and  $\mathbf{P} \in \mathbb{R}^{V_p \times (T + 1)}$  is the one-hot encoded absolute position sequence.  $V_{p}$  is the maximum absolute position. Therefore, the hidden state at column  $j$  is:
$$
\boldsymbol {h} _ {j} ^ {(0)} = \boldsymbol {e} _ {j} + \boldsymbol {p} _ {j} \tag {9}
$$
where  $e_j \in \mathbb{R}^d$  is the word embedding of token  $x_j$  and  $p_j \in \mathbb{R}^d$  is the positional embedding for position  $j$ . Then, the dot product for the first layer in Equation (6) is computed as:
$$
\begin{array}{l} \langle \boldsymbol {q} _ {t}, \boldsymbol {k} _ {i} \rangle = \langle \boldsymbol {W} _ {Q} \boldsymbol {h} _ {t} ^ {(0)}, \boldsymbol {W} _ {K} \boldsymbol {h} _ {i} ^ {(0)} \rangle \\ = \left\langle \boldsymbol {W} _ {Q} \left(\boldsymbol {e} _ {t} + \boldsymbol {p} _ {t}\right), \boldsymbol {W} _ {K} \left(\boldsymbol {e} _ {i} + \boldsymbol {p} _ {i}\right) \right\rangle \\ = \left(\boldsymbol {W} _ {Q} \left(\boldsymbol {e} _ {t} + \boldsymbol {p} _ {t}\right)\right) ^ {\intercal} \left(\boldsymbol {W} _ {K} \left(\boldsymbol {e} _ {i} + \boldsymbol {p} _ {i}\right)\right) \\ = e _ {t} ^ {\intercal} \boldsymbol {W} _ {Q} ^ {\intercal} \boldsymbol {W} _ {K} \boldsymbol {e} _ {i} + e _ {t} ^ {\intercal} \boldsymbol {W} _ {Q} ^ {\intercal} \boldsymbol {W} _ {K} \boldsymbol {p} _ {i} \\ + \boldsymbol {p} _ {t} ^ {\intercal} \boldsymbol {W} _ {Q} ^ {\intercal} \boldsymbol {W} _ {K} \boldsymbol {e} _ {i} + \boldsymbol {p} _ {t} ^ {\intercal} \boldsymbol {W} _ {Q} ^ {\intercal} \boldsymbol {W} _ {K} \boldsymbol {p} _ {i} \tag {10} \\ \end{array}
$$
In the learned variant of APE,  $\pmb{p}_j \in \mathbb{R}^d$  is learned during training. In the sinusoidal variant,  $\pmb{p}_j$  is calculated using a non-parametric function. Specifically,  $\pmb{p}_j$  is computed as:
$$
\boldsymbol {p} _ {j} = \left[ \sin (\omega_ {1}. j), \cos (\omega_ {1}. j), \sin (\omega_ {2}. j), \cos (\omega_ {2}. j), \dots , \sin (\omega_ {d / 2}. j), \cos (\omega_ {d / 2}. j) \right] ^ {\intercal} \tag {11}
$$
where  $\omega_{i} = \frac{1}{10000^{2i / d}}$
T5's Relative PE The Relative bias in T5 is a type of relative positional encoding that initially calculates the relative distance  $(t - i)$  between tokens at positions  $t$  and  $i$ . This distance is then transformed into a scalar bias value  $b$  and is incorporated into the dot product between the query and key.  $b$  is learned during training. Thus, 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} + b _ {\text {b u c k e t} (n - m)} \tag {12}
$$
$$
\operatorname {b u c k e t} (n) = \left\{ \begin{array}{l l} n & \text {i f} n <   \frac {\mathcal {B}}{2} \\ \frac {\mathcal {B}}{2} + \left\lfloor \frac {\log \left(\frac {n}{\mathcal {B} / 2}\right)}{\log \left(\frac {\mathcal {D}}{\mathcal {B} / 2}\right)} \times \frac {\mathcal {B}}{2} \right\rfloor & \text {i f} \frac {\mathcal {B}}{2} \leq n <   \mathcal {D} \\ \mathcal {B} - 1 & \text {i f} n \geq \mathcal {D} \end{array} \right.
$$
This function maps the relative distance  $d$  to a bucket index, which will be used to look up the weight corresponding to that bucket.  $\mathcal{B}$  is the number of buckets,  $\mathcal{D}$  is the maximum distance. It assigns half of the buckets to distances smaller than  $\frac{\mathcal{D}}{2}$  with linear spacing and the other half to distances larger than  $\frac{\mathcal{D}}{2}$  with logarithmic spacing. The weight for distances larger than  $\mathcal{D}$  is the same. This is to facilitate generalization to unseen distances. In the original implementation of T5,  $\mathcal{B} = 32$  and  $\mathcal{D} = 128$ . Following shows an example of the bucket function with  $\mathcal{B} = 5$  and  $\mathcal{D} = 6$ :