AdaRoPE: Not All Attention Heads Should Rotate and Scale Equally

2026-06-05

AI Summary

Standard Rotary Position Embedding (RoPE) employs uniform rotation frequencies and scaling factors across all attention heads, disregarding their functional heterogeneity and thereby limiting long-context modeling capacity. This work proposes AdaRoPE, which for the first time reveals substantial inter-head variation in positional encoding requirements and introduces learnable, head-specific rotation frequencies and attention scaling factors to enable adaptive position encoding. Built upon the Transformer architecture, AdaRoPE is theoretically grounded and validated through simplified retrieval tasks. It consistently outperforms existing RoPE variants in both pre-trained language models and context extension benchmarks, significantly enhancing long-context performance while preserving short-context accuracy.

Abstract

Rotary Position Embedding (RoPE) is 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 attention 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 propose AdaRoPE, which equips each attention head with learnable rotation frequencies and attention scaling factors. Pretrained LLMs with AdaRoPE consistently outperform 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 the extrapolation setting and the long-context continued pretraining setting. These results highlight the importance of optimizing rotary position embedding at the level of individual attention heads.

Problem

• Rotary Position Embedding
• Attention Heads
• Frequency Schedule
• Attention Scaling
• Long-context Generalization

Innovation

• AdaRoPE
• Rotary Position Embedding
• attention heads
• frequency scheduling
• context extension