RoPE: How 2D Rotations Solved Transformer Long-Context

Rotary Position Embedding (RoPE), introduced by Su et al. in 2021, solves the Transformer long-context problem by rotating Query and Key vectors in 2D sub-planes, making attention depend only on relative token offsets. This approach avoids the memory overhead and GPU fusion issues of prior methods while enabling models to handle sequence lengths beyond their training window.

Attention requires two things from positional encodings: Original absolute positional encodings added static sine/cosine waves directly to input embeddings: This forced the model to burn parameter capacity un-mixing semantic meaning ei from positional location pi . Worse, if a model was trained on sequence lengths up to N=2048 , position 2049 presented an unseen vector p2049 , causing immediate generation collapse. Subsequent approaches tried adding relative bias terms bi,j directly into the attention matrix QKT+B . While functionally effective, modifying the N×N matrix broke kernel-level GPU fusions like FlashAttention and introduced huge memory overheads. Enter Rotary Position Embedding RoPE Su et al., 2021 . Instead of adding positional vectors or patching the attention matrix, RoPE rotates the Query and Key vectors in 2D sub-planes before computing attention. To make attention relative, we want an encoding function R x,m applied to a vector x at position m such that the dot product between Query at position m and Key at position n depends only on the relative offset m−n : How do you preserve vector norms while encoding an angle shift? Complex space rotations. In a 2D plane, rotating a vector x= x1,x2 T by an angle mθ is represented by the orthogonal rotation matrix: When you compute the dot product between a rotated Query at position m and a rotated Key at position n : Because RT m R n =R n−m , the absolute positions m and n cancel out completely. The resulting attention weight is strictly a function of the distance m−n . For a d -dimensional vector, RoPE splits the channels into d/2 pairs of 2D planes, applying a different rotation frequency θi to each pair: In practice, explicitly constructing full block-diagonal rotation matrices for every head and token is slow. We can implement RoPE efficiently using the complex number representation or the real-valued slice trick: where x~= −x2,x1,−x4,x3,… . Here is a clean, production-ready PyTorch implementation: python python import torch import torch.nn as nn class RotaryPositionalEmbedding nn.Module : """ Rotary Position Embedding RoPE for multi-head attention. Applies 2D plane rotations across d model / 2 feature pairs. """ def init self, dim: int, max seq len: int = 8192, base: float = 10000.0 : super . init self.dim = dim self.base = base Compute frequency bands theta i = 10000^ -2 i-1 /d inv freq = 1.0 / self.base torch.arange 0, dim, 2 .float / dim self.register buffer "inv freq", inv freq, persistent=False Precompute cosine and sine cache for max seq len self. build cache max seq len def build cache self, max seq len: int : t = torch.arange max seq len, dtype=self.inv freq.dtype Outer product: seq len, dim / 2 freqs = torch.outer t, self.inv freq Duplicate frequencies to match feature dimensions: seq len, dim emb = torch.cat freqs, freqs , dim=-1 self.register buffer "cos cached", emb.cos , persistent=False self.register buffer "sin cached", emb.sin , persistent=False def rotate half self, x: torch.Tensor - torch.Tensor: """Splits vector in half and negates/swaps halves: -x2, x1 """ x1 = x ..., : self.dim // 2 x2 = x ..., self.dim // 2 : return torch.cat -x2, x1 , dim=-1 def forward self, x: torch.Tensor, seq len: int - torch.Tensor: """ Args: x: Tensor of shape batch size, num heads, seq len, head dim seq len: Current sequence length Returns: Tensor with RoPE applied to Query or Key """ cos = self.cos cached :seq len, : .unsqueeze 0 .unsqueeze 1 1, 1, seq len, dim sin = self.sin cached :seq len, : .unsqueeze 0 .unsqueeze 1 1, 1, seq len, dim Apply elementwise rotation formula return x cos + self. rotate half x sin