CLEX: Continuous  Length Extrapolation for Large Language Models
Authors: Guanzheng Chen
Xin Li
Zaiqiao Meng
Shangsong Liang
Lidong Bing
Venue: ICLR 2024 poster (2024)
Keywords: Length Extrapolation
Long Context
Large Language Model (LLM)
Neural Ordinary Differential Equation (ODE)
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: Extrapolating the context length of LLMs to lengths 4x~ 8x times beyond training length via continual and trainable RoPE scaling.

Abstract: Transformer-based Large Language Models (LLMs) are pioneering advances in many natural language processing tasks, however, their exceptional capabilities are restricted within the preset context window of Transformer. Position Embedding (PE) scaling methods, while effective in extending the context window to a specific length, demonstrate either notable limitations in their extrapolation abilities or sacrificing partial performance within the context window. Length extrapolation methods, although theoretically capable of extending the context window beyond the training sequence length, often underperform in practical long-context applications. To address these challenges, we propose Continuous Length EXtrapolation (CLEX) for LLMs. We generalise the PE scaling approaches to model the continuous dynamics by ordinary differential equations over the length scaling factor, thereby overcoming the constraints of current PE scaling methods designed for specific lengths. Moreover, by extending the dynamics to desired context lengths beyond the training sequence length, CLEX facilitates the length extrapolation with impressive performance in practical tasks. We demonstrate that CLEX can be seamlessly incorporated into LLMs equipped with Rotary Position Embedding, such as LLaMA and GPT-NeoX, with negligible impact on training and inference latency. Experimental results reveal that CLEX can effectively extend the context window to over 4× or almost 8× training length, with no deterioration in performance. Furthermore, when evaluated on the practical LongBench benchmark, our model trained on a 4k length exhibits competitive performance against state-of-the-art open-source models trained on context lengths up to 32k. Our code is available at https://github.com/DAMO-NLP-SG/CLEX.

Published as a conference paper at ICLR 2024
CLEX: CONTINUOUS LENGTH EXTRAPOLATION FOR LARGE LANGUAGE MODELS
Guanzheng Chen 12,3, Xin Li Zaiqiao Meng"
Shangsong Liang Lidong Bing
Sun Yat-sen University DDMMO Academy, Alibaba Group
Hupan Lab, 310023, Hangzhou, China *University of Glasgow
Mohamed bin Zayed University of Artificial Intelligence guanzzh. chen @ com, {xinting. 1x, w.w.g.iinn.ll......... baabibiaabbbba com
zaiqiao. eegiao.mengee mennggglagoow uk, Liangshangssongegmail com.
ABSTRACT
Transformer-based Large Language Models (LLMs) are pioneering advances in many natural language processing tasks, however, their exceptional capabilities are restricted within the preset context window of Transformer. Position Em- bedding (PE) scaling methods, while effective in extending the context window to specific length, demonstrate either notable limitations in their extrapolation abilities or sacrificing partial performance within the context window. Length ex- trapolation methods, although theoretically capable of extending the context win- dow beyond the training sequence length, often underperform in practical long- context applications. To address these challenges, propose Continuous Length EXtrapolation (CLEX) for LLMs. We generalise the PE scaling approaches to model the continuous dynamics by ordinary differential equations scaling factor, thereby overcoming the constraints of current PE scaling methods designed for specific lengths. Moreover, by extending the dynamics to desired context lengths beyond the training sequence length, CLEX facilitates the length extrapolation with impressive performance in practical tasks. We demonstrate that CLEX can be seamlessly incorporated into LLMs equipped with Rotary Position Embedding, such as LLaMA and GPT-NeoX, with negligible impact on training and inference latency. Experimental results reveal that CLEX can effectively ex- tend the context window to over 4x or almost 8x training length, with no deteri- oration in performance. Furthermore, when evaluated on the practical LongBench benchmark, our model trained on a 4k length exhibits competitive performance against state-of-the-art open-source models trained context lengths up to 32k. Our code is available at httttts://ggthhti com -LP-MONN-PP---0 -
the length
INTRODUCTION
Transformer-base based large language models (LLMs), such as GPT-4 OpenAl 2023) and
LLaMA (Touvron et al., 2023a b). have now emerged as the tte models in various nat- ural language processing (NLP) tasks. However, these models grapple with the limitations inher- ent to the Transformer architecture mainly, a preset context window, beyond which performance plummets catastrophically (Press et al. (022). The quadratic complexity of the attention mecha- nism renders training LLMs with a larger context window extraordinarily resource-intensive Prior works (Dai et al. 2019 Beltagy al.) 2020 Bulatov et al. 2022) have proposed circumventing full context length access via hicrarchical architecture or sparse attention, albeit at the expense of forfeiting partial context information.
Recently, there have been two lines of methods aimed at efficiently extending the pre-trained con- text length of LLMs, both centred position embedding (PE). The first line of methods, known as PE scaling, are proposed to effectively extend the context window of LLMs integrated with Rotary
'This work was done during the internship of Guanzheng Chen at Alibaba DAMO Academy. 1 CCorrsponniig authors.

---

Published as a conference paper at ICLR 2024
FF1
24
Pl (c-16)
(t=16)
Yam-16k (t-32)
CLEX 16k (t-1G)
PPL
25
0.0
10k
20k
30k
40k
50%
60k
Eval uation Sequence Length
Figure 1: The PPLs of our CLEX and various baselines tested on 64k context length.
Position Embedding (RoPE) (Su ct al. 2022). They allow LLMs to access longer context by scal- ing either position indices (Chen et al. 2023) or frequency basis (Rozière al. 2023 Peng al.
2023] of RoPE, demonstrating remarkable performance in long-context applications. However, such methods are designed for extending the context length corresponding to fixed scaling factor, which either restricts their ability to extrapolate to longer sequences (when using small factors) or impairs performance even within the native context window (when using large factors) as shown in Fig-
ure On the other hand, length extrapolation methods (Press et al. 2022 Sun et al. 2023 Chi ctt al. 2022 2023). typified by ALiBi (Press et al. 2022). strive to achieve test-time context length extension (i.c., "training short, testing on long" ) by substituting position embeddings with addi- tional biases, where the biases encode positional information to the attention scores. Despite their impressive capability in language modelling, ALiBi-like methods usually struggle in the practical
tasks requiring long-context dependency (Pal et al. 2023) (also see
In this work, present Continuous Length EXtrapolation (CLEX), novel approach that effi ciently extrapolates the context window of LLMs through continuous PE scaling. Concretely, we propose unified view of PE scaling via generalising the PE scaling methods to the transition of frequency basis. Upon it, we formulate the PE scaling as continuou dynamical system, which models the transition of frequency basis through the continuous dynamics over the length scaling factor. We argue that previous PE scaling methods, training models using fixed (discrete) scaling fac- tors, overlook the progressively continuous dynamics over the gradually This ensnares themselves in the aforementioned dilemma between extrapolating the length and pre- serving the performance within shorter lengths. contrast, our CLEX exploits a neural ordinary differential equation (ODE) (Chen et al.) 2018), parameterised by up-and-down projection layer with lightweight parameters to learn these continuous dynamics, enabling fine mee-graiedd extending to long context. More essentially, by extending the dynamies beyond training length, CLEX empowers models to progressively extrapolate to longer contexts even when trained with short sequences.
process.
CLEX can serve drop-in component for RoPE-based LLMs, such LLaMA TTuvron al. 2023a b) and GPT-NeoX (Black et al., 2022), with negligible overhead in computation and param- size. We evaluate the performance of CLEX two datasets: (1) subset of RedPajama- Book (Computer) 2023) for long-context language modelling, and (2) LongBench (Bai et al.) 2023) for long-context practical tasks. Empirically, CLEX demonstrates remarkable length extrapolation ability in language modelling, which can extend the context window to more than 4x training length without any performance deterioration. For example, ILL.MA--7B trained with CLEX on 16k con- text length achieves comparable perplexities when testing on 16k and 64k tokens, respectively. By scaling the base model scale from 7B to 13B, CLEX exhibits an expanded extrapolation scope from 4x to almost 8x training length. To be complementary, we also conduct instruction tuning (Weil et al. 2022) with the proposed CLEX on the sequences of 4k length. The resulting model, when evaluated the LongBench benchmark, is on par with current state-of-the-an open-source mod- els trained on context lengths up to 32k. These findings underscore the effectiveness of CLEX in extrapolating context length, signifying its efficiency for developing long-context L.IMs.
2

---

Published as a conference paper at ICLR 2024
2
PRELIMINARIES
2.1 ROTARY POSITION EMBEDDING (RoPE)
Rotary Position Embedding (RoPE) (Su et 2022) has recently emerged as the most prevailing po- sitional encoding method in open-source LLMs like LLaMA. It integrates both absolute and relative positional information for Transformer models. Given position index m [1, L]. RoPE injects the absolute positional information into æ E Rd via the transformation f R - IR as:
(x, m, 0) Ra. 2;
where 0 E R [2///2] is the rotation frequency basis and 0, 10, 1..0000-500 2//d Romm E Rdd is block diagonal matrix formed by the elements
Re m) = cos mf, sin mQ, for i = 1, - Ld/2].
(2)
sin mf, cos m0;
The transformation in Eq. 1I is applied to the query and key vectors during self-attention. When calculating the attention score for the query vector qm Rd at position m and the key vector k, E Rdd at position n, we have
(Re (Ro 9m Ro
(3)
Hence, the relative positional information Rg.n-m is implicitly incorporated into the attention scores. However, even given the relative information, L.LMs trained with RoPE, egg. LLaMA, still cannot achieve reasonable performance beyond the pre-trained context length.
2.2 PE SCALING METHODS
To extend the context length L, several strategies are proposed to adjust the position embedding by scaling either the position index or frequency basis 0 in Eq. (I). Formally, we define t = L' / as the length scaling factor where T' denotes the desired extended length. While Chen et al. (2023) introduces scaling the index by Position Interpolation (PI)
m (a, m, 0) = f(æ, T 0).
(4)
This strategy maintains the position indices within the range [1. L], while effectively extending the processed range [1, L] by minimal fine-tuning steps on t L sequences. On the other hand, Peng et al. (2023) proposes Yarn, a.k.a. NTK-Aware Scaled RoPE, extends the context window by
frequency basis scaling (FBS). This strategy is similarly utilised by CodeLLaMA (Rozière et al.
2023 Formally, the FBS methods are denoted as
f!TS( (2, m, 0) f(x, m, 0c),
(5)
where 0, is the scaled frequency basis. Specifically, 0,,; = 0; (())((1)))))))) (1)-2//(d-2) in Yarn and Oz,i 0: 100 in CodeLLaMA.
3 METHODOLOGY
This section demonstrates the details of CLEX. We first generalise the PE scaling to continuous dynamical system unified manner (see $3.I). On top of the continuous dynamical system, CLEX employs the neural ODE, parameterised by an up-and-down projection layer, to adaptively learn the continuous dynamics during PE scaling (see $3.2). $3.3 we introduce the training strategy of CLEX that distributes the continuous dynamics beyond the training sequence length, thereby enabling the generalisation of continuous PE scaling achieve the length extrapolation.
3.1 POSITION EMBEDDING SCALING: UNIFIED VIEW
Given the various methods that extend models' context length through position indices scaling and frequency basis scaling, we first show that the transformations applied to position indices are essen- tially casting the frequency basis, which is formalised in Theorem
3

---

Published as a conference paper at ICLR 2024
4k
8k
12k
16k
8k
4k
16k
02
12k
1.0
2.0
3.0
4.0
1.0
2.0
3.0
4.0
t
t
Figure 2: The graphical model of discrete PE scaling (left) and our continuous PE scaling (right).
Theorem 1. For the transformation T to position index m, there exists equivalent transforma- tion't T to frequency basis 0 in Eq. (7) namely
f(x, m, 0) f(x m, T о 0),
where T = = and denotes the element-wise transformation.
Proof. From Eg. ( we have f(x,, T m, 0) Ro, m and f(x, m, T ④ O) RT a. For
any T = [
cos Tmo; sin T mm, [oos ml m 0.)
-sn m( +..
(Ro.Tm);
=RTooom))
sin T Tm8;
cos Tm0;
sin ml T о 02)
cos cosm( O 0;)
Hence, there is unified form for PE scaling that consistently projects the frequency basis by axt):
fz m, 0) = f (æ, m, a<t))
(8)
where (t) is sangle-variable transformation defined over the length scaling factor t. Through this unified formulation, (Chen et al.) 2023) and Yarn (Peng et al. 2023) can be viewed as the special
cases when plugging (1) [1/4]/ and a Yarm
2i/((d 20 d/2 into Eq. respectively.
Note that 0, (t) @ 0 denotes the scaled frequency basis at context length of t T. and 0 (namely ax1) = 1). As illustrated Figure 2 this indicates a progressive chain across discrete values that
a(t) (t) z(1) a(t) z(t 1) log a 1)'
where 2(t)= O,.
By continuizing the progressive chain, we can formulate the PE scaling as a continuous dynamical system, with the continuous dynamics of frequency basis (l)//dt
dz(t) dooo (tt) di
(10)
dt
In essence, recent PE scaling methods, concentrating on manually formulating the a<))), are equiva- lent to applying various dynamics for frequency basis that enable models to adapt to longer contexts.
3.2 CONTINUOUS PE SCALING VIA NEURAL ODE
Even given the continuous dynamics of frequency basis, previous methods are inherently designed for extending the context length at discrete values. For example, PI (Chen et al., 2023) fine-tunes the model on a specific scaling factor to extend the context window length to f L. One potential issue of these methods, as depicted in Figure I is that they are susceptible to overfitting to the speci- fied frequency basis, leading to either poor extrapolation ability to longer lengths beyond training or performance drops within short lengths, both in some cases. Therefore, CLEX aims to build a continuous PE scaling, which induces the model to adapt the frequency basis corresponding to a continuous scope of as illustrated in Figure (right).
4

---

Published as a conference paper at ICLR 2024
Recall that previous PE scaling, corresponding to manually defined axt),, implies the constant dynamics in Eq. (10). In method, we utilise variable function g: 12//2 - 12/2 to model the dynamics, hence towards a more general and flexible view as:
dz(() g(=(1),))
(11)
dt
By restricting the function to be associated with the latent states (tt), is capable of capturing the fine-grained changes of frequency basis during the length-extending process. However, it is non- trivial to manually define the =(w)-a function g. Here, we directly parameterise the function using the neural network Φ. Therefore, for any L' E [1, there neural ODE modelling the scaling of frequency basis as
x(t)= + S, ()(=,
(12)
where the frequency basis at the length I I. can derived by Or = exp(z(1')). More specifically, we adopt an up-and-down projection as the neural network, expressed as:
g (=)) Wdown a (Wap =()) + &t,
(13)
where Wap E x and Wownn E R Adx are the transformation matrices, of which the parameters are determined by the amplification factor X; is the SiLU activation function and & is the scalar embedding typifying the scaling procedure at factor of t. Here, we adopt the constant dynamics of Yarn as the & for speeding up convergence, namely
& dlog Yar dt
2i
d/2
(14)
(d - 2)
i=1
3.3 CONTINUOUS LENGTH EXTRAPOLATION: TRAIN ON SHORT, TEST ON LONG
Continuous PE scaling can serve more adaptive and flexible PE scaling method to extend the context length to given training length Train, Unlike the previous PE scaling methods built on larger scaling factor, which would lead inferior performance the lengths corresponding to smaller counterparts, the continuous PE scaling would enable non-destructively generalisation to larger scaling factors via adaptive continuous dynamics. Therefore, by simply extending the contin- uous dynamics beyond the factor t= Traain /L. during training (where we denote the desired scaling access the continuous length extrapolation (CLEX) method, which achieves
factor as {Train), the capability of "training on short, testing on long".
Moreover, learn the neural ODE in Eq. 12 for continuous I, randomly sample l' [1, [Train) for each training step, enabling the model to adapt to the broad scope frequency basis without over- fitting a specific one. Note that the frequency basis bound with the position index in Eq. (I. This reveals the aforementioned training involves inconsistency between the frequency basis and position indices: the frequency basis is varied corresponding the E (1,ETrain) while the position indices are fixed as (1,2..... I TTTain. }. Here, we propose the position extrapolation strategy to address this consistency. Contrary to PI, which shrinks the position indices into the context length, we enlarge the position indices 21,2.... ..., train } of the trained sequences up the range [1,'- for each train- ing step. The position indices can be acquired by uniformly scaling to {1-8, 2-8, L [.Itain..) s) where = L alternatively, by randomly sampling L Train of indices from [1, L). Empirically, we found that random sampling generally performs better. More discussions can be found in $4.2
During inference, the ideal scenario is to acquire the frequency basis corresponding to each sequence length. However, this approach is computationally demanding. To improve efficiency, we first cache some frequency basis derived from IS for K discrete t values as {tx lk E [1, K])- For each sequence employ the frequency basis corresponding the nearest
with length of Infer during inference, upper bound within 4. T. for k 1,..., K. Through this, our method introduces negligible time cost compared to naive inference of I.LMs.
4 EXPERIMENTS
In this section, we conduct a thorough evaluation of CLEX's performance in terms of handling long contexts and its extrapolation capabilities. We compare approach against other methods covering
5

---

Published as a conference paper at ICLR 2024
Truin
Evaluation Length
Methods
Length
4096
(4k)
8192
(8k)
16,384 (16k)
32,768
(32K)
65,536 (64k)
PPL.
ACC.
PPL
ACC.
PPL
ACC.
PPL
ACC.
PPL
ACC.
LLaMA-2
4k
6.04
58.18
20.54
44.50
>100
22.43
>- >1000
12.70
> 1000
10.64
CodeLLLMMA
16k
7.60
54.88
7.40
55.19
7.33
55.30
15.12
44.70
52.02
31.16
Naive FT
16k
5,98
58.83
5.93
58.91
5,91
58.58
18.31
43.04
>100
26.05
PI
16k
5,90
59.05
5.71
59.44
5.72
59.87
6.05
58.5
8.75
52.02
Yarn (t=16)
16k
6,50
57.28
5.71
59.57
5.73
59.87
5.99
58.13
8,51
52.62
Yarn (f=32)
16k
6.61
57.12
5.94
58.27
5.96
58.04
6.08
57.73
6.22
$7.98
Cl.-Scating
16k
24.99
37.84
5.86
59.08
5.87
59.05
10.56
50.47
41.09
34.16
6.34
58.01
6.39
57.8
57.78
57.47
56.44
RandomPos
4k
5.88
58.49
>> 1-00
34.23
>1000
18.27
>1000
9.31
>> 1000
7.44
4k 5.86
59.21
5.70
59.62
5.87
58,93
14,53
47,55
30,51
35,33
CLEX
8k
5,98
58.75
5.78
59.44
5.71
59.64
5.99
58.66
11.74
47.50
16k
5.88
59.21
5.68
59.73
5.52
60.28
5.55
60.06
5.64
59.94
Table 1: Perplexity (PPL) and next-token -ereteeeeeeeeeitionnnn accuracy (ACC.) on language modeling with evaluation lengths from 4k 64k. We train the LLaMA-2-7B using length extrapolation methods on 4k length and PE scaling methods on 16k length, while reporting the results of CLEX trained across 4k, 8k and 16k. CL-Scaling denotes training LLaMA-2- -2 with the scaling method of CodeI_LaMA but using training data. The training loss curves are depicted in Figure
both length extrapolation (i.c., ALiBi (Press et al. 2022) and random positions (denoted as Random- Pos) (Ruoss et al. 2023) ) and PE scaling methods (i.e., PI (Chen et al. 2023) and Yarn (Peng al.
2023 We primarily conduct experiments on the LLaMA-2-7B model. For the language mod-
elling, train our model and the baselines on 2B tokens extracted from Redpajama-Book (Com-
puter 2023), which is collected from Pile-Books3 (Gao et al.] (020) and PG-19 (Rae et al., 2019) datasets. The performance of the models is assessed based on perplexity and wext-tooen-prediction accuracy, with evaluation sequence lengths up to 64k. Furthermore. we conduct instruction tun- ing for LLaMA-2-7B using CLEX on the UltraChat dataset (Ding et al. (0236). The evaluation is performed on the LongBench benchmark (Bai et al. 2023), where we compare our model with GPT-3.5-ture and other LLaMA-2 based open-source models designed for handling long context.
Further details about baselines and training configuration will be discussed in Appx. A as well as more experimental results and ablations in Appx. S
4.1 LONG-CONTEXT LANGUAGE MODELLING
CLEX achieves length extrapolation. We first report the experimental results baselines and CLEX on language modelling, with the evaluation length from 4k to 64k. As shown in Table I our CLEX consistently demonstrates remarkable performance in length extrapolation, being able to extrapolate the context length to more than 4x training length without any performance drops. Taking CLEX-4k as an example, its PPL on 4k sequence (training length) is comparable to that on 16k sequence (5.86 vs. 5.87). When evaluated on the sequences within 16k, CLEX-4k is par with or even better than all of the compared methods trained on lengths up to 16k. Moreover, with the increase in training length, our CLEX not only exhibits promising generalisation capability to very long contexts (up to 64k) but also guarantees performance on short sequences.
We also found that discrete PE scaling methods (i.e., PI and Yarn) have self-extending property: training with scaled frequency basis equips the model with the ability to extrapolate to further- scaled counterparts (see Appx. (B.2 for more discussions.).
depicted in |Figure T] however, the extrapolation capability of these methods is limited, accompanied by significant performance decline even within the naive context length. This indicates the inherent challenge of achieving a delicate balance between extrapolation to longer lengths and performance maintenance within short lengths when using the discrete scaling factor. In contrast, CLEX tackles this issue via Iearnable continuous dynamics, providing more fine- -ne-gra extrapolation while preserving the performance for the internal context.
Note that ALiBi may extend further than CLEX trained on 4k sequences (though typically pro- ducing inferior results), our experiments reveal that these gains may come at the cost of long-term information, leading underperformance in long-context practical tasks (see $4.3 for more details).

---

Published as a conference paper at ICLR 2024
30
CLEX-7B-4
CLEX-7B-8K
CLEX-13B-4K
CLEX-7B-4k 16k
PPL
20
20
CLEX-7B-4k 32k
CLEX-13B-4K 16k
CLEX-13B-4K on 32k
10
10
4k
8k
16k
32k
64k
0.25
0.5
0.75
1.25
1.5
1.75
2
Evaluation Sequence Length
Training Dataset Size (Billion tokens)
Figure 3: Left: The PPLs of CLEX on different evaluation sequence lengths with 7B and 13B parameter sizes. Right: The PPLs of CLEX cross variable training data size with different parameter sizes and evaluation lengths.
Continuous Diserete
Sampling Strategy
À Factor
Continuos
Random
30
A=
Discrete
Uniform
2
2
PPl
sampling
20-
20
10
4k
16k
Evaluation Sequence Length
Figure 4: The ablation studies for continuous dynamics, sampling strategies and À factor in Eq.
The scaling law for the extrapolation ability of CLEX. To investigate the effectiveness of CLEX over the scale of the base model and training data size, we further port our method to LLaMA-2- 13B. As depicted in Figure 3 when trivially extending the base model scale from 7B to 13B, CLEX demonstrates an increased capacity to extrapolate to longer context lengths. Specifically, the extrapolation ability of CLEX-13B trained 4k length approaches that of CLEX-7B trained 8k. While the training data scale, more surprisingly. does not significantly impact the extrapolation capability of CLEX. Models trained with 0.25B 2B tokens with 4k sequence length achieve comparable PPLs when evaluating on 16k or 32k lengths in Figure 3] indicating the negligible margins from the larger training data size. This also implies that CLEX can efficiently extend the context length of LLMs through minimal training steps resembling PI and Yarn.
Based on these findings, we propose scaling law for CLEX: to scale the context length of LLMs to moderately desired lengths (e.g., 16k + 64k), one should proportionally enlarge the training sequence lengths (e.g.. 4k - 16k). For scaling the context length up to considerably long lengths (e.g., 200k), the parameter size of the base model should be correspondingly increased while maintaining the training length, since the contexts may take more footprints than model parameters. Note that scaling the training data does not directly affect the extrapolation ability of CLEX, but may be implicitly incorporated when scaling the base pre-trained LLMs.
4.2 ABLATION STUDY
We conduct three types of ablations to investigate the efficacy of the components in CLEX:
Continuous dynamics. To learn the continuous dynamies using neural ODE, we adopt distinct training approach that involves sampling the scaling factor for each data batch. Here we seek to explore if the exceptional extrapolation ability CLEX is solely derived from the variable rather than the continuous dynamics. We employ the discrete Yarn method with = 16, that undergoes the same training procedure of CLEX but removes the ODE parameters, serving as a discrete baseline. In Figure 4 (lef)), we discover that the discrete approach equipped with the icantly underperforms our CLEX, indicating the essentiality of the learnable continuous dynamics in CLEX for accessing the extrapolation ability.
signif-

---

Published as a conference paper at ICLR 2024
Single-Doc QA
CPT 1.
Longhhe
Code Completion
Multi-Doc QA
40-
Avg. All
30-
Synthetic Task
Summarization
20-
Few-shot Learning
4k
16k Trained Sequence Length
32k
GPT-3 5/Turbo
LongChen--11 5-78 Codel Jenak-BB--BBB 7B-t CLEX-7B-44 (Cuss)
V...maa1..--BB..kk
Beichran 138 (ALIBO
Left: the average scores for each domain of tasks in LongBench. Right: the average scores of all tasks corresponding to the training length of each model. Note that CLEX trained on 4k sequence length but directly tested on 16k context length without truncation.
Figure 5:
Position extrapolation. We adopt the position extrapolation strategy, which extends the scope of position indices in training sequences by sampling from a broader range, to reconcile the inconsis- tency between frequency basis and position indices during the training process. In this study, examine the impact of various sampling strategies (uniform or random) and contrast them with the naive position indices. The results in FFiuure 4 underscore the efficacy of position extrapolation in CLEX, without which the extrapolation ability of models declines significantly. Furthermore, ran- dom sampling slightly performs better than uniform sampling, so we adopt across all experiments.
The parameter scale of ODE. We also study the impact of parameter size of the neural ODE in CLEX. The parameter size is determined by the X, namely the amplification factor Eq. (3). In Figure 4| plot the results of CLEX with À 1,2, 4, where they achieve similar performance. Note that the parameter size of neural ODE in CLEX is quite small even when À = 4. as the dimension d in Eq. (13 usually equal 128. Although it is possible to enhance CLEX with larger X (e.g., 32), we set the x=1 in all experiments for the minimal effect on inference latency.
4.3 EVALUATION ON LONG-CONTEXT BENCHMARK
To ascertain the comprehensive performance of CLEX in real-world scenarios, we further conduct an evaluation on the zero-shot LongBench benchmark. This benchmark encompasses broad range summarization, and code completion, where the evaluation
of tasks, such as question-answerii length ranges from 5k 15k. We perform a pilot instruction tuning for LLaMA-2-77 by employing CLEX the UltraChat dataset, with sequence length of 4k. During inference, harness all models to tackle the context length of 16k, thereby ensuring the comprehensive exploitation of contextual information in the tasks. As depicted in Figure 3] present the average scores of each domain in LongBench for CLEX, in comparison to the GPT-3.5-Turbo- model and strong open- and CodeLLaMA-7B-16k.
source LLMs like LongChat-v1.5-7B-
Generally, when trained with sequences of 4k length, CLEX holds its own against any open-source LLMs that trained on lengths up 32k. In the specific domains of Summarization, Few-shot Learning, and Code Completion, CLEX on LLaMA-2-7B remains competitive with or even sur- passes the GPT-3.5-Turbo-16k. We note that the Baichuan-13B-44 pre-trained with ALiBi (Press) et al. 2022), demonstrates marked underperformance on the LongBench although with a larger parameter size. Additionally, similar poor results are achieved by ALiBi when applying it upon LLaMA-2-7B using the same training procedure as CLEX (see Appx. {B.5). This could likely be attributed to ALiBi's overemphasis on local context through the attention bias, which, while advan- tageous for language modelling, restricts access to long-context information in practical tasks, In contrast, CLEX directly extends the context length of LLMs without imposing any constraints on context, which consistently achieves superior extrapolation ability both language modelling and the LongBench. This substantiates the considerable potential of CLEX to serve as the state-of- of-the approach for extrapolating the context length of L.LMs to excel in long-context applications.
8

---

Published as a conference paper at ICLR 2024
In addition, highlight that our CLEX merely introduces minuscule inference latency. Given context length of 16k in LongBench with a generation length of 512, the generation throughput between our CLEX-7B and LLaMA-2-77 2-BB is comparable (27.8 tokens/s 28.3 tokens/s, in single A 1,0)), when using the cache mechanism introduced in $3.3
5 RELATED WORK
Hierarchical Architecture Sparse Attention. To overcome the quadratic complexity of atten- tion, Dai et al. 2019) proposes the Transformer-XL. that handles the long sequence at segment level by Transformer, with these segments interacting through a recurrence mechanism. The Recurrent Memory Transformer (Bulatov et al. 2022) refines this mechanism by incorporating special memory tokens into the recurrence, which is capable of scaling the context length to the millions (Bulatov) al. 2023). On the other hand, Child al. (2019); [Beltayy et al. (2020) proposed using the sparse attention to circumvent the full access to the long sequences, hence reducing the complexity. The sparse attention has been adopted by Ding et al. (2023a) to scale the context length of transformers into the billions. However, these methods sacrifice the utilisation of the entire sequence during at- tention, resulting in an inevitable loss of some contextual information. Additionally, modifications to the model architecture make these methods challenging to apply to existing pre-trained LLMs. Conversely, our CLEX serves as a drop-in component for LLMs, can efficiently extend the capacity of models to tack the entire long sequences without explicit drops of context information.
Length Extrapolation. Building on the foundation laid by ALiBi (Press et al. 2022 a series of works (Sun al. [023 Chi et al. 2022 2023) seek train the Transformer-base models on short length, while directly testing on longer counterparts. These methods substitute the position embedding with bias introduced into attention scores, thereby incorporating positional information. Notably, the bias typically gives higher profits to closer tokens. This mechanism intuitively am- plifies the local context for each token at the expense of distant information. Consequently, these eetgth--ttrapoittion methods encounter challenges in effectively handling long contexts in practi- cal applications (Pal et al., 2023). However, our CLEX demonstrates remarkable effectiveness in practical tasks such summanization, indicating the de facto extrapolation ability for applications.
Position Embedding (PE) Scaling. Recent research has sought to extend the context length of
Transformers through the scaling of the extensively employed RoPE. Specifically. Chen et al. (2023) proposed position interpolation, a method that efficiently extends the context window by scaling the position index within RoPE. In similar vein, Peng et al. (202);; Rozière et al. (2023) opted to scale the frequency basis, achieving superior performance. However, these methods necessitate training (or fine-tuning) on the desired extended length. As a result, they exhibit a limited ability to extrapo- late beyond the trained length and even suffer from performance drops within the shorter lengths. In CLEX, we generalise the discrete PE scaling to a continuous counterpart, hence uniformly extrapo- lating the context length of LLMs while preserving the performance within short lengths.
6 CONCLUSION
We have presented the Continuous Length EXtrapolation (CLEX), novel approach that efficiently extrapolates the context length of Large Language Models (LLMs) to over 4x the training (fine- tuning) length without any decline in performance. CLEX utilises the neural ODE to learn the continuous dynamics over the length scaling factor during PE scaling, hence enabling fine-grained extension for the frequency basis in the RoPE. We conduct thorough experiments to investigate the effectiveness of CLEX compared to a variety of strong LLMs, covering the language modelling task and LongBench benchmark. The experimental results have demonstrated the exceptional extrapola- tion ability of CLEX, where our CLEX trained with a sequence length of 4k holds the potential to remain competitive to any open-source long-conted LLMs (e.g., CodeLLaMA) trained on lengths up to 32k. These results highlight the potential of CLEX as a state-of-the-ar approach for effi- ciently extrapolating the context length of L.LMs, paving the way for advancements in long-context applications. By scaling the base model size up, we found CLEX can be correspondingly enhanced and subsequently is capable of extrapolating the model to a longer context length. This property indicates the tempting effectiveness of CLEX in the era of LLMs.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work was substantially supported by DAMO Academy through DAMO Academy Research Intern Program. Shangsong Liang was supported by the National Natural Science Foundation of China (Grant No. 61906219) and the Mohamed bin Zayed University of Artificial Intelligence, United Arab Emirates.
REFERENCES
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. 2023. URL https::/aaxtaii org/
abs/2308 14508
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05150, 2020. /RR.atttpss/aaxx...oor////////// 0000.//00000000000000. 05150
Sidney Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Ho- race He. Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, Usvsn Sai Prashanth, Shiv- anshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. GPT-NeoX- 20B: An open-source autoregressive language model. In Proceedings of BigScience Episode #5 Workshop on Challenges & Perspectives in Creating Large Language Models. Association for Computational Linguisties, May 2022. URL https: hataa/aaaannttoooooooooooo org/2022. bigscience-1...
Aydar Bulatov, Yuri Kuratov, and Mikhail Burtsev. Recurrent memory transformer. In Al- ice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neu- ral Information Processing Systems, 2022. URL https: Uyn. 3iPhksa.
Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Scaling transformer 1m tokens and beyond with rmt. 2023.
11062
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary dif- ferential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- ResacBianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Cur- Associates, Inc., 2018 URL httpppppppppccccagggss.... nww..nn.....euussrreeee.r..r cc/paprr files/ paper/2018/ / 66/666933 http:///e......oocooooeeeeee6663 pappa
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. 2023. URL https://azxiv.org
abs abs/2306.. 15595
Ta-Chung Chi, Ting-Han Fan, Peter J Ramadge, and Alexander Rudnicky. Ker- ple: Kernliized relative positional embedding for length extrapolation. Advances in Neural Information Processing Systems. Curran Associates, Inc., 2022. URL
Ta- Ta-Chung Chi, Ting-Han Fan, Alexander Rudnicky, and Peter Ramadge. Dissecting transformer length extrapolation via the lens of receptive field analysis. In Proceedings the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Associa- tion for Computational Linguistics, July 2023.
acl 10ll-oongg-- wwo.onn... 756
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. 2019. http:p/://aaxivv..tr////////0011 10509.
Together Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023.
http://pp../gggggtoooco/ooooo/cr jama
10

---

Published as a conference paper at ICLR 2024
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. Association for Compu- tational Linguistics, July 2019. URL htt pt:pp::t///t...mm ptss://aclanttoo
o wgw..rgg//9.111000
Tri Dao. FlashAttention--:: Faster attention with better parallelism and work partitioning. 2023.
URL
08691
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Process- ing Systems, 2022. /
Jiayu Ding, Shuming Ma. Li Dong, Xingxing Zhang, Shaohan Huang. Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. 2023a. URL
fttps://arxiv: org/ hbtgg//bs//3// 02486
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations. arXiv preprint arXiv:2305. 2023b. URL
wwwp:vvvorg/ao. /bs/
2305. 14233
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 di- verse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. URL https:
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gi- anna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William Sayed. Mixtral of experts. 2024. URL https:: awwiv. abs 04088
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. Yoshua
Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. URL http:
6980
OpenAl. GPT-4 Technical Report. 2023. URL
Arka Pal. Deep Karkhanis, Manley Roberts, Samuel Dooley, Arvind Sundararajan, and Siddartha Naidu. Giraffe: Adventures in expanding context lengths in 1lms. 2023. URL https://
arxiv.org/abs. 22008 10882
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole Yarn: Efficient context window extension of large language models, 2023. URL htRl:htss:://// /ttp///rxvv/../////////////ae.l/ wwwpwawwwww...////////0020
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:://ppenrevieee net forum (it.rrmmmid-BB0PPGGOG
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL
pettps://axxvv.orgge
05507
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usumier, Thomas Scialom, and Gabriel Synnaeve. Code llama: Open foundation models for code. 2023. URL https:://
arxiv. www./rgg.........ooo. /2008 12950
11

---

Published as a conference paper at ICLR 2024
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Ben- nani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. In Proceedings of the 6l Annual Meeting of the Association for Computational Linguistics ( Volume 2: Short Papers), Association for Computational Linguistics, July 2023. URL. https://accaatthie otg. w0w. -0w.0223...ll-hh.......
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. 2022. htlhhttss///xxiv.coo 21 09864
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaud- hary, Xia Song, and Furu Wei. A eeteth-cttraplltablle transformer. In Proceedings of the 61st Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). As- sociation for Computational Linguisties, July 2023. URL htt.hhh//acclleeee hntho ogy org/
2023.ac1 616.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. 2023a. /tttt::s/////xxiiii wttp://wwww...r///////////2 13971
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fermandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Ilama 2: Open foundation and fine-tuned chat models. 2023b. htLhhtts:// /ttpppp:://axi...... wwweeiiv.orggabee. wttp/rra/abb./ooo/eee 09288
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe V Le. Finetuned language models are zero-shot learners. In Interna- fional Conference Learning Representations, 2022. URL / net./
Forum?id
12

---

Published as a conference paper at ICLR 2024
A EXPERIMENTAL DETAILS
A.1 BASELINES
We compare our CLEX with a variety of baselines in language modelling, all of which are trained with the subset of 2B tokens from Pile-Books3 and PG-19 except LLaMA-2 and LaMA that are evaluated by the open-source pre-trained checkpoints. We train the PE scaling and length extrapolation methods with 16k and 4k sequence lengths, respectively. Specifically, we use the scaling factor = 4 to train models using 16k sequence length for PI following Chen et al. (2023). While for Yam, we train models with the scaling factor of 16 and 32 (but also using the 16k sequence length), respectively, following the settings in Peng et al. (023). For the ALiBi, we directly remove the RoPE L- and train models with the attention bias. For the random positions, we sample the position indices from [1, 16k] while train models with 4k sequence length. We also utilise the PE scaling method from CodeLLaMA to train LLaMA-2-7B using our training data, which is denoted CL-Scaling.
A.2 TRAINING DETAILS
We use a subset from the Redpajama- (Computer 2023) as the training data for language modelling, which is collected from Pile-Books3 and PG- 19 datasets and the subset comprises ap- proximately 2B tokens. We set the learning rate of 2e-5 for all models, which are optimised by Adam optimiser (Kingma & Ba 2015). The batch size set to 64k tokens for 7B models and 16k tokens for 13B models. The maximum desired f during training in CLEX (namely & in $3.3) is set as 16 for LLaMA-2. We utilise the Flash Attention (Dao et al.) 2022 Dao 2023) to support efficient training. The amplification factor of ODE layer X is set as for all 7B models and 2 for 13B models. For the instruction tuning, we train our model using the UltraChat dataset for epoch, starting with the checkpoint after the training of language modelling. The training procedure of CLEX is shown in Alg. 1
Algorithm Training Procedure of CLEX
1: repeat
2: Given a batch of sequences B with length of 7, Train;;
3:
sample 1' ~ [
[min);
4:
calculate =(t') by Eq.
5:
Or = ex=exp((2(( (t'));
6: sample (randomly or uniformly) t' L Train position indices from [1, t Train I Train]
7:
calculate RoPE with O2 and sampled position indices;
8: train the model with RoPE B with language modelling objective.
9:
until converged
A.3 EVALUATION DETAILS
For language modelling, evaluate all models with similar subset the training set but con- taining 20 million tokens. For different evaluation lengths, we group the test set into corresponding lengths of sequences. We found fixing the scaling factor the training one would hamper the general performance and extrapolation ability of baselines, therefore, we proportionally enlarge the scaling factor when the evaluation length is beyond the training length. For example, when evaluat- ing on 32k length, we set the scaling factor of PI as 8 rather than the training factor of 4. We also adopt the log scaling trick[] for all baselines and our model, where we scale the attention scores by the max ( 11 log Tuis [ }. We found it can improve the performance when evaluating on lengths beyond training length as shown in Figure 8 For the LongBench, we follow the setups from Bai et al. (2023), where the decoding method is set to greedy search. We set the 4 in $3.3 as {1, 3, 4J. to cache the frequency basis corresponding to (4k, 8k; 12k, 16k).
https://spaces, hwt.p//cc.e....c//ccccc.c...
13

---

Published as a conference paper at ICLR 2024
Base= 1e4 (original)
Sinusoidal Frequency
0.8
CLEX (t 4)
CLEX (t 8)
CLEX (t 16)
0.6
Base=1e6
0.4
0.2
0.0
20
40
60
Dimension
Figure 6: The sinusoidal frequency over the dimension of the frequency basis. The lower dimension denotes the higher frequency.
Pt: 16k (t=4)
24
PPL
23
10k
60k
Eval ualion Sequence Length
Figure 7: The results of PI-16k (trained with 4) 64k sequence length, while evaluating with t t= 4 and = 16.
B ADDITIONAL RESULTS
B. VISUALISATION OF LEARNED FREQUENCY BASIS
In Figure 6 we show the frequency basis according to different sequence lengths after the training of CLEX LLaMA-2 with 4k sequence length. The plot reveals that CLEX tends upscale the high frequencies at certain dimensions while simultancously downscaling some others. More surprisingly, we have observed that the frequency bases associated with different / values in CLEX exhibit an isotropic behaviour, that the dimensions where downscaling and upscaling occur are sim- ilar across different é values, with larger values resulting in further scaling. This finding may contribute to the design of heuristics PE scaling methods without training.
B.2 THE SELF-EXTENDING PROPERTY OF PE SCALING METHODS
Previous PE scaling methods usually evaluate models with a specified scaling factor during training. We found this would significantly hamper the performance evaluation beyond the training length. Hence, we adjust the scaling factor the previous PE scaling method corresponding to the evalu- ation lengths. In Figure 7, P1-16k (trained with i = achieves non-trivial results when evaluating
14

---

Published as a conference paper at ICLR 2024
Train
Evaluation Length
Models
Length
32K
64K
128K
256K
CLEX-Phi-2
32K
5.11
5.17
6.55
>10
CLEX-MixtraI-8x7B
32K
2.56
2.53
2.57
3.78
Table 2: Perplexity (PPL) of CLEX-Phi-2 and CLEX-Mixtra-8XX7B on language modelling with evaluation lengths from 32k to 256k.
log Scaling
Constant Dynamic &t.
30
Scaling
W/ &
30
No Scaling
W/O
PPL
20
20
10
10
4k
8k
16k
32k
64k
4k
8k
16k
32k
64k
Evaluation Sequence Length
Figure 8: The ablation of log scaling (left) and constant dynamic & (right) for CLEX trained on 4k sequence length.
with ł 16 64k length, which significantly outperforms the naive evaluation using = 4. Note that Yarn is naively trained with larger so it is not necessary to adjust the scaling factor for Yarn during evaluation. When the evaluation lengths become longer, however, self-extending property would benefit the Yarn as well. Regarding these, the extrapolation ability from the self-extend previous PE scaling methods is still far weaker than CLEX as illustrated in Figure
believe the
B.3 RESULTS OF MORE MODELS
We also extend CLEX to more models beyond LLaMA-2, i.e., Phi--and Mixtral-8x7B (Jiang et al. 2024). The results in Table further support the general effectiveness of CLEX, where the context lengths extrapolated to 4x training length across various LLMs.
B.4 ABLATION ON SCALAR EMBEDDING
The learnalle dynamic in Eq. 13] based the constant dynamic &. from Yarn. To conduct ablation on the impact of & we remove & in Eq. (13) and train the ILLMMA2-7B on 4k sequence length that undergoes the same training procedure. From Figure
found the training of CLEX without & would not hamper the extrapolation ability (or even better). However, as depicted in Fig- ure 10 training with & can speed up the convergence and produce more stable loss curve. This is essential for training larger LLMs on longer sequence lengths, we integrate the & into the learnable dynamics.
B.5 RESULTS OF LONGBENCH
We list the numerical results of our CLEX trained on 4k compared with variety of baselines trained up to 32k, covering single-document question answering (Table 3), multi-document question an- swering Table 4), summarization Table 5), few-shot learning (Tbble 6). synthetic (Table 7). and
httts::/huggingfaaee ccw.ac/c...ccc.ccooommmmooo /trr
15

---

Published as a conference paper at ICLR 2024
Naive-FT-18
4.0
PI-16k
Yarn-16k (t=16)
CL-Scaling 16k
3.5
ALiBi-4k
RandomPos-44
Loss
CLEX-4K
3.0
CLEX-8k
CLEX-16K
2.5
2.0
0.0
0.2
0.4
0.6
0.8
1.0
Epoch
Figure 9: The training loss curves of various baselines and our CLEX models Table all models upon LLaMA-2-7B on 2B tokens for one epoch.
We train
4.5
CLEX-4k (w/o E,)
CLEX-4k (w/ Ei)
4.0
3.5
Loss
3.0
2.5
2.0
0.0
0.2
0.4
0.6
0.8
1.0
Epoch
Figure 10: The training loss curves of our CLEX-4k trained on LLaMA-2-7B with or without the constant dynamics (E.) in Eq. (14).
code completion Table 8) tasks. Note that the average sequence length of most tasks in Long- Bench ranges from 5k to 15k, so the samples would be truncated into the sequence length within the supported context length (if necessary) for baselines, except Baichuan-13B-o
ALiBi-7B-4F and CLEEX-7B-4K which enables length extrapolation are evaluated with a context window of 16k. Baichuan-1334 13B-44k is a model pre-trained with ALiBi on sequence length of 4k, while ALiBi-7B- 4k is the model trained with ALiBi that undergoes the same training procedure our CLEX. We found that both ALiBi-based models significantly underperform in the LongBench benchmark when evaluated with sequence length (16k) beyond their training length, indicating the challenges of their extrapolation ability in practical tasks. Our CLEX, however, achieves well perf-porn results across most tasks trained with 4k but tested with 16k length, which further reveals the superior length extrapolation ability of our method.
16

---

Published as a conference paper at ICLR 2024
Models
NarrativeQA
Qasper
MultiFieldQA-zhh MultiFieldQA-en
AVG.
GPT-3.5-Turbo-16k
16k
23.6
43.3
52.3
61.2
45.1
Llama2-7B-chat-4k
18.7
19.2
36.8
11.9
21.65
LongChat-v1.5-7B-22
16.9
27.7
41.4
29.1
28.78
Codel.LaMA-7He
22.93
30.69
43.37
31.76
32.19
XGen-7B-8k
18
18.1
37.7
14.8
22.15
InternL--7B..
12.1
16.7
23.4
33.6
21.45
Vicuna-v (iama-vv.5-7B-... 1.5-7B-16k
19.4
26.1
38.5
43
31.75
Baichuan-1338 133-44
0.07
17.55
17.28
38.55
18.36
ALiBi--BB-4k
0.04
8.13
17.87
2.89
7.23
CLEX-7B-4k
18.05
23.68
44.62
31.15
29.38
Table 3: Single-document question answering tasks in LongBench.
Model
HotpotQA
2WikiMQA
Musique
DuReader
AVG.
GPT-3.5-Turbo-16k
51.6
37.7
26.9
28.7
36.23
Llama2-7B-chat-44
25.4
32.8
9.4
5.2
18.2
LongChat-v1.5-7BB32
31.5
20.6
9.7
19.5
20.33
CodeLLLaMA-7B-116k
33.05
27.93
14.2
10.78
21.49
XGen-7B-8k
29.7
21.1
10.3
11
18.02
IntemLM-7B-8k
28.7
22.8
11.1
17.9
Vicuna-vl.5-7B-166
25.3
20.8
9.8
19.3
18.8
Baichuan-13B-4k
3.29
15
0.1
8.77
6.79
ALiBi-7B-4k
2.73
1.33
11.87
5.98
CLEX-7B-44
28.44
19.53
9.15
23.21
20.08
Table 4: Multi-document question answering tasks in LongBench.
Model
GovReport
QMSum
MultiNews
VCSUM
AVG.
GPT-3.5-Turbo-16k
29.5
23.4
26.7
16
23.9
Llama2-7B-chat-4k
27.3
20.8
25.8
0.2
18.525
LongChat-vl.5-7BBB22
30.8
22.7
26.4
9.9
22.45
CodeL.LaMA-7B-166
28.43
24.18
26.84
0.79
20.06
XGen-7B-8k
27.3
20.5
26.2
2.2
19.05
InternLM-7B--8k
9.7
15.9
22.8
12.4
15.2
Vicuna-v1.5-7B-
27.9
22.8
27.2
15.1
23.25
Baichuan-13B-4k
6.8
1.71
23.1
8.09
9.925
ALiBi-7B-4k
5.31
1.64
19.38
3,25
7.395
CLEX-7B-4k
32.52
22.9
25.55
12.03
23.25
Table 5: Summarization tasks in LongBench.
17

---

Published as a conference paper at ICLR 2024
Model
TREC
TriviaQA
SAMSum
LSHT
AVG.
GPT-3.5-Turbo-16k
68
91.4
41.7
29.2
57.575
Llama2-7B-chat-4k
61.5
77.8
40.7
19.8
49.95
Lom-Chat-v1.5-7B-32k
63.5
82.3
34.2
23.2
50.8
XGen-7B-8k
65.5
77.8
25.3
20.5
47.275
CodeLLaMA-7B-16k
70
84.97
43.43
32.5
57.725
InternLM--B-88k
52
77.8
21.2
15.2
41.55
Vicuma-v1.5-7B-166
71.5
86.2
40.8
28.8
56.825
Baichuan-13B-4k
20.05
20.06
5.77
1
11.72
ALiBi-7B-4k
9.25
8.83
4.67
0
5.6875
CLEX-7B-4k
68
84.92
42.82
28.35
56.0225
Table 6: Few-shot learning tasks in LongBench.
Passage Count
PussageRetrieval-enn
PussageRetrreval-zh
AVG.
GPT-33--Turbo-166
4.5
77.5
51
Llama2-7B-chat-4k
2.1
9.8
0.5
4.13
LongChat-vl.5-7BBB3
1
30.5
7.6
13.03
CodeLLLaMA-7B-16k
2
13.5
11.25
8.92
XGen-7B-8k
8.5
3.5
4.7
InterL-M-7B-88
3
0.9
3.3
Sccuma-vl55-7B--
6.5
4.5
5
5,33
0.06
0.5
5
1.85
ALiBi--B--kk
0
1.27
0.75
0.67
CLEX-7B-4k
0
11.5
17.5
9.67
Table 7: Synthetic tasks in LongBench.
LCC
RepoBench-P
AVG
GPT-3.5-Turbo-166
54.7
53.6
54.15
JJama2-7B-chat-44
52.4
43.8
48.1
LongChat-vl.5-77-2
53
55.3
54.15
CodeLLaMA-7BBI6k
64.35
55.87
60.11
XGen-7B-8k
38.6
38.6
38.6
InternLM-7B-8k.
44.1
28.8
36.45
Becuna-v1.57B-B1
43.5
47.25
Baichuan -3attaann-1B--4--
47.98
16.58
32.28
ALiBi-7B-4k
46.69
18.54
32.61
CLEX-7B-4K
59.01
56.87
57.94
Table 8: Code completion tasks in LongBench.
18


Reviews:
[{"review_id": "v87GFNqqNZ", "reviewer": "Reviewer_ZBA8", "summary": "The paper studies the length extrapolation problem of large language models, i.e., training on short sequences while testing on long sequences. The work is built upon RoPE. Continuous PE scaling is introduced as a RoPE embedding scaling method.", "questions": "- The work is built upon RoPE, which limits its application to other models that don't use RoPE. How to use the proposed method for other PE methods?\n\n- According to Table 1, the models still do not perform \"real\" length extrapolation. The PPL results become worse when the length is increased. If PPL becomes worse, why not directly use window-based methods in practice? The real-world value of the proposed method is questionable.\n\n- Straightforward method (such as https://arxiv.org/abs/2309.16039) works well in practice. It also challenges the value of research on length extrapolation, as long as we finetune the models. The proposed method can be integrated into the above pipeline, which provides more valuable evaluation metrics.\n\n- Fig 5 indicates that different models perform similarly across tasks, despite GPT. The significance of the method is not clearly demonstrated.\n\n- The lines in the right subfigure of Fig 5 are not correctly shown. The figure can be updated.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality:\nContinuous PE scaling is introduced as a RoPE embedding scaling method.\n\n\nClarity:\nThe paper is easy to follow and understand.\n\nSignificance: \nLong-sequence modeling is important for many downstream applications.", "weaknesses": "- The work is built upon RoPE, which limits its application to other models that don't use RoPE.\n\n- According to Table 1, the models still do not perform \"real\" length extrapolation. The PPL results become worse when the length is increased. If PPL becomes worse, why not directly use window-based methods in practice? The real-world value of the proposed method is questionable.\n\n- Straightforward method (such as https://arxiv.org/abs/2309.16039) works well in practice. It also challenges the value of research on length extrapolation, as long as we finetune the models. So the evaluation setting can be improved.\n\n- Fig 5 indicates that different models perform similarly across tasks, despite GPT. The significance of the method is not clearly demonstrated.", "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": "f6l7I0yaNS", "reviewer": "Reviewer_nREr", "summary": "The paper proposes a new positional embedding scaling to be used for using a model with different context lengths than seen during training.\nThe idea is an extension of rotary positional embedding, for which the frequencies used are dynamically updated depending on the desired context length. The method used to actually update those frequencies is through a neural ODE whose parameters are also trained.", "questions": "The paper is mostly interesting and I guess that anyone working on the topic would have a few questions\n* the random sampling method you propose look like a strong and nice ingredient of your approach. Could you just make it clear for me whether the _order_ of the samples is maintained within the sequence?\n* Your method definitely allows some extrapolation as per your experiments. However, I somehow feel that it could also shine for \"superresolution\"/\"interpolation\", i.e. infilling missing data within a sequence. This feeling comes from your random sampling idea. It looks like you are basically simulating \"missing data\".\n* \"Unlike the previous PE scaling methods built on a larger scaling factor would lead to inferior performance on the lengths corresponding to smaller counterparts, the continuous PE scaling would enable non-destructively generalisation to larger scaling factors via adaptive continuous dynamics\". This would be great, but at this point in the paper, I don’t see why the proposed scaling method would _necessarily_ enable it. Maybe you could rephrase that in a more humble way", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the proposed method has some value. It sticks to a well established PE scheme, and then proposes a way to update its parameters that is _trained_ to be good, instead of just wishing it will be based on some assumptions. For this reason, the paper definitely deserves consideration in my opinion.\n\nThe proposed method moreover seems to be providing good performance for extrapolation, which was the intent.", "weaknesses": "* Not much details is provided in the main text regarding how we train such a beast. I must say this looks quite daunting to me how I would train a NODE along my transformer model. I guess it would help to have some explanations to it.\n* I am missing some exploration of what the model is producing regarding the frequencies for ROPE. As I understand, it boils down to being able to produce a new set of frequencies for ROPE to use for any input lengths. This would have been feasible to actually display that. Since many people have played with the idea of manually setting such parameters, I am curious whether a trained method could give us insights as to what good frequencies actually look like. Are we observing high frequencies to disappear to favor long term dependencies? Such things.", "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": "tQOrD3xgEy", "reviewer": "Reviewer_pxa3", "summary": "The paper introduces CLEX, a method to efficiently extend the context window of LLMs without compromising performance. Traditional methods either have length limitations or suffer performance drops. CLEX overcomes these by modeling the relation of sequence length and frequency in RoPE during position extrapolation as a continuous system. Specifically, it utilizes Neural ODE as a tool to do this. In tests, CLEX extends context windows to over 4x the training length without performance loss and beat several popular methods both in long sequence modeling and downstream tasks.", "questions": "1: By my understanding, the core idea of this paper is to do position extrapolation with an appropriate frequency for different sequence lengths. Is it necessary to utilize Neural ODE? In another word, can we use a regular NN? As it seems that a regular NN can do the same thing. (Not quite sure, for I'm not an expert of Neural ODE)\n\n2: Could you please provide more details about the training/fine-tuning? Did you train all the baseline models with the same number of tokens, the same batch size as well as the same steps? If so, for PI, its paper mentioned that they only fine-tuned the LLM for ~ 1000 steps, while for some other baselines such as replacing RoPE with Alibi for Llama-2, the tuning steps should definitely be much larger. With different required number of training steps is the performance comparison fair enough? \n\n3: Also, I'm wondering if there's any explanation to the poor performance on LongBench's synthetic tasks.\n\n4: Compared to Random Position, the main difference is that CLEX add adaptive frequency for different sequences, it that correct?\n\n5: Still about Random Position, its original paper and some blogs (https://kexue.fm/archives/9444, it's in Chinese, you may translate it to English first) showed that it shows good length generalization ability. But in  Table-1, Random Position does not work at all ( trained on 4k, and can only keep low PPL at 4k), do you have any thoughts about it?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1: The theory part is closely combined with practical part. And the performance of the proposed method also aligns with the theoretical derivation. \n2: The performance of the proposed method is good. And the experiments are comprehensive. \n3: Besides the main results, this paper also provided some insightful observations about LLMs' length generalization.", "weaknesses": "Please check the question section", "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": "EwKZv5BgPu", "reviewer": "Reviewer_jLqk", "summary": "This paper proposes CLEX, a method that enables length extrapolation on Rotary position embedding (RoPE) by finetuning on a dataset. Prior work has found that finetuning with a position embedding scaling (PE scaling) on either the position or the frequency values of RoPE can enhance the extrapolation of a pre-trained language model. Based on the PE scaling, CLEX models the position-frequency values of RoPE with a neural ODE and the YaRN model, with an aim to learn the continuous dynamics over the length scaling factor during PE scaling. The experiments demonstrate a strong extrapolation compared with prior PE scaling methods (Position Interpolation, YaRN, CodeLLaMA) and non-PE scaling methods (ALiBi, RandomPos).\n\n**Post Rebuttal Update**\n\nI acknowledge the authors' efforts in addressing the questions. The new experiments have effectively addressed my concerns. As a result, I have raised my score.", "questions": "* In Table 1, the authors provided numbers for CLEX with training length 4k, 8k, and 16k. However, most of the other PE scaling methods (PI and YaRN) are trained only on 16k. I wonder how PI and YaRN perform when finetuning with 4k and 8k context length.\n\n* For other questions, see the Weakness section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* CLEX demonstrates strong length extrapolation results (evaluation length > train length) compared with prior PE scaling methods such as Position Interpolation,YaRN, and CodeLLaMA.\n\n* When the evaluation length is smaller than the train length (i.e. finetuning length), CLEX exhibits a better performance compared with prior PE scaling methods.\n\n* The ablation includes a few useful topics such as continuous vs. discrete dynamics, sampling strategy, and log-scaling.", "weaknesses": "* Despite a strong extrapolation performance, the motivation for adopting continuous modeling is a bit unclear. It seems that the continuous model has to be somehow discretized on a few points (e.g. evaluating the integral of equation 12). If this is true, doesn't this imply an equivalent discrete modeling?\n* CLEX is adopting YaRN in equation 13, so it seems some part of the performance of CLEX is due to YaRN. An ablation of CLEX without YaRN is needed.\n* CLEX is based on PE scaling, which requires a finetuning dataset. However, non-PE scaling methods (e.g., ALiBi and RandomPos) don't require finetuning. So it doesn't seem fair to compare CLEX with non-PE scaling methods.\n* The author mentioned that CLEX is computationally demanding due to the evaluation on the integral. Maybe the author can comment more on the training time.\n* The author claimed that AliBi-like methods (attention biasing) struggle in practical tasks requiring long-context dependency; however, the cited evidence is on AliBi. Among the author-cited AliBi-like methods, there are attention-biasing methods that achieve better long-context dependency than Alibi. Maybe the author can clarify on this.\n* The notations are confusing sometimes. For example, $\\lambda$ is supposed to be an amplification factor but is missing in equation 13.", "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)