RoPE

RoPE, rotary positional embedding, injects positional information to a model. Early LLM models uses learned absolute embeddings, which adds a position embedding to token embeddings before they go through training:
$$\text{Input}=\text{TokenEmbed}+\text{PositionEmbed}$$ The issue with the this approach is 1, the additional learning the model has to go through to learn the position information; 2, the position information learned this way is aboslute positions, “this word appears in the 5th position in a sequence”, where as what probably matters more are relative positions, how a word positions in relation to other words.

At a high level, RoPE pairs planes in query ($Q$) and key ($K$) vectors into many 2-d vectors before they go through the attention computation. Each vector then rotates by a slightly different angle. There is no additional embeddings added to the vectors. See the diagram below for a comparison between RoPE and absolute position embedding.

Absolute position embedding:
┌───────────────┐
│ token vectors │
└───────┬───────┘
        ▼
  + position vector
        ▼
┌───────────────┐
│ Wq, Wk → q, k │
└───────┬───────┘
        ▼
   score: q · k
 
RoPE:
┌───────────────┐
│ token vectors │
└───────┬───────┘
        ▼
┌───────────────┐
│ Wq, Wk → q, k │
└───────┬───────┘
        ▼
   R(q, mθ), R(k, nθ)
        ▼
   score: q · k

Given the fact that we want to focus on the relative position of the vectors, not absolute values, let’s see how RoPE help us do that. To understand how RoPE encodes relative positions, we need to prove that attention score, which is $q\cdot k$, only depends on the difference between the two positions.

Let $q$ points at angle ${\psi }_q$ and $k$ points at angle ${\psi }_k$. Because vectors’ rotation angles are position based in RoPE, we also introduce the original position indices of $q$, $m$; and the original position of $k$, $n$. Assuming vectors rotate by $originPos\times \theta$, knowing the dot product of two vectors are $a\cdot b=\Vert a\Vert \Vert b\Vert \cos \alpha$ where $\alpha$ is the angle between the two, it is easy to show that after rotation, we have
$$q_m\cdot k_n=\Vert q_m\Vert \Vert k_n\Vert \cos ({\psi }_q-{\psi }_k+(m-n)\theta )\text{\hspace{1em}(1)}$$ We can see that

By rotating them, we introduce the relative positions $m-n$ into the attention score.

Also, note that the absolute positions of the vectors doesn’t matter in attention calculation, only their difference $m-n$. For example, a query at position 5 and a key at position 3 get exactly the same positional effect as a query at 105 and a key at 103, because both pairs are 2 apart.

Implementation

Now to implement RoPE, we start with a simple, 2-D scenario. At the end of the day, RoPE is about rotating vectors by a certain agle. It does not need to know whether the vector in query is a query or key projection. And 2-d rotation is trivial:
$$R(\theta )=\left(\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right)$$

And the implementation is straight forward:

from math import cos, sin
import torch
 
def rotate(theta, x):
	(a, b) = x
	return (a * cos(theta) - b * sin(theta), a * sin(theta) + b * cos(theta))

The innovation of RoPE lies in how it handles high dimensionality. The issue with rotation in high dimensionality is there are many possible ways of rotating. What RoPE does is for each tensor, it divides the head dimension into two halve, and 1 dimension from each half to form many 2-d pairs and rotate each individually.

Now there are many rotations to perform, we also want to rotate each 2-d pair by a different angle so that the positions are encoded differently. By itself, position-based rotation doesn’t care about which dimensionality. Positions in $(1)$ are indices of each vector in a sequence like vector “is” has an index of 1 in sequence “Here is an apple”. To register dimensionality. We use another trick, a dimension-dependent formula to derive the rotation angle:

$${\theta }_{i,k}=\frac{i}{{\Theta }^{(2k-2)/d}}\text{\hspace{1em}(2)}$$

Here $i$ is the token position just like $m$ and $n$, but we introduce $k\in \left\{1,...,d/2\right\}$ where $d$ is the head dimension. So k is an index for dimensionality. Since head dimensions are paired up, we only need to define half of the head dimensions. $\Theta$ here is a constant. Because now the rotation angle is dimension and index-based, we no longer need to supply it externally.

Putting it together, we have:

def rotate_tensor(in_feature: torch.Tensor, pos: torch.Tensor):
	shape = in_feature.shape # [batch, heads, seq_len, head_dim]
	head_dim = in_feature.shape[-1]
	half_idx = torch.arange(head_dim//2) # [head_dim//2]
	theta = torch.pow(10_000, -2 * half_idx / head_dim) # Theta is 10_000, and shape of theta [head_dim//2]
	angles = pos.unsqueeze(-1)* theta # [..., seq_len, head_dim//2]
	angles = angles.repeat_interleave(2, dim=-1)
 
	a_cos, b_cos = (in_feature * angles.cos()).reshape(-1, 2).unbind(dim = 1)
	a_sin, b_sin = (in_feature * angles.sin()).reshape(-1, 2).unbind(dim = 1)
	first, second = a_cos - b_sin, a_sin + b_cos
	return torch.stack((first, second), dim=1).reshape(shape)
 

In this implementation, we are pairing up head dimension $i-1$ with $i$, as seeing in repeat_interleave(2), which duplicate each element twice. This is different from the canonical implementation. Which pairs $i$ with $i+d/2$. In this implementation, the first half of the function implements the angle formula stated in $(2)$. The second half implements the actual rotation, similar to the original rotate function but in a higher dimension. There, we pull (a, b) out of the tensor and calculate their $sin$ and $cos$ value first before recombining them into the familiar shape of a_cos - b_sin, a_sin + b_cos and put them back to the original vector.

Finally, we perform several checks on our rotation:

1. Position 0 doesn’t rotate
2. Rotation keep each vector’s length
3. Attention scores depend only relative positions (same vector in position 1 and 3 get the same results after rotation)

torch.manual_seed(42)
 
batch_size = 2
num_heads = 2
seq_len = 4
d_model = 16
head_dim = 8
 
# Token representations entering attention
hidden = torch.randn(batch_size, seq_len, d_model)
 
# Separate query and key projection weights
W_q = torch.randn(d_model, num_heads * head_dim) / d_model**0.5
W_k = torch.randn(d_model, num_heads * head_dim) / d_model**0.5
 
# [batch, heads, seq_len, head_dim] = [2, 2, 4, 8]
q = (hidden @ W_q).reshape(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
k = (hidden @ W_k).reshape(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
token_positions = torch.arange(seq_len)  # Shared across batches and heads
q_rot =rotate_tensor(q, token_positions)
# 1. poisition 0 doesn't rotate
torch.testing.assert_close(q_rot[..., 0, :], q[..., 0, :])
print("position 0 is unchanged")
 
# 2. Rotation preserves each vector's length
torch.testing.assert_close(
q_rot.norm(dim=-1), q.norm(dim=-1)
)
print("rotation preserves vector length")
 
k_rot = rotate_tensor(k, token_positions)
 
q_shift = rotate_tensor(q, token_positions + 3)
k_shift = rotate_tensor(k, token_positions + 3)
 
# 3. advance a vector by 3 indices doesn't change the attention score
torch.testing.assert_close(
	q_rot @ k_rot.transpose(-2, -1),
	q_shift @ k_shift.transpose(-2, -1),
	rtol=1e-5,
	atol=1e-5,
)
print("attention scores depend only on relative position")

position 0 is unchanged
rotation preserves vector length
attention scores depend only on relative position