Relative Positional Information in Permutation-Equivariant Attention
Standard Multi-Head Self-Attention is permutation-equivariant: for any permutation matrix , . To represent sequential word order, Transformers must inject explicit positional inductive biases.
Early positional encoding strategies present fundamental mathematical and operational trade-offs:
- Absolute Sinusoidal Encodings (Vaswani et al., 2017) add fixed frequency vectors to input embeddings (), failing to preserve clean relative distance invariance in inner product operations .
- Learned Absolute Embeddings (BERT, GPT-2) introduce static lookup tables , hard-capping maximum sequence length and failing to generalize to unseen context lengths ().
- Learned Relative Position Biases (T5, Raffel et al., 2020) add scalar bias matrices , introducing memory lookups and breaking compatibility with hardware-fused attention kernels (FlashAttention).
Rotary Position Embedding (RoPE) (Su et al., 2021) resolves these trade-offs by parameterizing absolute positions via block-diagonal orthogonal rotation matrices in the complex plane, provably inducing exact relative distance dependencies in the self-attention inner product with zero additional trainable parameters and zero KV-cache memory overhead.
Positional Encoding Paradigms Comparison
| Positional Paradigm | Mathematical Operation | Relative Distance Invariance | Extrapolation Capacity | Learnable Parameters | FlashAttention Kernel Fused |
|---|---|---|---|---|---|
| Absolute Sinusoidal (Vaswani) | Additive: | Approximate (Trig expansion) | Moderate | Yes (Input Add) | |
| Learned Absolute (GPT-2) | Additive: | None | Catastrophic () | Yes (Input Add) | |
| T5 Relative Bias (Raffel et al.) | Logit Bias: | Exact (Discretized buckets) | Moderate | No ( Matrix Add) | |
| ALiBi (Press et al.) | Linear Logit Bias: $-m \cdot | m-n | $ | Exact (Linear decay penalty) | Strong |
| RoPE (Su et al.) | Multiplicative Rotation: | Exact (Algebraic Rotation) | High (via Base Scaling) | Yes (In-Place Q/K Rotate) | |
| YaRN (Peng et al.) | Frequency-Split Scaled RoPE | Exact (Multi-Scale Interpolated) | Frontier () | Yes (Precomputed Angles) |
Mathematical Foundations
flowchart LR
Q_IN["Unrotated Query Vector q_m in R^d\n(Layer Norm -> Linear Projection)"] --> SPLIT["Partition into d/2 Orthogonal 2D Subspaces\n(q_2i, q_2i+1) in R^2"]
subgraph ROTATION_MANIFOLD["Complex 2D Plane Rotation"]
COMPLEX["Map to Complex Number:\nz_i = q_2i + i * q_2i+1"]
PHASE["Compute Phase Rotation Factor:\nexp(i * m * theta_i) where theta_i = base^(-2i/d)"]
MULT["Complex Multiplication:\nz_rot = z_i * exp(i * m * theta_i)"]
COMPLEX --> MULT
PHASE --> MULT
end
SPLIT --> COMPLEX
MULT --> UNPACK["Unpack Real Coordinates:\n[ q_2i * cos(m*theta_i) - q_2i+1 * sin(m*theta_i),\n q_2i * sin(m*theta_i) + q_2i+1 * cos(m*theta_i) ]"]
UNPACK --> Q_OUT["Rotated Query Vector q_tilde_m in R^d"]
Figure 1: RoPE vector transformation decomposing d-dimensional embeddings into d/2 independent 2D plane rotations by position angle .
1. The 2D Complex Rotation Formulation
Let and represent 2-dimensional query and key vectors at sequence positions and . We seek transformation functions and whose inner product depends strictly on the relative displacement :
Treating 2D vectors as complex numbers and , the rotation operator evaluates to multiplication by the unitary complex exponential :
In matrix form, this corresponds to the orthogonal rotation matrix :
2. Relative Distance Invariance Proof
Computing the Euclidean inner product between the rotated query and key vectors:
\langle \mathbf{R}_{\Theta, m}\mathbf{q}_m, \mathbf{R}_{\Theta, n}\mathbf{k}_n \rangle &= (\mathbf{R}_{\Theta, m}\mathbf{q}_m)^\top (\mathbf{R}_{\Theta, n}\mathbf{k}_n) \\ &= \mathbf{q}_m^\top \left( \mathbf{R}_{\Theta, m}^\top \mathbf{R}_{\Theta, n} \right) \mathbf{k}_n \end{aligned}$$ By the group properties of planar rotation matrices $\mathbf{R}(\alpha)^\top \mathbf{R}(\beta) = \mathbf{R}(-\alpha)\mathbf{R}(\beta) = \mathbf{R}(\beta - \alpha)$: $$\mathbf{R}_{\Theta, m}^\top \mathbf{R}_{\Theta, n} = \mathbf{R}_{\Theta, n - m} = \begin{pmatrix} \cos((n-m)\theta) & -\sin((n-m)\theta) \\ \sin((n-m)\theta) & \cos((n-m)\theta) \end{pmatrix}$$ $$\boxed{\langle \mathbf{R}_{\Theta, m}\mathbf{q}_m, \mathbf{R}_{\Theta, n}\mathbf{k}_n \rangle = \mathbf{q}_m^\top \mathbf{R}_{\Theta, n - m} \mathbf{k}_n = g(\mathbf{q}_m, \mathbf{k}_n, n - m)}$$ The inner product is mathematically invariant to absolute sequence shifts $m \to m + k, n \to n + k$, depending strictly on relative displacement $n - m$. --- ### 3. Generalization to $d$-Dimensional Vector Spaces For a $d$-dimensional embedding space (where $d$ is even), RoPE constructs a block-diagonal rotation matrix $\mathbf{R}_{\Theta, m}^d \in \mathbb{R}^{d \times d}$ spanning $d/2$ orthogonal 2D subspaces: $$\mathbf{R}_{\Theta, m}^d = \operatorname{diag}\left( \mathbf{R}_{\theta_0, m}, \mathbf{R}_{\theta_1, m}, \dots, \mathbf{R}_{\theta_{d/2 - 1}, m} \right)$$ where frequencies $\theta_i$ follow an exponentially decaying geometric progression: $$\theta_i = b^{-2i / d}, \quad i \in \left\{0, 1, \dots, \frac{d}{2} - 1\right\}$$ The base constant $b$ defaults to $b = 10{,}000$ (RoFormer, LLaMA-1/2) and scales up to $b = 500{,}000$ in LLaMA-3/3.1 for long-context stability. The inner product over $d$ dimensions expands to: $$\langle \tilde{\mathbf{q}}_m, \tilde{\mathbf{k}}_n \rangle = \sum_{i=0}^{d/2 - 1} \operatorname{Re}\left[ \mathbf{q}_m^{(i)} \overline{\mathbf{k}_n^{(i)}} e^{i(m - n)\theta_i} \right] = \sum_{i=0}^{d/2 - 1} \left( a_i \cos((m-n)\theta_i) + b_i \sin((m-n)\theta_i) \right)$$ where $a_i = q_{2i}k_{2i} + q_{2i+1}k_{2i+1}$ and $b_i = q_{2i}k_{2i+1} - q_{2i+1}k_{2i}$. --- ### 4. Multi-Dimensional Axial RoPE (2D / 3D Vision & Video) In Diffusion Transformers (DiT) and video generation architectures (LTX-2, Sora), RoPE generalizes to spatial-temporal coordinates $\mathbf{p} = (p_x, p_y, p_t) \in \mathbb{N}^3$ by partitioning the head dimension $d$ across independent coordinate axes: $$\mathbf{q}_{\text{axial}} = \left[ \mathbf{R}_{\Theta_x, p_x} \mathbf{q}_{0:d_x}, \; \mathbf{R}_{\Theta_y, p_y} \mathbf{q}_{d_x:d_x+d_y}, \; \mathbf{R}_{\Theta_t, p_t} \mathbf{q}_{d_x+d_y:d} \right]$$ The composite multi-dimensional self-attention kernel factors into independent multiplicative components: $$\langle \tilde{\mathbf{q}}_{\mathbf{p}}, \tilde{\mathbf{k}}_{\mathbf{p}'} \rangle = \rho_x(p_x - p_x') \cdot \rho_y(p_y - p_y') \cdot \rho_t(p_t - p_t')$$ --- ### 5. Long-Context Scaling Dynamics (PI, NTK-Aware, YaRN) When evaluating models at context lengths exceeding training bounds ($L_{\text{eval}} > L_{\text{train}}$), high-frequency rotary phases cycle beyond observed empirical bounds, causing attention collapse. ```mermaid flowchart TD subgraph BASE["Original RoPE (L_train)"] BASE_FREQ["theta_i = base^(-2i/d)\nPhases wrap rapidly at L_eval > L_train"] end subgraph PI["Linear Position Interpolation (PI)"] PI_FREQ["theta_i' = theta_i / s\nUniform down-scaling across all dimensions\nDegrades high-frequency local resolution"] end subgraph NTK["NTK-Aware Scaling"] NTK_FREQ["base' = base * s^(d / (d - 2))\nPreserves high frequencies (local tokens)\nStretches low frequencies (long-range tokens)"] end subgraph YARN["YaRN Multi-Scale Scaling"] YARN_FREQ["Dynamic blend: PI for low-freq, NTK for mid-freq, Unscaled for high-freq\n+ Attention Temperature Scaling: sqrt(1 / (0.1 * ln(s) + 1))"] end BASE --> PI BASE --> NTK BASE --> YARN ``` #### 1. Position Interpolation (PI; Chen et al., 2023): Uniformly compresses position indices $m \to m / s$, where $s = L_{\text{eval}} / L_{\text{train}}$: $$\theta_i' = \frac{\theta_i}{s}$$ #### 2. NTK-Aware Scaling (bloc97, 2023): Scales the base constant $b \to b'$ rather than position indices directly, preserving high-frequency local positional resolution while expanding low-frequency ranges: $$b' = b \cdot s^{\frac{d}{d - 2}}$$ #### 3. YaRN (Yet another RoPE extensioN; Peng et al., 2023): Applies ramped piecewise interpolation across dimension ratios $r_i = \frac{2\pi}{\theta_i L_{\text{train}}}$ and scales attention entropy: $$\tilde{\theta}_i = (1 - \alpha_i) \frac{\theta_i}{s} + \alpha_i \theta_i, \quad \alpha_i = \operatorname{clamp}\left(\frac{r_i - \beta_{\text{low}}}{\beta_{\text{high}} - \beta_{\text{low}}}, 0, 1\right)$$ $$\operatorname{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \operatorname{softmax}\left( \frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d} \cdot t_{\text{scale}}} \right)\mathbf{V}, \quad t_{\text{scale}} = \sqrt{0.1 \ln(s) + 1}$$ --- ## Production PyTorch Implementation Below is a complete, modular PyTorch implementation of the **RotaryEmbedding** module with precomputed cosine/sine buffers, vectorized in-place rotation kernels, and multi-head attention integration with **KV-cache handling**. ### Step 1: Vectorized Rotary Embedding Module ```python title="src/rotary_embedding.py" from __future__ import annotations import torch import torch.nn as nn def rotate_half(x: torch.Tensor) -> torch.Tensor: """Rotates half the hidden dimensions: [-x2, x1].""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Applies RoPE rotation to query and key tensors.""" # cos, sin: [1, seq_len, 1, head_dim] or [seq_len, head_dim] if position_ids is not None: cos = cos[position_ids].unsqueeze(2) # [batch, seq_len, 1, head_dim] sin = sin[position_ids].unsqueeze(2) else: cos = cos.unsqueeze(0).unsqueeze(2) # [1, seq_len, 1, head_dim] sin = sin.unsqueeze(0).unsqueeze(2) # In-place complex rotation formulation: (x * cos) + (rotate_half(x) * sin) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed.to(q.dtype), k_embed.to(k.dtype) class RotaryEmbedding(nn.Module): """Precomputes and caches rotary frequencies for LLM attention.""" def __init__( self, dim: int, max_position_embeddings: int = 8192, base: float = 10000.0, scaling_factor: float = 1.0, scaling_type: str = "none", # "none", "linear", "ntk" ) -> None: super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base self.scaling_factor = scaling_factor self.scaling_type = scaling_type # Adjust base for NTK scaling if scaling_type == "ntk" and scaling_factor > 1.0: self.base = base * (scaling_factor ** (dim / (dim - 2))) self._set_cos_sin_cache(max_position_embeddings) def _set_cos_sin_cache(self, seq_len: int) -> None: self.max_seq_len_cached = seq_len inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) if self.scaling_type == "linear" and self.scaling_factor > 1.0: inv_freq = inv_freq / self.scaling_factor t = torch.arange(seq_len, dtype=torch.float32) freqs = torch.outer(t, inv_freq) # [seq_len, dim / 2] # Duplicate across contiguous pairs: [theta_0, theta_0, theta_1, theta_1, ...] emb = torch.cat((freqs, freqs), dim=-1) # [seq_len, dim] self.register_buffer("cos_cached", emb.cos(), persistent=False) self.register_buffer("sin_cached", emb.sin(), persistent=False) def forward(self, seq_len: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len) return ( self.cos_cached[:seq_len].to(device), self.sin_cached[:seq_len].to(device), ) ``` --- ### Step 2: Attention Layer with RoPE and KV-Cache Support ```python title="src/attention_rope.py" from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from rotary_embedding import RotaryEmbedding, apply_rotary_pos_emb class MultiHeadAttentionWithRoPE(nn.Module): """Multi-Head Attention integrating RoPE and dynamic KV-caching.""" def __init__( self, d_model: int = 4096, num_heads: int = 32, num_kv_heads: int = 8, # Grouped-Query Attention (GQA) max_seq_len: int = 8192, rope_base: float = 500000.0, ) -> None: super().__init__() self.d_model = d_model self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = d_model // num_heads # 128 self.num_queries_per_kv = num_heads // num_kv_heads self.q_proj = nn.Linear(d_model, num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(num_heads * self.head_dim, d_model, bias=False) self.rotary_emb = RotaryEmbedding( dim=self.head_dim, max_position_embeddings=max_seq_len, base=rope_base, ) def forward( self, hidden_states: torch.Tensor, position_ids: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: batch_size, seq_len, _ = hidden_states.shape # 1. Project Q, K, V q = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) # 2. Compute RoPE Rotations cos, sin = self.rotary_emb(seq_len=position_ids.max().item() + 1, device=hidden_states.device) q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids) # 3. KV-Cache Update (Keys are stored already rotated) if kv_cache is not None: past_k, past_v = kv_cache k = torch.cat([past_k, k], dim=-2) v = torch.cat([past_v, v], dim=-2) current_kv_cache = (k, v) # 4. GQA Expansion if self.num_queries_per_kv > 1: k_expanded = k.repeat_interleave(self.num_queries_per_kv, dim=1) v_expanded = v.repeat_interleave(self.num_queries_per_kv, dim=1) else: k_expanded, v_expanded = k, v # 5. Scaled Dot-Product Attention is_causal = (seq_len > 1) and (kv_cache is None) attn_out = F.scaled_dot_product_attention( q, k_expanded, v_expanded, is_causal=is_causal ) attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) return self.o_proj(attn_out), current_kv_cache ``` --- ## Empirical Benchmark Evaluation Context extrapolation fidelity, effective sequence length, and retrieval pass rates across RoPE configurations: | Model Backbone | Base Context ($L_{\text{train}}$) | RoPE Strategy | Base Constant ($\theta_{\text{base}}$) | Max Effective Context | Needle-in-a-Haystack Pass Rate | Perplexity (32K Context $\downarrow$) | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | **LLaMA-2 (7B)** | $4\text{K}$ | Vanilla RoPE | $10{,}000$ | $4\text{K}$ | $100\%$ ($4\text{K}$) / $12\%$ ($8\text{K}$) | $> 100.0$ (Diverged) | | | $4\text{K}$ | Position Interpolation (PI) | $10{,}000$ | $16\text{K}$ | $94.2\%$ | $8.42$ | | | $4\text{K}$ | NTK-Aware Scaling | $80{,}000$ | $32\text{K}$ | $96.8\%$ | $6.84$ | | | $4\text{K}$ | **YaRN ($s=8$)** | **$10{,}000$** | **$32\text{K}$** | **$99.4\%$** | **$5.92$** | | **LLaMA-3 (8B)** | $8\text{K}$ | Extended $\theta$ Base | $500{,}000$ | $16\text{K}$ | $98.1\%$ | $5.12$ | | **LLaMA-3.1 (8B)**| $128\text{K}$ | **Dual-Band RoPE + FT**| **$500{,}000$** | **$128\text{K}$** | **$100.0\%$ (All 128K)** | **$4.18$** | | **Qwen-2.5 (7B)** | $32\text{K}$ | **Dual-Chunk RoPE** | **$1{,}000{,}000$** | **$128\text{K}$** | **$99.8\%$** | **$4.06$** | --- ## Troubleshooting Common RoPE Implementation Faults ### 1. Re-Applying RoPE to Cached Keys During Autoregressive Decoding - **Symptom**: Model generates coherent tokens for the prompt prefill phase, then outputs complete gibberish after generation step 1. - **Root Cause**: Passing the entire accumulated key cache through `apply_rotary_pos_emb` during single-token decoding steps, compounding phase rotations $\mathbf{k}_t \cdot e^{i 2t \theta}$. - **Remedy**: Apply RoPE strictly to the newly generated key token $\mathbf{k}_{\text{new}}$ at position $t$, concatenating it to the pre-rotated cache tensor: `past_k = torch.cat([past_k, k_new_rotated], dim=-2)`. ### 2. Applying RoPE Before QKV Linear Projections - **Symptom**: Training loss fails to converge, or downstream fine-tuning destroys model instruction-following capabilities. - **Root Cause**: Rotating hidden states before linear projection ($\mathbf{W}_q \mathbf{R}_{\Theta}\mathbf{x}$), which violates orthogonal cancellation in inner products ($\mathbf{W}_q^\top \mathbf{W}_k \neq \mathbf{I}$). - **Remedy**: Always apply RoPE after RMSNorm and linear projection: $\mathbf{q} = \mathbf{R}_{\Theta, m}(\mathbf{W}_q \operatorname{RMSNorm}(\mathbf{x}))$. ### 3. Numerical Precision Underflow in Low Frequencies - **Symptom**: Long-context inference ($L > 64\text{K}$) exhibits NaN loss spikes or degraded retrieval accuracy in FP16 precision. - **Root Cause**: High-index frequencies $\theta_i = b^{-2i/d} \ll 10^{-6}$ lose significant floating-point bits when multiplied by large position indices $m > 65{,}536$. - **Remedy**: Precompute and cache the frequency table in `torch.float32`, casting to `torch.bfloat16` only immediately prior to tensor addition/multiplication. --- ## References 1. Su, J., Lu, Y., Pan, S., Wen, B., & Liu, Y. (2021). *RoFormer: Enhanced Transformer with Rotary Position Embedding*. Neurocomputing 2024 / arXiv:2104.09864. 2. Chen, S., Wong, S., Chen, L., & Tian, Y. (2023). *Extending Context Window of Large Language Models via Positional Interpolation*. arXiv:2306.15595. 3. Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). *YaRN: Efficient Context Window Extension of Large Language Models*. ICLR 2024. 4. Press, O., Smith, N. A., & Lewis, M. (2022). *Train Short, Test Long: Attention with Linear Biases Enables Input Length Generalization (ALiBi)*. ICLR 2022. 5. Touvron, H., et al. (2023). *LLaMA: Open and Efficient Foundation Language Models*. arXiv:2302.13971.