Extending Context Window of Large Language Models from a Distributional Perspective (2024)

Abstract: Scaling the rotary position embedding (RoPE) has become a common method for extending the context window of RoPE-based large language models (LLMs). However, existing scaling methods often rely on empirical approaches and lack a profound understanding of the internal distribution within RoPE, resulting in suboptimal performance in extending the context window length. In this paper, we propose to optimize the context window extending task from the view of rotary angle distribution. Specifically, we first estimate the distribution of the rotary angles within the model and analyze the extent to which length extension perturbs this distribution. Then, we present a novel extension strategy that minimizes the disturbance between rotary angle distributions to maintain consistency with the pre-training phase, enhancing the model’s capability to generalize to longer sequences. Experimental results compared to the strong baseline methods demonstrate that our approach reduces by up to 72% of the distributional disturbance when extending LLaMA2’s context window to 8k, and reduces by up to 32% when extending to 16k. On the LongBench-E benchmark, our method achieves an average improvement of up to 4.33% over existing state-of-the-art methods. Furthermore, Our method maintains the model’s performance on the Hugging Face Open LLM benchmark after context window extension, with only an average performance fluctuation ranging from -0.12 to +0.22.

## introduction given the remarkable capabilities of transformerbased large language models (llms) in addressing a wide range of natural language processing tasks @xcite @xcite 2024), modeling arbitrarily long textual sequences remains a significant challenge. on the one hand, llms trained on short sequences often encounter out-of-distribution (ood) issues when applied to the longer ones @xcite . on the other hand, training an llm with extremely long context windows (i.e., the maximal sequence length) from scratch is expensive and inefficient. currently, the most popular approach is pre-training a large language model, such as llama, qwen2 @xcite , with a limited context window and the rotary position embedding @xcite ). during the inference stage, the context window is dynamically extended via fine-tuning or tuning-free position interpolation strategies @xcite @xcite on the rotary position embedding. however, these position interpolation strategies primarily rely on intuition and are developed from an empirical perspective, resulting in a lack of interpretability @xcite and sub-optimal performance for context extension. for example, pi @xcite equally stretches all dimensions of the rope with the context extension ratio. yarn @xcite observes that heuristically utilizing different strategies for different dimensions yields better performance. however, the reasons behind this phenomenon have not been thoroughly investigated, resulting in it likely not achieving the best results. moreover, the optimal hyperparameters determined experimentally in yarn potentially hinder its generalization to new model settings. to bridge the gap between experiments and theoretical analysis, we tackle context window extension from the view of rotary angle distribution. hence, we propose a method for length extension strategy selection, which has the potential to be theoretically optimal by minimizing the perturbation to the rotary angle distributions of the pre-trained language model. specifically, we first compare the pre-training rotary angle distribution with the distributions introduced by interpolation and extrapolation. as illustrated in figure 1 (a), interpolation can introduce too many ood angles that have a frequency of 0 in the pre-training distribution, indicating a significant disturbance to the original distribution and posing a challenge for the model to adapt to the new distribution. while direct extrapolation may have a negligible impact on the distribution. contrarily in another dimension demonstrated in figure 1(b) , direct extrapolation introduces numerous ood angles in this situation, causing a severe distribution disturbance, whereas interpolation performs better. from such distributional view, we find that the consistency between the pre-training rotary angle distribution and the extension distribution varies across different dimensions. thus, we propose to employ different extension strategies in different dimensions according to the rotary angle distribution. we first approximate the distributions of rotary angles by calculating the frequency of angles in minimal discrete intervals. then, we estimate the disturbance introduced by different extension strategies by computing the distance between the interpolated or extrapolated distribution and the original one. finally, we determine the most appropriate extension strategy for each rotary angle dimension independently. experiments across llms of different sizes and various long-context tasks demonstrate the effectiveness of our distributional approach. we outperform the strong extension baselines pi @xcite and yarn @xcite on longbench-e @xcite , achieving a new state-of-the-art. besides, our method achieves 100% accuracy on passkey retrieval (mohtashami and jaggi, 2023) and matches the performance of original llms on short-text tasks in the huggingface open llm leaderboard @xcite . in summary, our contributions are as follows: ## method in this section, we first introduce how to estimate the rotary angle distribution. then, we propose a novel approach that extends the context window of llms by minimizing the disturbance of the rotary angle distribution. ## experiments in this section, we evaluate our distribution-based method on both long-and short-context benchmarks. the results show that models employing our method outperform existing context window extension methods, indicating a better context window extension of rope-based llms while maintaining their original short-context capabilities. ## analysis in this section, we analyze the impact of distributional disturbance on model performance. moreover, we analyze the selection of different interpolation dimension numbers in eq. ( 9 ) and the impact of the number of intervals in eq. ( 5 ). all analyses are based on the task of extending the context window of llama2-13b from 4k to 8k. ## related works long-sequence modeling is a crucial issue in the application of llms. recent efforts focus on improving position embedding to enable llms have larger context window. currently, the most popular relative position embedding are alibi @xcite and rope @xcite . alibi @xcite but only generalizes to limited lengths on downstream tasks @xcite . rope @xcite cannot generalize to lengths beyond its pre-training length. some works have been done to overcome such limitation. ruoss et al. (2023) 1 )). chen et al. (2023) achieves proposed to extend the context window by interpolating positions, using a scaling factor @xmath0, and fine-tuning on a small amount of data to extend the model's context window. bloc97 (2023b,a) base on the neural tangent kernel (ntk) theory, they scale lower dimensions less and higher dimensions more, this is also referred to as adjusted base frequency (abf). @xcite achieves an effect similar to ntk by modifying the base of rope. yarn @xcite improved ntk by dividing rope dimensions into three frequency-based groups and applying different strategies to each group. low frequency (r @xmath1, dimensions that fall in-between employs the ntk. yarn achieved good performance, but lacked interpretability, the hyperparameters α and β were also empirically chosen, making it hard to obtain the optimal results. different from these empirical method, our work initially highlights the consistency of rotary angle distribution as a theoretical guidance for extending the context window. ## conclusion in this work, we proposed to study the context window extension from a distributional perspective and demonstrated that the consistency of rotary angle distributions has a significant impact on extending the context window of llms based on the rotary position embedding. we designed a framework to select scaling strategies with the guidance of minimizing the disturbance of rotary angle distributions. experimental results demonstrated the effectiveness and superiority of our approach. although our approach is limited by the rotary position embedding, we believe that our distributional perspective has the potential to inspire future work.