RoPE \u2014 Rotary Position Embedding

Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding", 2021

Learning Objectives

After completing this section, you will be able to:

1. Explain why encoding position through rotation is fundamentally different from adding position to embeddings or biasing attention scores, and what mathematical advantage this provides.
2. Derive and interpret the 2D rotation formula that transforms Query and Key vectors, understanding each symbol and matrix operation involved.
3. Understand the multi-frequency design that gives RoPE both local and global position sensitivity — and why this allows generalization to unseen sequence lengths.
4. Implement RoPE from scratch in NumPy and PyTorch, applying it to the shared "the cat sat on the mat" example.
5. Connect RoPE to modern systems including LLaMA, GPT-NeoX, PaLM, CodeLlama, and understand how it interacts with Flash Attention, KV-cache, and context length extension techniques like YaRN and NTK-aware scaling.

Where RoPE appears: LLaMA 1/2/3 (Meta), GPT-NeoX (EleutherAI), PaLM/Gemini (Google), CodeLlama, Mistral, Qwen, Falcon, Phi — virtually every major open-weight LLM since 2022 uses RoPE. It has become the default positional encoding for modern transformers.

The Real Problem

Transformers process all tokens in parallel — unlike RNNs, they have no built-in notion of sequential order. Without positional information, the sentence "the cat sat on the mat" would be indistinguishable from "mat the on sat cat the." The attention mechanism computes $Q \cdot K^T$, which depends only on what the tokens are, not where they are. This is the position encoding problem, and it has been solved in several ways, each with different tradeoffs.

Position as Addition — The Limitation

The original Transformer (Vaswani et al., 2017) solved this by adding a position vector to each token embedding: $x_m = \text{token}_m + \text{PE}(m)$. This is simple and elegant, but it has two fundamental problems.

First, the position signal gets diluted as it passes through layers. After multiple rounds of self-attention, layer normalization, and feed-forward transformations, the additive positional component can lose its influence. The model must "remember" the position it was told about many layers ago.

Second, the dot product $(x_m + \text{PE}(m))^T(x_n + \text{PE}(n))$ expands into four terms: $x_m^T x_n + x_m^T \text{PE}(n) + \text{PE}(m)^T x_n + \text{PE}(m)^T \text{PE}(n)$. Only the last term captures purely positional interaction. The cross terms $x_m^T \text{PE}(n)$ mix content and position in ways the model must learn to disentangle — this is wasted capacity.

Position as Score Bias — Better but Still External

Relative position biases (Chapter 7) improve on this by adding a learned bias directly to the attention score: $\text{score}(m, n) = Q_m \cdot K_n + b(m - n)$. This cleanly separates content matching from position matching. However, the bias $b(m-n)$ is typically a learned lookup table with a fixed maximum distance. When the model encounters a sequence longer than what it was trained on, the bias has no entry for large $|m - n|$ values, and generalization breaks down.

From Intuition to Rotation

The Rotation Insight

Jianlin Su and his coauthors at ZhuiYi Technology asked a beautiful question: is there a transformation $f(x, m)$ that we can apply to both Q and K such that the inner product $\langle f(Q_m, m), f(K_n, n) \rangle$ depends only on $Q_m$, $K_n$, and the relative position $(m - n)$?

The answer comes from a well-known property of 2D rotations. If you rotate a vector $\mathbf{a}$ by angle $\alpha$ and another vector $\mathbf{b}$ by angle $\beta$, their dot product depends only on the difference of angles $(\alpha - \beta)$:

$R(\alpha)\mathbf{a} \cdot R(\beta)\mathbf{b} = \mathbf{a}^T R(\alpha - \beta) \mathbf{b}$

This is the key identity. If we set $\alpha = m\theta$ and $\beta = n\theta$ for some frequency $\theta$, then the dot product after rotation depends on $(m - n)\theta$ — exactly the relative position, scaled by frequency. This is not an approximation. It is an exact mathematical consequence of rotation in 2D.

The idea is simple: rotate Q by position $m$ and rotate K by position $n$ before computing their dot product. The dot product will automatically encode relative position.

The Multi-Frequency Design

A single rotation frequency can only encode one scale of position. To encode both "this word is 2 positions away" (local) and "this word is 500 positions away" (global), RoPE uses multiple frequencies, one for each pair of embedding dimensions. This is directly inspired by the sinusoidal positional encoding from the original Transformer, which also uses a geometric sequence of frequencies.

With embedding dimension $d$, RoPE splits the vector into $d/2$ pairs. Each pair $i$ gets its own frequency $\theta_i = 1/10000^{2i/d}$. The first pair ($i=0$) has $\theta_0 = 1.0$ — it rotates by 1 radian per position (fast, local). The last pair has a tiny $\theta$ and barely rotates even at position 1000 (slow, global). Together, these form a multi-scale positional representation, analogous to wavelets or Fourier decomposition.

The Mathematical Definition

Symbol-by-Symbol Breakdown

RoPE transforms a vector $\mathbf{x}$ at position $m$ by applying a 2D rotation to each consecutive pair of dimensions. For dimension pair $i$ (pairing the $i$-th first-half dimension with the $i$-th second-half dimension), the rotation is:

$x_i' = x_i \cos(m\theta_i) - x_{i+d/2} \sin(m\theta_i)$

$x_{i+d/2}' = x_i \sin(m\theta_i) + x_{i+d/2} \cos(m\theta_i)$

Let us define every symbol precisely:

The Frequency Formula

The frequency for dimension pair $i$ is: $\theta_i = \frac{1}{10000^{2i/d}}$

This creates a geometric sequence from $\theta_0 = 1.0$ down to $\theta_{d/2-1} \approx 0.0001$ for typical $d = 128$. The base 10000 was chosen by Vaswani et al. (2017) for the original sinusoidal encoding and reused by Su et al. for RoPE.

For our example with $d = 4$:

• Pair 0 ($i = 0$): $\theta_0 = 1/10000^{0/4} = 1/10000^0 = 1.000000$ — fast rotation, 1 radian per position (57.3° per step)
• Pair 1 ($i = 1$): $\theta_1 = 1/10000^{2/4} = 1/10000^{0.5} = 1/100 = 0.010000$ — slow rotation, 0.01 radian per position (0.57° per step)

Pair 0 rotates 100 times faster than pair 1. This means pair 0 distinguishes nearby positions (local context), while pair 1 distinguishes distant positions (global context). In a real model with $d = 128$, you would have 64 pairs spanning five orders of magnitude in frequency.

Why the Dot Product Depends Only on Relative Position

Consider a single dimension pair. After rotating Q at position $m$ and K at position $n$, the contribution of this pair to the dot product is:

$(q_1 \cos m\theta - q_2 \sin m\theta)(k_1 \cos n\theta - k_2 \sin n\theta) + (q_1 \sin m\theta + q_2 \cos m\theta)(k_1 \sin n\theta + k_2 \cos n\theta)$

Expanding and using the identity $\cos(A)\cos(B) + \sin(A)\sin(B) = \cos(A - B)$:

$= q_1 k_1 \cos((m-n)\theta) - q_1 k_2 \sin((m-n)\theta) + q_2 k_1 \sin((m-n)\theta) + q_2 k_2 \cos((m-n)\theta)$

Every term depends on $(m - n)\theta$, the relative distance scaled by frequency. The absolute positions $m$ and $n$ have vanished. This is the mathematical guarantee that RoPE provides: the attention score is a function of content and relative position only.

Interactive: 2D Rotation Visualizer

The visualizer below shows how RoPE rotates a single dimension pair. Select a token to see how its position determines the rotation angle. Notice that "The" (position 0) has zero rotation, while later tokens rotate progressively more.

Interactive: Multi-Frequency Explorer

This explorer shows all dimension pairs simultaneously. Move the position slider and observe how low-index pairs (red) oscillate rapidly while high-index pairs (purple/pink) barely move. This multi-scale structure is what enables RoPE to encode both local and global position.

Step-by-Step Calculation

We use the same shared example from every chapter: tokens = ["The", "cat", "sat", "on", "mat"] with $d = 4$ and the standard Q, K, V matrices. The implementation pairs the first half of dimensions with the second half: pair 0 = (col 0, col 2), pair 1 = (col 1, col 3).

Step 1: Compute Frequencies

With $d = 4$ and base = 10000, we have 2 dimension pairs:

• $\theta_0 = 1/10000^{0/4} = 1.000000$ (pair 0: dims 0 and 2)
• $\theta_1 = 1/10000^{2/4} = 0.010000$ (pair 1: dims 1 and 3)

Rotation angles = position $\times$ frequency:

Step 2: Rotate All Q Vectors

For each token, we split Q into first-half $x_1 = Q[\text{:,\,:2}]$ and second-half $x_2 = Q[\text{:,\,2:}]$, then apply the rotation:

"The" (pos=0): angle = 0 for both pairs, so cos=1, sin=0. $Q_{\text{rot}}$[The] = Q[The] = [1.0000, 0.0000, 1.0000, 0.0000] — no change.

"cat" (pos=1): Q[cat] = [0.0, 2.0, 0.0, 1.0], so $x_1$ = [0.0, 2.0], $x_2$ = [0.0, 1.0]

• Pair 0 (angle=1.0): $0.0 \times \cos(1.0) - 0.0 \times \sin(1.0) = 0.0000$, $0.0 \times \sin(1.0) + 0.0 \times \cos(1.0) = 0.0000$
• Pair 1 (angle=0.01): $2.0 \times 0.9999 - 1.0 \times 0.0100 = 1.9899$, $2.0 \times 0.0100 + 1.0 \times 0.9999 = 1.0199$

$Q_{\text{rot}}$[cat] = [0.0000, 1.9899, 0.0000, 1.0199]

"sat" (pos=2): Q[sat] = [1.0, 1.0, 1.0, 0.0], angle pair 0 = 2.0 (cos = −0.4161, sin = 0.9093)

• Pair 0: $1.0 \times (-0.4161) - 1.0 \times 0.9093 = -1.3254$, $1.0 \times 0.9093 + 1.0 \times (-0.4161) = 0.4932$
• Pair 1 (angle=0.02): $1.0 \times 0.9998 - 0.0 \times 0.0200 = 0.9998$, $1.0 \times 0.0200 + 0.0 \times 0.9998 = 0.0200$

$Q_{\text{rot}}$[sat] = [−1.3254, 0.9998, 0.4932, 0.0200]

Full rotated Q matrix:

Step 3: Rotate All K Vectors

Same rotation angles applied to K:

Step 4: Compute Scaled Scores

Raw scores = $Q_{\text{rot}} \times K_{\text{rot}}^T$, then divide by $\sqrt{4} = 2.0$:

Notice how "on" (pos=3) gives its highest score to itself (1.0000) and to "mat" (pos=4, score 0.8058) — nearby tokens — while giving negative scores to "cat" (pos=1, score −0.6627) and "sat" (pos=2, score −0.4257), which are farther away. RoPE is creating a positional preference for nearby tokens.

Step 5: Apply Softmax

Softmax converts scaled scores to probabilities (each row sums to 1.0). Detailed computation for "The" (row 0):

• Scaled scores: [0.0000, 0.5403, 0.2466, −0.5656, −0.6794]
• Max = 0.5403, shifted = [−0.5403, 0.0000, −0.2937, −1.1059, −1.2197]
• exp = [0.5826, 1.0000, 0.7455, 0.3309, 0.2953], sum = 2.9543
• Softmax: [0.1972, 0.3385, 0.2523, 0.1120, 0.1000]

"The" attends most strongly to "cat" (0.3385), which is the next token — the positional rotation has boosted the score for nearby tokens.

Step 6: Compute Output

Output for "The" = weighted sum of V (unrotated):

• 0.1972 $\times$ V[The] + 0.3385 $\times$ V[cat] + 0.2523 $\times$ V[sat] + 0.1120 $\times$ V[on] + 0.1000 $\times$ V[mat]
• = [0.1972, 0.0000, 0.0000, 0.0000] + [0.0000, 0.3385, 0.0000, 0.0000] + [0.0000, 0.0000, 0.2523, 0.0000] + [0.0000, 0.0000, 0.0000, 0.1120] + [0.0500, 0.0500, 0.0500, 0.0500]
• = [0.2472, 0.3885, 0.3023, 0.1620]

Full Attention Weights and Output

Attention Weights — RoPE ($5 \times 5$)

Output Matrix — RoPE ($5 \times 4$)

Interactive: RoPE vs Standard Attention

Compare RoPE attention weights with standard (position-free) attention. Toggle between views and hover over cells to see exact values. The "Difference" view highlights where RoPE increases (green) or decreases (red) attention compared to the standard mechanism.

Applications Across Domains

RoPE was designed for language modeling but its mathematical properties make it valuable across many domains:

Connection to Modern Systems

Flash Attention Compatibility

RoPE is applied before the attention computation: you rotate Q and K, then pass them to standard scaled dot-product attention. This means RoPE is fully compatible with Flash Attention (Chapter 13) — the rotation happens outside the memory-efficient kernel. In practice, most implementations apply RoPE as a preprocessing step, then call $\texttt{flash\_attn\_func(Q\_rot, K\_rot, V)}$.

KV-Cache and Incremental Rotation

During autoregressive generation, K and V are cached to avoid recomputation. With RoPE, each K vector is rotated by its position once when generated, then stored in the cache already rotated. This means the KV-cache works identically to standard attention — no recomputation of rotations is needed when the cache grows.

Context Length Extension

Because RoPE uses a mathematical formula (not a learned table), it can be extended to longer sequences by modifying the frequency base:

• NTK-aware interpolation (bloc97, 2023): scales the base from 10000 to a larger value, effectively stretching the wavelengths so existing RoPE representations cover more positions without retraining.
• YaRN (Peng et al., 2023): combines NTK scaling with an attention temperature adjustment. Extends a 4K context model to 128K+ with a small amount of fine-tuning.
• Dynamic NTK: adjusts the scaling factor at inference time based on actual sequence length, providing automatic adaptation.

Multi-Head Latent Attention (MLA)

DeepSeek-V2 and V3 (Chapter 15) combine RoPE with low-rank key-value compression. They split the query and key into a RoPE portion (carrying positional information) and a content portion (compressed via low-rank projection). This decouples position from content at the architectural level, achieving massive KV-cache savings while preserving RoPE's positional properties.

Complexity Analysis

RoPE adds zero parameters to the model — the frequencies are computed from a formula, not learned. The computational overhead is $O(Nd)$ for the rotation step, which is negligible compared to the $O(N^2 d)$ attention computation. In practice, the rotation is fused into the Q/K projection kernel and adds less than 1% to wall-clock time.

Python Implementation

Full NumPy implementation with class structure. Click any line to see the execution trace with actual matrix values from "the cat sat on the mat."

        θ₀       θ₁
The  0.0000   0.0000
cat  1.0000   0.0100
sat  2.0000   0.0200
on   3.0000   0.0300
mat  4.0000   0.0400

        θ₀       θ₁
The  0.0000   0.0000
cat  1.0000   0.0100
sat  2.0000   0.0200
on   3.0000   0.0300
mat  4.0000   0.0400

          pair0     pair1
The    1.000000  1.000000
cat    0.540302  0.999950
sat   -0.416147  0.999800
on    -0.989992  0.999550
mat   -0.653644  0.999200

          pair0     pair1
The    0.000000  0.000000
cat    0.841471  0.010000
sat    0.909297  0.019999
on     0.141120  0.029996
mat   -0.756802  0.039989

     d0   d1
The  1.0  0.0
cat  0.0  2.0
sat  1.0  1.0
on   0.0  0.0
mat  1.0  0.0

     d2   d3
The  1.0  0.0
cat  0.0  1.0
sat  1.0  0.0
on   1.0  1.0
mat  0.0  1.0

         d0       d1
The   1.0000   0.0000
cat   0.0000   1.9899
sat  -1.3254   0.9998
on   -0.1411  -0.0300
mat  -0.6536  -0.0400

         d0       d1
The   1.0000   0.0000
cat   0.0000   1.0199
sat   0.4932   0.0200
on   -0.9900   0.9996
mat  -0.7568   0.9992

         d0       d1       d2       d3
The   1.0000   0.0000   1.0000   0.0000
cat   0.0000   1.9899   0.0000   1.0199
sat  -1.3254   0.9998   0.4932   0.0200
on   -0.1411  -0.0300  -0.9900   0.9996
mat  -0.6536  -0.0400  -0.7568   0.9992

        The      cat      sat       on      mat
The   0.0000   1.0806   0.4932  -1.1311  -1.3589
cat   3.0098   0.0000   2.0099   0.9598   0.4698
sat   1.0198   1.0806   2.0000  -0.3112  -0.1796
on    0.9696  -1.3254  -0.8515   2.0000   1.6116
mat   0.9592  -0.8489  -0.4361   1.8414   1.5000

        The      cat      sat       on      mat
The   0.0000   0.5403   0.2466  -0.5656  -0.6794
cat   1.5049   0.0000   1.0049   0.4799   0.2349
sat   0.5099   0.5403   1.0000  -0.1556  -0.0898
on    0.4848  -0.6627  -0.4257   1.0000   0.8058
mat   0.4796  -0.4244  -0.2181   0.9207   0.7500

        The      cat      sat       on      mat
The   0.1972   0.3385   0.2523   0.1120   0.1000
cat   0.4052   0.0900   0.2457   0.1454   0.1138
sat   0.2116   0.2181   0.3454   0.1088   0.1162
on    0.2095   0.0665   0.0843   0.3508   0.2889
mat   0.2098   0.0849   0.1044   0.3260   0.2749

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

        d0       d1       d2       d3
The  0.2472   0.3885   0.3023   0.1620
cat  0.4620   0.1468   0.3026   0.2023
sat  0.2697   0.2762   0.4035   0.1668
on   0.3540   0.2109   0.2287   0.4952
mat  0.3472   0.2224   0.2418   0.4635

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

         d0       d1       d2       d3
The   1.0000   0.0000   1.0000   0.0000
cat   0.0000   1.9899   0.0000   1.0199
sat  -1.3254   0.9998   0.4932   0.0200
on   -0.1411  -0.0300  -0.9900   0.9996
mat  -0.6536  -0.0400  -0.7568   0.9992

         d0       d1       d2       d3
The   0.0000   1.0000   0.0000   1.0000
cat  -0.3012   0.0000   1.3818   0.0000
sat  -0.4161   0.9998   0.9093   0.0200
on   -0.1411  -0.0300  -0.9900   0.9996
mat  -0.2752  -0.0200  -1.0836   0.4996

        The      cat      sat       on      mat
The   0.0000   1.0806   0.4932  -1.1311  -1.3589
cat   3.0098   0.0000   2.0099   0.9598   0.4698
sat   1.0198   1.0806   2.0000  -0.3112  -0.1796
on    0.9696  -1.3254  -0.8515   2.0000   1.6116
mat   0.9592  -0.8489  -0.4361   1.8414   1.5000

        The      cat      sat       on      mat
The   0.0000   0.5403   0.2466  -0.5656  -0.6794
cat   1.5049   0.0000   1.0049   0.4799   0.2349
sat   0.5099   0.5403   1.0000  -0.1556  -0.0898
on    0.4848  -0.6627  -0.4257   1.0000   0.8058
mat   0.4796  -0.4244  -0.2181   0.9207   0.7500

        The      cat      sat       on      mat
The   0.1972   0.3385   0.2523   0.1120   0.1000
cat   0.4052   0.0900   0.2457   0.1454   0.1138
sat   0.2116   0.2181   0.3454   0.1088   0.1162
on    0.2095   0.0665   0.0843   0.3508   0.2889
mat   0.2098   0.0849   0.1044   0.3260   0.2749

        d0       d1       d2       d3
The  0.2472   0.3885   0.3023   0.1620
cat  0.4620   0.1468   0.3026   0.2023
sat  0.2697   0.2762   0.4035   0.1668
on   0.3540   0.2109   0.2287   0.4952
mat  0.3472   0.2224   0.2418   0.4635

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

        d0       d1       d2       d3
The  0.2472   0.3885   0.3023   0.1620
cat  0.4620   0.1468   0.3026   0.2023
sat  0.2697   0.2762   0.4035   0.1668
on   0.3540   0.2109   0.2287   0.4952
mat  0.3472   0.2224   0.2418   0.4635

1import numpy as np
2import math
3
4class RotaryPositionAttention:
5    """
6    RoPE — Rotary Position Embedding (Su et al., 2021)
7
8    Encodes position by ROTATING Query and Key vectors
9    before computing their dot product. The rotation makes
10    Q[m] · K[n] depend only on (m - n), the relative offset.
11
12    Values V are NOT rotated — only Q and K receive position.
13    """
14
15    def __init__(self, d_model: int, base: int = 10000):
16        self.d_model = d_model
17        self.base = base
18        self.scale = math.sqrt(d_model)
19        # Precompute frequencies: one theta per dimension pair
20        self.theta = np.array([
21            1.0 / (base ** (2 * i / d_model))
22            for i in range(d_model // 2)
23        ])
24
25    def _softmax(self, x: np.ndarray) -> np.ndarray:
26        """Numerically stable softmax along last axis."""
27        x_shifted = x - np.max(x, axis=-1, keepdims=True)
28        exp_x = np.exp(x_shifted)
29        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
30
31    def compute_angles(self, N: int) -> np.ndarray:
32        """Compute rotation angles for positions 0..N-1."""
33        positions = np.arange(N)
34        return np.outer(positions, self.theta)
35
36    def apply_rotation(self, X: np.ndarray) -> np.ndarray:
37        """Apply RoPE rotation to a matrix X (N, d_model)."""
38        N, D = X.shape
39        angles = self.compute_angles(N)
40        cos_a = np.cos(angles)
41        sin_a = np.sin(angles)
42        x1 = X[:, :D // 2]
43        x2 = X[:, D // 2:]
44        rotated_first = x1 * cos_a - x2 * sin_a
45        rotated_second = x1 * sin_a + x2 * cos_a
46        return np.hstack([rotated_first, rotated_second])
47
48    def compute_scores(self, Q_rot: np.ndarray, K_rot: np.ndarray):
49        """Raw dot-product scores after rotation."""
50        return Q_rot @ K_rot.T
51
52    def scale_scores(self, scores: np.ndarray) -> np.ndarray:
53        """Divide by sqrt(d_model)."""
54        return scores / self.scale
55
56    def compute_weights(self, scaled: np.ndarray) -> np.ndarray:
57        """Apply softmax to get attention weights."""
58        return self._softmax(scaled)
59
60    def compute_output(self, weights: np.ndarray, V: np.ndarray):
61        """Weighted sum of value vectors (V is NOT rotated)."""
62        return weights @ V
63
64    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
65        """
66        Full RoPE attention forward pass.
67
68        Args:
69            Q: Query matrix  (N, d_model)
70            K: Key matrix    (N, d_model)
71            V: Value matrix  (N, d_model) — NOT rotated
72
73        Returns:
74            output:  Context vectors  (N, d_model)
75            weights: Attention matrix (N, N)
76        """
77        Q_rot = self.apply_rotation(Q)
78        K_rot = self.apply_rotation(K)
79        raw_scores = self.compute_scores(Q_rot, K_rot)
80        scaled_scores = self.scale_scores(raw_scores)
81        weights = self.compute_weights(scaled_scores)
82        output = self.compute_output(weights, V)
83        return output, weights
84
85    def explain(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray,
86                tokens: list, query_idx: int = 0):
87        """Print a detailed trace for a specific query token."""
88        Q_rot = self.apply_rotation(Q)
89        K_rot = self.apply_rotation(K)
90        token = tokens[query_idx]
91
92        print(f"\n=== RoPE trace for '{token}' (pos {query_idx}) ===")
93        print(f"Q[{token}] original = {Q[query_idx]}")
94        print(f"Q[{token}] rotated  = {np.round(Q_rot[query_idx], 4)}")
95
96        raw = self.compute_scores(Q_rot, K_rot)
97        scaled = self.scale_scores(raw)
98        w = self.compute_weights(scaled)
99        out = self.compute_output(w, V)
100
101        for j, t in enumerate(tokens):
102            print(f"  score[{t}] = {raw[query_idx,j]:.4f}"
103                  f" -> scaled = {scaled[query_idx,j]:.4f}"
104                  f" -> weight = {w[query_idx,j]:.4f}")
105        print(f"  output = {np.round(out[query_idx], 4)}")
106
107
108# ── Shared Example (used in every chapter) ──
109tokens = ["The", "cat", "sat", "on", "mat"]
110
111Q = np.array([
112    [1.0, 0.0, 1.0, 0.0],   # The
113    [0.0, 2.0, 0.0, 1.0],   # cat
114    [1.0, 1.0, 1.0, 0.0],   # sat
115    [0.0, 0.0, 1.0, 1.0],   # on
116    [1.0, 0.0, 0.0, 1.0],   # mat
117])
118
119K = np.array([
120    [0.0, 1.0, 0.0, 1.0],   # The
121    [1.0, 0.0, 1.0, 0.0],   # cat
122    [1.0, 1.0, 0.0, 0.0],   # sat
123    [0.0, 0.0, 1.0, 1.0],   # on
124    [1.0, 0.0, 0.5, 0.5],   # mat
125])
126
127V = np.array([
128    [1.0, 0.0, 0.0, 0.0],   # The
129    [0.0, 1.0, 0.0, 0.0],   # cat
130    [0.0, 0.0, 1.0, 0.0],   # sat
131    [0.0, 0.0, 0.0, 1.0],   # on
132    [0.5, 0.5, 0.5, 0.5],   # mat
133])
134
135# ── Run ──
136rope = RotaryPositionAttention(d_model=4, base=10000)
137output, weights = rope.forward(Q, K, V)
138
139print("RoPE Attention Weights (5x5):")
140print(np.round(weights, 4))
141
142print("\nRoPE Output (5x4):")
143print(np.round(output, 4))
144
145# Detailed trace
146rope.explain(Q, K, V, tokens, query_idx=0)
147rope.explain(Q, K, V, tokens, query_idx=1)

PyTorch Implementation

Production-ready PyTorch version with GPU support, register_buffer for frequencies, and F.softmax. Click any line for the execution trace.

        θ₀       θ₁
The  0.0000   0.0000
cat  1.0000   0.0100
sat  2.0000   0.0200
on   3.0000   0.0300
mat  4.0000   0.0400

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class RotaryPositionAttention(nn.Module):
7    """
8    RoPE — Rotary Position Embedding (Su et al., 2021)
9    PyTorch implementation with GPU support.
10
11    Rotates Q and K by position before computing attention.
12    V is NOT rotated. Supports automatic differentiation.
13    """
14
15    def __init__(self, d_model: int, base: int = 10000):
16        super().__init__()
17        self.d_model = d_model
18        self.base = base
19        self.scale = math.sqrt(d_model)
20
21        # Register frequencies as a buffer (not a parameter)
22        theta = torch.tensor([
23            1.0 / (base ** (2 * i / d_model))
24            for i in range(d_model // 2)
25        ])
26        self.register_buffer("theta", theta)
27
28        # In production: learned projection matrices
29        # self.W_q = nn.Linear(d_model, d_model)
30        # self.W_k = nn.Linear(d_model, d_model)
31        # self.W_v = nn.Linear(d_model, d_model)
32        # self.W_o = nn.Linear(d_model, d_model)
33
34    def apply_rotation(self, X: torch.Tensor) -> torch.Tensor:
35        """Apply RoPE rotation. X shape: (N, d_model)."""
36        N, D = X.shape
37        positions = torch.arange(N, device=X.device, dtype=X.dtype)
38        angles = positions.unsqueeze(1) * self.theta.unsqueeze(0)
39        cos_a = torch.cos(angles)
40        sin_a = torch.sin(angles)
41        x1 = X[:, :D // 2]
42        x2 = X[:, D // 2:]
43        rotated_first = x1 * cos_a - x2 * sin_a
44        rotated_second = x1 * sin_a + x2 * cos_a
45        return torch.cat([rotated_first, rotated_second], dim=-1)
46
47    def forward(
48        self,
49        Q: torch.Tensor,
50        K: torch.Tensor,
51        V: torch.Tensor,
52    ) -> tuple[torch.Tensor, torch.Tensor]:
53        """
54        Args:
55            Q: (N, d_model)  query matrix
56            K: (N, d_model)  key matrix
57            V: (N, d_model)  value matrix (NOT rotated)
58        Returns:
59            output:  (N, d_model) context vectors
60            weights: (N, N) attention weights
61        """
62        Q_rot = self.apply_rotation(Q)
63        K_rot = self.apply_rotation(K)
64        scores = Q_rot @ K_rot.T / self.scale
65        weights = F.softmax(scores, dim=-1)
66        output = weights @ V
67        return output, weights
68
69
70# ── Shared Example ──
71tokens = ["The", "cat", "sat", "on", "mat"]
72
73Q = torch.tensor([
74    [1.0, 0.0, 1.0, 0.0],
75    [0.0, 2.0, 0.0, 1.0],
76    [1.0, 1.0, 1.0, 0.0],
77    [0.0, 0.0, 1.0, 1.0],
78    [1.0, 0.0, 0.0, 1.0],
79])
80
81K = torch.tensor([
82    [0.0, 1.0, 0.0, 1.0],
83    [1.0, 0.0, 1.0, 0.0],
84    [1.0, 1.0, 0.0, 0.0],
85    [0.0, 0.0, 1.0, 1.0],
86    [1.0, 0.0, 0.5, 0.5],
87])
88
89V = torch.tensor([
90    [1.0, 0.0, 0.0, 0.0],
91    [0.0, 1.0, 0.0, 0.0],
92    [0.0, 0.0, 1.0, 0.0],
93    [0.0, 0.0, 0.0, 1.0],
94    [0.5, 0.5, 0.5, 0.5],
95])
96
97# ── Run ──
98rope = RotaryPositionAttention(d_model=4, base=10000)
99output, weights = rope.forward(Q, K, V)
100
101print("RoPE Weights (PyTorch):")
102print(weights.detach().numpy().round(4))
103
104print("\nRoPE Output (PyTorch):")
105print(output.detach().numpy().round(4))

Key Takeaways

1. Rotation encodes position intrinsically. Unlike additive position embeddings or score biases, RoPE makes the dot product $Q_m \cdot K_n$ depend on relative position $(m - n)$ as a mathematical consequence, not an approximation.
2. Multi-frequency design captures multiple scales. Fast-rotating pairs encode local position (is this word the immediate neighbor?), while slow-rotating pairs encode global position (is it in the same paragraph?). This is why RoPE generalizes to longer sequences.
3. Zero additional parameters. Frequencies come from a formula, not a learned lookup table. This means no out-of-vocabulary positions and no parameter overhead.
4. V is NOT rotated. Only Q and K carry positional information through rotation. V provides the content that gets aggregated.
5. Compatible with all optimizations. RoPE works seamlessly with Flash Attention, KV-cache, multi-head/multi-query/grouped-query attention, and context length extension techniques like YaRN.
6. De facto standard for modern LLMs. LLaMA, Mistral, PaLM, Qwen, Falcon, Phi, and virtually every major open-weight model since 2022 uses RoPE.

Exercises

1. Verify the relative position property: Using the code above, create two identical vectors $\mathbf{x} = [1, 1, 1, 1]$ and place them at positions (0, 2) and (3, 5). Verify that the dot product of their rotated versions is the same in both cases, confirming that the score depends only on the distance 2.
2. Frequency analysis: For $d = 128$ and base = 10000, compute all 64 frequencies. What is the wavelength (in positions) of the fastest and slowest pair? How many full rotations does pair 0 complete over a 4096-token sequence?
3. NTK scaling experiment: Modify the code to use base = 100000 instead of 10000. Recompute the attention weights. How does this affect which tokens "on" attends to? Why would this help with longer sequences?
4. Compare with Chapter 7: Run both the relative position bias attention (Chapter 7) and RoPE attention on the same input. Which produces stronger local preference? What happens when you increase the sequence length to 20 tokens?
5. 2D RoPE for images: Design a 2D RoPE scheme for a 4×4 image patch grid (16 tokens). Each token has a (row, col) position. Apply separate rotations for row-position and column-position using different frequency bands. What properties does this have compared to 1D RoPE?

References

1. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). "RoFormer: Enhanced Transformer with Rotary Position Embedding." arXiv:2104.09864.
2. Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). "Attention Is All You Need." NeurIPS 2017.
3. Touvron, H., Lavril, T., Izacard, G., et al. (2023). "LLaMA: Open and Efficient Foundation Language Models." arXiv:2302.13971.
4. bloc97. (2023). "NTK-Aware Scaled RoPE Allows LLaMA Models to Have Extended (8k+) Context Size Without Any Fine-tuning and Minimal Perplexity Degradation." Reddit/r/LocalLLaMA.
5. Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). "YaRN: Efficient Context Window Extension of Large Language Models." arXiv:2309.00071.
6. Black, S., Biderman, S., Hallahan, E., et al. (2022). "GPT-NeoX-20B: An Open-Source Autoregressive Language Model." arXiv:2204.06745.
7. Liu, Z., Oguz, B., Zhao, C., et al. (2024). "World Model on Million-Length Video and Language with RingAttention." arXiv:2402.08268.