Frayed RoPE and Long Inputs: A Geometric Perspective

TL;DR

This paper provides a geometric analysis of Rotary Positional Embedding (RoPE), revealing how long inputs cause performance issues by disrupting key/query clustering, and proposes RoPE-ID to enable better generalization to longer inputs.

Contribution

It introduces a geometric perspective on RoPE behavior, identifying the cause of long input degradation, and proposes RoPE-ID, a simple modification for improved long-input performance.

Findings

RoPE damages key/query clustering on long inputs.

RoPE-ID improves long-input handling in large transformers.

Empirical results show RoPE-ID outperforms standard RoPE on benchmarks.

Abstract

Rotary Positional Embedding (RoPE) is a widely adopted technique for encoding position in language models, which, while effective, causes performance breakdown when input length exceeds training length. Prior analyses assert (rightly) that long inputs cause channels to rotate ``out of distribution,'' but it is not clear how extra rotation relates to or causes pathological behavior. Through empirical and theoretical analysis we advance a unified geometric understanding of attention behavior with RoPE. We find that attention induces tight clustering of separated key and query latent point clouds, allowing for creation of sink tokens: placeholders that allow attention heads to avoid token mixing when not required. RoPE applied to longer inputs damages this key/query cluster separation, producing pathological behavior by inhibiting sink token functionality. From this geometric perspective, we propose RoPE-ID (In Distribution), a straightforward modification that allows attention layers to generalize to longer inputs out of the box: apply RoPE with high frequency to a subset of channels. We demonstrate the effectiveness of RoPE-ID for extended inputs using 1B and 3B parameter Transformers on the LongBench and RULER information retrieval benchmarks.