Beyond Real: Imaginary Extension of Rotary Position Embeddings for Long-Context LLMs

TL;DR

This paper introduces an extension to Rotary Position Embeddings that re-incorporates the imaginary component of the complex-valued attention scores, significantly improving long-context modeling in Large Language Models.

Contribution

It proposes a novel method that leverages the full complex-valued representation in RoPE, enhancing long-context dependency modeling in LLMs.

Findings

Improved performance on long-context benchmarks

Enhanced modeling of positional information

Benefits increase with longer contexts

Abstract

Rotary Position Embeddings (RoPE) have become a standard for encoding sequence order in Large Language Models (LLMs) by applying rotations to query and key vectors in the complex plane. Standard implementations, however, utilize only the real component of the complex-valued dot product for attention score calculation. This simplification discards the imaginary component, which contains valuable phase information, leading to a potential loss of relational details crucial for modeling long-context dependencies. In this paper, we propose an extension that re-incorporates this discarded imaginary component. Our method leverages the full complex-valued representation to create a dual-component attention score. We theoretically and empirically demonstrate that this approach enhances the modeling of long-context dependencies by preserving more positional information. Furthermore, evaluations on a suite of long-context language modeling benchmarks show that our method consistently improves performance over the standard RoPE, with the benefits becoming more significant as context length increases. The code is available at https://github.com/OpenMOSS/rope_pp.