PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training
Authors: Dawei Zhu
Nan Yang
Liang Wang
Yifan Song
Wenhao Wu
Furu Wei
Sujian Li
Venue: ICLR 2024 poster (2024)
Keywords: context window extension
efficiency
positional skip-wise training
Primary area: representation learning for computer vision, audio, language, and other modalities

Abstract: Large Language Models (LLMs) are trained with a pre-defined context length, restricting their use in scenarios requiring long inputs. Previous efforts for adapting LLMs to a longer length usually requires fine-tuning with this target length (Full-length fine-tuning), suffering intensive training cost. To decouple train length from target length for efficient context window extension, we propose Positional Skip-wisE (PoSE) training that smartly simulates long inputs using a fixed context window. This is achieved by first dividing the original context window into several chunks, then designing distinct skipping bias terms to manipulate the position indices of each chunk. These bias terms and the lengths of each chunk are altered for every training example, allowing the model to adapt to all positions within target length. Experimental results show that PoSE greatly reduces memory and time overhead compared with Full-length fine-tuning, with minimal impact on performance. Leveraging this advantage, we have successfully extended the LLaMA model to 128k tokens using a 2k training context window. Furthermore, we empirically confirm that PoSE is compatible with all RoPE-based LLMs and position interpolation strategies. Notably, our method can potentially support infinite length, limited only by memory usage in inference. With ongoing progress for efficient inference, we believe PoSE can further scale the context window beyond 128k.

Published as a conference paper at ICLR 2024
POSE: EFFICIENT CONTEXT WINDOW EXTENSION OF LLMS VIA POSITIONAL SKIP-WISE TRAINING
Dawei Zhu Nan Yang
Liang Wang
Yifan Song
Wenhao Wu
Furu Wei 9
Sujian Li 9
School of Computer Science, Peking University
National Key Laboratory for Multimedia Information Processing, Peking University
Microsoft Corporation
ABSTRACT
Large Language Models (LLMs) are trained with a pre-defined context length, restricting their use in scenarios requiring long inputs. Previous efforts for adapting LLMs to a longer length usually requires fine-tuning with this target length (Full- length fine-tuning), suffering intensive training cost. To decouple train length from target length for efficient context window extension, we propose Positional Skip-wisE (PoSE) training that smartly simulates long inputs using a fixed context window. This is achieved by first dividing the original context window into several chunks, then designing distinct skipping bias terms to manipulate the position indices of each chunk. These bias terms and the lengths of each chunk are altered for every training example, allowing the model to adapt to all positions within target length. Experimental results show that PoSE greatly reduces memory and time overhead compared with Full-length fine-tuning, with minimal impact on per- formance. Leveraging this advantage, we have successfully extended the LLaMA model to 128k tokens using a 2k training context window. Furthermore, we empir- ically confirm that PoSE is compatible with all RoPE-based LLMs and position interpolation strategies. Notably, our method can potentially support infinite length, limited only by memory usage inference. With ongoing progress for efficient inference, we believe PoSE can further scale the context window beyond 128k.
INTRODUCTION
Large Language Models (LLMs) have revolutionized language modeling and demonstrated impres- sive abilities to perform various tasks (Brown et al., 2020). However, even with their remarkable capacity, these LLMS remain restricted by pre-defined context window sizes, suffering from notable performance decline when input tokens exceeds these limits. Nevertheless, numerous application scenarios demand extremely long input sequences, including long document summarization (Huang al., 2021), in-context learning with numerous examples (Li al., 2023), and long document retrieval (Zhou et al., 2022), etc. This naturally poses a significant challenge of context window extension: Extending the context window of pre-trained LLM to accommodate longer sequences.
Naively fine-tuning LLMs on inputs of target length for window extension has received limited success due the large disruption introduced by new position indices (Chen et al., 2023a; Han et al., 2023). Addressing this, Position Interpolation (Chen et al., 2023a; kaiokendev, 2023; Peng et al., 2023) propose to down-scale the position indices to match the original window size, yielding improved results for context extension. However, these methods still rely on Full-length fine-tuning, i.e., fine- tuning with context of target length, which is memory and time-intensive due to the computational complexity that increases guadratically with input length. For example, Chen et al. (2023a) use 32 A100 GPÜs to extend LLaMA models from 2k to 8k context, and 128 A100 GPUs for even larger context. These overhead has made it impossible to extend context window to extreme lengths.
Work done during Dawei's internship at MSRA. Sujian Li is the corresponding author.

---

Published as a conference paper at ICLR 2024
Full-length: 8192 tokens
target / original context size 8192 2048
8191
PoSE: 2048 tokens
skip
Train Example
Relative Positions
0
511
6560
8095
#
[1,1535]u[6049,805]
skip
#
2
[1,1024]u[252944575]
1023
3552
4575
Figure I: Position indices of Full-length fine-tuning v.s. PoSE fine-tuning for extending the context window size from 2,048 to 8,192. each iteration, the former directly takes 8, 191 tokens for fine-tuning, while PoSE manipulates the position indices of 2,048 tokens to simulate longer inputs. For example, we partition the original context window 2,048 tokens into two chunks, and adjust the position indices of the second chunk by adding distinct skipping bias term. These bias terms, as well as the length of each chunk, are altered for each training example, so that the model can adapt to all relative positions of the target context window through fine-tuning.
In this paper, we introduce Positional Skip-wisE (PoSE) fine-tuning to decouple the fine-tuning length from the target context window length, unleashing the possibility of efficiently extending context window to extreme size. The key idea of PoSE is to simulate long inputs by manipulating position indices within a fixed context window. As depicted in Figure we partition the original context window into several chunks, and adjust the position indices of each chunk by adding distinct skipping bias term. These bias terms, as well as the length of each chunk, are altered for each training example, that the model can adapt all positions (including both absolute and relative) within the target context window through fine-tuning. Meanwhile, by maintaining continuous position indices within each chunk, PoSE bears close resemblance to pre-training. As result, the model's pre-trained capacity for language modeling and comprehension is retained to the greatest degree.
The advantages of our PoSE are threefold: 1) Memory and Time Efficiency: By only requiring the original context size for fine-tuning, PoSE circumvents the quadratic increase in computational complexity with respect target length during the fine-tuning stage, thereby significantly reducing memory and time overhead. 2) Potential for Extremely-Long Context: We manage to extend the context window of LLaMA (Touvron al., 2023a) by up to times (2k 128k, k=1,024) while pre- serving decent ability of language modeling and understanding. 3) Compatible with all RoPE-based LLMs and PI strategies: The effectiveness of PoSE has been empirically validated across several representative RoPE-based LLMs, including LLaMA. LLaMA2 (Touvron et al., 2023b), GPT-J (Wang & Komatsuzaki, 2021), ), and Baichuan (Baichuan, 2023). Additionally, PoSE has been demonstrated to be compatible with variety of position interpolation methods, including Linear (Chen al., 2023a), NTK (Peng & Quesnelle, 2023), and YaRN (Peng et al., 2023) interpolation.
Notably, by decoupling the fine-tuning and target length, PoSE theoretically extend context window to an infinite length. The only constraint is the memory usage during the inference phase. Hopefully, with the continuous advancements in efficient inference techniques, including Flash Attention (Dao et al., 2022: Dao, 2023), xFormers (Lefaudeux et al., 2022), vLLM (K won et al., 2023), etc, we believe PoSE can promisingly push the context window size to a even larger scale.
RELATED WORK
Training Length-Exxen Extrapolatable Models. Length extrapolation requires the model to handle continually increasing input tokens, even beyond the context window size used for training (Press et al., 2021). To this end, a series of positional embedding schemes have been proposed, including ALibi (Press et al., 2021), xPos (Sun al., 2023), NoPos (Haviv al., 2022), etc.
Similar to our work, Ruoss et al. (2023) also attempted to simulate longer sequences during training time to mitigate out-of-distribution lengths. They proposed randomized positional encoding
2

---

Published as a conference paper at ICLR 2024
(RandPos), which randomly selected an ordered subset of position indices from longer sequences. Our proposed method, PoSE, diverges from their approach in several key aspects: First, RandPos is positional embedding scheme designed to pre-train encoder-only models from scratch for length extrapolation. contrast, PoSE a fine-tuning method aiming at efficiently extend the context window of pre-trained L.Ms,, which are majorly decoder-only models. Second, in RandPos, the position indices between adjacent tokens are not continuous. However, in PoSE, the position indices within each chunk are intentionally made continuous resemble the pre-training phase, therefore reducing the risk of disrupting the language modeling abilities learned during pre-training.
Fine-tuning LLMs for Longer Context. Differing from length extrapolation, which primarily involves training model from scratch to support lengths exceeding those it was initially trained for, context window extension focuses on extending the context window of a pre-trained L.M.. Directly fine-tuning an existing LLM with longer context window has been shown to progress slowly (Chen al., 2023a). To expedite and stabilize training. Chen et al. (2023a) first down-scaled position indices to match original context size through Linear Position Interpolation. Subsequently, range of Positional Interpolation (PI) strategies have been introduced, including NTK (Peng & Quesnelle, 2023) and YaRN (Peng al., 2023). More recently, LongLora (Chen et al., 2023b) propose shift short attention to approximate full attention. However, all these methods require Full-length fine-tuning, suffering computational cost that grows with target context size. By contrast, method managed to decouple train / target length, requiring only the original context size for fine-tuning.
Memory Transformers. An alternative strategy for extremely long input sequences involves memory mechanisms. Typically, there are two lines of research for utilizing memory: the recurrence- based approach (Dai al., 2019; Bulatov al., 2022) and the retrieval-basee approach (Wu et al., 2022: Wang et al., 2023; Tworkowski et al., 2023). The former segments long inputs and reuses the hidden states of preceding segments as memory, suffering from information loss and limited capacity for random access. The latter encodes prior sequences as (key, value) pairs and utilizes a memory retriever and reader to extract previously encoded information, primarily limited by the lack of interaction between discrete memory segments. More recently, Mohtashami & Jaggi (2023) introduced landmark attention to facilitates random access to any chunk of the input. In contrast, our method achieves full access to the entire input without any modifications to the attention mechanism.
3
METHODOLOGGY
3.1 PRELIMINARIES
Rotary Position Embedding (RoPE). The use of RoPE (Su al., 2021) has become pervasive
in contemporary LLMs, including LLaMA (Touvron et al., 2023a), GPT-J (Wang & Komatsuzaki, 2021), etc. It encodes position information of tokens with rotation matrix that naturally incorporates explicit relative position dependency. To elucidate, given a hidden vector h [ho, hı .... ha-
where d the hidden dimension, and a position index m, RoPE operates as follows:
ho h
cos m&g
-h ho ha
sin m00 sin m00
cos m00 m00 COS m0
h2
sin m0,
f(h,
hig
h2
sin m01
(I)
ha-2
COS ml 2/00//211
-ha-1 sin m0a/2-1 sin m0a/2-1/
ha-1
COS m&
h
where 0, 10000 1000022//d_.. {0. 1}. Unlike previous absolute position encodings that are directly applied to the input vector æ, RoPE is employed on the query and key vectors at each layer. Given a query q at position m and a key le at position n, attention score a(g, k)) is defined as: a(g. k) m). f),((knn)) f(k,n)
d/2-1
Σ 1=0
1112kk2 + 12i+1k2j+1 cos (m - n)00 + (g2gk25+1 - 42j+1k22) sin (m - n)),
o(g. k. 0 m n)
(2)
3

---

Published as a conference paper at ICLR 2024
Hence, RoPE encodes position information in relative manner, the attention score depends on the relative distances between positions rather than their absolute position values.
Problem Formulation. Given Large Language Model pre-trained with context window size of Le. our objective is to extend this context size to a target length L, so that the model maintains good performance when processing input sequences containing maximum of L tokens.
Position Interpolation (PI). In contrast to directly extending the position indices to L, - when dealing with input text {Io, 21...... Ft, J. position interpolation down-scales the position indices to align with the original context window size Le- This approach effectively mitigates the risk of encountering extreme values and has been empirically demonstrated to enhance stability during fine-tuning Various interpolation strategies have been proposed, with = L. denoting the scaling factor:
Linear Interpolation. As described by Chen et al. (2023a) and kaiokendev (2023), linear interpola- tion involves a proportional down-scaling of the position index m to m/a. Consequently, the atten- tion score between a query q at position m and key k at position becomes g(g, k. 0, (m-n)/aa) as defined in Equation 2. Theoretical analysis has substantiated that the interpolated attention score exhibits significantly greater stability compared the extrapolated counterpart.
Neural Tangent Kernel (NTK) Interpolation, In contrast to linear interpolation, NTK Interpola- tion alters the base of RoPE, effectively modifying the rotational "speed" of each dimension of RoPE (Peng& Quesnelle, 2023). Specifically, the original 0g (0000-25/d] : E {0, 1,.., d/2-1]
in RoPE is transformed into 07 = (10000A) -2//d, where À = 4/d-2, It is noteworthy that the value of À is chosen to ensure that m6' d/2-1 = (m/a))a./22 /a/22-1
(2000/0)-21/4,
YaRN Interpolation. Different from Linear and NTK interpolation that treat each dimension of RoPE equally, YaRN (Peng et al., 2023) employs a ramp function to combine Linear and NTK interpolation at varying proportions across different dimensions. Simultaneously, it introduces a temperature factor to mitigate distribution shift of attention matrix caused by long inputs.
3.2 PROPOSED APPROACH: POSITIONAL. SKIP-WISE TRAINING (PoSE)
Although position interpolation effectively addresses -itto---iitritttion position indices, extending to an extreme length by fine-tuning on context window of this size remains impractical, owing to the quadratic growth in computational complexity of attention as sequence length increases. Instead, explore to train within the original context window L. and achieve context window extension via manipulating position indices simulate longer inputs.
There are two designing desiderata for this endeavor: First, to avoid nut-of-distribution positions during inference, the relative distance of manipulated position indices should comprehensively cover the range of (L.... L, 1}. Second, fine-tuning with the manipulated position indices should not harm the original abilities of LLMs, so the structure of manipulated position indices should closely adhere to the original structure to the greatest extent possible.
Initially, we randomly divide the original context window Le into N chunks cu, ci,..., CN-1. each with lengths lo. l,...., IN...., where Σ 4 = Le. We introduce the starting index st, for each chunk c; which facilitates the formulation of its position indices follows:
(os(c;) = (st;, st. + 1,..., st, + + - 1}, st, = Σ54 1
1=0
Subsequently, we employ the discrete uniform distribution U(S) to sample a skipping bias term u4 ~ ((([u-1,.-..; Ia T.r)) for each chunk c,. This bias term is applied the corresponding chunk to transform the original position indices into:
PoSE(c;) {w + sti, + st, 1,...,, U4 + st, 44 - 1}
Note that the constraint of u; 2 u is applied prevent position index overlaps between chunks.
Intuitively, the introduction of skipping bias terms exposes model to a more diverse range of relative positions. To achieve comprehensive coverage of the target context window, re-sample both the
4

---

Published as a conference paper at ICLR 2024
length and skipping bias term of every chunk for each training example. Moreover, the continuity of position indices within each chunk closely resembles the structure employed during pre-training Consequently, fine-tuning the model on these position indices for language modeling does not compromise its original capabilities.
Concerning the text contained within each chunk, similar procedure is followed to select continuous spans of tokens from the input text x = = #1 21,,..,,,,, 1. To elaborate, we begin by sampling a bias term 72 ~ U((v:=1;.... Liz Le) followed by assigning the content of chunk e below:
G2 ar + 844 : 04 + 844 +
Notably, we have also explored other assigning strategy of vs, including scenarios where o; which results in genuinely continuous content for the chunks, or vú ws, aligning the manipulated position indices with actual positions in the original text. However, we observe that these variations have relatively little impact on the outcomes of fine-tuning.
After position indices and content for each chunk are settled, we perform position interpolation for stabilized fine-tuning For simplicity, We set the initial bias terms uo and to O. In terms of chunk number N, we view as an trade-off between efficiency and effectiveness. Because an increase the number of chunks will further deviates from the position structure of pre-training, which may harm the ability acquired during pre-training. Hence, this paper we set N to 2, exposing the models to a wider range of relative positions, while adhering as close to the original position structure as possible. (See Appendxi and B for further discussion of v and N.)
EXPERIMENTS
In this section, conduct experiments to verify the effectiveness of PoSE for context window extension. Our method demonstrates impressive results on context lengths of both 16k and 32k for language modeling as well passkey retrieval. Other advantages of PoSE are discussed in Section 5.
4.1 SETUPS
Training Procedure. For each setting in the main experiments, we train LLaMA-7B with the next token prediction objective. This training process comprises 1,000 steps, employing global batch size of 64 on 8 V GPUs using Deepspeed ZeRO stage 3 (Rajbhandari et al., 2020). We use learning rate -5 and linear scheduler, with 10 warmup steps. We use AdamW optimizer with its default hyperparameters setup. The fine-tuning dataset is sourced from The Pile (Gao et al., 2020), with a minimum length requirement of 2, 2,048 tokens. Our default choice for interpolation strategies linear interpolation. For evaluation, we use single A100 GPU. Flash Attention V2 (Dao, 2023) is applied, making it possible to evaluate long documents of up to 128k tokens (k=1,024)
Evaluation Tasks and Datasets. We examine the ability of long text modeling on two tasks: language modeling and passkey retrieval. The language modeling task a fundamental task that reflects the overall capability of a model in handling long text. Passkey retrieval, on the other hand, can effectively measure the maximum distance that a token can attend to during the inference stage. We evaluate language modeling GovReport (Huang et al., 2021) and Proof-pile (Zhangir et al., 2022) datasets. For passkey retrieval, we follow Mohtashami & Jaggi (2023) to construct synthetic prompts for evaluation.
Baseline Methods.
We compare our PoSE training method against following baselines:
Full-length fine-tuning takes input tokens of target length for fine-tuning. For this method, com- putation complexity scales quadratically with target context window size. Following Chen et al. (2023a) and Peng et al. (2023), we perform PI before fine-tuning LLMs on inputs of target length.
RandPos (Ruoss et al., 2023) is initially designed to train an encoder-only model from scratch for length extrapolation. However, since it shares similar idea of simulating longer sequences via changing position indices, we include it for a comprehensive comparison. Given the original / target context window length L. / L,, it uniquely samples L. positions from the set {O...., L, 1). arranges them in ascending order, and employs them as new position indices for training. For fair comparison, we also apply PI for this method.
5

---

Published as a conference paper at ICLR 2024
Table 1: Perplexity of models trained with different methods. We conduct evaluation on the GovReport and Proof-pile datasets, varying evaluation context window size from 2k to 32k. Our POSE, with fixed training window size of 2k, effectively extended to a target context size of 16k / 32k for inference while receiving only minimal performance degradation compared to Full-length.
Context size
GovReport
Proof-pile
Method
Train / Target
2k
4k
8k
16k
32k
2k
4k
8k
16k
32k
Original
-/-
4.74
10
>
10
10
> 10a
2.83
10a
>
10a
10
10
Full-length
16k 16k
4.87
4.70
4.61
4.59
2.93
2.71
2.58
2.53
2k / 16k
11.63
11.17
11.54
15.16
7.26
6.83
6.76
7.73
RandPos
2k /3kk
93.43
95.85
91.79
93.22
97.57
60.74
63.54
60.56
63.15
66.47
/ 16k
4.84
4.68
4.60
4.60
2.95
2.74
2.61
2.60
PoSE (Ours)
2k 32k
4.91
4.76
4.68
4.64
4.66
3.01
2.78
2.66
2.60
2.59
4.2
LANGUAGE MODELING
First, we investigate the impacts of different fine-tuning methods on long sequence language modeling using the GovReport and Proof-pile datasets. GovReport a summarization dataset comprising 19,402 reports published by the Congress and the U.S. Goverment, with an average document length 7 7,866 tokens. We randomly select reports containing more than 32,768 tokens for evaluation. Similarly, Proof-pile is 13GB mathematical dataset of long mathematical documents. In line with the approach taken for GovReport, we choose samples from Proof-pile that contain more than 32,768 tokens for evaluation.
Table I presents the results of scaling to 16k and 32k using Full-length, RandPos, and PoSE training method, each with linear interpolation (See Appendix C for results of NTK and YaRN). For each scaled model, as well as the Original LLaMA model, we report perplexity scores at various evaluation context window sizes, ranging from 2k to 32k, employing the sliding window approach proposed by Press et al. (2021). For evaluation efficiency, we set the stride of the sliding window to 1,024.
First, observe overall decreasing trend of perplexity for both models scaled to 16k and 32k via PoSE as evaluation context window size increases, proving their abilities to leverage longer context. Second, with significantly shorter context length during fine-tuning, our PoSE achieves comparable results with Full-length, consolidating its effectiveness. Third, our method achieves much stronger results than RandPos. We suppose because our manipulated position indices closely resembles that of pre- pre-training, hereby preserving the pre-true pre-trained language modeling ability to the greatest extent.
We also notice that all the scaling methods suffers certain performance degradation as the supported context length increases. We perceive this as a trade-off between the quantity of tokens the model can process and the level of granularity in the attention the model can pay to each individual token.
4.3 PASSKEY RETRIEVAL FOR EFFECTIVE CONTEXT WINDOW
To effectively measure the maximum distance that a token can attend to during the inference stage, we adopt the passkey retrieval test proposed by Mohtashami & Jaggi (2023). In this test, models tasked with recovering random passkey hidden within lengthy document. Prompt template used for this task presented in Figure 2a.
Specifically, we compare the original LLaMA model with the PoSE-extended versions for 16k and 32k context. For each model, we vary the prompt length from 2k to 32k. For each length, we conduct the passkey retrieval test for 50 times, with a random passkey of 5 digits generated and placed at random position inside the prompt. We also include results from Full-length, RandPos, and PI-only (position interpolation without fine-tuning). Figure 2b illustrates the results. For the Original, Pl-only, and RandPos models, their retrieval accuracy rapidly drop to 0 when the context exceeds 2k. In contrast, both PoSE-16k / 32k models managed maintain a high retrieval accuracy (2 90%) within their respective target context window, comparable to Full- Full-length. This indicates that models trained via PoSE genuinely possess the capability to attend to all tokens within the extended context windows.

---

Published as a conference paper at ICLR 2024
There important info hidden inside lot of irrelevant text. Find and memorize them. I will quiz you about the important information there.
100
80
Original
The grass green. The sky is blue. The sun is yellow. Here we go. There and back again. (repeat times)
Accuracy
60
PoSE-32k
PoSE-16k
Pl-only-1od
The pass key B15o1. Rememberit. it. 8a5os is the pass key.
40
ReedPos-
Full-langth-6 16k
The grass green. The sky blue. The sun is yellow. Here we go. There and back again. (repeat times)
20
0-
What the pass key? The pass ks
2k
16k
24k
32k Tooeess
(a)
(b)
Figure (a) Prompt template used for passkey retrieval; (b) Retrieval accuracy for the PoSE-extended 16k / 32k models, compared with other baselines. Both PoSE-extended models maintain a high retrieval accuracy (2 90%) within their respective context window.
5 ANALYSIS
In this section, we analyze the advantages of POSE, including 1) memory and time efficiency; 2) compatibility with all RoPE-based LLMs and diverse interpolation strategies; 3) potential for extremely-long context. In Section 5.4, We also verify that model performance within the original context window only receives minimal degradation.
5.1 MEMORY AND TIME EFFICIENCY
We study the memory and time efficiency of PoSE compared with Full-length fine-tuning. For each method, we scale I.LMAA-BB to 4k 16k through 1,000 training steps with a global batch size of 16 on 8 V100 GPUs Experiment results are demonstrated in Figure 3. Figure 3(a) and (b) respectively illustrates memory and time consumption for 1,000 steps of Full-length versus POSE. While the training cost of Full-length increases rapidly with target window length, PoSE only requires fixed quota of memory and time for context extension, which is significantly lower. Figure 3(c) further compares model perplexity of the two training methods different steps GovReport. Notably, both models achieve relatively low perplexity levels within the initial 100 training steps. Moreover, requiri
each step, our proposed POSE, while requiring only a training context size of 2k tokens, exhibits
very close language modeling ability to Full-length fine-tuning, which requires an extended training context of 16k. We did not experiment with context window of 32k or above, because V100 machines cannot afford full fine-tuning of these lengths. But it can be expected that the overhead ration between Full-leng and PoSE will become more exaggerated as target length increases. Consequently, we can confidently assert that our proposed approach is both memory and time-efficient.
5.2
COMPATIBILITY WITH ROPE-BASED LLMS AND DIVERSE INTERPOLATION STRATEGIES
We also delve into the effectiveness of PoSE when applied to different RoPE-based LLMs, as well various interpolation strategies. Specifically, we employ PoSE four distinct models: LLaMA-7B, LLaMA2-7B, GPT-J-6B, and Baichuan2-7B, all of which encompasses RoPE in their architectures. The original context size of LLaMA-7B and GPT-J-6B is 2k, while that of LLaMA2-7B and Baichuan2-7B is 4k. For each model, we examine the integration with Linear, NTK, and YaRN interpolation, as well the original version for comparative purposes. The same GovReport dataset as described in Section 4.2 is utilized, The test set is truncated to the first 1k to 16k tokens for plotting the perplexity curve, as depicted Figure 4. First, is evident that PoSE is effective across all four models and three interpolation strategies, by all 12 combinations in comparison to the 4 original model. Second, we observe that NTK and YaRN interpolation generally yields superior results compared to Linear interpolation. However, it is noteworthy that NTK exhibits a significant increase in perplexity after a certain turning point, which occurs prior to reaching the target context length. This behavior is consistent with previous findings, indicating that for a given scaling factor a, NTK cannot genuinely expand the context window by a times (Peng& & Quesnelle, 2023; Quesnelle, 2023; Peng al., 2023).
evidenced by the low perplexities achieved
7

---

Published as a conference paper at ICLR 2024
Mcmory (GB)
Time (h)
Perplexity
CCM
Full-length
Full-lenghh
30
POSE
PoSE
Full-length
10
PoSE
5.1
10
0
4.7
4%
8k
16k
4k
8k
16k 10
100
300
1000
Steps
(8)
(b)
(c)
Figure 3: Full-length fine-tuning PoSE in terms of (a) Memory and (b) Time consumption for extending LLaMA-7B from to 8k 16k context, each finishing 1000 training steps. (c) Perplexity of both 16k-context models at every training steps. We show that PoSE takes constantly reduced time and memory for context extension, while attaining a comparable level of PPL performance with Full-length fine-tuning at each step.
LLaMA-7B
LLaMA2-7B
GPT-J-6B
Baichman2-7B
54
5.2
9.5
7.3
Pepleciy
5.2
9.0
6.8-
8.5
5.0
6.3
8.0
4.8
7.5
5.8
4x
8k
16k
1k
2k
16%
2k
4k
16k
1k
8k
16k
Original
-
PoSE Linear
POSE-NIK
- PoSE-YaRN
Figure 4: Perplexity of LLaMA-7B, LLaMA2-7B, GPT-J-6B, Baichuan2-7B extended 16k via PoSE with Linear NTK YaRN interpolation, along with the Original model. The consistently low perplexity observed across all nine combinations serves as an indication of the effectiveness of our method across RoPE-based LLMs and diverse interpolation strategies.
5.3 POTENTIAL FOR EXTREMELY-LONG CONTEXT
Because PoSE only takes a fixed context window at training stage to extend to target context window size, we can promisingly extend LLMs to support infinite input lengths using this method. In this section, we extend context window size to 96k and 128k to explore PoSE's potential for extreme context window extension. Given the need to evaluate extremely long documents, have opted to employ two book datasets, namely Books3 (Presser, 2020) and Gutenberg (PG-19) (Rae et al., 2019). Both of these datasets consist of extensive collections of literary works, rendering them well-suited subjects for the assessment of long-range modeling. For our evaluation, we randomly selected 20 books from each dataset, each containing more than 128k tokens.
Fine-tuning LLaMA models using POSE, we experimented with Linear NTK YaRN interpolation for both the 96k and 128k models. To calculate perplexity, we adhere the sliding window strategy adopted in Section 4.2, with an increased sliding window step of 16k to enhance evaluation efficiency, The outcomes of these experiments are detailed in Table It is observe that, PoSE successfully extends the model's context window to 96k when coupled with Linear interpolation, and further extends the context window 128k when paired with YaRN. These promising results consolidates the effectiveness of PoSE for extreme context window extension.
5.4 EVALUATION OF CAPABILITY ON ORIGINAL CONTEXT WINDOW
In this section, we examine the capabilities of the PoSE-extended models on the original context win- dow using standard benchmarks. We combine the Hugging Face Open L.M Leaderboard (Face, 2023)
8

---

Published as a conference paper at ICLR 2024
Table 2: Perplexity of models extended extreme context size via PoSE on PG-19 and Books3. We show that our training method can effectively extend context window size to 128k when combined with YaRN interpolation.
Gutenberg (PG-19)
Books3
Model
32k
64k
96k
128k
32k
64k
96k
128k
FoSE-Lincar-96.
10.18
11.11
13.57
9.98
10.90
13.42
PoSE-NTK-966
7.98
20.39
38.73
8.29
20.82
40.39
PoSE- -aR-99K
8.31
8.65
9.36
8.90
9.40
10.38
PoSE-Linear- 128k
16.90
22,47
26.77
31.18
26.20
43.62
57.08
70.87
POSE-NTK-128k 8.04
14.84
29.48
34.80
8.34
16.04
31.42
37.00
PoSE- YaRN- 9.32
10.36
10.77
11.33
10.56
12.30
13.07
13.81
Table 3: Performance of PoSE--xteeded LLaMA model on standard benchmarks in comparison with
Full-Iength fine-tunie tinettuning and the original LLaMA. We show that PoSE-extended models exhibit only marginal performance degradation compared with Full-length fine-tuning and the original version.
Zero-Shot
Few-Shot
Model
BoolQ
PIQA
WinoGrande
TruthfulQA
ARC-C
HellaSwag
Original I.aMMA
75.11
78.67
69.85
34.08
51.19
77.75
Full-L.neerr-166 16k
70.95
77.64
69.06
31.89
48.55
74.19
Full-NTK-16k
75.80
78.08
68.98
33.83
48.81
76.57
Full- Full-YaRN-16k
73.88
77.64
68.15
34.12
50.60
77.18
PoSE-Linear-164.
74.50
78.13
68.59
32.05
48.29
75.56
POSE-NTK-16s
74.28
78.24
68.90
33.89
49.83
76.82
PoSE-YaRN-16k
74.28
78.02
69.06
34.00
49.23
77.04
PoSE-Lincar-IP 128k
67.71
76.22
67.56
36.16
39.93
66.04
PoSE-NTK-
75.35
78.18
68.98
32.71
49.66
76.19
DOSE-YaRN- 128k
73.61
77.80
70.01
34.47
48.46
75.54
with subset of LLaMA benchmarks to assess zero-shot and few-shot performance. For zero-shot evaluation, employ BoolQ (Clark et al., 2019), PIQA (Bisk al., 2020), WinoGrande (Keisuke et al., 2019), and TruthfulQA (Lin et al., 2022). For few-shot evaluation. we utilize 25-shot ARC- Challenge (Clark al., 2018) and 10-shot HellaSwag (Zellers al., 2019). Our evaluation metries are Denchhark-speciffee for BoolQ. PIQA, and WinoGrande, we report accuracy; for TruthfulQA, we report mc2; and for ARC-C and HellaSwag, we report normalized accuracy.
Table 3 summarizes the results. It is observed that, PoSE-extended models exhibit only marginal performance degradation compared with Full-length fine-tuning and the original LLaMA, with the only exception of the 128k model employing linear interpolation. This indicates that while extending context window size, PoSE effectively preserves original language comprehension ability.
CONCLUSION
In this paper, we introduce Positional Skip-wisE (PoSE) training to efficiently extend the context window of Large Language Models. PoSE simulates long inputs by manipulating position indices, thereby requiring only the original context window for fine-tuning, successfully decoupling train length and target length. Experiments have shown that, compared with fine-tuning the full length, PoSE greatly reduces memory and time overhead. Taking advantage of this, we have managed to extend LLaMA model 128k on 8 V100 GPUs, observing only minimal performance degradation standard benchmarks. We have also empirically verified that PoSE is compatible with all RoPE-base I.LMs and position interpolation strategies.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGEMENT
We thank all the anonymous reviewers for their helpful comments on this paper. We thank Xueguang Ma. Yang Ouyang, Pengyun Yue, Hanyu Li, Fangwei Zhu for the thoughtful discussion. This work was partially supported by the Okawa Research Grant.
REFERENCES
Baichuan. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:2309.10305. 2023.
URL https://arxis: org/ abs//309. 10305.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Thirty-Fourts AAAI Conference on Artificial Intelligence, 2020.
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.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35: 1079-11091, 2022.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. ArXiv, abs/2366.15595 2023a.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu. Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309. 12307, 2023b.
Christopher Clark, Kenton Lee, Ming-Was Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In NAACL, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint
2018.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xt: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 2978 2988, 2019.
Tri Dao. FlashAttention-2:: Faster attention with better parallelism and work partitioning. 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with (O-awareness. Advances Neural Information Processing Systems, 2022.
Hugging Face. Open Ilm leaderboard. https://hugginnfaaee http:/aaccccccpccccaceeeesscccc http://g...g..g...o..././//c//ar llm_leaderboard, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint 1rXiv:2101.00023 2020.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308. 16137, 2023.
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 Computa- tional Linguistics: EMNLP 2022, pp. 1382-1390 Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguisties. doi: 10.18663/v 1/2022. www.fiinnngeeenl........
10

---

Published as a conference paper at ICLR 2024
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pp. 1419-1436, Online, June 2021, Association for Computational Linguistics. doi: 10. 10.18653/v1/2011 maacl--min. 112.
kaiokendev. Things i'm learning while training superhot. https://kaiokeneee .irleeeeiiiittti gith... io/til# extending- conteet-tooottttttooooooot tontt--o----h 8k, 2023.
Sakaguchi Keisuke, Le Bras Ronan, Bhagavatula Chandra, and Choi Yejin. adversarial winograd schema challenge at scale. 2019.
Winogrande: An
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang. Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang. and Ion Stoica. Efficient memory management for large language model serving with pagedattention, 2023.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. xformers: A modular and hackable transformer modelling library. https://gi thub. com/ 2022.
bacebokkresearchh/ffooeee
Mukai Li, Shansan Gong, Jiangtao Feng. Yiheng Xu, Jun Zhang, Zhiyong Wu, and Lingpeng Kong. In-context learning with many demonstration examples. arXiv preprint arXiv:2302.04933 2023.
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3214-3252 2022.
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random- aacesss infinite context
length for transformers, 2023.
Bowen Peng and Jeffrey Quesnelle. Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degra- dation. https://wwee reddit. sope_allows__ae aamllow__lmma_mddellllllllllll tooeeee__eeee___eoeeeeeeeeee 2023.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models, 2023.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021.
Shawn Presser. ht https://tiee tter. wto. httm://hhhw///ttuuuuseeses 220aaaauuss111122222222222238333 2020.
Jeffrey Quesnelle. Dynamically scaled rope further increases performance of long context Iaaaa with zero fine-tuning. https://www. reddit. com/ http:///ocaalLaa//oooeettttt /mrrgrr/
2023.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEF, 2020.
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 Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 1889-1903, Toronto, Canada, July 2023, Association for Computational Linguisties. doi: 00.18653/v1/2033. act-soot...
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, arXiv preprint arXiv:2 / (44009844 2021.
11

---

Published as a conference paper at ICLR 2024
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A cength-extrapolattble transformer. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 14590-14604, Toronto, Canada, July 2023. Association for Computational Linguistics. doi:
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 preprint arXiv:2307.0922 2023b.
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miloš. Focused transformer: Contrastive training for context scaling. arXiv preprini arXiv:2307.03170, 2023.
Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://githun
May 2021.
Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Augmenting language models with long-term memory. arXiv preprint arXiv:2306.077 07700, 2023.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers.
arXiv preprint arXiv:2203.cmm 08913, 2022.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, 2019.
Azerbayev Zhangir, Ayers Edward, and Bartosz Piotrowski. Proof-pile. https://github. com/ aaangir-azerbaye///rooof
2022.
Yucheng Zhou, Tao Shen, Xiubo Geng, Chongyang Tao, Guodong Long, Can Xu, and Daxin Jiang. Fine-grained distillation for long document retrieval. arXiv preprint arXiv:2212.10423, 2022.
12

---

Published as a conference paper at ICLR 2024
Table 4: Comparison of different methods for choosing v,. We report perplexity with evaluation context window ranging from 2k to 16k. We show that these variations have relatively little impact on the outcomes of fine-tuning.
GovReport
Proof-pile
Method
2k
4k
8k
16k
2k
4k
8k
16k
U; ~ ((...)
4.84
4.68
4.60
4.60
2.95
2.74
2.61
2.60
U/
4.85
4.72
4.64
4.68
2.96
2.75
2.63
2.61
U2 4.84
4.68
4.60
4.60
2.95
2.73
2.60
2.56
A ABLATION OF TEXT CONTAINED WITHIN EACH CHUNK
PoSE divide the original context window into several chunks, and modify the position indices of each chunk to cover a wider range of relative positions in a fixed window. However, does not impose particular constraint on the text contained within each chunk. Recall that in Equation 5, we assign the content of chunk e; below:
= x: U1 + st. + st, 4.]
1) sampling from uniform distribution, i ~ U((ui_1,..... Tog I.,), which the one used in PoSE; 2) o = O, which results genuinely continuous content for the chunks; 3) U2 Ug aligning the manipulated position indices with actual positions in the original text. We use the same test setting as Section 4.2, extending LLaMA-7B from 2k to 16k context. As can be seen in Table 4, we show that these variations have relatively little impact on the outcomes of fine-tuning.
In this section, we explore several strategies for determining
B ANALYSIS OF CHUNK NUMBER N
1.0
Original
0.8
2 chumks
Probability
3 chumks RandPos
0.6
0.4
0.2
0.0
2500
5000
7500
10000
12500
15000
Relative Position
Figure 5: Coverage probability for each relative position in single training example (2k -> 16k). Utilizing multiple chunks reduces coverage probability within the original [0, 2, 048) context window, while enhancing the coverage likelihood of relative positions in the range of [2, 048, 16, 383). Proba- bility of coverage increases with the number of chunks. Pushing the chunk number to the limit is RandPos, utilizing 2048 chunks, capable of covering every relative position in each training example by expectation.
PoSE achieves coverage of all positions within the target context window by randomly sampling the chunk sizes and skipping bias terms for each training example. In this section, we explore the probability of each relative position being covered by a training example, using a context extension of 2,048 to 16,384 position within 2048 being covered 1, and the probability of a relative position above 2,048 being the covered is O. For the cases where the number of chunks is 2, 3, 2,048 (i.e., RandPos),
example. For the unextended original version, the probability of relative
13

---

Published as a conference paper at ICLR 2024
pssio_ 1ist (p.aaraa([[]] Lt)
visit @isit_proh_lit list [p.array([[] Lt)
proh, np.array([0] LE)
iter steest 18800
stve_ - 18900
iter. tines 180
for in nange(iter timos):
for in range(iter times):
for in rango(Ster tines]:
10 randinta Lc-1)
18 randos, (l,dint), Lc-16-1)
LC-2)
randos. randin, Lt-Lc)
tot list 1ist(range(L))
nes_pos__ist www.oe.sapeeeetotpooooooo
rander.rand Lt-Lc)
rogt set(rarge(s, (1)))
randee.nendint(us. L8-1e)
rog2 set(range(utss,, u1+))
distance_ mag sat()
mgl set(range(), 11,11)
for 1 in range(0, leninew_pos, list))))
rog mgl mg2
set/range(uirl, uiv10+113) set(rangeuu2-uttl, 02- ul-1]+12])
m2g
for 1 in rangetiel, sen(new pes list)):
eng1
distance, aettann__nngeeeeee add(new_pos__ssaee Histj]]
for x in rng:
nen, rok_pos__ie los__is
visit prob Hist[s] e*
eng4 set(nange(u24114, u2aic))
mga mg4
for distance visit_prot__istt
prop_ iter_tires
rog:
visit prob list[s]
visit /- Ster times
vislt_prob__te Hsst / tires
Figure 6: Python Code used for calculating coverage probability of each relative position in Figure 5.
Table 5: Comparison of different chunk numbers. We report perplexity with evaluation context window ranging from 2k to 16k. By increasing chunk number, relative positions in [2, 048, 16, 383] receive increased chance of being trained, rendering better results for context extension. However, extremely large chunk number also damages model performance.
Proof Priof-pile
Chunk number
2k
4k
8k
16k
2.83
> 10a
10a
> 103
2.95
2.74
2.61
2.60
2.93
2.72
2.60
2.59
2048
7.26
6.83
6.76
7.73
Monte Carlo method to estimate this coverage probability. The code used is demonstrated in Figure 6. The estimated results are shown in Figure It can be seen that PoSE reduces the coverage probability of positions within the original context window, while all relative positions in [2, 048, 383) receives certain increase in chance of being covered, and the probability of coverage increases as the number of chunks increases. For the case where the number of chunks equal to 2,048, the probability of each relative position being covered is close to 1. With this observation, we further compare the impact of chunk number on language modeling capability, as presented in Table Increasing chunk number efficiently renders better results for context extension. However, extremely large chunk number also damages model performance, due to the severe deviation from the position encoding structure used in pre-training phase. We believe that the choice of the number of chunks is a trade-off between training efficiency and performance.
C SLIDING WINDOW PPL FROM LINEAR NTK/ / YARN INTERPOLATION
Evaluation results Table are based on Linear interpolation. this section, we comprehensively provide experiment results with all three PI strategies, and compare four scenarios for each: Full- length fine-tuning, POSE, PI-only, and Original. Note that PI-only means we only apply position interpolation without fine-tuning, while Original means the original LLaMA model with neither PI nor fine-tuning. For testing data and sliding window stride, we use the same setup used for Table 1. From Table 6, We can that for all strategies, the performance follows the same trend: Full-length & PoSE > PI-only >> Original. We also notice that the NTK method suffers from a significant increase in ppl 16k, mainly because NTK cannot effectively extend the context window by a scaling factor a (Peng & Quesnelle, 2023; Quesnelle, 2023; Peng et al., 2023). YaRN alleviates this issue, achieving progressively decreasing ppl the context window grows.
14

---

Published as a conference paper at ICLR 2024
Table 6: Perplexity of models trained with different methods using Linear NTK YaRN interpolation. is observed that for all interpolation strategies, the performance follows same trend: Full-length PoSE > Pl-only >> Original.
Context size
GovReport
Proof-pile
Method
Train / Targge
2k
4k
8k
16k
2k
4k
8k
16k
Original
-/-
4.74
10
10
10
2.83
10
10
10
Linear Interpolation
PI-only Full-length PoSE (Ours)
/ 16k
43.80
43.35
45.89
54.33
25.32
24.20
24.88
29.59
16k 16k
4.87
4.70
4.61
4.59
2.93
2.71
2.58
2.53
2k / 16k
4.84
4.68
4.60
4.60
2.95
2.74
2.61
2.60
NTK Interpolation
PI-only Full-Iength
-/ 16k
5.62
5.61
5.80
550
3.27
3.1.
3.19
517
16k / 16k
4.78
4.63
4.57
7.24
2.93
2.71
2.61
5.66
PoSE (Ours)
2k 16k
4.79
4.63
4.57
7.24
2.92
2.71
2.60
4.37
YaRN Interpolation
Pi-only Full-length
-/ 16k
5.57
5.51
5.57
5.83
3.17
2.97
2.87
2.89
16k 16k
4.78
4.62
4.54
4.53
2.90
2.68
2.56
2.52
PoSE (Ours)
2k 16k
4.79
4.63
4.55
4.55
2.91
2.69
2.57
2.53
15


Reviews:
[{"review_id": "aJHjuyo3G3", "reviewer": "Reviewer_UVnU", "summary": "The paper proposes a method called Positional Skip-wise Training (PoSE) to extend the context window size of large language models (LLMs) during fine-tuning. PoSE simulates long inputs by manipulating position indices within a fixed training context window, decoupling the training length from the target length.", "questions": "Please see the weakness.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The proposed PoSE approach is effective in extending the context window size of various RoPE-based LLMs to up to 128k tokens with minimal performance degradation compared to full-length fine-tuning.\n\n2. PoSE significantly reduces the memory and time complexity compared to full-length fine-tuning, achieving up to 64x speedup.\n\n3. PoSE is compatible with various RoPE-based language models and positional interpolation strategies.\n\n4. The idea of manipulating position indices within chunks to simulate longer inputs is novel and effective as shown in the experimental results.", "weaknesses": "1. The theoretical analysis of PoSE is limited. More discussion of how manipulating position indices allows the model to generalize to longer contexts could strengthen the paper.\n\n2. The experimental setups and hyperparameters are not thoroughly described. More details would increase reproducibility.\n\n3. It would be better that if the authors could provide some failure cases that this method might leads to some missing details or inferior answers to the standard fine-tuned models. Because it should be inevitable that the speedup has a cost on performance. Only the perplexity comparison is a bit unclear for understanding. It would be more clear that the perplexity gap lies in retrieval, summarization, or some other abilities.", "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": "J9kRcTFmjO", "reviewer": "Reviewer_T2GT", "summary": "The paper proposed Positional Skip-wisE (PoSE) training for extending the context length of LLMs. The idea is to divide the input into multiple chunks and add randomly sampled bias values to later chunks. This effectively simulates the positions the model will see during the inference stage. Experiments with LLaMA show that PoSE is effective in extending RoPE-based LLMs, and is compatible to other length extension techniques like linear, NTK, and YaRN,", "questions": "How did you do full-length finetuning? Have you used length expansion before finetuning on the target sequence length?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of PoSE is neat and novel. PoSE is able to reduce the GPU memory requirement for training long-context LLMs.\n2. Experiments show that PoSE is effective in extending models trained on 2k context length to 32k. Also, the author explored extending models trained on 2k to 128k in Table 2 and demonstrated the potential of PoSE.", "weaknesses": "1. Finetuning the model with PoSE will degrade its performance on standard benchmarks, according to Table 3. It is unclear if it can be mitigated from the paper. Also, the perplexity of PoSE is worse than finetuning with the target sequence length (shown in Table 1). Thus, if the user really cares about the performance, finetuning with the target sequence length still gives the best result.\n2. From Table 1, I haven't seen the comparison between PoSE and applying NTK directly without training. In addition, the author has not provided enough details in how it conducted full-length finetuning. Is it to first extend the sequence length to 16K via linear or NTK and finetune the model on the 16k sequence length?", "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": "uErCVsMzFk", "reviewer": "Reviewer_VfnN", "summary": "In this paper, the authors introduce a novel technique known as Positional Skip-wise Training (PoSE) for extending the context window of Large Language Models (LLMs). While existing methods like Positional Interpolation (PI), YaRN, and NTK Interpolation have proven effective in extending the context window of LLMs, they require extensive fine-tuning of the LLM with the entire context window. This fine-tuning process imposes a high computational cost due to the quadratic complexity of LLMs relative to the length of the context window.\n\nTo address this challenge, the authors propose a strategic approach during training by selectively skipping certain positions in the middle for each batch. As illustrated in Figure 1, this technique allows the model to effectively perceive a context length of 8192 tokens while using a training-time context window of just 2048 tokens. Consequently, this method significantly reduces the training cost. When evaluated on prominent long-context datasets, including GovReport, Proof-pile, and PG19, the authors demonstrate that PoSE achieves comparable performance to full-context fine-tuning, all while delivering a substantial speedup (as shown in Figure 3).", "questions": "I believe that PoSE demonstrates proficiency in both PPL and key retrieval evaluations. I look forward to your response to my inquiry regarding the real extension of the context window, as mentioned in the Weaknesses section.\n\nIn addition, I am curious about the performance of the proposed method when the first trunk does not initiate from position 0. Recent concurrent works, such as LM-Infinite and StreamingLLM, have highlighted the significance of retaining the first few tokens in the KV cache for zero-shot context extension, without requiring fine-tuning. I am intrigued to know if the design of $u_0=v_0=0$, which resembles the approach proposed in these two papers, holds the same relevance when applied to a fine-tuned LLM. It would be greatly appreciated (though not obligatory) if the authors could include a discussion on the efficacy of fine-tuning in comparison to these zero-shot methods. For example, does the PPL on long context benchmarks improve with PoSE? Does PoSE unlock new capabilities which zero-shot methods do not have? I will update the score to 8 if the authors provide interesting discussions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and offers clear, accessible explanations.\n\n2. The results presented in Table 1 demonstrate the effectiveness of the proposed method in extending the context of Language Model (LLM). Furthermore, the performance of PoSE is comparable to that of full-length fine-tuning, highlighting the evident benefit of skipping contiguous token chunks.\n\n3. PoSE excels in both memory reduction and computation time, as illustrated in Figure 3.\n\n4. PoSE proves to be adaptable across various positional encoding interpolation methods, including linear, NTK, and YaRN, thereby enhancing its versatility and applicability.", "weaknesses": "I am uncertain about whether PoSE has genuinely extended the contextual window. While the current PPL evaluation results are undoubtedly promising, there's a possibility that methods demonstrating low PPL on PG19/GovReport might not effectively comprehend longer contexts. It's possible that the proposed method excels in consistently generating fluent text over extended periods but falls short in truly grasping the nuances of extended context. If that is the case, the potential impact of PoSE could be rather limited, since there are training-free methods (StreamingLLM, LM-Infinite) that can achieve this goal.\n\nTo address this concern, I recommend that the authors consider conducting additional experiments using benchmarks like LongBench and ZeroSCROLLS. Additionally, employing benchmarks utilized by Llama 2 Long could provide valuable insights into the capabilities of PoSE in handling more extensive contextual information.", "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": "XJojzgPGVL", "reviewer": "Reviewer_NXZX", "summary": "the paper presents pose, a novel mechanism to extend the context window of LLMs that decouples train length from target length. It achieves this by simulating longer inputs instead of directly training on longer inputs.", "questions": "Please address the weakness above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. the paper is well written; the method is well illustrated.\n2. the problem is important, and the method can greatly reduced the memory and computation requirements.", "weaknesses": "1. section 3.1 and 3.2 are better to be placed in a different background section, or in related work.\n2. the evaluation for long context is limited - they are mostly all perplexity based. The only available non-perplexity experiment is the one with passkey retrieval, but the baselines are not comprehensive. Can the authors either provide more non-perplexity based measure, or include more baseline for Figure 2?", "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)