PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training

1 INTRODUCTION

• Full-length fine-tuning for context extension is memory- and time-intensive because computational complexity grows quadratically with input length.The passage notes that Chen et al. (2023a) used 32 A100 GPUs to extend LLaMA from 2k to 8k context, and 128 A100 GPUs for larger contexts.Source: §1 · p. 1
• PoSE divides the original context window into chunks and varies their skipping biases and lengths across examples to simulate longer inputs and expose the model to target-window positions.Position indices remain continuous within each chunk, resembling pre-training and helping retain the model’s pre-trained language-modeling and comprehension capacity.Source: §1 · p. 2, Figure 1 · p. 2
• PoSE extends LLaMA from 2k to 128k while retaining decent language-modeling and understanding ability, and is compatible with RoPE-based LLMs and several position-interpolation strategies.The reported compatible models include LLaMA, LLaMA2, GPT-J, and Baichuan; the interpolation strategies include Linear, NTK, and YaRN.Source: §1 · p. 2
• PoSE uses only the original context size for fine-tuning, avoiding the quadratic increase in computation with target length and reducing memory and time overhead.The paper describes this as one of PoSE’s three main advantages.Source: §1 · p. 2
• PoSE can theoretically extend context length indefinitely, with inference memory usage as the stated constraint.The authors connect further scaling potential to continued progress in efficient inference techniques.Source: §1 · p. 2

Training Length-Extrapolatable Models

• Length extrapolation methods aim to handle inputs beyond training length, using positional schemes such as ALiBi, xPos, and NoPos.The passage describes this line of work as training models from scratch for increasing input lengths.Source: §2 · p. 2
• RandPos selects an ordered subset of positions from longer sequences, whereas PoSE fine-tunes pretrained decoder-only LLMs and keeps positions continuous within chunks.The passage contrasts RandPos’s encoder-only pretraining from scratch with PoSE’s context-extension objective for pretrained LLMs.Source: §2 · pp. 2–3

Fine-tuning LLMs for Longer Context

• Position Interpolation, NTK, YaRN, and LongLoRA address context extension, but these methods still require full-length fine-tuning and its target-length-dependent computational cost.PoSE instead uses the original context size for fine-tuning.Source: §2 · p. 3

Memory Transformers

• Memory-transformer approaches use recurrence or retrieval, with the passage noting information loss and limited random access for recurrence and limited interaction between discrete segments for retrieval.Landmark attention is described as supporting random access to input chunks; PoSE is contrasted as providing full input access without modifying attention.Source: §2 · p. 3

3.1 PRELIMINARIES

• RoPE applies rotations to queries and keys to encode position; position interpolation rescales indices, with linear, NTK, and YaRN strategies differing in how they modify positions.Linear interpolation proportionally scales indices; NTK changes RoPE's base, while YaRN blends the approaches across dimensions and adds a temperature factor.Source: §3.1 · pp. 3–4

3.2 PROPOSED APPROACH: POSITIONAL SKIP-WISE TRAINING (POSE)

• Fine-tuning at target length is impractical for extreme context extension because attention computation grows quadratically with sequence length.PoSE instead trains within the original context window and manipulates position indices to simulate longer inputs.Source: §3.2 · p. 4
• PoSE partitions the original context into chunks and adds sampled, nondecreasing skipping biases to their position indices, targeting broad relative-distance coverage without overlaps.The method also aims to keep manipulated positions close to the pre-training structure.Source: §3.2 · p. 4
• PoSE selects continuous text spans for chunks, and alternative content-offset assignments had relatively little impact on fine-tuning outcomes.The explored alternatives used zero offsets for continuous content or aligned manipulated positions with positions in the original text.Source: §3.2 · p. 5
• Using two chunks balances efficiency with effectiveness by exposing models to a wider range of relative positions while staying close to pre-training structure.The paper notes that increasing the chunk count further departs from the pre-training position structure and may harm previously acquired abilities.Source: §3.2 · p. 5

4.1 SETUPS

• Evaluation covers GovReport and Proof-pile language modeling plus synthetic passkey retrieval, which measures the maximum token-attention distance at inference.The reported setup trains LLaMA-7B for 1,000 steps with a global batch size of 64; the fine-tuning data comes from The Pile and requires at least 2,048 tokens.Source: §4.1 · p. 5

4.2 LANGUAGE MODELING

• PoSE achieves comparable long-sequence language-modeling results to Full-length with shorter fine-tuning context and outperforms RandPos.Perplexity decreases as evaluation context grows for PoSE models scaled to 16k and 32k.Source: §4.2 · p. 6, Table 1 · p. 6
• All scaling methods show some performance degradation as supported context length increases.The authors describe this as a trade-off between the number of processable tokens and attention granularity.Source: §4.2 · p. 6

4.3 PASSKEY RETRIEVAL FOR EFFECTIVE CONTEXT WINDOW

• PoSE-16k and PoSE-32k maintain ≥90% passkey retrieval accuracy within their respective context windows, comparable to Full-length.Original, PI-only, and RandPos retrieval accuracy rapidly drops to 0 beyond 2k tokens.Source: §4.3 · p. 6, Figure 2 · p. 7

5.2 COMPATIBILITY WITH ROPE-BASED LLMS AND DIVERSE INTERPOLATION STRATEGIES

• PoSE achieved low perplexity across four RoPE-based models and three interpolation strategies; NTK and YaRN generally outperformed Linear interpolation.The models were LLaMA-7B, LLaMA2-7B, GPT-J-6B, and Baichuan2-7B; NTK sometimes showed a pre-target-length perplexity increase.Source: §5.2 · p. 7, Figure 4 · p. 8

5.1 MEMORY AND TIME EFFICIENCY

• PoSE's memory and time costs stayed fixed and far below full-length fine-tuning as LLaMA-7B scaled from 2k to 4k, 8k, and 16k, with comparable 16k GovReport perplexity at each step.Each method used 1,000 steps; PoSE trained with a 2k context while full-length fine-tuning used 16k for the perplexity comparison.Source: §5.1 · p. 7, Figure 3 · p. 8

5.3 POTENTIAL FOR EXTREMELY-LONG CONTEXT

• PoSE extended LLaMA to 96k with Linear interpolation and to 128k with YaRN on PG-19 and Books3.The evaluation used 20 books from each dataset, each longer than 128k tokens.Source: §5.3 · p. 8, Table 2 · p. 9

5.4 EVALUATION OF CAPABILITY ON ORIGINAL CONTEXT WINDOW

• PoSE-extended models showed only marginal degradation versus full-length fine-tuning and the original LLaMA on standard benchmarks, except for the 128k Linear model.The authors interpret this as preservation of original language-comprehension ability while extending the context window.Source: §5.4 · p. 9, Table 3 · p. 9

• PoSE simulates long inputs by manipulating position indices, decoupling fine-tuning length from target context length.It requires only the original context window for fine-tuning.Source: §6 · p. 9
• Compared with full-length fine-tuning, PoSE greatly reduced memory and time overhead and extended LLaMA to 128k on 8 V100 GPUs with minimal standard-benchmark degradation.Source: §6 · p. 9
• Experiments verified PoSE's compatibility with all RoPE-based LLMs and position interpolation strategies.Source: §6 · p. 9

• PoSE changes chunk position indices to cover wider relative positions in a fixed window, without constraining the text assigned to each chunk.Source: A ABLATION OF TEXT CONTAINED WITHIN… · p. 13
• Uniformly sampled chunk starts, zero starts yielding continuous text, and starts aligned to original text positions had relatively little impact on fine-tuning outcomes.The ablation extended LLaMA-7B from 2k to 16k context.Source: A ABLATION OF TEXT CONTAINED WITHIN… · p. 13

B ANALYSIS OF CHUNK NUMBER N

• With more chunks, coverage probability falls within [0, 2,048] and rises across [2,048, 16,383]; 2,048 chunks give each position near-1 coverage probability.RandPos uses 2,048 chunks and can cover every relative position in each training example by expectation.Source: Figure 5 · p. 13, B ANALYSIS OF CHUNK NUMBER N · p. 14
• PoSE estimates position coverage using Monte Carlo sampling for chunk counts of 2, 3, and 2,048 when extending 2,048 to 16,384 tokens.The unextended model covers positions within 2,048 with probability 1 and positions above 2,048 with probability 0.Source: B ANALYSIS OF CHUNK NUMBER N · p. 13
• Increasing chunk number improves context-extension results, but extremely large chunk counts damage performance as position encodings deviate from pretraining.The authors frame chunk number as a trade-off between training efficiency and performance.Source: Table 5 · p. 14, B ANALYSIS OF CHUNK NUMBER N · p. 14

C SLIDING WINDOW PPL FROM LINEAR / NTK / YARN INTERPOLATION

• NTK shows a significant perplexity increase at 16k, while YaRN perplexity decreases progressively as the context window grows.The passage attributes the NTK increase mainly to difficulty extending the context window by a scaling factor α.Source: C SLIDING WINDOW PPL FROM LINEAR / NTK… · p. 14
• Across Linear, NTK, and YaRN, perplexity follows Full-length ≈ PoSE > PI-only ≫ Original.Table 6 compares all three interpolation strategies across Full-length fine-tuning, PoSE, PI-only, and Original.Source: C SLIDING WINDOW PPL FROM LINEAR / NTK… · p. 14, Table 6 · p. 15