AdaRoPE: Not All Attention Heads Should Rotate and Scale Equally

Rotary Position Embeddings (RoPE) are widely adopted in Transformers to encode positional information, yet standard implementations enforce a uniform frequency schedule and scaling across all attention heads. Using simplified retrieval tasks and length generalization scenarios, we show—both empirically and theoretically—that heads with different functional roles require distinct frequency ranges and scaling factors to operate effectively. Ignoring this structure leads to suboptimal utilization of embedding dimensions and degraded performance, particularly under long-context settings. To address these limitations, we proposeAdaRoPE, which equips each attention head with learnable rotation frequencies and scaling factors. Pretrained LLM with AdaRoPE consistently outperforms existing RoPE variants, including partial-RoPE and NoPE baselines. For context extension, we further show that uniform frequency and attention scaling, used in methods such as YaRN, are suboptimal. By applying head-specific scaling, AdaRoPE enables better context extension while better preserving short-context performance in both extrapolation setting and long context continued pretrain setting. These results highlight the importance of optimizing rotary position embeddings at the level of individual attention heads.

Large language models need a way to understand the order and distance of words in a long piece of text. A common method gives every attention head in the model the same rule for handling position. However, different heads often serve different purposes: some focus on nearby words, while others gather information from much farther away. Treating all heads the same can waste model capacity and make long-context performance worse. This paper introduces AdaRoPE, a simple modification that lets each attention head learn its own way of handling position and adjusting attention as the input becomes longer. This allows local heads to keep focusing on nearby tokens, while long-range heads can better use distant information. Across several language models and long-context tasks, AdaRoPE improves both general language performance and the ability to handle longer inputs. The method adds very few parameters and remains compatible with standard training and inference systems.