Loading your learning…
Loading your learning…
Attention compares queries and keys with dot products. A dot product only sees the two vectors’ contents, so on its own attention cannot tell which token came first or how far apart two tokens are. RoPE rotates vector pairs by angles determined by token position. Input and output are [B, H, T, D]: batch, attention heads, tokens, and features per head. D must be even. positions [T] supplies one position per token, shared across this batch.
import torchfrom torch import nn def apply_rope(x, positions, base=10_000): # x: [B, H, T, D]; positions: [T]. D must be even. if x.shape[-1] % 2: raise ValueError("RoPE needs an even head dimension") half = x.shape[-1] // 2 frequencies = base ** (-torch.arange(half, device=x.device).float() / half) angles = positions.to(device=x.device, dtype=torch.float32)[:, None] * frequencies cos = angles.cos().to(x.dtype)[None, None, :, :] sin = angles.sin().to(x.dtype)[None, None, :, :] first, second = x[..., :half], x[..., half:] return torch.cat([first * cos - second * sin, first * sin + second * cos], dim=-1) if __name__ == "__main__": x = torch.tensor([[[[1., 0., 0., 1.], [1., 0., 0., 1.]]]]) positions = torch.tensor([0, 1]) rotated = apply_rope(x, positions) print("Shape:", rotated.shape) # [1, 1, 2, 4] print("Position zero:", rotated[0, 0, 0]) print("Lengths:", x.norm(dim=-1), rotated.norm(dim=-1)) # Same vectors, same gap between positions: the score doesn't change. query = torch.tensor([[[[1., 0., 0., 1.]]]]) # [1, 1, 1, 4] key = torch.tensor([[[[0., 1., 1., 0.]]]]) for m, n in [(2, 0), (7, 5), (0, 2)]: score = (apply_rope(query, torch.tensor([m])) * apply_rope(key, torch.tensor([n]))).sum() print(f"Query at {m}, key at {n}: {score:.3f}")With D=4, this implementation pairs coordinates (0, 2) and (1, 3): one coordinate from each half. Splitting the last axis gives first and second, each [B, H, T, D/2]. Pairing conventions must match for queries and keys and any pretrained weights. This half-split convention differs from adjacent-coordinate pairing.
if x.shape[-1] % 2: raise ValueError("RoPE needs an even head dimension") half = x.shape[-1] // 2 frequencies = base ** (-torch.arange(half, device=x.device).float() / half) angles = positions.to(device=x.device, dtype=torch.float32)[:, None] * frequencies cos = angles.cos().to(x.dtype)[None, None, :, :] sin = angles.sin().to(x.dtype)[None, None, :, :] first, second = x[..., :half], x[..., half:] return torch.cat([first * cos - second * sin, first * sin + second * cos], dim=-1)Pair i uses frequency base^(-i/(D/2)). For D=4 and base=10,000, frequencies are [1, 0.01] radians per token. Multiplying positions [T, 1] by frequencies [D/2] gives angles [T, D/2]. Fast and slow rotations carry position information at different scales.
raise ValueError("RoPE needs an even head dimension") half = x.shape[-1] // 2 frequencies = base ** (-torch.arange(half, device=x.device).float() / half) angles = positions.to(device=x.device, dtype=torch.float32)[:, None] * frequencies cos = angles.cos().to(x.dtype)[None, None, :, :] sin = angles.sin().to(x.dtype)[None, None, :, :]