Positional Encoding — infra.bea.ai

Self-attention is permutation-equivariant — shuffling input tokens produces shuffled outputs with the same values. Without positional information, the model cannot distinguish “the cat sat on the mat” from “mat the on sat cat the.” Positional encoding solves this.

1. Why Position Matters

Given tokens $(x_1, x_2, x_3)$ and any permutation $\pi$ , self-attention satisfies:

\text{Attn}(x_{\pi(1)}, x_{\pi(2)}, x_{\pi(3)}) = \pi(\text{Attn}(x_1, x_2, x_3))

The attention scores depend only on the content of Q and K, not their positions. We need to inject position information explicitly.

2. Absolute Positional Encoding

Sinusoidal (Original Transformer)

Add position-dependent vectors to token embeddings:

PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d}}\right), \qquad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d}}\right)

Each dimension oscillates at a different frequency. Low dimensions change slowly (capture long-range patterns), high dimensions change rapidly (capture nearby positions).

Properties: Fixed (no learnable parameters), theoretically enables extrapolation to longer sequences, but in practice extrapolation is poor.

Learned Positional Embeddings (GPT-2)

Simply learn a $(S_{\max}, d)$ embedding table — one vector per position. Simple and effective, but hard-coded to a maximum sequence length. Cannot extrapolate beyond training length.

Limitations of Absolute Encoding

Both approaches encode position as a fixed offset from the start. Problems:

Relative distances between tokens are hard to learn from absolute positions
No generalization to sequence lengths beyond training
Position 1000 and position 1001 have very different absolute encodings despite being adjacent

3. Rotary Position Embedding (RoPE)

RoPE is the dominant position encoding in modern LLMs (LLaMA, Qwen, Mistral, DeepSeek). It encodes relative position directly into the attention score computation.

Core Idea

Instead of adding position to embeddings, RoPE rotates query and key vectors by a position-dependent angle:

\mathbf{q}_m = R_{\Theta, m} \mathbf{q}, \qquad \mathbf{k}_n = R_{\Theta, n} \mathbf{k}

Where $R_{\Theta, m}$ is a rotation matrix at position $m$ . The attention score becomes:

\mathbf{q}_m^\top \mathbf{k}_n = \mathbf{q}^\top R_{\Theta, m}^\top R_{\Theta, n} \mathbf{k} = \mathbf{q}^\top R_{\Theta, n-m} \mathbf{k}

The score depends only on the relative distance $n - m$ , not the absolute positions.

Why this works mathematically: Rotation matrices have the property $R_a^\top R_b = R_{b-a}$ . So the product of the two rotations collapses into a single rotation by the difference. This means two tokens at positions (100, 105) produce the same attention score as tokens at (0, 5)
same relative distance, same score.

Implementation

RoPE operates on pairs of dimensions. For a $d_k$ -dimensional vector, group dimensions into $d_k/2$ pairs, each rotated by angle $\theta_i \cdot m$ :

\theta_i = 10000^{-2i/d_k}

def apply_rope(x, positions):
    """Apply RoPE to query or key tensor.
    x: (B, Nh, S, dk), positions: (S,)
    """
    d = x.shape[-1]
    freqs = 1.0 / (10000 ** (torch.arange(0, d, 2).float() / d))  # (dk/2,)
    angles = positions.unsqueeze(-1) * freqs.unsqueeze(0)           # (S, dk/2)

    cos = angles.cos()  # (S, dk/2)
    sin = angles.sin()  # (S, dk/2)

    x1, x2 = x[..., ::2], x[..., 1::2]  # split even/odd dimensions
    out_even = x1 * cos - x2 * sin
    out_odd  = x1 * sin + x2 * cos
    return torch.stack([out_even, out_odd], dim=-1).flatten(-2)

Infrastructure implication: RoPE is applied to Q and K after projection but before the attention score computation. It’s typically fused into the attention kernel for efficiency — separate application would add an extra memory read/write pass.

4. ALiBi (Attention with Linear Biases)

ALiBi adds a position-dependent bias directly to attention scores instead of modifying embeddings:

\text{score}_{ij} = \mathbf{q}_i^\top \mathbf{k}_j - m \cdot |i - j|

Where $m$ is a head-specific slope. Closer tokens get higher scores; distant tokens are penalized linearly.

Advantages: No parameters, trivially extends to any sequence length. Disadvantage: The linear penalty may be too aggressive for some tasks. Used by BLOOM and some earlier models, but largely superseded by RoPE.

5. Long-Context Extension

Training with a 4K context but serving at 128K requires extending position encoding. Key techniques:

Position Interpolation

Scale positions to fit within the training range:

\theta'_i = \theta_i / s, \qquad s = \frac{L_{\text{target}}}{L_{\text{train}}}

Requires a small amount of fine-tuning to adapt. Used by Code LLaMA (16K → 100K).

NTK-aware Scaling

Instead of uniformly scaling all frequencies, adjust the base:

\theta'_i = \left(\text{base} \cdot s^{d/(d-2)}\right)^{-2i/d}

This preserves high-frequency (local position) information while extending low-frequency (long-range) capacity.

YaRN

Combines NTK scaling with an attention temperature factor, achieving better quality on long contexts. Used by Qwen and other recent models.

Infrastructure Impact

Long-context extension is critical for AI infra:

KV cache scales linearly with context: 128K context = 32× the KV cache of 4K
Attention compute scales quadratically: 128K context = 1024× the attention FLOPs of 4K
RoPE computation must support dynamic context lengths at serving time

Self-Check

Explain why self-attention needs positional encoding
Explain the difference between absolute and relative positional encoding
Explain how RoPE makes attention scores depend on relative distance
Implement basic RoPE in PyTorch
Explain Position Interpolation and NTK scaling for long-context extension
Explain the infrastructure impact of long context (KV cache, attention FLOPs)