S110 · Positional Encoding — Sinusoidal, Learned, RoPE

# positional_encoding.py — sinusoidal, learned, and RoPE from scratch.
# Run: uv run positional_encoding.py
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
def sinusoidal_pe(max_len: int, d_model: int) -> torch.Tensor:
    """Standard 2017 Transformer positional encoding. Returns [max_len, d_model]."""
    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # [T, 1]
    # div_term shape [d_model/2]: 10000^(2i/d) for i=0..d/2
    div_term = torch.exp(
        torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
    )
    pe[:, 0::2] = torch.sin(position * div_term)   # even indices
    pe[:, 1::2] = torch.cos(position * div_term)   # odd indices
    return pe
 
 
class LearnedPE(nn.Module):
    """Learned absolute positional embeddings — one vector per position."""
 
    def __init__(self, max_len: int, d_model: int):
        super().__init__()
        self.pe = nn.Embedding(max_len, d_model)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, T, d_model]
        T = x.size(1)
        positions = torch.arange(T, device=x.device)  # [T]
        return x + self.pe(positions).unsqueeze(0)     # broadcast
 
 
def precompute_rope_freqs(d_head: int, max_len: int, base: float = 10000.0) -> tuple[torch.Tensor, torch.Tensor]:
    """Precompute cos/sin tables for RoPE. Called once at model init."""
    # d_head must be even.
    theta = 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head))  # [d_head/2]
    positions = torch.arange(max_len).float()                               # [max_len]
    freqs = torch.outer(positions, theta)                                   # [max_len, d_head/2]
    cos = freqs.cos()
    sin = freqs.sin()
    return cos, sin
 
 
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply rotary position embedding to a Q or K tensor of shape [B, T, N, d_head]."""
    # Split last dim into two halves: even indices and odd indices.
    x_even = x[..., 0::2]  # [B, T, N, d_head/2]
    x_odd = x[..., 1::2]   # [B, T, N, d_head/2]
 
    # cos, sin: [T, d_head/2]. Broadcast against [B, T, N, d_head/2].
    T = x.size(1)
    cos_t = cos[:T].unsqueeze(0).unsqueeze(2)  # [1, T, 1, d_head/2]
    sin_t = sin[:T].unsqueeze(0).unsqueeze(2)
 
    # 2D rotation applied to each (even, odd) pair.
    rotated_even = x_even * cos_t - x_odd * sin_t
    rotated_odd = x_even * sin_t + x_odd * cos_t
 
    # Interleave back into a tensor of the original shape.
    out = torch.empty_like(x)
    out[..., 0::2] = rotated_even
    out[..., 1::2] = rotated_odd
    return out
 
 
class AttentionWithRoPE(nn.Module):
    """Multi-head attention with RoPE applied to Q and K before the dot product."""
 
    def __init__(self, d_model: int, n_heads: int, max_len: int = 2048):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        cos, sin = precompute_rope_freqs(self.d_head, max_len)
        # Non-learnable buffers — travel with the model on .to(device).
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_head)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_head)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_head)
 
        # Apply RoPE to Q and K only (NOT to V — value doesn't need positional encoding).
        Q = apply_rope(Q, self.cos, self.sin)
        K = apply_rope(K, self.cos, self.sin)
 
        # Standard MHA from here on.
        Q = Q.transpose(1, 2)  # [B, N, T, d_head]
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_head)
        weights = F.softmax(scores, dim=-1)
        out = (weights @ V).transpose(1, 2).contiguous().view(B, T, D)
        return self.W_o(out)
 
 
def visualise_sinusoidal() -> None:
    pe = sinusoidal_pe(max_len=50, d_model=16)
    print(f"Sinusoidal PE shape: {pe.shape}")
    print(f"\nFirst 3 dims across positions 0..9:")
    for pos in range(10):
        print(f"  pos {pos}: [{pe[pos, 0]:+.3f}, {pe[pos, 1]:+.3f}, {pe[pos, 2]:+.3f}, ...]")
    print(f"\nSimilarity between positions (dot product):")
    for p in [0, 1, 5, 20, 49]:
        print(f"  PE[0] · PE[{p}] = {(pe[0] * pe[p]).sum().item():+.3f}")
 
 
def demonstrate_rope_relative() -> None:
    """Prove that RoPE-transformed Q·K depends only on the position difference."""
    torch.manual_seed(0)
    d_head = 8
    cos, sin = precompute_rope_freqs(d_head, 100)
 
    # Two random vectors, at different absolute positions but same difference.
    q = torch.randn(1, 1, 1, d_head)  # Q at position 0 (pre-rotation)
    k = torch.randn(1, 1, 1, d_head)  # K at position 0 (pre-rotation)
 
    print("\nTest: rotate Q at pos m, K at pos n. Check that Q·K depends only on (n-m).")
    for m, n in [(0, 5), (10, 15), (50, 55), (0, 10), (20, 30)]:
        q_ = q.clone()
        k_ = k.clone()
        # Rotate at absolute positions.
        q_rot = apply_rope(q_, cos[m:m+1], sin[m:m+1])
        k_rot = apply_rope(k_, cos[n:n+1], sin[n:n+1])
        dot = (q_rot * k_rot).sum().item()
        print(f"  m={m:2d}, n={n:2d}, diff={n-m:2d} → q_rot·k_rot = {dot:+.4f}")
    print("(Notice: same difference → same dot product, regardless of absolute positions.)")
 
 
def main() -> None:
    print("=== Sinusoidal PE ===")
    visualise_sinusoidal()
 
    print("\n=== Learned PE ===")
    x = torch.randn(1, 5, 16)
    learned = LearnedPE(max_len=100, d_model=16)
    out = learned(x)
    print(f"Learned PE: input {x.shape} → output {out.shape}")
    print(f"Params: {sum(p.numel() for p in learned.parameters()):,}")
 
    print("\n=== RoPE ===")
    demonstrate_rope_relative()
 
    # Full attention layer with RoPE.
    print("\n=== AttentionWithRoPE end-to-end ===")
    attn = AttentionWithRoPE(d_model=64, n_heads=4)
    x = torch.randn(2, 10, 64)
    out = attn(x)
    print(f"Input: {x.shape} → Output: {out.shape}")
 
 
if __name__ == "__main__":
    main()