Effective Length Extrapolation via Dimension-Wise Positional Embeddings Manipulation

TL;DR

This work tackles the challenge of extending context windows in large language models without costly retraining. It introduces Dimension-Wise Positional Embeddings Manipulation (DPE), a training-free framework that identifies and selectively modifies key RoPE dimensions by their effective relative distances and 2-norm attention contributions, then scales the corresponding position indices per dimension group. Empirically, DPE achieves state-of-the-art extrapolation on NIAH, RULER, and InfiniteBench, enabling models like Llama3-8B to reach 128K context while outperforming strong baselines; it also boosts performance within training-length contexts and, in some cases, rivals GPT-4-128K. The approach integrates with Flash Attention 2 and incurs negligible overhead, illustrating the practical value of dimension-aware RoPE manipulation for long-context reasoning and document-processing tasks.

Abstract

Large Language Models (LLMs) often struggle to process and generate coherent context when the number of input tokens exceeds the pre-trained length. Recent advancements in long-context extension have significantly expanded the context window of LLMs but require expensive overhead to train the large-scale models with longer context. In this work, we propose Dimension-Wise Positional Embeddings Manipulation (DPE), a training-free framework to extrapolate the context window of LLMs by diving into RoPE's different hidden dimensions. Instead of manipulating all dimensions equally, DPE detects the effective length for every dimension and finds the key dimensions for context extension. We reuse the original position indices with their embeddings from the pre-trained model and manipulate the key dimensions' position indices to their most effective lengths. In this way, DPE adjusts the pre-trained models with minimal modifications while ensuring that each dimension reaches its optimal state for extrapolation. DPE significantly surpasses well-known baselines such as YaRN and Self-Extend. DPE enables Llama3-8k 8B to support context windows of 128k tokens without continual training and integrates seamlessly with Flash Attention 2. In addition to its impressive extrapolation capability, DPE also dramatically improves the models' performance within training length, such as Llama3.1 70B, by over 18 points on popular long-context benchmarks RULER. When compared with commercial models, Llama 3.1 70B with DPE even achieves better performance than GPT-4-128K.

Figures (7)

• Figure 1: Detecting the effective relative distance across different dimension groups. We show the NIAH accuracy on Llama3-8b in Figure \ref{['fig:distance_detect_acc']} and rank the results in Figure \ref{['fig:distance_detect_rank']}. When the NIAH accuracy is the same, we prioritize ranking based on larger relative distances.
• Figure 2: 2-norm Attention Contribution for different heads and layers of Llama3-8b. We select top-$k$ dimensions as the key dimensions for extrapolation. For example, the top-$1$ dimension is selected in the red dashed line.
• Figure 3: NIAH Accuracy on Llama3-8b. Only top-$k$ dimensions' position indices are scaled.
• Figure 4: An illustrative example of DPE with three main procedures. (a) We detect the effective length $e_{i}=4,e_{j}=5$ for dimension group $g_{i},g_{j}$. (b) Dimensions $29$ and $47$ are identified as key dimensions, and the corresponding effective length is obtained from the detection result. (c) We set $W=2$, and all the position indices are scaled within the effective length, thereby avoiding the impact of OOD position indices for every dimension.
• Figure 5: Ablation study on the rank of effective length and key dimension selection.