Rotary Positional Embeddings

Encoding Relative Position

In the previous section, we implemented scaled dot-product attention, the core mechanism that allows tokens to interact. However, attention itself is permutation-invariant. It treats the input as a bag of tokens, blind to their order. To fix this, we inject positional information.

Absolute positional encodings, like the sinusoidal functions used in the original Transformer, add a unique vector to each token embedding based on its absolute position $m$ . While effective, this approach has limitations. The model must learn relative relationships from absolute signals, which can be inefficient. More importantly, it struggles to generalize to sequence lengths longer than those seen during training.

Rotary Positional Embeddings (RoPE) offer a more elegant solution. Instead of adding a positional vector, RoPE rotates the query and key vectors based on their absolute position. This seemingly simple change has a profound effect: the dot product between a query at position $m$ and a key at position $n$ becomes dependent only on their relative distance, $m-n$ . This encodes relative position directly into the self-attention mechanism.

The Mathematics of Rotation

The goal is to define a function $f$ that transforms query $q_m$ and key $k_n$ (at positions $m$ and $n$ respectively) such that their inner product depends only on the tokens themselves ( $q$ , $k$ ) and their relative position ( $m-n$ ).

$⟨f(q, m), f(k, n)⟩ = g(q, k, m-n)$

To achieve this, RoPE leverages the properties of complex numbers. A vector $x \in \mathbb{R}^d$ can be reinterpreted as a set of $d/2$ complex numbers. The position $m$ is encoded as a rotation. For a 2D vector $[x_1, x_2]$ , a rotation by angle $m\theta$ is a linear operation.

\begin{pmatrix} x'_1 \\ x'_2 \end{pmatrix} = \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix}

This extends to a $d$ -dimensional embedding by pairing up adjacent dimensions and applying a rotation to each pair. Each pair $(x_{2i}, x_{2i+1})$ gets its own rotation frequency $\theta_i$ . By applying this rotation to both the query and key vectors, the dot product becomes invariant to absolute position.

Let $q'_m = R_m q$ and $k'_n = R_n k$ where $R$ is the rotation matrix. The dot product is:

$(R_m q)^T (R_n k) = q^T R_m^T R_n k = q^T R_{n-m} k$

This works because the transpose of a rotation matrix is its inverse ( $R_m^T = R_{-m}$ ), and a rotation by $n$ followed by a rotation by $-m$ is equivalent to a single rotation by $n-m$ . The attention score now explicitly contains relative positional information.

PyTorch Implementation

We first pre-compute the rotary embeddings for all possible positions up to a maximum context length. The frequencies $\theta_i$ are chosen to form a geometric progression, inspired by sinusoidal encodings.

import torch

def precompute_rope_embeddings(embedding_dim: int, max_seq_len: int, theta: float = 10000.0):
    # Frequencies for each dimension pair (d/2)
    inv_freq = 1.0 / (theta ** (torch.arange(0, embedding_dim, 2).float() / embedding_dim))
    
    # All possible positions [0, 1, ..., max_seq_len-1]
    t = torch.arange(max_seq_len, device=inv_freq.device, dtype=inv_freq.dtype)
    
    # Calculate angles for each position and frequency: (max_seq_len, embedding_dim/2)
    freqs = torch.einsum("i,j->ij", t, inv_freq)
    
    # Concatenate to get both sine and cosine components
    # This creates a tensor of shape (max_seq_len, embedding_dim)
    emb = torch.cat((freqs, freqs), dim=-1)
    
    # Convert to complex numbers for rotation: (max_seq_len, embedding_dim/2)
    freqs_complex = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_complex

With the complex-valued frequencies pre-computed, we can apply the rotation to our query and key tensors. The key insight is to view the last dimension of the Q and K tensors (the embedding dimension) as a series of $d/2$ complex numbers and then perform element-wise complex multiplication.

def apply_rotary_embeddings(x: torch.Tensor, freqs_complex: torch.Tensor):
    # Reshape x to view last dim as complex numbers
    # x shape: (batch_size, seq_len, num_heads, head_dim)
    # x_complex shape: (batch_size, seq_len, num_heads, head_dim/2)
    x_complex = x.float().reshape(*x.shape[:-1], -1, 2).view_as_complex()
    
    # Reshape freqs_complex for broadcasting
    # freqs_complex shape: (seq_len, head_dim/2)
    # freqs_complex shape after reshape: (1, seq_len, 1, head_dim/2)
    freqs_complex = freqs_complex.unsqueeze(0).unsqueeze(2)

    # Apply the rotation via complex multiplication
    # The multiplication broadcasts across batch and head dimensions.
    x_rotated = x_complex * freqs_complex
    
    # Reshape back to the original real number format
    # x_out shape: (batch_size, seq_len, num_heads, head_dim)
    x_out = torch.view_as_real(x_rotated).flatten(3)
    
    return x_out.type_as(x)

# --- Example Usage within an attention block ---
# freqs_complex = precompute_rope_embeddings(self.head_dim, self.max_seq_len)
# q_rot = apply_rotary_embeddings(q, freqs_complex)
# k_rot = apply_rotary_embeddings(k, freqs_complex)
# # ... then proceed with attention calculation using q_rot and k_rot

Performance Considerations

RoPE offers significant advantages over traditional absolute positional encodings, but it's not a free lunch. The choice involves trade-offs in performance, implementation complexity, and model behavior.

Feature	RoPE (Rotary)	Absolute (Sinusoidal/Learned)
Position Type	Relative	Absolute
Extrapolation	Excellent; generalizes to unseen sequence lengths	Poor; performance degrades beyond training length
Parameters	Zero (parameter-free)	Potentially millions (if learned)
Locality	Strong local dependency decay	Uniform; no inherent distance decay
Computation	Element-wise multiplication on Q/K	Addition to input embeddings
Compatibility	Works well with linear attention variants	Less compatible with some linear attentions

RoPE's main advantage is its ability to generalize to longer sequences. Because it encodes relative distance, the patterns it learns for a relative distance of +5 are applicable regardless of whether it's between tokens 1 and 6 or tokens 1001 and 1006. Absolute encodings must learn these relationships from scratch for different absolute positions.

Furthermore, RoPE's parameter-free nature makes it more memory-efficient than learned absolute embeddings, especially for models with very large context windows. The computational overhead of applying the rotations is generally minimal compared to the main matrix multiplications in the attention block.

Test your understanding of RoPE's implementation and properties.

Quiz Questions 1/5

What is the primary problem that positional encodings like RoPE are designed to solve in the self-attention mechanism of Transformers?

Quiz Questions 2/5

Instead of adding a positional vector to the token embedding, how does Rotary Positional Embedding (RoPE) encode positional information?

By incorporating relative position directly into the attention mechanism, RoPE provides a robust and scalable method for ordering tokens, making it a cornerstone of many modern LLMs like Llama and GPT-NeoX.

Architecting Transformers from Scratch in Python