Functional Interpolation for Relative Positions improves Long Context Transformers
Authors: Shanda Li
Chong You
Guru Guruganesh
Joshua Ainslie
Santiago Ontanon
Manzil Zaheer
Sumit Sanghai
Yiming Yang
Sanjiv Kumar
Srinadh Bhojanapalli
Venue: ICLR 2024 poster (2024)
Keywords: Transformers
positional encoding
long context
length generalization
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We propose a novel functional relative position encoding with progressive interpolation to improve Transformer generalization to longer contexts.

Abstract: Preventing the performance decay of Transformers on inputs longer than those used for training has been an important challenge in extending the context length of these models. Though the Transformer architecture has fundamentally no limits on the input sequence lengths it can process, the choice of position encoding used during training can limit the performance of these models on longer inputs. We propose a novel functional relative position encoding with progressive interpolation, FIRE, to improve Transformer generalization to longer contexts. We theoretically prove that this can represent some of the popular relative position encodings, such as T5's RPE, Alibi, and Kerple. We next empirically show that FIRE models have better generalization to longer contexts on both zero-shot language modeling and long text benchmarks.

Published as a conference paper at ICLR 2024
FUNCTIONAL INTERPOLATION FOR RELATIVE POSI- TIONS IMPROVES LONG CONTEXT TRANSFORMERS
Shanda Li'; Chong You2, Guru Guruganesh? Joshua Ainslie?, Santiago Ontanon² Manzil Zaheer", Sumit Sanghai?, Yiming Yang', Sanjiv Kumar", Srinadh Bhojanapalli 'Carnegie Mellon University GGoogle Research shandal @cs. cmu. edu
'Google DeepMind
ABSTRACT
Preventing the performance decay of Transformers on inputs longer than those used for training has been an important challenge in extending the context length of these models. Though the Transformer architecture has fundamentally no limits on the input sequence lengths it can process, the choice of position encoding used during training can limit the performance of these models on longer inputs. We propose novel functional relative position encoding with progressive interpolation, FIRE, to improve Transformer generalization to longer contexts. We theoretically prove that this can represent some of the popular relative position encodings, such as T5's RPE, Alibi, and Kerple. We next empirically show that FIRE models have better generalization longer contexts on both zero-shot language modeling and long text benchmarks.
INTRODUCTION
Transformer based Language Models have demonstrated state-of-the-art zero-shot performance on many natural language processing tasks (Brown et al., 2020), enabling increasingly longer context applications such chat bots (Roller al., 2021; Zhang al., 2020b) and long document summarization and question answering (Zhang et al., 2020a; Guo et al., 2022; Ainslie et al., 2023). However, the accuracy of these models usually drops quickly for inputs longer than the during training (Press et al., 2022: Anil et al., 2022; Deletang et al., 2023) - which are usually relatively short (e.g. 2048 for LLaMM (Touvron et al., 2023a:b)) avoid the expensive quadratic attention cost during training. This has led to a significant interest in improving length generalization of Transformers where we train the model using shorter inputs (e.g. 2048) and test the models performance on longer inputs (e.g. 8192) (Press et al., 2022; Anil et al., 2022; Chi et al., 2022; 2023; Chowdhury & Caragea, 2023; Chen et al., 2023).
used
Transformers are fundamentally permutation equivariant, and are agnostic to input sequence order- ing (Vaswani al., 2017: Yun et al., 2 They rely on position encodings to learn the ordering of input tokens. Popular position encodings such as Absolute Positional Encoding (APE) (Vaswani et al., 2017) and more recent Rotary Positional Encoding (RoPE) (Su et al., 2021) do not generalize to longer contexts than seen during training (Kazemnejad al., 2023). T5's relative positional encoding (Raffel et al., 2019) generalizes to longer contexts by using the same representation for all out of distribution (OOD) sequence lengths, but suffers from slow vector operations modern accelera- tors (Press et al., 2022). Another line of recent work promotes length generalization by encoding specific inductive biases how attention should decay with sequence length (Press et al., 2022; Chi et al., 2022; 2023). More recently, Kazemnejad et al. (2023) show that having no position encodings in decoder-only models can have better length generalization, albeit for small-scale synthetic tasks.
In this work we take a functional approach to learn the relative position biases?, instead of having hard coded inductive biases, towards training language models with length generalization (focusing on decoder-only models). We propose FIRE (Functional Interpolation for Relative Positional Encoding) method that i) uses a learnable function to map the input positions biases, and ii) uses progressive
WWrk done during internship at Google Research.
Note that decoder-only models can infer position from the causal attention mask (Haviv et al., 2022). We consider relative position encodings for their superior performance over absolute position encodings (Raffel et al., 2019; Chen al., 2021).

---

Published as a conference paper at ICLR 2024
interpolation technique, which ensures bounded input for the position encoding function for all input
sequence lengths, thereby enabling length generalization.
C4 language modeling (base model)
A functional approach to learn the biases allows the model to adapt to the given task instead of always having the same inductive bias, e.g. bias towards nearby tokens in (Press et al., 2022; Chi et al., 2022; 2023). In particular we use an MLP to learn these biases, which we theoretically prove can represent several popular methods such as T5's RPE, Alibi, and Kerple in parameter efficient manner. In fact, all our experiments use tiny MLP with a hidden size of 32, which is also pacelerator-friendle unlike T5's RPE. Next, our progres- sive interpolation technique normalizes the query-key relative distance by the query position. Since for causal attention in language models the relative distance is always between O and the query position, progressive interpolation results in an output that is always bounded between [0, 1]. This results in a bounded input to the position encoding function for all input sequence lengths, leading to better generalization performance. As result, with increasingly longer sequence lengths, the positional inputs will form progressively finer grids, interpolating the positional encoding function on [o, 1]-
Validation 18 perplex
3.5
RoPE
NoPE
Alibi
Karple
T5's
3.3
YeRN
3.1
Validation sequence length (x10')
Figure I: Language modeling per- plexity on C4 with varying evalua- tion sequence lengths. Models are trained on length 2048.
Inspired by the existing methods, we incorporate the following two transformations into FIRE, which we find helpful to improve the model quality. i) To encourage locality bias in FIRE, we apply the popular log transformation (Raffel et al., 2019; Chi et al., 2022) the relative distance before feeding it to the MLP, which amplifies the input differences for local tokens. ii) Next we modify progressive interpolation with learnable threshold in the normalizer to yield exact distances for shorter contexts. Note that both these transformations do not limit the ability of the model to learn arbitrary biases. In fact we show that FIRE learns to pay more attention far away contexts in some attention heads.
We conduct an extensive empirical study to demonstrate the effectiveness of FIRE for length general- ization. We benchmark FIRE as well as other positional encoding approaches on a wide range of real-world language modeling (C4, arXiv, and Github), long text benchmark (SCROLLS), zero-shot long-context question answering (NarrativeQA). and natural language understanding benchmarks (GLUE/SuperGLUUE) Our empirical results show the strong length generalization performance and long text modeling capability of FIRE. Our experiments on standard natural language understanding benchmarks show that FIRE is competitive on short sequence tasks well. We further visualize the learned positional encoding of FIRE showing that it learns diverse patterns, beyond just locality bias.
The main contributions of our paper are summarized below:
We propose FIRE. a new functional relative positional encoding method. Using progres- sive interpolation, FIRE is able to transform arbitrary input lengths into bounded domain, followed by a learned mapping.
We theoretically prove that FIRE can represent popular position encodings such as T5's RPE, Alibi, and Kerple, thereby unifying a class of existing position encoding approaches.
We empirically show strong length generalization behavior of FIRE, significantly improving over existing methods in zero-shot and finetuning settings benchmarks. For instance, consistently delivers strongest performance on C4 language modeling across various sequence lengths, outperforming the best baseline by 2.28 perplexity points (Fig. 1). On SCROLLS long text benchmark, FIRE surpasses all the competing methods on average by over point (Table 1).
wide range of datasets and
We present visualization of learned position embeddings of FIRE model showing that it can learn both local and anti-local position biases.
2 POSITIONAL ENCODINGS AND LENGTH GENERALIZATION
We are interested in building Transformer models with length generalization ability, i.e., we expect that the model can be trained on sequences of length Ltrain and be directly applied to sequence length Lusst without performance degradation for Ltest > Lirain (Press et al., 2022). Length generalization requires Transformers generalize unseen positions during training, and designing better position
2

---

Published as a conference paper at ICLR 2024
encodings is active line of research towards improving the length generalization (Chi et al., 2022; 2023; Kazemnejad et al., 2023; Chen et al., 2023). this section, we review existing positional encoding approaches with emphasis on their length generalization abilities. More discussions related work can be found in Appendix D.
2.1 ABSOLUTE POSITIONAL ENCODING
The Transformer paper (Vaswani et al., 2017) proposes Absolute Positional Encoding (APE) to endow Transformers with positional information. In particular, (learnable or fixed sinusoidal) real-valued embedding e; E IR is assigned to each position i, leading to an Absolute Positional Encoding matrix E [i, .... en en]", which will be added the input sequence. Though simple and straightforward, APE- based Transformers usually generalize poorly to longer sequences (Press et al., 2022).
2.2 RELATIVE POSITIONAL ENCODING
Relative Positional Encoding (RPE) is an increasingly popular way to encode positional information for Transformers. Shaw al. (2018) are the first to introduce RPE Transformers and their proposed method adds position encodings to the key (and optionally the value) in the attention layer, instead of the input. Raffel al. (2019) simplify the vector representations of relative positions to scalars and use them as a bias term added to the pre-softmax attention logits. They further map any OOD sequence lengths to the same position, resulting in length generalization. This form of additive RPE has proven to be highly effective in many applications (Dai et al., 2019; Liu et al., 2021; Ying al., 2021). Following this, multiple additive RPE methods have been proposed to improve both length generalization and efficiency, such Alibi (Press et al., 2022), Kerple (Chi et al., 2022), and Sandwich (Chi et al., 2023).
For most of these additive RPE methods, the computation of the (pre-softmax) attention logits can be unified using the following formula:
Additive RPE.
ARPE X) = XWo(XWx)) + B,
where the bias matrix B E R=x* is induced by the position encoding function :: N + R. Let the j-th entry of be ji. Different formulations and parameterizations of lead different RPE variants. A few examples that support arbitary sequence length include:
T5's RPE (Raffel al., 2019): b(i. (i-n(i-- where K is ayper-parameter and (Ty {r are learnable scalars."
Alibi (Press et al., 2022): b(i,j) -rlė jز where r >> 0 is a cypeer-arameter.
Kerple (Chi et al., 2022): b(i.j) = -r1 log( (1 1 ji) (logarithmic variant) or -rili-j (power variant), where r1, 12 > are learnable scalars.
Sandwich (Chi et al., 2023): b(i.j) rı Exe cos ( (i- (11000/0000 where rı and are hyper-parameters.
The above methods can be applied to longer sequences than training, but they also have several limitations. T5's RPE uses the same attention bias for all query-key pairs with distance greater than K, lacking representational power to distinguish between different positions in long sequences. Furthermore, it relies on vector operations that are not serelerraoo-freenllyy making its training and inference relatively slow (Press et al., 2022). Alibi, Kerple, and Sandwich significantly bias towards local attention, making harder to attend to more distant query-key pairs (Chi et al., 2023). This property can prevent the model from capturing long-range dependencies and lead to performance degradation on some tasks. In the subsequent section, we will present our method to overcome these limitations.
Rotary Positional Encoding. In addition to the aforementioned methods, there are also several non-additive RPE variants. Among them, the most popular one large language models is Rotary Position Encoding (RoPE) (Su et al., 2021; Chowdhery et al., 2022; Touvron et al., 2023a). RoPE
In practice, T5's RPE segments relative distances into distinct buckets with logarithmic scale, each associated with unique parameter. Refer to Appendix A.1 for further details.
3

---

Published as a conference paper at ICLR 2024
rotates the query and key vectors with angle proportional to their absolute positions before the dot product attention, which results in attention being a function of the relative distance between the tokens, capturing the relative positional information.
Press et al. (2022); Kazemnejad et al. (2023) find that RoPE-based language models have poor length generalization. To address this, Chen al. (2023) propose RoPE with position interpolation, and show this allows better length generalization of these models. Such interpolation techniques ((Chen al.. 2023) for RoPE and (Dosovitskiy et al., 2021) for APB), usually requires I) knowing the target sequence length priori, which may not be feasible in practical generative applications, 2) finetuning the model at the new target sequence length, which can be challenging for larger scale models. In contrast, our proposed approach uses progressive interpolation technique that does not require any prior information of the target sequence length. This property is appealing since the maximum sequence length can be hard to predict for auto-regressive language models. Further, our experiments show that the proposed approach does not require any additional finetuning achieve strong zero-shot length generalization behavior.
2.3 No POSITIONAL ENCODING
While encoder-only Transformer models (c.g.. BERT (Devlin et al., 2019)) are permutation equivariant without positional encoding, Haviv al. (2022) show that decoder-only Transformers with causal attention masks can learn positional information even without any explicit positional encoding. Recently, Kazemnejad et al. (2023) show that the no positional encoding (NoPE) model shows strong length generalization on small scale synthetic tasks.
METHOD
In this section, we formally introduce FIRE (Functional Interpolation for Relative Positional Encoding), a new relative positional encoding approach for improving length generalization of Transformers.
3.1 FUNCTIONAL. POSITION ENCODING WITH PROGRESSIVE INTERPOLATION
Our proposed approach FIRE uses learnable continuous function to map input positions to biases. We implement the function using an MLP fo R + R, where denotes the MLP parameters. This avoids hard coding specific inductive biases and lets the position encoding be learnt jointly with the task at hand. standard approach would be to feed the relative query-key distance as the input to the MLP. However this suffers from generalization issues when the inputs (the relative distances) outside the training domain of the MLP.
We propose Progressive Interpolation address this challenge. Instead of using the query-key relative distance as the input to the MLP, we normalize the distance by the query position index. Formally, we consider the following positional encoding function:
b(i.j) = fo (" where fo (x) = vg (V2o(v)))) 0 (U1, V2, v3).
Here is the ReLU activation function; and denote the query and key positions respectively. Note that in causal attention, the relative distance satisfies < i-jj j < i. Therefore, the normalized relative distance constrained to be [0, 1] regardless of the sequence length. In particular, with increasingly longer sequence lengths, the positional inputs will form progressively finer grids, interpolating the positional encoding function on [0, 1]. Hence, this technique aligns inference domain with training domain for any sequence lengths, leading to better length generalization.
Discussion on the choice the normalizer. FIRE uses the query position i to normalize the relative distance and implement interpolation. For auto-regressive generation with causal attention,
4 'erre we focus on a single attention head. Generally, with H heads, FIRE learns an ML.P fo -+ R uses different attention biases for different heads. U1 ny E R². V2 R*x* and s denotes the hidden size. Bias terms omitted for brevity.
4

---

Published as a conference paper at ICLR 2024
the query position index corresponds to the length of current context. Another possible choice is to use some pre-defined max context length as the normalizer. this case, the model will still suffer from unfamiliar (large) distances when the texts exceed the pre-defined max lengths, making such choice suboptimal. Using the query position index as the normalizer avoids this issue.
3.2 ADDITIONAL TRANSFORMATIONS
Inspired by existing methods, introduce two transformations FIRE for further improvement. We note that these transformations do no limit the expressive power of FIRE to learn arbitrary biases.
Amplifying the differences among local positions. Existing works show that RPE attention biases change more rapidly for the local tokens than for the distant tokens (Khandelwal et al., 2018; Wang et al., 2021). Thus, it's appealing consider some monotonically increasing transformation /b : + R. with a monotonically decreasing slope (i.e., a concave function) to the relative distance, so that more modeling capacity can be allocated learn RPE for local positions:
(i (
b(i,j) fa
(b(1)
For example, in our experiments, we use v0 : - log(cr + 1) where c > 0 is a learnable parameter. This transformation & amplifies the differences among local positions. Note that, the log transformation is applied to both the relative distance and the normalizer. Thus, the MLP inputs are still constrained to [0, 1] for any sequence lengths as long as monotonically increasing.
Thresholding the normalizer for better short sequence modeling, While the progressive in- terpolation technique offers robust length generalization capabilities, our preliminary experiments indicate marginal degradation in model performance for shorter sequences. We posit that it's because the actual relative distances are important in RPE of short sequences, while the normalization in progressive interpolation obfuscates this information. To address this, we introduce an adaptive thresholding mechanism, activating the progressive interpolation technique only for larger query position indices, i.e., long contexts. Specifically, we define learnable threshold L and only apply progressive interpolation when L. For short sequences with less than tokens, we use (L) to normalize the relative distance.
Based on the above, the positional encoding function of FIRE can formulated as
wi - j) (maxx (L,
bFIRE(i.j)
(4)
where v/r N + R4 monotonically increasing and L > a learnable scalar. Our main
experiments of FIRE are based Eq. (4) with v + log(ez + 1). We present experiments ablating these design choices in Appendix B.
3.3 EXPRESSIVENESS OF FIRE
In this subsection, we theoretically prove that FIRE can represent all the existing additive RPE approaches discussed in Sec. 2.2. This expressiveness allows FIRE to learn suitable position encoding functions from the data. We state this formally in the theorem below. The proof can be found in Appendix A.
Theorem 3.1. Let by be the positional encoding function of T5's RPE, Alibi, Kerple, Sandwich as defined in Sec. 2.2. Consider FIRE function bi Eq. (4). Given any sequence length Lo E N., there exist some transformation v. threshold L, and MLP configuration (weights 9 and activation function
σ) functtio such that bFIRE(i.)
(o(i)) for any < S Lo.
We point out that our proof is constructive, and does not leverage the universal approxi- mation property of MLP, i.e., the MLP does not need to be extremely wide or deep. In fact, FIRE parameter efficient in the sense that it represents T5's RPE, Alibi, and Kerple with nearly the same number parameters (up to a constant factor). Further, in all our experiments with FIRE, we show that small MI. with hidden size of 32 suffices for strong performances.
Remark.
5

---

Published as a conference paper at ICLR 2024
C4 lanquage modeling (large model)
arXiv lanquage modeling Darge modell
Github language modeling (large model)
Validation log perplesty
RePE
NoFE
Valldation log perplextt
4.0-
RoFE
NoPE
3.2
Albi
Albi
Alei
Kerple
Kerple
Kerpfe
3.1
T5's RPE
TE%s RPE
15% RPE
3.0
FIRE
Volidation log
FIRE (ours)
3.0
FIRE tours)
2.9
2.5
Validation sequence length (x110))
Validation sequence longth (x110')
Validation sequence longth (x300)
Figure 2: Language modeling perplexity with varying evaluation sequence lengths for large models trained on sequence length 2048.
4 EXPERIMENTS
In this section we present experimental results comparing our proposed unified relative encoding method FIRE with T5's RPE (Raffel et al., 2019), Alibi (Press et al., 2022), and Kerple (Chi et al., 2022), showing that the proposed approach significantly improves long context generalization while not sacrificing short context performance. We also include comparisons to other popular methods Rotary Positional Encoding (RoPE) (Su et al., 202.) and no positional encoding (NoPE) (Kazemnejad et al., 2023). We use hidden size of 32 for the MLPs in FIRE for all our experiments.
We consider language models trained on the C4 dataset (Raffel et al., 2019) with 2048 input length, with different positional encoding methods. We first compare the zero-shot perplexity values inputs with different lengths (512 to 8192) from various datasets, comparing the long context generalization ability of different position encoding methods (Sec. 4.1). Later, we present finetuning results on both longer inputs of length 8192 on SCROLLS (Shaham et al., 2022) and shorter inputs of length 1024 on GLUE/SuperGLUE (Wang et al., 2019b;a) (Sec. 4.2 & 4.4). In addition, we conduct experiments on zero-shot long contect question answering on NarrativeQA (Kočiský et al., 2018) with different context lengths from 512 to 32768 (Sec. 4.3). Appendix B, we present some ablation experiments studying the design choices of FIRE. The complete experimental setup along with the hyper-parameters for each of the tasks and hardware details is provided in Appendix C.
4.1 LANGUAGE MODELING WITH LENGTH GENERALIZZATIO
Following Brown et al. (2020), use the causal LM objective to pretrain decoder-only Transformers with different position encodings on C4 dataset (Raffel et al., 2019). We experiment with two model size settings, base (125M parameters) and large (350M parameters). The evaluation metrics validation log perplexity on C4, arXiv, and Github (Raffel et al., 2019; Gao et al., 2020). We pretrain the models on sequence length to 2048, and evaluate their zero-shot perplexity on sequence lengths (512, 1024, 2048, 4096, 8192]. For base- base-sized models, we additionally compare our method with a concurrent work, YaRN (Peng et al.. 2024), which improves length generalization of RoPE-based Transformer models." Model and training configurations are detailed in Appendix C.I.
The results are shown in Fig. 1, 2, & 7. We first notice that FIRE consistently achieves lower perplexity across different model sizes, validation sequence lengths, and datasets. In comparison to existing approaches, the performance gain is particularly significant for validation sequences that are longer than training sequences nuut-ofddistribution sequence lengths), showing better length generalization behavior. For example, for base models trained on sequence length 2048 and evaluated on sequence length 8192, FIRE outperforms the best baseline method, Kerple, by 2.28 points (21.24 v.s. 23.52 perplexity). Methods such as RoPE achieve strong performance for in-distribution sequence lengths, but their performances quickly degrade with longer inputs. YaRN requires knowledge of target sequence length and further finetuning, but we can see from Fig. & that it underperforms FIRE on long sequences and sacrifices model quality on short sequences (c.g., length 512). Note that in all our experiments, perplexity is computed in a single forward pass for given input, and we do not use any sliding window tricks during inference (Press al., 2022).
'Wiile finetuning not the same as the zero-shot long-contex generalization. it still measures ability of the pre-trained model to adapt to longer inputs in the downstream applications. We note that YaRN needs additional tuning on long sequences. All the other methods in this subsection, including FIRF, are evaluated long context without any tuning.

---

Published as a conference paper at ICLR 2024
Table 1: Experimental results on SCROLLS benchmark. Abbreviations for dataset names:
Qasper (Qas), ContractNLI (CNLI), QMSum (QMS), NarrativeQA (NQA), SummScreenFD (SumS), GovReport (GovR), and QuALITY (QuAL). We provide the evaluation metrics, the median sequence lengths in each dataset (Ainslie et al., 2023), and detailed results for base/large models. RoPE-PI refers to the RoPE interpolation (Chen et al., 2023). Best results are highlighted in bold.
QAS
CNLI
QMS
NQA
SumS
GovR
QuAL
Average
Metric
F1
EM
Rgm
F1
Rgm 9046
Rgm 8841
EM
Median length
5472
2148
14197
57829
7171
Base models
NoPE
10.98
72.90
14.36
5.90
15.44
16.24
22.10
22.56
RoPE
10.44
71.75
14.90
8.71
14.40
15.72
6.71
20.38
RoPE-PI
15.41
71.94
13.12
9.21
15.77
16.86
20.33
23.23
Alibi
8.38
67.21
5.48
4.24
3,49
6.96
9.68
15.06
Kerple
11.67
75.99
14.39
9.24
15.73
16.42
25.36
24.11
T5's RPE
12.80
74.93
16.12
9.00
15.37
15.96
24.83
24.14
FIRE (ours)
16.24
82.93
14.58
9.55
15.87
16.31
24.02
25.64
Large models
NoPE
15.34
74.25
15.79
7.56
16.60
16.66
24.16
24.34
RoPE
11.01
79.94
15.13
9.40
15.84
15.50
9.92
22.39
RoPE-PI
17.02
84.28
14.05
10.14
16.72
17.03
23.01
26.04
Alibi
8.20
68.95
5.81
4.91
4.34
11.58
12.27
16.58
Kerple
18.93
77.24
15.09
9.97
17.14
16.85
24.83
25.72
T5's RPE
17.51
75.70
16.17
9.62
16.68
16.76
24.45
25.27
FIRE (ours)
19.47
85.15
15.10
10.27
17.27
16.83
25.26
27.05
4.2 FINETUNING ON LONG TEXT BENCHMARK
To further test the models' capability of learning and modeling long sequences, we conduct finetuning experiments SCROLLS, long text benchmark (Shaham al., 2022) which contains different datasets. We initialize the models with the C4 checkpoints pretrained on sequence length 2048, and finetune them on sequence length 8192 for each individual task. In addition position encoding methods in Sec. 4.1, we also experiment with RoPE with positional interpolation (RoPE-PI) (Chen al., 2023), which extends the context window of RoPE-based pretrained models given downstream maximum sequence length. Following existing works by Shaham et al. (2022); Ainslie et al. (2023), we use three different evaluation metrics (Rgm, F1, and EM scores) for different datasets. We also compute the average score across different datasets as done in the SCROLLS benchmark. Detailed descriptions of the datasets and evaluation metrics are provided in Appendix C.2.
The results on SCROLLS benchmark are shown in Table 1. We first notice that FIRE attains the best average score, outperforming existing approaches by over 1.0 point on both model sizes. Even at the individual task level, FIRE achieves the best performances models. RoPE-PI significantly improves RoPE as expected, but lags behind FIRE. One drawback though is that RoPE-PI requires the knowledge of maximum input sequence length beforehand, which not always known in practice for decoder-only models.
4/5 out of 7 tasks among the base/large
4.3 ZERO-SHOT LENGTH GENERALIZATION ON NARRATIVEQA
We next evaluate the zero-shot length generalization capabilities of the finetuned models on the downstream NarrativeQA dataset. We use the NarrativeQA dataset (Kočiský et al., 2018) with different input context lengths to test the model's ability to leverage long context zero-shot learning settings. We use the base-sized model checkpoints pretrained on C4 (sequence length 2048) and finetuned on NarrativeOA (sequence length 8192). We evaluate the models on context lengths (512, 2048, 4096, 8192, 16384, 24576, 32768) without any further tuming the target con- text lengths. For RoPE with position interpolation (Chen et al., 2023), we consider two variants with max sequence lengths set to 8192 32768. We unigram overlap (F1) the evaluation metric.
7

---

Published as a conference paper at ICLR 2024
Ferformance on NarrativeQA
GLUE/SuperGLUEEbase)
GLUE/SuperGLWE
Average accura 6
76
Average accuracy
76
RoPE-Plaizz
72
score
RoPE-PPaz7ee
Kerple
T5's
64
FIRE (ours)
60
56
56
512
2048
4096
8192
16384
24576
32768
NoPE
RoPE
RoPE PL
Albbi
Kerple.
T5's RPE
FIRE
NoPE
ROPE
RoPE- -L.
Alibi
Kerple
T5's RPE
FIRE
Context length
Figure 3: Left: Comparisons on NarrativeQA with different context lengths. "RoPE-Plszey and
"RoPE-PIs27es" refers RoPE interpolation with max sequence lengths 8192 and 32768 respectively.
Right: Results on GLUE and SuperGLUE benchmarks. We report the average accuracy across all the tasks on these two benchmarks.
Layer 2
Layer 8
Layer 12
FIRE positional bias
1st head
2nd head
4th head
-2
-5
6th head
7th head
-4
-1
9th head
10th head
30
60
90
120
30
60
90
120
30
60
120
12th head
Relative distance
Relative distance
Relative distance
Figure 4: Visualization of FIRE learned position biases for the 128th query position with key
positions between 1 and 128. We notice that FIRE learns both local and anti-local position patterns.
We compare FIRE with the most competitive baselines in the left panel of Fig. Detailed results (including omitted baselines) can be found in Table 1. We notice that FIRE achieves top perfor- mances consistently across different sequence lengths. The plot also shows sensitivity of RoPE-PI to the max sequence length parameter in this zero-shot length generalization setting. Setting the max sequence length to small value (8192) results in good performance until 8192, but with a steep drop for longer contexts. On the other hand, using a larger value for max sequence length (32768) gets rid of the steep drop for long contexts; but results worse performance across all sequence lengths. In contrast, FIRE using progressive interpolation is able to generalize across all sequence lengths.
4.4 FINETUNING GLUE/SUPERGLUE
We next evaluate the C4 pre-trained models on GLUE/SuperGLUE benchmarks, to test these methods on shorter sequence lengths. finetune on standard natural language understanding benchmarks, GLUE (Wang et al., 2019b) and SuperGLUE (Wang al., 2019a), with shorter sequence lengths (1024) to evaluate the general quality of the models. use the average accuracy/exact match across all the tasks
main evaluation metric. Detailed experiment results can be found in Table 12.
The results are shown in the right of Fig. 3. Among the baseline approaches, NoPE and Alibi slightly lag behind, while RoPE, Kerple, and T5's RPE all achieve similarly good accuracies. FIRE is on par with these approaches, demonstrating good performance on GLUE and SuperGLUUE tasks. These results show that although FIRE is designed to enhance length generalization of Transformers, it does not sacrifice the accuracy on downstream tasks with shorter sequence lengths.
4.5 VISUALIZATION OF FIRE
In this subsection we present visualization of learned position encoding biases from a FIRE model pretrained on C4. We plot the learned position encoding bias for the query token at the 128th position, for all the attention heads from selected layers in Fig. 4. We notice that, in different attention heads, FIRE learns both local and "anti-locall" attention patterns that emphasize far away keys more, showing the advantage of functional approach, as opposed to a fixed local inductive bias (Press et al., 2022; Chi et al., 2022; 2023).
8

---

Published as a conference paper at ICLR 2024
Table 2:
Comparing FIRE with/without positional encoding function sharing across layers. FIRE and FIRE-S refer to models without and with sharing, respectively.
C4 log perplexity with varying lengths
GLUE & SuperGLUE
512
1024
2048
4096
8192
Average accuracy
FIRE
3.15
3.08
3.05
3.05
3.06
71.14 71.04
FIRE-S
3.22
3.14
3.10
3.09
3.10
SCROLLS benchmark
Qas
CNLI
QMS
NQA
SumS
GovR
QuAL
Average
FIRE FIRE-S
16.24
82.93
14.58
9.55
15.87
16.31
24.02 24.11
25.64
17.93
75.22
15.05
9.22
16.02
16.25
24.83
4.6 LAYERWISE SHARING
Another important factor beyond length generalization is compu- tational cost of these approaches. Most of FIRE's computation is based on matrix multiplication, which is more accelerator-frient friendly than the vector operations used in T5's RPE. To further improve the computational efficiency of FIRE, we consider FIRE-S, weight-sharing version which uses the same position encoding bias for all the layers. This way the position encoding bias only needs to be computed once, and the cost is amortized over all the layers. Note that sharing position encoding across layers a common inductive bias in many existing methods (Su et al., 2021; Press et al., 2022; Luo et al., 2022).
Efficiency evaluation
Ø 2.7
2.6
2.5
Φ 2.4
2.3
2.2
NoPE
RoPE
Alibi
Kerple
T5's RPE
FIRE
FIRE S
Figure 5: Inference time com- parisons for different methods. The reported results are averaged over 10 runs for each method.
We conduct experiments to evaluate FIRE-S (with layerwise shar- ing) on C4 language modeling, SCROLLS long text benchmark, and GIJUF/SuperGI.Ur. We also measure the inference speed of different methods. Experimental details are provided C.6.
Model quality. Table 2 compares the accuracy of FIRE-S and the standard FIRE. The results show that sharing position encoding function across layers only leads to slight performance degradation. FIRE-S still outperforms other baselines in the long sequence regime. For example, on C4 language modeling with sequence length 8192, it outperforms Kerple, the best baseline Fig. 1 (3.10 v.s. 3.16 log perplexity). On SCROLLS, its average score outperforms all the strong baseline methods including T5's RPE, RoPE with positional interpolation, and Kerple.
Inference speed. Fig. 5 compares the model speed of FIRE/FIRE-S with baselines. We first notice that FIRE and FIRE-S are both faster than T5's RPE while achieving stronger performances. Moreover, FIRE-S significantly improve the efficiency of FIRE and is faster than all the baselines but NoPE (no positional encoding). In conclusion, the experiments show that FIRE-S demonstrates good speed-accuracy trade-off.
5 CONCLUSION
We propose functional interpolation for relative position encoding (FIRE) to improve Transformer's ability to generalize to longer contexts, and present theoretical and empirical results showing its effectiveness. We prove that FIRE unifies many existing additive RPE methods, while being adaptive enough to learn diverse position encoding biases in long context settings. Empirical results show strong length generalization behavior pushing the paradigm of train short test long. Our work does suffer from some limitations. 1) We only study decoder models. 2) We do not analyze the role of other components of Transformer and other training components (data, optimizer) in length generalization. These questions interesting directions for future exploration.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work is supported in part by the United States Department of Energy via the Brookhaven National Laboratory under Contract No. 384608.
REFERENCES
Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo, James Lhe-Th Yi Tay, al. Colt5: Faster long-range transformers with conditional computation. arXiv preprint arXiv:2303.09752, 2023.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 35:38546-38556, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
2020.
Mirelle Candida Bueno, Carlos Gemmell, Jeff Dalton, Roberto Lotufo, and Rodrigo Nogueira. Induced natural language rationales and interleaved markup tokens enable extrapolation in large language models. In Proceedings of the 1st Workshop on Mathematical Natural Language Processing (MathNLP), pp. 17-24, 2022.
Pu-Chin Chen, Henry Tsai, Srinadh Bhojanapalli, Hyung Won Chung, Yin-Wen Chang, and Chun- Sung Ferng. A simple and effective positional encoding for transformers. arXiv preprint arXiv:2104.08690 2021.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306..153 15595, 2023.
Ta- Ta-Chung Chi, Ting- Ting-Han Fan, Peter Ramadge, and Alexander Rudnicky. Kerple: Kerneiieed relative positional embedding for length extrapolation. Advances in Neural Information Processing Systems, 35:8386-8399, 2022.
Ta-Chung Chi, Ting-Han Fan, Alexander Rudnicky, and Peter Ramadge. Dissecting transformer length extrapolation via the lens of receptive field analysis. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 13522- 13537, 2023.
Krzysztof Marcin Choromanski, Shanda Li, Valerii Likhosherstov, Kumar Avinava Dubey, Shengjie Luo, Di He, Yiming Yang. Tamas Sarlos, Thomas Weingarten, and Adrian Weller. Learning fourier transform for linear relative positional encodings in transformers. arXiv preprint arXiv:2302.01925, 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023111 2022.
Jishnu Ray Chowdhury and Cornelia Caragea. Monotonic location attention for length generalization. arXiv preprint arXiv:2305.20011 2023.
Xiangxiang Chu, Zhi Tian, Bo Zhang, Xinlong Wang, and Chunhua Shen. Conditional positional encodings for vision transformers. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: httpt:://peeeeeee/ http///weweeeeeeeeeeeeeeeeeeeeep n forum? i wewddd-nnnn-RRBa
Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self- attention and convolutional layers. arXiv preprint urXiv:1911.03584 2019.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le. and Ruslan Salakhutdinov. Transformer-xl xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978-2988, 2019.
10

---

Published as a conference paper at ICLR 2024
Gregoire Deletang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, and Pedro A Ortega. Neural networks and the chomsky hierarchy. In The Eleventh International Conference on Learning Representations. 2023. URL https:t////..mm..m openreeeevvv wewerreeiew....... net/ forum? dd-odddtaeeeeeeeee
Jacob Devlin, Ming- Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186, 2019.
Chao Dong, Chen Change Loy, Kaiming Hc, and Xiaoou Tang. Image super-resoluties using deep convolutional networks. IEEE transactions on pattern analysis and machine intelligence, 38(2): 295-307, 2015.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021. URL https: //te:://oooeee openreviee..eeeeeeeeeee net/ forumf i dibniiiddTYyyyyyyy...
Yann Dubois, Gautier Dagan, Dieuwke Hupkes, and Elia Bruni. Location attention for extrapolation to longer sequences. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 403 413,441,,2 2020.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, al. Faith and fate: Limits of transformers on compositionallty. arXiv preprint 2rXiv:2305.18653 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00022 2020.
Mandy Guo, Joshua Ainslie, David C Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. In Findings of the Association for Computational Linguistics: NAACL 2022, pp. 724-736, 2022.
Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. G-mixup: Graph data augmentation for graph classification. In International Conference on Machine Learning, pp. 8230-8248. PMLR, 2022.
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 1382-1390, 2022.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
wwXiv:1606.08415. 2016.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part II 14, pp. 694-711. Springer, 2016.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. arXiv preprini arXiv:2305. 19466, 2023.
Guolin Ke, Di He. and Tie- Yan Liu. Rethinking positional encoding in language pre-training. In International Conference on Learning Representations, 2021. URL htttps://pperee//// /pwp//eeee.......eomeeee. net/ fortm?orum? -
Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, and Bing Liu. Continual pre- pre-training of language models. The Eleventh International Conference on Learning Represen- fations, 2022.
11

---

Published as a conference paper at ICLR 2024
Urvashi Khandelwal, He He, Peng Qi, and Dan Jurafsky. Sharp nearby, fuzzy far away: How neural language models use context. In Proceedings of the 56th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 284-294, 2018.
Nikita Kitaev and Dan Klein. Constituency parsing with
encoder. arXiv preprint
arXiv:1805.01052. 2018.
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018.
Shanda Li, Xiangning Chen, He, and Cho-Jui Hsieh. Can vision transformers perform convolution? arXiv preprint arXiv:2111.0135. 2021.
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36, 2024.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models long contexts. arXiv preprint arXiv:2307.03172 2023.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103. 14030, 2021.
Antoine Liutkus, Ondřej Cifka, Shih-Lun Wu, Umut Simsekli, Yi-Hsuan Yang, and Gael Richard. Relative positional encoding for transformers with linear complexity. In International Conference on Machine Learning, pp. 7067-7079. PMLR, 2021.
Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3431-3440, 2015.
Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tin-Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances in Neural Information Processing Systems, 34, 2021.
Shengjie Luo, Shanda Li, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. Your transformer may not be as powerful as you expect. Advances in Neural Information Processing Systems, 35: 4301-4315 2022.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference on machine learning (ICML-10), pp. 807-814, 2010.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, 2024. URL https: http///pppweeeeeeeeeeeeeeeeeeeep net wid=
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL
https: http:////peeeeeeeeneeeeeee net http:/rwwwwweeeeeeeooooiiiiiieee htt::/ffmmmio-sppppoooo
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv: 15/1v 1.06434, 2015.
12

---

Published as a conference paper at ICLR 2024
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint 19X0v:I910.100
2019.
Stephen Roller, Emily Dinan, Naman Goyal, Da Ju, Mary Williamson, Yinhan Liu, Jing Xu, Myle Ott, Eric Michael Smith, Y-Lan Boureau, et al. Recipes for building an open- doen-nmmiin chatbot. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pp. 300-305, 2021.
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. In 61 Annual Meeting of the Association for Computational Linguistics, 2023.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, et al. Scrolls: Standardized comparison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. pp. 12007-12021, 2022.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pp. 464-468, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10.18653/v1/
N18-2074. URL https: beattps:/aallanthtllll hat/aaaaaanthooggyy..oo.o org N188- -0774.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprini arXiv:2307 230.09988,, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 6000-6010, 2017.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019a.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2019b. URL httptt:tttppsteeeeeee /ppenrevie.. net/ fortm/o 44d=:44k22555
Benyou Wang, Lifeng Shang, Christina Lioma, Xin Jiang, Hao Yang, Qun Liu, and Jakob Grue Simon- sen. On position embeddings in [bert). In International Conference on Learning Representations, 2021. URL https: : / wwwp//peeeeeeeeeeeeeeee net/ forum? onxoVA XMW.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837, 2022.
Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie Yan Liu. Do transformers really perform badly for graph representation? Advances Neural Information Processing Systems, 34, 2021.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are trans- formers universal approximators of sqquence-to-eqquence functions? In International Conference Learning Representations, 2019.
13

---

Published as a conference paper at ICLR 2024
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018. URL httpa: / ///ppperevew..../ net wwtp://eee//oumuteeeeeeee Forum?iddrlldse
Jingqing Zhang, Yao Zhao, Mohammad Saleh, and Peter Liu. Pegasus: Pre-training with extracted gap-sentences for abstractive summarization. In International Conference on Machine Learning, pp. 11328-11339, PMLR, 2020a.
Yizhe Zhang, Siqi Sun, Michel Galley, Yen-Chun Chen, Chris Brockett, Xiang Gao, Jianfeng Gao, Jingjing Liu, and William В Dolan. Dialogpt: Large-scale generative pre-training for conversational response generation. Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, pp. 270-278, 2020b.
Wenting Zhao, Mor Geva, Bill Yuchen Lin, Michihiro Yasunaga, Aman Madaan, and Tao Yu. Complex reasoning in natural language. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 6: Tutorial Abstracts), pp. 11-20, Toronto, Canada, July 2023. Association for Computational Linguisties. doi: 10.18653/v atuttooas.. -tttorills..
2. URL ht attps: / acl anthology wrg//02/000 2023. acl-tutorials.2
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization. arXiv preprint arXiv:2310.16028 2023.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but robustly. arXiv preprint arXiv:2402.09371, 2024.
14

---

Published as a conference paper at ICLR 2024
A
OMITTED PROOF
In this section, first provide more general formulation of T5's positional encoding function as mentioned in Sec. 2.2. Then we provide the proof of Theorem 3.1.
A.1 T5's RPE WITH BUCKETING
In Sec. 2.2, we use a simplified description for T5's RPE. In practice, T5's RPE does not assign different position bias for all different relative positions. Instead, all possible relative distances partitioned into several buckets, and the relative distances in one bucket share a (learnable) attention bias. Formally speaking, T5's RPE pre-defines 0 s0 81 attention bias as
8k-1 SK, and computes the
rk 5k Si -i-j < sk+li == , KK-1
b(t
(5)
Tx i - 2 SK
It's easy to that the formulation in Sec. 2.2 is special case of Eq. (5) by setting k. In the official T5 implementation" the buckets are defined based on "log binning". With K K+1 buckets and pre-defined distance L1, the attention bias is calculated as (assuming K + is even)
ri-j rx =
0 K+1 2
b(i,))
log
K=T
log
2hg KF1
K+1 Si IA. 2
(6)
Fk
i-j L
This is also special case of Eq. (5). In the proof of Theorem 3.1, we will be working on the most general formulation (Eq, (5)), so that the proof works for any specific instances.
A.2 PROOF OF THROREM
Proof. For each RPE variant (T5's RPE, Alibi, Kerple, and Sandwich), provide constructions in which FIRE represent each the target by for < j S i Ly-
T5's RPE. We consider the general T5's RPE formulation with bucketing in Eq. (5). The target positional encoding function can be rewritten
(i,))) = To +o + Σ (rs rk-1) ((1-22)))
(7)
Consider a two-layer MLP with activation a(x) = 1 (x20) and hidden neurons:
fe(x) = u2 0 b1) + +b2.
(8)
Let v Lo1 (where denotes an all-one vector), b = [-91,, .., U2 [ri
T1.-.. Tx - - -x= and b2 = To-
In the positional encoding function of FIRE (Eq. (4)), we set the transform to be the identity mapping + and the threshold L to Lo-
"https: tttb//ttthhhhhcccc http:///ww...ooooooeeoeooeeeoena
15

---

Published as a conference paper at ICLR 2024
Then for any kisis Lo.
( i--
brine(i))=)) =ʃ0 Lo
(9)
i j 81
i- 52
[r]
PK rk-1jo
+ To
(10)
j 8K
(1-12813
(1-1202
[r1 -1
TK TK-1]
+
(11)
(1-12xx]
K Σ (rk - (1- +
(12)
k=1
Thus, we have (FIRE(i, = bo(i.) for any 0 <jsis Lo.
Alibi. The target positional encoding function is bu j) -rii 5) (note that focus the setting where 2 f). Consider a one-layer MLP with identity activation and no bias term (which degrades to linear mapping) fo(x) vy and let f1 Lo In the positional encoding function of FIRE (Eq. (4)), we set the transform vp to be the identity mapping x + x and the threshold /. to
Lo- Then for any 0 < S Lo-
i (1-j) -r(i-- 5) = bi(
(F.R = fo ( Lo
(13)
which concludes the proof.
Kerple (logarithmic variant). The target positional encoding function is bu( -1 log(1 - (2 j)) (note that we focus on the setting where 2 j). Consider a one-layer MLP with identity activation and no bias term (which degrades to linear mapping) fa(x)
vrz. and let in -ri log(1 + r2 La). the positional encoding function FIRE (Eq. (4)), we set the transform to be the log transform + log(r2x + 1) and the threshold L to Lo. Then for any 0 < jjiss Lo.
(FIRE(i, fo loog(1 + r2(1 +r2(i-33)) (()(-))) -ri log(1 + + - 5)) = (i(
(14)
log(1 r2
which concludes the proof.
Kerple (power variant). The target positional encoding function is (i (1 j) (note that we focus the setting where i j). Consider two-layer MLP with activation a (xx) x² one hidden neuron, and no bias term: fe (x) = (2((12) Let 01 = ' and - In the positional encoding function of FIRE (Eq. (4)). we set the transform e be the identity mapping - x and the threshold I, to Lag Then for any 0 < Sis To.
i (i-3) = - -1-
(i.)r fo
j) = bo(i.)),
(15)
To
which concludes the proof.
Sandwich. The target positional encoding function
d'
(o(i.) = cos
(i--)/1100005
(16)
k=1
Consider two-layer MLP with cos activation, d' hidden neurons, and no bias term:
So (x) ng cos(v137).
(17)
16

---

Published as a conference paper at ICLR 2024
Let 1a flo/10000 Lo/ 100000 10000000000000000000000000000000 and v2 c1. In the positional encoding
function of FIRE (Eq. (4)), we set the transform vp to be the identity mapping - x and the threshold to Lo. Then for any < iis Lo.
( (1-2)
bFIRE(i,j) =fo
(Lo Lo
(18)
(i - (o--j)11000 10/00000 +
cos
(i j) 100100
=/c
d
cos
(19)
cos
((i 5)/ 10000")
d'
C
cos
j// (i-=)/110003)
(i
(20)
k=1
Thus, we have bF)RE(i,J) = (i( for any 0 j S i S Lo.
B ABLATION STUDY
The positional encoding function of FIRE can be viewed as composition of position transformation and a function approximator b(i.j) = fo(g(i - j. i). The position transformation takes the relative distance i and the query position i as the input and produces "normalized" distance. For example, in Eq. (4), the position transformation choices of & leads different position transformation g. The function approximator fo should be in an expressive function class parametrized by 0, which transforms the normalized distances into attention biases. For example, we use a wwo-hidden-layer MLP with 32 neurons in each hidden layer and ReLU activation by default, as discussed in Appendix C.I.
(i j, + v6(i - (max/ (1, L}). Different
In this section we ablate our design choices for both the position transformation and the function approximator. We also conduct ablation experiments to test the length generalization performances on different training sequence lengths. All the ablation experiments are based on base-sized models.
B. THE LOG TRANSFORM AND THRESHOLDING IN POSITION TRANSFORMATIONS
In Sec. 3.2, we propose two modifications, the log transformation and thresholding operation, as additional transformations to the relative distance. We conduct experiments to ablate these design choices and demonstrate their effectiveness. We experiment with base-sized models and compare FIRE variants with or without the additional transformations in Sec. 3.2. Specifically, we consider three variants with the following positional encoding functions:
Without log transform//trresoolding (i(i.) - fo (-=2)
(21)
With log transform but without thresholding: (2().))) So [log(c(i -
(22)
log(ci +
With log transform and thresholding: bg(i.)) = fo ( log(c(i - +)+ +1)
(23)
[og((c +1
For all the three variants (Eq. (21-23)), fo is parameterized as two-hidden-layer MLP with 32 neurons in each hidden layer and ReLU activation to ensure a fair comparison. Eq. (23) is the standard FIRE positional encoding function used in Sec. 4. We experiment on C4 language modeling and GLUEESuperGLUE benchmark using the settings and evaluation metrics described in Appendix C. The experimental results are shown in Table 3. From the language modeling results, we can that both the log transformation and the thresholding operation improve the language modeling quality for all the lengths, and the standard FIRE positional encoding function in Eq. (23) is the best variant. In particular, the log transformation largely improve the performance long sequences,
17

---

Published as a conference paper at ICLR 2024
indicating that amplifying the differences among local positions helps in the long sequence regimes. We further study the effectiveness the thresholding operation on GLUE/SupprGL benchmark which contains relatively short sequences. The results show that the thresholding operation leads to 0.72 point performance gain on average Gi.UEE/uperGLL.U accuracy, verifying its effectiveness on improving short sequence modeling.
Table 3: Ablation study the position transformation. We compare FIRE variants with or without the additional transformations See. 3.2. For log transform, indicates (bx)) = x, i.c. no log transform; while / indicates (xx)) = log(cr cr + 1), i.e., applying log transform for the relative distance. For thresholding, X indicates using (2) to normalize the relative distance, i.e., thresholding operation; while / indicates (max L}) to normalize the relative distance with L being learnable threshold.
Method
C4 log perplexity with varying lengths
Log transform
Thresholding
Formula
512
1024
2048
4096
8192
X
Eq. (21)
3.194
3.128
3.099
3.216
3.334
Eq. (22)
3.161
3.093
3.062
3,057
3.085
Eq. (23)
3.149
3.083
3.054
3.046
3.056
Method
GLUE/SuperGLUE
Log transform
Thresholding
Formula
Average accuracy
X
Eq. (21)
69.06
X
Eq. (22)
70.42
Eq. (23)
71.14
Additional discussions on the thresholding operation. We note that even FIRE without thresh- olding outperforms all the baselines (including RoPE, T5's RPE, etc) on all the sequence lengths on C4 language modeling. Detailed comparisons are in Table 4.
In all the experiments presented in the paper, the threshold of FIRE Eq. (23) a learnable parameter. For the base-sized model pretrained on sequence length 2048, the learned parameter L is between 1200 to 1600 across different layers. Setting to fixed value also a viable option. In our preliminary exploration, FIRE with either fixed or learnable outperforms all the baselines, while the learnable variant leads to better performances. The fixed variant introduces one more hyper-parameree and may require more tuning. Thus, FIRE uses learnable threshold L the default choice.
Table 4: Comparing FIRE variants with baselines. We present additional comparisons between existing methods and FIRE variants with or without thresholding.
C4 log perplexity with varying lengths
Method
512
1024
2048
4096
8192
NoPE
3.206
3.14
3.111
3.287
3.410
RoPE
3.178
3.102
3.070
3.375
3.519
Alibi
3.320
3.248
3.216
3.438
3.537
Kerple
3.326
3.217
3.170
3.156
3.158
T5's RPE
3.164
3.095
3.064
3.095
3.181
FIRE without thresholding (Eq. (22))
3.161
3.093
3.062
3.057
3.085
FIRE (Eq. (23))
3.149
3.083
3.054
3.046
3.056
B.2 EFFECTS OF THE FUNCTION APPROXIMATOR CAPACITY ON THE PERFORMANCES
We experimentally study the impact of the function approximator (fø) capacity on the model perfor- mance. We compare linear layer, soe-hhdden-lyyee MLP, and two-hidden-layer MLP. The MLPs both have 32 neurons in the hidden layers and use ReLU (Nair & Hinton. 2010) activation function. Two-hidden-llaer MLP is the defualt choice for FIRE in Sec. 4. We experiment C4 language
18

---

Published as a conference paper at ICLR 2024
modeling and evaluate the models varying sequence lengths using the settings and evaluation metrics described in Appendix C.I and present the experiment result 5. The result shows that linear layer is not experssive enough and leads suboptimal performance on C4 language modeling. Introducing non-linearty and parametrizing fo as a mee/two-hidden--aye MI.P leads to much better results. In particular, using one-hidden-layer MLP has largely improve the overall performances especially in the long sequence regimes. For example, it outperforms a linear fo by 0.24 point log perplexity on sequence length 8192. Morcover, using an MLP with larger capacity (two hidden layers v.s. one hidden layer) can further brings performance gains. That being said, the MLP is still very tiny (with only 32 hidden neurons) and we believe it's the non-linearty that helps.
Table 5: Ablation study on the capacity of the function approximator (fo). We compare FIRE variants with different activation functions in MLP.
C4 log perplexity with varying lengths
Parametrization of fe
512
1024
2048
4096
8192
Linear
3.21
3.14
3.11
3.20
3.32
One-hidden-Rase MLP (32 hidden neurons)
3.17
3.10
3.07
3.06
3,08
Two-hidden-layer MLP (32 hidden neurons)
3.15
3.08
3.05
3.05
3.06
B.3 CHOICE OF THE MLP ACTIVATION FUNCTION
We study the impact of the MLP activation function on the model performance. We experiment on C4 language modeling and evaluate the models on varying sequence lengths using the settings and evaluation metrics described in Appendix C.I. We compare ReLU (Nair & Hinton, 2010) and GeLU (Hendrycks & Gimpel, 2016) activation functions and present the experiment result in 6. The result shows that the model performance is not sensitive the choice of activation function in the length generalization setting, while ReLU works better on normal sequence lengths. Thus, we use ReLU as default activation function.
Table 6: Ablation study on the MLP activation. We compare FIRE variants with different activation functions MLP.
C4 log perplexity with varying lengths
512
1024
2048
4096
8192
ReLUU
3.15
3.08
3.05
3.05
3.06
GeLU 3.36
3.26
3.06
3.05
3.06
B.4 CHOICE OF FINAL ACTIVATION OF MLP OUTPUT
In our main experiments, we focus on MLPs of the form fo(x) (0 o(-(---001 where a is the activation function. In this implementation, the MLP ends with a linear layer and no activation function is applied to the MLP final output. A slightly different choice is to consider fe(x) = (v, (-.- (ac.--o0,,aa (v)x))) where final activation applied to the MLP output. We compare these two choices by experimenting on C4 language modeling and evaluating the models on varying sequence lengths. We use haddeed---- MLP with 32 hidden neurons and the ReLU (Nair & Hinton, 2010) activation function in both model variants. The results are presented in Table 7. We find that MLP without final activation leads to better performances on long sequences and use it as our default choice.
B.5 FIRE STILL STRONG WHEN TRAINED ON SEQUENCE LENGTH 512
In most of our pretraining experiments, the training sequence length is set to 2048 (see Appendix C.I). In this experiment we train models with different positional encodings on C4 with training sequence length 51 512 to confirm that the overall performance trends not sensitive to the pretraining
19

---

Published as a conference paper at ICLR 2024
Table 7: Ablation study on final activation of MLP output. We compare FIRE variants using MLP
with/without final activation to its output.
C4 log perplexity with varying lengths
512
1024
2048
4096
8192
With final activation
3.16
3.10
3.07
3.09
3.19
Without final activation
3.17
3.10
3.07
3,06
3.08
sequence length. Other experimental settings are te same as those in Appendix C.I. We evaluate the models on varying sequence lengths and report the log perplexity in Fig. 6. It's clear that FIRE still achieves the strongest overall performance compared with the other baselines. The results in Fig. 1 & 6 demonstrate that FIRE can robustly deliver higher modeling quality regardless of the training sequence lengths.
C4 language modeling (base model trained on length 512)
Valida Validation log perplexii explexity
5.00
NoPE
4.75
RoPE
Alibi
4.50
Kerple
4.25
TS's RPE
4.00
FIRE (ours)
3.75
3.50
3.25
512
1024
2048
4096
8192
Validation sequence length
Figure 6: Language modeling perplexity evaluated on varying sequence lengths on C4 validation set. The plots are base-sized models with training sequence length 512.
C EXPERIMENT SETTINGS & ADDITIONAL RESULTS
C.I LANGUAGE MODELING WITH LENGTH GENERALIZATION
Model configurations. this experiment, we train decoder-only Transformer language models with different positional encoding variants while keeping all the other configurations the same. For T5's RPE, we follow Raffel et al. (2019) and use 64 position bucket for each attention head. For Alibi, follow Raffel al. (2019) to set the hyperparameters in the positional encoding function in each attention head. For our FIRE method, we use the positional encoding function defined in Eq. (4). In Eq. (4), we let (xx log(ez: + 1) where is learnable parameter; is parametrized two-hidden-layer MLP with 32 neurons in each hidden layer and ReLU activation.
We experiment with model size settings, base (125M parameters) and large (350M parameters). The model configurations follow (Brown et al., 2020) and are presented in Table 8.
Training recipe. Following Brown et al. (2020), we use the causal LM objective to pretrain decoder-only Transformers with different position encodings. We use the C4 dataset (Raffel et al., 2019) as the pretraining corpora. We set pretraining sequence lengths to 2048, and evaluate the zero-shot perplexity sequence lengths (512, 1024, 2048, 4096, 8192). We truncate documents with length greater than 2048 to multiple sequences of length 2048 during training; similar trucation done to construct the validation sets of different sequence lengths. Our training recipe follows (Brown et al., 2020) and is presented in Table 9.
Additional results. We evaluate language modeling log perplexity with varying lengths on C4, arXiv, and Github datasets (Raffel et al., 2019; Gao al., 2020) for both base and large models. The
20

---

Published as a conference paper at ICLR 2024
Table 8: Model configurations for language model pretraining.
Small model
Large model
Training sequence length
2048
2048
Number of layers
12
24
Attention heads
12
16
Hidden layer size
768
768
Head dimensions
64
64
FFN activation
GeLU
GeLU
Number of parameters
125M
350M
results of base models on C4 are presented in Fig. 1. The results of large models on all the three datasets are presented in Fig. 2. In Fig. 7, additionally present the results of base models on arXiv and Github. All the results show similar trends and FIRE consistently demonstrate strong length generalization behavior.
arXiv language modeling (base model)
Github language modeling (base model)
NoPE
NoPE
Validation log perplexity
Validation 19 perp perplext
6.5
RoPE
4.0
RoPE
6.0
Albb
Albi
Karple
Karple
5.5
T5 's RPE
3.5
T5 's RPE
5.0
YaRN
YaRN
FIRE [ours]
FIRE (ours)
4.5
3.0
4.0
3.5
2.5
8
Validation sequence length (x10²)
Validation sequence length (x10²)
Figure 7: Language modeling perplexity evaluated on varying sequence lengths on arXiv (left) and Github (right) validation set. The plots base-sized models with training sequence length 2048.
Table 9: Training recipe for language model pretraining.
Small model Training sequence length
Large model
2048
2048
Batch size
256
256
Numer of iterations
600k
600k
Dropout prob.
0.0
0.0
Attention dropout prob.
0.0
0.0
Optimizer
Adam W
Adam W
Learning rate
6c
3e
Hardware (TPUv4 chips)
128
256
C.2 FINETUNING ON LONG TEXT BENCHMARK
Datasets and evaluation metrics. We use SCROLLS long text benchmark (Shaham et al., 2022) to further test the models' capability of learning and modeling long sequences. SCROLLS benchmark includes question-answering datasets Qasper, NarrativeQA, and QuALITY natural language inference datasets ContractNLI; and summarization datasets QMSum, SummScreenFDD, and GovReport. Following existing works Shaham et al. (2022); Ainslie et al. (2023), three different evaluation metrics are used for different datasets: Rgm score (the geometric mean of ROUGE- 1,2.L) for GovReport, SummScreenFD, and QMSum, unigram overlap (F1) for Qasper and NarrativeQA, and exact match (EM) for ContractNLI and QuALITY. We also compute the average score across different datasets done in the SCROLLS benchmark.
21

---

Published as a conference paper at ICLR 2024
Model and training configurations. We finetune the checkpoints pretrained C4, the model configurations are the same as those in Table 8. We use the same set of hyperparameters for all the models and all the tasks, and report the best results on the validation set. Table 10 presents our finetuning configurations.
Table 10: Finetuning configurations for SCROLLS benchmark.
Batch size
128
Numer of iterations
25k
Dropout prob. Attention dropout prob. Optimizer Learning rate Hardware (TPUv4 chips)
0.1
0.1
A
le - 5
128
C.3 ZERO-SHOT LENGTH GENERALIZATION ON NARRATIVEQA
Datasets and evaluation metrics. We use the NarrativeQA dataset (Kočiský et al., 2018) with different input context lengths to test the model's ability to leverage long context in zero-shot learning settings. We use the base-sized model checkpoints pretrained on C4 (sequence length 2048) and finetuned on NarrativeQA (sequence length 8192). We evaluate the models on context lengths (512, 2048, 4096, 8192, 16384, 24576, 32768) and use unigram overlap (F1) as the evaluation met- ric.
Detailed results. We provide detailed performances of all the tested models in Table 11. The result shows that FIRE is consistently outperforming all the baselines across all different context lengths.
Table 11: Detailed performance comparisons on NarrativeQA with varying context lengths. "RoPE-PI,," refers RoPE interpolation with max sequence lengths Lo. Best performances highlighted in bold.
Context length
512
2048
4096
8192
16384
24576
32768
Average
NoPE
2.245
4.070
4.277
5.661
4.770
4.716
3,930
4.238
RoPE
1.546
1.482
2.060
8.737
1.071
0.190
0.132
2.174
RoPE-Pls192
5.241
4.639
6.070
8.301
0.565
0.728
0.623
3.738
RoPE-Ply2768
4.092
5.912
5.769
5,459
5.677
5.446
6.767
5.589
Alibi
4.036
4.339
4.190
4.251
4.144
4.086
3,899
4.135
Kerple
5.590
7.832
8.001
9.249
9.483
9.204
9.010
8.338
T5's RPE
4.595
5.557
6.528
8.983
3.872
2.226
1.757
4.788
FIRE (ours)
6.232
8.076
8.178
9.581
9.581
9.868
9.417
8.705
C.4 FINETUNING ON GLUEESUPPERGI.UE
Datasets, evaluation metrics, and configurations. GLUE and SuperGL.UF. are widely-used bench- marks to evaluation the natrual language understanding capability of neural language models (Wang et al.. 2019b;a). We finetune the models on mixture of the tasks in GLUE and SuperGLUE for simplicity. We evaluate the model on each task separately. We use the macro average accura- cy/exact match across all the tasks as our main evaluation metric. Table 12 presents our finetuning configurations.
Detailed results. For reference, we present detailed results for all the models on each individual dataset in Table 13. In general, FIRE achieves decent performances. Thus, FIRE's strong perfor- mances on long sequences does not come at the price of sacrificing model quality on short sequences and standard tasks.
22

---

Published as a conference paper at ICLR 2024
Table 12: Finetuning configurations for GLUE/SuperGLUE benchmark.
Batch size
256
Numer of iterations
25k
Dropout prob. Attention dropout prob. Optimizer
0.1
0.1
AdamW
Learning rate Hardware (TPUv2 chips)
le le-3 5
Table 13: Detailed performances on GLUE and SuperGLUE tasks. The evaluation metrics are
EM (exact match) for Multire & Record; and accuracy for the remaining tasks.
Base models
Boolq
Cb
Cola
Copa
Mnli
Mrpc
Qnli
Qqp
NoPE
72.51
73.21
69.42
67.00
79.72
75.98
84.70
88. 72
RoPE
75.78
80.36
74.78
60.00
83.11
79.17
87.70
90.03
RoPE-PI
75.72
80.36
72.87
64.00
82.87
80.64
86.89
89.93
Alibi
69.76
76.79
69.32
58.00
78.02
76.72
83.97
88.14 14
Kerple
77.31
82.14
74.11
61.00
82.69
80.64
87.66
90.22
T5's RPE
76.30
83.93
71.33
61.00
82.10
81.37
87.61
89.87
FIRE (ours)
76.76
83,93
73.63
59.00
83.01
80.39
87.83
89.97
Rte
Sst2
Wic
Wnli
Multirc
Record
Wsc
NoPE
71.84
91.17
58.78
63.38
16.89
35.50
67.31
RoPE
73.65
92.89
66.93
61.97
23.19
46.57
71.15
RoPE-PI
71.48
91.51
65.05
60.56
22.46
45.96
70.19
Alibi
68.23
88.76
57.05
61.97
12.70
29.34
63.46
Kerple
69.68
92.43
64.89
53.52
22.56
47.74
66.35
T5's RPE
73.65
92.20
63.79
60.56
20.57
45.71
69.23
FIRE (ours)
75.81
92.66
64.58
60.56
25.81
46.89
66.35
Large models
Boolg
Cb
Cola
Copa
Mnli
Mrpc
Qnli
Qqp
NoPE
79.27
83.93
78.24
61.00
84.39
79.90
89.79
90.74
RoPE
79.66
91.07
80.54
63.00
85.67
81.86
90.87
91.04
RoPE-PI
79.45
92.86
80.54
63.00
85.31
81.62
90.52
91.05
Alibi
74.77
80.36
71.05
58.00
81.72
79.41
86.18
89.75
Kerple
80.70
92.86
79.29
65.00
85.63
80.88
90.56
90.86
T5's RPE
79.88
87.50
78.33
65.00
84.80
83.58
89.77
90.71
FIRE (ours)
79.60
85.71
79.10
65.00
84.93
81.13
90.37
90.84
Rte
Sst2
Wic
Wnli
Multirc
Record
Wsc
NoPE
77.26
93.69
62.70
59.16
26.65
51.18
70.19
RoPE
79.42
94.38
69.59
60.56
30.64
58.23
72.12
RoPE-PI
79.06
94.61
70.69
56.34
31.17
56.69
68.27
Alibi
72.56
91.97
60.35
50.70
22.77
40.79
66.35
Kerple
79.06
94.61
67.24
53.52
31.17
58.55
71.15
T5's RPE
79.78
92.89
64.58
54.93
29.80
52.54
69.23
FIRE
80.87
93.92
67.71
59.16
31.90
54.67
72.12
Kerple
79.06
94.61
67.24
53.52
31.17
58.55
71.15
T5's RPE 79.78
92.89
64.58
54.93
29.80
52.54
69.23
FIRE (ours)
80.87
93.92
67.71
59.16
31.90
54.67
72.12
23

---

Published as a conference paper at ICLR 2024
C.5
VISUALIZATION
We present another visualization of learned FIRE biases for query at position 8192 in Figure 8.
Layer
Layer
Layer 10
FIRE positional bias
1st head
0
2nd head
-10
-5-
-10
4th head
5th head
10
7th head
-20
-20
-20
10th head
-30-
2000
8000
6000 Relative distance
2000
12th head
Relative distance
Relative distance
Figure 8: Visualization of FIRE learned position biases for the 8192nd query position with key positions between and 8192. We notice that FIRE models learn both local and anti-local position patterns.
C.6 EFFICIENCY AND
For FIRE-S (FIRE with layerwise sharing), we experiment with the base-sized model (125M parame- ters), and keep all the configurations and training recipes the same as those in previous subsections. The models are pretrained on C4 with sequence length 2048. The finetuning sequence lengths are 8192/1024 for SCROLLS and GLUE/SuperGLUE, respectively.
For the inference time evaluation, we test the forward time of base-sized model with different positional encodings on sequence length 2048. We measure the forward time on 4 TPUv2 chips for all the models, and report the average result over 10 runs.
D RELATED WORKS
In the main body of the paper, we cover the most relevant works to our paper (Sec. 2). In this section, we provide more discussions on related works.
Length generalization. Many existing works show the length generalization failure of standard Transformer models (Press et al., 2022; Anil et al., 2022; Deletang et al., 2023; Liu et al., 2024). Recently, there have been growing interests in long-context applications such as multi-step reason- ing (Wei et al., 2022; Dziri et al., 2023; Zhao et al., 2023) and document/hook understanding (Kočiský et al., 2018; Ke et al., 2022; Guo et al., 2022; Ainslic et al., 2023; Liu et al., 2023). Designing length- generalizable Transformers is appealing for these applications. Dubois et al. (2020); Chowdhury & Caragea (2023) introduce location attention for length generalization on synthetic tasks. Bueno et al. (2022) show that generating step-by-step rationales and using marker tokens positional guides helps length generalization. Studying positional encoding approaches for length generalization is main direction in this line of research. Press et al. (2022): Chi al. (2022; 2023) propose new relative positional encoding methods which emphasize recency bias and improve language modeling on longer sequences. Chu al. (2023) propose Conditional Positional Encodings enhance Vision Transformer length generalization. The most relevant to our work is concurrent paper by Chen al. (2023). propose Position Interpolation (PI) for Rotary Positional Encoding (RoPE), which extends the context window of RoPE-based pretrained models given downstream max sequence length. However, this requires additional finetuning on longer sequence data, albeit for much fewer steps than original training. By contrast, our proposed FIRE does not require pre-defined max sequence length, and can be directly applied to length generalization setting without tuning. We provide extensive experimental comparisons in Sec. 4. More recently, Zhou et al. (2024) show that standard Transformers can generalize to a sequence length that is 2.5x the training input length on integer addition using FIRE (and other techniques (Ruoss al., 2023; Zhou et al., 2023)).
24

---

Published as a conference paper at ICLR 2024
Positional encoding in Transformers. Positional encoding is critical component of Transformers. Vaswani et al. (2017) propose sinusoidal Absolute Positional Encoding (APE) to encode positional information in the sequential input. Shaw al. (2018) are the first to propose Relative Positional Encoding (RPE) for Transformers, and many follow-up works explore different RPE strategies (Dai al., 2019; Raffel al., 2019). There also many works that study positional encoding from different perspectives, including the disentanglement of positional and content information (Kitaev Klein, 2018; Ke al., 2021), the representational power of attention modules and Transformers (Cordonnier et al., 2019; Chen et al., 2021; Li et al., 2021; Luo et al., 2022), computational efficiency al., 2023), and length generalization (Press et al., 2022; Chi et al., 2022; 2023; Kazemnejad et al., 2023). Our work is based on unified formulation of existing additive relative positional encoding approaches, and proposes new RPE variant aimed at improving length generalization.
(Su al., 2021; Liutkus et al., 2021: Luo et al., 2021; Choromanski
Interpolation techniques in deep learning. Interpolation techniques successfully applied to many deep learning applications, especially in computer vision. Long et al. (2015) employ bilinear interpolation in up-sampling layers of convolutional neural networks for dense visual prediction. Dong et al. (201 5); Johnson et al. (2016) employ bicubic interpolation for image super-resolution. Radford al. (2015) probe generative models by interpolation in the latent space. Zhang al. data
(2018); Han et al. (2022) use interpolating between pairs of examples and their labels augmentation method. Recently, Dosovitskiy al. (2021) propose perform 2D interpolation the pre-trained APE for Vision Transformer to apply the model to higher resolution images. In contrast, our interpretation is applied
relative position encoding functions. Besides, we are focused on causal attention setting where "global" information such as the total sequence length is unknown, while Dosovitskiy et al. (2021) work on encoder-only Transformers with fixed input lengths.
E IMPLEMENTATION
In this section, we present the implementation of proposed FIRE module in PyTorch (Paszke et al., 2019).
import torch
import torch.nn nn
4 class FIRE (nnn.Mdule))
def _init_ (self, num_heaaee san_hheds=122 mip_width-3 32, init_c-0.1, init_I-S12., eps-le-6)
PIRE attention bias module.
9
10
Args:
11
nun_heads: number of attention heads. mip_width, Width of MLP.
12
13
init_c: initial value of log transformation parameter
14
init_L: initial value of thresholding parameter
15
eps: small constant for numerical stability
16
17
super (FIRE, self) _init__iitt__
18
19
Define the MLP layers self mll nn. Sequential nn Li (1, mlp_width), Re.ReLU (), nn. Ln.ein (mip_width, num_heads)
20
22
23
24
25
26
Initialize (log transformation parameter)
27
self.c = nn. Parameeet (torch. tensor (esoo(init_cc)) (init_c))
28
25

---

Published as a conference paper at ICLR 2024
29
# Initialize L (threshold)
30
self.init_L an.Parameter (torch.tenscr (init_I)
31
# Learn a multiplier to requires_grad-Falsee
32
33
self.I_multte melf.i_ml
Parameter (torch (1.0))
34
35
self.eps eps
36
37
def forwarddse1f x: torch. Tensor)
38
Compute FIRE attention bias.
40
Args:
42
input sequence, shape [bsz, num_heads seq_len, hidden_dim]
43
44
45
Returns:
46
attention bias,
47
shape [1, num_heads, seq_len, seq_len]
48
.*.
49
seq length x.size (2)
50
positions torch.arange (seq_length,
51
syype=torch.flottt
52
device-x.device
53
rel_distance positions(:; None] positions [None, :]
55
A Thresholding the normalizer
threshold torch. (self.L. [teef.._mmltppliiiee self.init L)
57
pos_normal limmmaliieeeee izer torch. max (positions, threshold)
58
_is_normalizes = pos_normalizer(:: None]
59
60 # Amplifying differences among local positions with log transform
62
rel_distance = torch.log(
torch.abe (self.c * rel_distance)
64
pos_nermalizer torch.loge
66
torch.abs (self.c pos normalizer)) + self.eps
68
69
# Progressive interpolation
70
dismanc rel_distance pos_normalizer
71
fire_bias self.mlp haapppooomliiieedissstaaaa unsqueezee ( (-1))
72
fire_bias fire_bias, unsqueeze (0) permute (0, 3, 1, 2)
73
return fire bias
26


Reviews:
[{"review_id": "k0iPoCVGMe", "reviewer": "Reviewer_VXMT", "summary": "The paper proposes to use a small neural network that computes the positional bias used in self attention. The authors show that such a positional encoding is a generalization of several popular relative position encodings. Importantly, when combined with progressive interpolation (normalizing by the current position), they showcase superior performance in length generalization when using the proposed positional encoding. The authors perform extensive experiments on zero-shot length generalization as well as fine-tuning with longer sequences and show that the proposed method outperforms a host of other positional encodings.", "questions": "Asked in the weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is very well written. The related work is clearly presented as is the proposed method.\n\nGiven the simplicity of the contribution, extensive experiments are required to ensure the universality of the proposed positional encoding. The authors perform experiments at two model scales, evaluating zero-shot generalization as well as generalization with fine-tuning. Moreover the authors provide extensive ablation studies for all components of the method.", "weaknesses": "It is not very clear why ALiBi performs so poorly while both in the Kerple paper and ALiBi paper, ALiBi shows great zero-shot length generalization.\n\nAdditionally, a simple baseline is missing, namely interpolating a learnable positional bias using $\\frac{\\psi(i - j)}{\\psi(i)}$. This would show whether the nugget of the method is the progressive interpolation or the neural network or both. From the visualization of the learned positional biases the functions learned by the neural network are, as expected, not too complicated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HLxkh0g0kr", "reviewer": "Reviewer_WBMp", "summary": "In this work, the authors propose FIRE: functional interpolation for relative positional encoding. The idea is to first project relative position to a real number between 0 and 1, and this projected number is subsequently fed to a MLP to generate positional bias. FIRE enables language models to train on a short sequence length, and then generalize to longer sequence length during testing. Extensive experiments are conducted to demonstrate the effectiveness of the proposed method.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed method is very intuitive and easy to understand. The presentation is very clear. The main motivation is that many existing positional encoding methods cannot extrapolate to unseen sequence length, limiting their practicality. The idea of projecting relative positions to a real number between $[0,1]$ is very natural.\n\n* The authors conduct extensive experiments to demonstrate the effectiveness of the proposed method. Models pre-trained with different positional encoding approaches are evaluated under different settings: fine-tuning on long and short context taks, and zero-shot generalization to unseen sequence length.", "weaknesses": "* Some of the design choices are ad-hoc, especially the thresholding parameter for the normalizer. From Table 3, it seems model performance considerably degrades on short sequence length without this thresholding parameter. I’m wondering about the model performance of using Eq. 21 on GLUE/SuperGLUE.\n\n* Also regarding the thresholding parameter, it is mentioned that the parameter $L$ is learnable. Could the authors demonstrate what the learned parameter $L$ looks like after training? Also, will model performance significantly change if we set $L$ to a fixed value?\n\n* Another weakness is that FIRE is not better than existing methods when the sequence length is short. Even with the thresholding parameter and the slower inference speed, it seems FIRE is not better than RoPE on GLUE/SuperGLUE.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S5bsI6TkPJ", "reviewer": "Reviewer_hDCR", "summary": "Preventing performance degradation of Transformers on inputs longer than the training sequence lengths has posed a significant challenge in expanding the context length of these models. While the Transformer architecture inherently has no limitations on the input sequence lengths it can handle, the choice of position encoding during training can limit their performance on longer inputs. To address this, the author propose a novel approach called Functional Interpolation with Relative Encoding (FIRE), which aims to enhance Transformer's generalization to longer contexts. The author provide theoretical evidence that FIRE can effectively represent popular relative position encodings like T5's RPE, Alibi, and Kerple. Furthermore, The author empirically demonstrate that the FIRE models exhibit improved generalization to longer contexts in both zero-shot language modeling and long text benchmarks.", "questions": "The same as weaknesses part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The author introduces a new relative position encoding and demonstrates its superior extrapolation capabilities.", "weaknesses": "Although the performance of the model presented in this paper is impressive, I believe there are several shortcomings:\n\n1. Heavy reliance on handcrafted input features: The authors attempt to use a neural network to learn a relative position encoding. However, the input features of this network still heavily rely on manual design. Although some ablation experiments are conducted in section B.1, I believe the baseline should be a simpler approach like b(i, j) = f(i - j). For example, [1], [2] has achieved good results by learning relative positional relationships using a similar approach.\n\n2. Trade-off between efficiency and performance: Compared to Alibi and RoPE, the learnable FIRE raises the question of how much slower it makes the training process for language modeling. However, this aspect is not mentioned in the paper. Therefore, the authors need to evaluate and discuss the training speed.\n\n3. Lack of certain baselines: For instance, the paper fails to compare its approach with NTK-RoPE[3], YaRN[4], and NTK-ALiBi[5]. These methods have also made advancements in relative position encoding. The absence of comparison with these baselines is a drawback.\n\n\nCitations:\n[1] Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations (ICLR), 2023\n\n[2] Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficient long convolutions for sequence modeling. CoRR, abs/2302.06646, 2023\n\n[3] https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/\n\n[4] https://github.com/jquesnelle/yarn\n\n[5] https://github.com/keezen/ntk_alibi/blob/main/readme_en.md", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)