PSC : Extending Context Window of Large Language Models via Phase Shift Calibration (2024)

Abstract: Rotary Position Embedding (RoPE) is an efficient position encoding approach and is widely utilized in numerous large language models (LLMs). Recently, a lot of methods have been put forward to further expand the context window based on RoPE. The core concept of those methods is to predefine or search for a set of factors to rescale the base frequencies of RoPE. Nevertheless, it is quite a challenge for existing methods to predefine an optimal factor due to the exponential search space. In view of this, we introduce PSC (Phase Shift Calibration), a small module for calibrating the frequencies predefined by existing methods. With the employment of PSC, we demonstrate that many existing methods can be further enhanced, like PI, YaRN, and LongRoPE. We conducted extensive experiments across multiple models and tasks. The results demonstrate that (1) when PSC is enabled, the comparative reductions in perplexity increase as the context window size is varied from 16k, to 32k, and up to 64k. (2) Our approach is broadly applicable and exhibits robustness across a variety of models and tasks.

## introduction large-scale language models (llms) have shown impressive results across a variety of natural language processing (nlp) applications. for instance, openai has shown that @xcite can perform at a level comparable to humans in a range of professional tasks. additionally, open-source models such as llama2 @xcite and mistral @xcite have made significant contributions to the advancement and practical application of llms in both research and industry. however, one significant challenge that llms face is handling tasks that require processing long context, such as responding to questions based on multiple documents and summarizing lengthy texts such as books. in these scenarios, the perplexity of the responses can increase substantially, leading to a notable decrease in the performance of llms. therefore, equipping llms with long-range ability has become a critical and pressing issue for both academic and commercial sectors. an intuitive method is to fine-tune a pre-trained transformer with a longer context length. nevertheless, there are two limitations: first, models trained in this manner adapt to long context lengths very slowly @xcite ; second, finetuning updates all model parameters is memoryinefficient which prevents the model from adapting to a large context length @xcite . optimizing position encodings is another major direction for extending the context window of llms @xcite . the original transformer @xcite that serves as the core component of llms uses sinusoidal functions of various frequencies to enhance the model's extrapolate capability. it could be regarded as an absolute position encoding mechanism. since then, relative positional encoding techniques such as rope @xcite and alibi @xcite have further increased the length extrapolation of transformers. despite the effectiveness, many existing pre-trained llms that use these positional encoding methods exhibit weak extrapolation capabilities. for example, llama @xcite with 2048 predefined context size explodes perplexity metric when the input texts length is larger than 4096 @xcite . recently, new positional encoding schemes have been proposed to overcome such limitations. @xcite and @xcite show that the effective context size could be extended by modifying rope via position interpolation, which has a much smaller upper bound than the extrapolated method and is more stable @xcite . neural tangent kernel (ntk) theory shows that it's difficult for multilayer perceptron (mlp) to learn high-frequency information in a low-dimensional domain. therefore, ntk-based methods take the high-frequency information into account @xcite . furthermore, yarn hypothesizes that previous methods lead to a closer embedding distribution and remedy the issue by using different interpolating schemes at different frequencies @xcite . the shared characteristic of previous methods is that they utilize predefined frequency rescale factors. some algorithms leverage optimal methods to estimate optimal frequencies directly, such as lon-grope @xcite and clex @xcite . however, due to the exponential search space complexity, it is challenging for those methods to estimate an optimal frequency; they also need heavy searching cost, for instance, it costs longrope nearly 3 days to search an optimal frequency for a 256k context window using an a100 gpu. while existing techniques for encoding positional information are adept at handling long-range dependencies, they often depend on fixed patterns or necessitate extensive searches within large parameter spaces. as a result, adapter-based approaches such as lora @xcite have been utilized to further enhance performance. nonetheless, these methods still face limitations, primarily due to the low rank of the adapter weights @xcite and the inherently high-rank nature of long-context tasks. in this work, we introduce phase shift calibration to assist position encoding methods to improve their long-range capabilities. the main idea is that we propose a module to calibrate the predefined frequency to approximate the optimal frequency. to this end, we first present that there is a rotary transformation between the actual frequencies and the optimal frequencies. the transformation can be represented as a block diagonal matrix. it is full-rank if the predefined frequencies are far from the optimal ones. hence, it is challenging for lowrank adapter methods such as lora to learn the transformation. to remedy this issue, we introduce a calibration module into the base model, which approximates the rotary transformation matrix and helps calibrate the predefined frequencies to the ideal position. we conduct extensive experiments across different llms, position encoding schemes, and various long-context tasks. the results demonstrate the effectiveness of our methods. ## preliminaries and related work rotary position embedding (rope). transformer models leverage positional information to exploit the order of tokens within texts. in our research, we concentrate on rotary position embedding (rope) @xcite and its derivatives. rope acts as the positional encoding technique used across various large language models (llms), such as the llama @xcite and the mistral model @xcite . given a sequence of n word embeddings {x i } n @xmath0, where x i is a d-dimensional vector and d is the dimension of the embedding. rope applies a rotary transformation to each query/key embedding in a pairwise manner. take @xmath1, rope converts each vector into the query vector and key vector via a transformation in a complex space: the rotary transformation introduces an m-n term in the attention score: hence, rope possesses the capability of encoding relative positional information via absolute positional encoding. for a general form where @xmath2, rope divides the d-dimensional space into d/2 2d complex sub-spaces: in matrix form, the rotary-transformed query and key can be expressed as: ## experiments we demonstrate that phase shift calibration successfully realizes the context window extension of large language models by using rope extensions as its position encoding schemes. furthermore, our approach is compatible with a broad range of position encoding techniques, including search-based methods (longrope), position interpolation (pi), and the combination of interpolation and extrapolation techniques (yarn). ## conclusion in this work, we present psc: phase shift calibration, an approach for calibrating the existing extended position encoding methods. we first present that there is a rank inconsistency issue when the predefined frequencies are not optimal. a phase shift calibration module is designed to remedy this issue. we conduct extensive experiments on various tasks, and the results show that psc is compatible with various context extension methods, including interpolation, mixing of interpolation/extrapolation, and search-based techniques. with psc, the longrange abilities of llms can be further enhanced. moreover, our method only introduces a few more parameters (< 1%), which is parameter-efficient. this work thus supports many natural language processing tasks that require long-range capabilities. we discuss several promising future works in the appendix.