Rotary Positional Embeddings (RoPE) - Transformer Architecture

Rotary Positional Embeddings (RoPE)

What you'll build

Apply RoPE to queries and keys: rotate each pair of dimensions by an angle proportional to the token's position. LLaMA, Mistral, Qwen, GPT-NeoX and most modern LLMs use it.

Background

Sinusoidal encodings are added to the embeddings once, at the input. RoPE instead rotates q and k inside every attention layer. Rotating q at position m by m·θ and k at position n by n·θ makes their dot product depend only on the angle difference (m - n)·θ, so attention scores see relative positions, with no learned parameters and no change to the vectors' norms. Values V are not rotated.

Formula

For position pos (row index, starting at 0), d = d_model and pair index i = 0 .. d/2 - 1:

theta_i = 1 / 10000^(2i / d)
a       = pos * theta_i

x_rot[2i]     = x[2i] * cos(a) - x[2i+1] * sin(a)
x_rot[2i + 1] = x[2i] * sin(a) + x[2i+1] * cos(a)

This is the 2-D rotation [[cos a, -sin a], [sin a, cos a]] applied to the pair (x[2i], x[2i+1]). Apply it to every row of Q and of K.

Inputs

• Q: float array (seq_len, d_model), d_model even
• K: float array (seq_len, d_model)

Expected Return Format

A tuple (Q_rot, K_rot), each with the shape of its input.

Edge cases and conventions

• Interleaved pairs (0,1), (2,3), ... as in the RoFormer paper. Hugging Face LLaMA code pairs dimension i with i + d/2 ("rotate_half"); it is the same idea with permuted dimensions, but it gives different numbers, so use the interleaved layout here.
• Position 0 is unchanged (angle 0). Don't modify the inputs in place.

Constraints

NumPy only, vectorised over positions and pairs (strided slicing x[:, 0::2], x[:, 1::2]).

Follow-up

• Prove <R(m) q, R(n) k> = <q, R(n - m) k>.
• Context extension (NTK scaling, YaRN) changes the base 10000. Which pairs are affected most, and why?

Input: apply_rope(np.array([[1.0, 0.0, 1.0, 0.0]] * 3), K)[0]
Output: [[1.0, 0.0, 1.0, 0.0], [0.5403, 0.8415, 1.0, 0.01], [-0.4161, 0.9093, 0.9998, 0.02]]

Input: apply_rope(np.ones((2, 2)), np.ones((2, 2)))[0]
Output: [[1.0, 1.0], [-0.3012, 1.3818]]

import numpy as np

def apply_rope(Q, K):
    """
    Apply Rotary Positional Embeddings to Q and K (interleaved pairs).

    Args:
        Q: Queries (seq_len, d_model), d_model even
        K: Keys (seq_len, d_model)

    Returns:
        (Q_rot, K_rot): pair (2i, 2i+1) of row pos rotated by pos * theta_i,
                        theta_i = 1 / 10000^(2i / d_model)
    """
    # Your code here
    pass