Frayed RoPE and Long Inputs: A Geometric Perspective

Davis Wertheimer, Aozhong Zhang, Derrick Liu, Penghang Yin, Naigang Wang

Keywords: RoPE, context length extension, sink tokens, clustering, attention, long context, transformer, language model

TL;DR: We advance a geometric understanding of latent attention dynamics, use it to explain RoPE's failure to generalize to long contexts, and mitigate that failure with a straightforward architectural adjustment

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.

Primary Area: generative models

Submission Number: 10128