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

TL;DR

This work tackles the challenge of extending the effective context window of large language models by refining RoPE-based positional encodings. It introduces Phase Shift Calibration (PSC), a lightweight module that learns a small phase-shift to bring predefined RoPE frequencies closer to an optimal $\theta^*$, improving extrapolation when using extensions such as PI, YaRN, and LongRoPE. The authors demonstrate across multiple models and long-context tasks that PSC yields consistent perplexity improvements, enhanced passkey retrieval, and competitive standard benchmarks, with gains growing as the context length increases. PSC is parameter-efficient, adding less than $1\%$ more parameters, and is compatible with a range of RoPE-based extensions, making it a practical approach for robust long-context modeling in real-world LLM deployments.

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. The code can be found at https://github.com/WNQzhu/PSC.

Figures (9)

• Figure 1: Phase shift leads to the $\sin/\cos$ values deviating from their optimal positions. The $\theta^*$ is assumed to be an optimal frequency.
• Figure 2: The embeddings are calibrated to an ideal position, and then existing position encode methods are adopted.
• Figure 3: A comparison of passkey retrieval accuracy for context-augmented Large Language Models (LLMs). "Extend YaRN" indicates that the model incorporates YaRN without undergoing fine-tuning. "FT" denotes that the models have been fine-tuned using LoRA (r=8), while "PSC" signifies that the models have been fine-tuned with the phase shift calibration module activated. (The graphs for LLaMA-2 7B (PI PSC) and LLaMA-2 (YaRN PSC) coincide as they exhibit the same results: with 100% accuracy up to 34k.)
• Figure 4: An ablation study on the fine-tuning process utilizing phase shift calibration. The perplexity is assessed with a context length of 32k.
• Figure 5: GPU memory consumption by LoRA and PSC.