2.5 Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi¶

Take the attention equation from the previous chapter and stare at it for a moment:

Now permute the rows of the input — swap token 2 with token 17, reverse the whole sequence, shuffle it like a deck of cards. What happens to the output? The output permutes in exactly the same way and nothing else changes. Attention is permutation-equivariant: it treats its input as a set, not a sequence. The dot product \(q_i \cdot k_j\) depends on what token \(i\) and token \(j\) are, never on where they sit. To attention, “dog bites man” and “man bites dog” are the same bag of three vectors. The MLPs that follow are applied independently per position and do not mix anything across the sequence either. So a bare Transformer is, astonishingly, blind to word order.

That is a disaster for language, code, and basically every sequence we care about. The cure is positional encoding: we inject information about each token’s position so the model can tell where every token is and, more importantly, how far apart any two tokens are. This chapter is the complete tour of how that is done — from the original sinusoidal scheme, through learned tables, to Rotary Position Embedding (RoPE) and ALiBi, the two designs that dominate modern LLMs — and then the practical art of stretching a model trained on 4K tokens to run at 128K via position interpolation, NTK-aware scaling, and YaRN. We build RoPE entirely from scratch, twice, and verify its defining mathematical property numerically.

This chapter assumes the attention mechanism and the embedding pipeline. It pairs tightly with Multi-Head Attention, MQA, GQA & MLA (RoPE is applied per head) and feeds directly into Long-Context Pretraining & Context Extension. Everything here is cashed out concretely in the capstone: Stack-100M ships exactly this RoPE (with NoPE on every 4^th layer) in The Stack-100M Architecture, and its context is stretched from 2048 to 8192 tokens by RoPE base rescaling in Mid-Training: Quality Annealing, Long-Context Extension & Capability Injection.

Why Attention Needs Position, and the Two Ways to Add It¶

Let us nail down the permutation-equivariance claim, because it is the entire reason this chapter exists. Let \(P\) be a permutation matrix that reorders the rows of the input \(X \in \mathbb{R}^{n \times d}\). Self-attention computes \(Q = XW_Q\), \(K = XW_K\), \(V = XW_V\), then \(\operatorname{softmax}(QK^\top/\sqrt{d_k})V\). Permute the input to \(PX\) and every projection permutes: \(Q \to PQ\), \(K \to PK\), \(V \to PV\). The score matrix becomes \(PQ(PK)^\top = PQK^\top P^\top\), the softmax (applied row-wise) commutes with the row/column permutation, and the final output is \(P \cdot \operatorname{Attention}(Q,K,V)\). In words: shuffle the input, and the output is the same vectors shuffled the same way. There is nowhere in this computation for absolute or relative order to enter. We must add it ourselves.

There are two fundamentally different places to inject position, and the entire chapter is organized around this fork:

Approach A — modify the inputs. Give every position \(i\) a vector \(p_i\) and fold it into the token representation, classically by addition: \(\tilde{x}_i = x_i + p_i\). The model then sees a representation that encodes position, and attention can learn to use it. Sinusoidal and learned absolute encodings live here.

Approach B — modify the attention scores. Leave the token representations alone and instead add a position-dependent term directly to the pre-softmax scores \(S_{ij}\). Because attention’s only cross-token interaction is that score, this is the most surgical place to inject relative position \((i - j)\). ALiBi and T5’s relative bias live here.

RoPE is the brilliant hybrid. It modifies \(q_i\) and \(k_j\) (Approach A’s locus — the inputs to the score) but does so with a position-dependent rotation engineered so that the resulting dot product \(q_i \cdot k_j\) depends only on the relative offset \(i - j\) (Approach B’s goal). It gets the implementation simplicity of touching \(Q\) and \(K\) with the semantics of a relative scheme. We will see exactly how.

Three properties separate good schemes from bad ones, and we will grade every method against them:

1. Relative awareness. Language cares about distance (“the adjective modifies the next noun”), not absolute index. A scheme that makes the score depend on \(i-j\) rather than on \(i\) and \(j\) separately is structurally aligned with language.
2. Extrapolation / length generalization. If we train on sequences of length 4096 and run on 32768, does the model fall apart? Schemes with a learned table per position cannot extrapolate at all (there is no row for position 5000). Smooth functional schemes can, to varying degrees.
3. Efficiency and KV-cache friendliness. The encoding must be cheap and must not break the KV cache: the key for token \(j\) must be the same whether the sequence currently has length 100 or 100000, so we can compute it once and reuse it forever.

Sinusoidal and Learned Absolute Encodings¶

The original sinusoidal encoding¶

The 2017 Transformer of Vaswani et al. used a fixed, parameter-free encoding. For position \(\text{pos} \in \{0, 1, \dots\}\) and embedding dimension index \(i \in \{0, \dots, d-1\}\), define

Each pair of dimensions \((2k, 2k+1)\) is a sine/cosine at its own frequency. The frequencies form a geometric progression: dimension pair \(k=0\) oscillates fastest (wavelength \(2\pi\)), and the wavelengths grow geometrically up to \(2\pi \cdot 10000\) for the last pair. The intuition is a binary-clock for continuous space: just as the bits of an integer toggle at frequencies \(1, 2, 4, 8, \dots\), the dimensions of \(PE_{\text{pos}}\) oscillate at a spectrum of frequencies, so the full vector encodes the position with high resolution (fast dims) and long range (slow dims) simultaneously.

The clever part is why sinusoids specifically. For any fixed offset \(\Delta\), \(PE_{\text{pos}+\Delta}\) is a linear function of \(PE_{\text{pos}}\) — a rotation, in fact. Using the angle-addition formulas,

The shift-by-\(\Delta\) operator is a \(2\times 2\) rotation that does not depend on \(\text{pos}\), only on \(\Delta\). So a linear attention projection can in principle learn to detect relative offsets even though the encoding is absolute. (Hold onto this rotation matrix — RoPE is exactly this idea, moved from the input embedding into \(q\) and \(k\).)

import numpy as np

def sinusoidal_encoding(seq_len: int, d_model: int, base: float = 10000.0) -> np.ndarray:
    """Classic Vaswani et al. (2017) sinusoidal positional encoding.

    Returns an (seq_len, d_model) matrix PE where PE[pos] is added to the
    token embedding at position `pos`. Parameter-free and deterministic.
    """
    pe = np.zeros((seq_len, d_model), dtype=np.float64)
    position = np.arange(seq_len)[:, None]                 # (seq_len, 1)
    # 10000^(2k/d) for k = 0,1,...,d/2-1 -> the per-pair wavelengths.
    div_term = base ** (np.arange(0, d_model, 2) / d_model)  # (d_model/2,)
    pe[:, 0::2] = np.sin(position / div_term)              # even dims: sin
    pe[:, 1::2] = np.cos(position / div_term)              # odd  dims: cos
    return pe

PE = sinusoidal_encoding(seq_len=6, d_model=8)
print(PE.round(3))
# Add to embeddings: x_tilde = token_embedding + PE[:seq_len]

Grading sinusoidal. Relative awareness: implicit and weak — the structure is there (the rotation property) but the model must learn to exploit it, and empirically it does so only partially. Extrapolation: poor — the sinusoids are defined for any position so there is no crash, but trained attention heads simply have not seen the high-position patterns and generalize badly past the training length. Efficiency: excellent — zero parameters, precompute once. It was a fine first design; the field has largely moved past it.

Learned absolute encodings¶

GPT, GPT-2, BERT, and most early large models replaced the fixed sinusoids with a learned lookup table: a trainable matrix \(W_{\text{pos}} \in \mathbb{R}^{L_{\max} \times d}\) where row \(i\) is the embedding for position \(i\), learned by gradient descent exactly like a token embedding. You pick a maximum context length \(L_{\max}\) (e.g., 1024 for GPT-2, 512 for BERT), allocate that many rows, and add \(W_{\text{pos}}[i]\) to the token embedding at position \(i\).

import torch
import torch.nn as nn

class LearnedAbsolutePositionalEmbedding(nn.Module):
    """GPT-2 style learned position embeddings: a trainable (L_max, d) table."""
    def __init__(self, max_len: int, d_model: int):
        super().__init__()
        self.pos_emb = nn.Embedding(max_len, d_model)  # one learned vector per position
        self.max_len = max_len

    def forward(self, token_embeddings: torch.Tensor) -> torch.Tensor:
        # token_embeddings: (batch, seq_len, d_model)
        seq_len = token_embeddings.size(1)
        assert seq_len <= self.max_len, (
            f"seq_len={seq_len} exceeds max_len={self.max_len}: "
            "learned absolute encodings CANNOT extrapolate — there is no row for this position."
        )
        positions = torch.arange(seq_len, device=token_embeddings.device)
        return token_embeddings + self.pos_emb(positions)  # broadcast over batch

Learned tables are flexible — the model carves out whatever positional geometry minimizes loss — and they were the workhorse for years. But they have a fatal flaw for the long-context era, made explicit by the assert above: they cannot extrapolate even one token past \(L_{\max}\). There is simply no parameter for position 1025 in a model with 1024 rows. Worse, they tend to overfit absolute index and underperform on relative reasoning. This is the reason the field abandoned them.

Relative Position and the Road to RoPE¶

The insight that reorganized the field: what attention actually needs is the relative offset \(i - j\), not the absolute indices. A relative scheme makes the score between query \(i\) and key \(j\) a function of \((i-j)\), which is automatically translation-invariant (“two tokens apart” means the same thing at the start and end of a document) and naturally extrapolates (offset \(-3\) is the same operation everywhere).

The first influential realization, from Shaw et al. (2018) and refined in Transformer-XL (Dai et al., 2019), added learned vectors indexed by the clipped relative distance into the keys and values. T5 (Raffel et al., 2019) simplified this to a pure scalar bias on the score (Approach B):

where \(\text{bucket}(\cdot)\) maps relative distances into a modest number of learned buckets (logarithmically spaced, so nearby offsets get fine resolution and far ones get coarse, shared buckets), and \(b\) is a learned scalar per bucket per head. T5’s scheme is genuinely relative, parameter-light, and extrapolates moderately (far distances share a bucket). Its weaknesses: the bias is added to every layer’s scores (a small but real cost), and the bucketing is a discretization with hand-tuned hyperparameters.

This sets the stage. We want a relative scheme that is (a) continuous, not bucketed; (b) needs no extra learned parameters; © is cheap; and (d) plays perfectly with the KV cache. RoPE is that scheme.

RoPE: Rotary Position Embedding From Scratch¶

RoPE (Su et al., RoFormer, 2021) is the dominant positional encoding in modern LLMs, and it deserves a careful derivation because the idea is genuinely elegant. The goal is stated as a constraint: find functions \(f_q(x, m)\) and \(f_k(x, n)\) that encode position into the query at position \(m\) and the key at position \(n\) such that their inner product depends only on the relative position \(m - n\) and the original vectors:

The 2-D derivation¶

Start in two dimensions (\(d=2\)). Treat a 2-vector \(x = (x_0, x_1)\) as a complex number \(x_0 + i x_1\). Define the encoding as multiplication by a unit complex number whose angle is proportional to the position:

This is just rotating the 2-D vector by an angle \(m\theta\) that grows linearly with position. Now compute the inner product of a rotated query (position \(m\)) with a rotated key (position \(n\)). In complex form, the real inner product of two complex numbers \(a\) and \(b\) is \(\operatorname{Re}(a \bar{b})\), so

The position dependence collapsed to a single factor \(e^{i(m-n)\theta}\) — it depends only on \(m - n\). That is the defining property, achieved exactly, with no learned parameters and no approximation. Equivalently, in matrix form, \(R(m\theta)^\top R(n\theta) = R((n-m)\theta)\) because rotation matrices satisfy \(R(\alpha)^\top R(\beta) = R(\beta - \alpha)\). The two absolute rotations combine into one relative rotation.

Lifting to \(d\) dimensions¶

A head dimension \(d\) is much larger than 2, so RoPE partitions the \(d\)-dimensional vector into \(d/2\) consecutive pairs and rotates each pair by its own frequency. Pair \(k\) (dimensions \(2k, 2k+1\)) gets angular frequency

with \(\text{base} = 10000\) by default — the same geometric frequency ladder as the sinusoidal encoding, which is no coincidence. Low-index pairs rotate fast (encode fine, local position); high-index pairs rotate slowly (encode coarse, long-range position). The base is a real design knob, not a constant: models pretrained for long context now ship a much larger rope_theta from the start (Llama 3 uses 500,000 and the Qwen 2.5/3 families use 1,000,000, versus GPT-NeoX-era 10,000), because a bigger base slows every dimension down and keeps the slow pairs inside a well-trained angular range at long positions. There is a genuine floor here — the base needed to support a target context grows with that context, a point argued directly in the “base of RoPE bounds context length” line of work (Men et al., 2024) — so a model pretrained at base 10,000 has less headroom to be stretched later than one pretrained at 500,000. The full RoPE transform is a block-diagonal rotation matrix:

and because each \(2\times 2\) block individually satisfies the relative property, the whole thing does: \(R_\Theta(m)^\top R_\Theta(n) = R_\Theta(n - m)\). The score becomes

a function of \(q_m\), \(k_n\), and the relative offset \(n - m\) alone. RoPE is applied to \(q\) and \(k\) after the QKV projection and before the dot product, independently in every head and every layer; the value vectors \(V\) are not rotated (only the similarity needs position; the payload does not).

Implementation: the rotate-half trick¶

We never materialize the \(d \times d\) matrix \(R_\Theta(m)\) — it is \(99.6\%\) zeros. Instead we precompute per-position \(\cos\) and \(\sin\) vectors and apply the rotation elementwise. There are two equivalent conventions for which dimensions form a pair:

• Interleaved (original RoPE paper): pairs are adjacent dims \((0,1), (2,3), \dots\).
• rotate_half (GPT-NeoX / Llama / HuggingFace): pairs are dim \(k\) with dim \(k + d/2\), i.e., the first half is paired with the second half. This permits a vectorized rotate_half that is friendlier to hardware.

Both produce mathematically equivalent models as long as you are consistent (you can permute one into the other). The Llama-style rotate_half is by far the most common in practice, so we implement that, with the interleaved version as a from-scratch reference.

import torch

def build_rope_cache(seq_len: int, head_dim: int, base: float = 10000.0,
                     device=None, dtype=torch.float32):
    """Precompute cos/sin tables for RoPE (Llama / rotate_half convention).

    Returns cos, sin of shape (seq_len, head_dim). The first half of the
    frequency ladder is *duplicated* into the second half so that the same
    cos/sin vector multiplies x and rotate_half(x) elementwise.
    """
    assert head_dim % 2 == 0, "RoPE needs an even head dimension."
    # inv_freq[k] = base^(-2k/d) for k = 0..d/2-1  (the angular frequencies)
    half = head_dim // 2
    inv_freq = 1.0 / (base ** (torch.arange(0, half, device=device).float() / half))
    pos = torch.arange(seq_len, device=device).float()           # (seq_len,)
    freqs = torch.outer(pos, inv_freq)                           # (seq_len, d/2): angle = pos * theta_k
    emb = torch.cat([freqs, freqs], dim=-1)                      # (seq_len, d): duplicate the half
    return emb.cos().to(dtype), emb.sin().to(dtype)

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Map (x1, x2) -> (-x2, x1) where x1,x2 are the two halves of the last dim."""
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([-x2, x1], dim=-1)

def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply RoPE to x of shape (batch, n_heads, seq_len, head_dim).

    cos, sin: (seq_len, head_dim). We unsqueeze to broadcast over batch & heads.
    The elementwise formula  x * cos + rotate_half(x) * sin  is exactly the
    per-pair 2x2 rotation, written without ever forming the dense matrix.
    """
    cos = cos[None, None, :, :]   # (1, 1, seq_len, head_dim)
    sin = sin[None, None, :, :]
    return x * cos + rotate_half(x) * sin

# --- Wire it into a single attention head ---
torch.manual_seed(0)
B, H, T, D = 2, 4, 16, 64          # batch, heads, seq_len, head_dim
q = torch.randn(B, H, T, D)
k = torch.randn(B, H, T, D)
v = torch.randn(B, H, T, D)

cos, sin = build_rope_cache(T, D)
q_rot = apply_rope(q, cos, sin)   # rotate queries by their positions
k_rot = apply_rope(k, cos, sin)   # rotate keys   by their positions
# Now do ordinary scaled-dot-product attention with q_rot, k_rot, and *unrotated* v:
scores = (q_rot @ k_rot.transpose(-2, -1)) / D ** 0.5
# ... add causal mask, softmax, multiply by v (see Chapter 2.3) ...
print("q_rot shape:", q_rot.shape)   # torch.Size([2, 4, 16, 64])

Let us verify the defining property directly — that two absolutely rotated vectors yield a dot product depending only on the relative offset.

import torch

def rope_single(x, pos, base=10000.0):
    """Apply RoPE (rotate_half convention) to a single vector x at position `pos`."""
    d = x.shape[-1]
    half = d // 2
    inv_freq = 1.0 / (base ** (torch.arange(0, half).float() / half))
    ang = pos * inv_freq                       # (d/2,)
    emb = torch.cat([ang, ang])                # (d,)
    cos, sin = emb.cos(), emb.sin()
    x1, x2 = x[:half], x[half:]
    rot = torch.cat([-x2, x1])
    return x * cos + rot * sin

torch.manual_seed(7)
d = 64
q = torch.randn(d); k = torch.randn(d)
m, n = 25, 10                                  # query at 25, key at 10  -> offset 15
lhs = rope_single(q, m) @ rope_single(k, n)    # absolute positions 25 and 10
rhs = rope_single(q, m - n) @ rope_single(k, 0)  # relative: offset 15 and 0
print(float(lhs), float(rhs))                  # the two numbers match
assert torch.allclose(lhs, rhs, atol=1e-4), "RoPE relative property failed!"
print("RoPE relative-position property verified.")

This runs and prints two essentially identical numbers: the score from rotating \(q\) at position 25 and \(k\) at position 10 equals the score from rotating \(q\) at offset 15 and \(k\) at offset 0. The absolute positions cancel; only the offset \(15 = 25 - 10\) survives.

Positions are an input, not arange¶

The code above quietly assumed every token’s position is its index in the batch, arange(T). That assumption is wrong in two of the three situations you will actually hit, and both failures are silent — the model still runs, it just gets worse. Build your RoPE application to take an explicit position_ids tensor and index the tables with it; that single indirection is what makes packing, incremental decoding, and context extension all work with no further changes.

import torch
# Reuses build_rope_cache() and rotate_half() defined earlier in this chapter.

def apply_rope_at(x, position_ids, cos_tab, sin_tab):
    """RoPE with EXPLICIT positions.

    x:            (B, H, T, D) queries or keys
    position_ids: (B, T) int64 — each token's TRUE absolute position
    cos_tab/sin_tab: (max_seq, D) tables from build_rope_cache
    """
    cos = cos_tab[position_ids].unsqueeze(1)   # (B, 1, T, D): broadcast over heads
    sin = sin_tab[position_ids].unsqueeze(1)
    return x * cos + rotate_half(x) * sin

B, H, D, MAX_SEQ = 1, 4, 64, 512
cos_tab, sin_tab = build_rope_cache(MAX_SEQ, D)

# (a) Training on PACKED documents: positions restart at each document boundary,
#     matching the block-diagonal mask so document 2 does not think it starts at 4000.
pos_packed = torch.tensor([[0, 1, 2, 3, 0, 1, 0, 1]])          # three documents in one window
q_packed = apply_rope_at(torch.randn(B, H, 8, D), pos_packed, cos_tab, sin_tab)

# (b) Incremental decoding with a KV cache: the ONE new token sits at position
#     past_len, NOT 0. Rotating it at 0 is the single most common RoPE bug —
#     the model degenerates into treating every generated token as the sequence start.
past_len = 137
q_step = apply_rope_at(torch.randn(B, H, 1, D),
                       torch.tensor([[past_len]]), cos_tab, sin_tab)
print(q_packed.shape, q_step.shape)   # torch.Size([1, 4, 8, 64]) torch.Size([1, 4, 1, 64])

Three practical rules fall out of this. (1) Rotate the key once, at write time, and store the rotated key in the cache — never re-rotate cached keys, and never rotate them by an offset relative to the current query. (2) Under GQA, apply RoPE to the narrow key tensor (n_kv_heads) before expanding it to n_heads; the rotation is per-position and head-count-agnostic, so doing it after expansion is pure wasted work. (3) RoPE happens before the attention kernel — FlashAttention and torch.nn.functional.scaled_dot_product_attention take already-rotated \(q,k\) and know nothing about position. The capstone’s model does exactly this and derives position_ids from the packed token stream itself (The Stack-100M Architecture).

The libraries you will actually call¶

You will implement RoPE from scratch exactly once — to understand it. In production it comes from a library, and the two you will meet are HuggingFace transformers (training/reference) and a fused CUDA kernel (serving).

# --- HuggingFace Transformers: the same math, config-driven ---
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
print(sorted(ROPE_INIT_FUNCTIONS))
# e.g. ['dynamic', 'linear', 'llama3', 'longrope', 'proportional', 'yarn']
#      (plus a default/no-scaling path; the exact set grows with the version)

from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, rotate_half
# Inside LlamaAttention the flow is exactly ours:
#   cos, sin = self.rotary_emb(value_states, position_ids)   # tables built from config
#   q, k     = apply_rotary_pos_emb(q, k, cos, sin)          # x*cos + rotate_half(x)*sin
# `rotary_emb` picks its table-builder from ROPE_INIT_FUNCTIONS[<the config's rope_type>],
# which is how one config dict switches between linear / yarn / llama3 / longrope.

Every scaling method in the next section is therefore a config edit, not a code change. For throughput, flash_attn.layers.rotary.apply_rotary_emb (from the flash-attn package) fuses the rotation into one kernel pass and supports interleaved=True for the original RoFormer pairing — note it expects half-width cos/sin of shape (seqlen, rotary_dim/2) rather than the duplicated full-width tables above, a mismatch worth checking against your own implementation with torch.allclose. vLLM and SGLang ship their own fused rotary kernels (vLLM’s get_rope(...) factory mirrors the HF config keys), and llama.cpp implements RoPE as a GGML op configured by rope_freq_base / rope_scaling_type stored in the GGUF header. All of them compute the same rotation; they differ only in the pairing convention and where the multiply happens.

Grading RoPE. Relative awareness: exact, by construction. Extrapolation: good and tunable (the basis for the scaling tricks below). Efficiency: excellent — no parameters, a cheap elementwise op, and it is fully KV-cache compatible because the rotation applied to key \(n\) depends only on \(n\), never on the current sequence length, so cached keys never need recomputation. This combination is why RoPE won.

ALiBi and NoPE: Biasing the Score, or Nothing at All¶

ALiBi: a linear distance penalty¶

ALiBi (Attention with Linear Biases; Press, Smith, Lewis, 2021) takes the radical Approach-B route: add no positional information to the embeddings or to \(q\)/\(k\) at all, and instead subtract a penalty from each attention score that grows linearly with the distance between query and key:

Here \(|i - j|\) is the relative distance and \(m_h\) is a fixed, non-learned slope assigned per head. The slopes form a geometric sequence; for a model with \(H\) heads they are typically \(m_h = 2^{-8h/H}\) for \(h = 1, \dots, H\) — so some heads have a steep penalty (sharply local, attend only to very recent tokens) and others have a shallow penalty (nearly global). The bias is added to the score matrix and otherwise the attention is ordinary.

import torch

def alibi_slopes(n_heads: int) -> torch.Tensor:
    """Geometric slopes m_h = 2^(-8h/H), the standard ALiBi schedule (power-of-two heads)."""
    start = 2 ** (-8 / n_heads)             # ratio so that head 1 gets 2^(-8/H), head H gets 2^-8
    return torch.tensor([start ** (i + 1) for i in range(n_heads)])

def alibi_bias(seq_len: int, n_heads: int) -> torch.Tensor:
    """Per-head additive bias of shape (n_heads, seq_len, seq_len): -m_h * |i - j|."""
    slopes = alibi_slopes(n_heads)                                  # (n_heads,)
    i = torch.arange(seq_len)[:, None]                              # query index (rows)
    j = torch.arange(seq_len)[None, :]                              # key   index (cols)
    distance = (j - i).clamp(max=0).abs().float()                  # causal: |i-j| for j<=i, else 0
    causal = (j <= i)                                              # boolean causal mask
    bias = -slopes[:, None, None] * distance[None]                # (n_heads, seq_len, seq_len)
    bias = bias.masked_fill(~causal[None], float("-inf"))         # forbid attending to the future
    return bias

bias = alibi_bias(seq_len=8, n_heads=4)
# Usage: scores = q @ k.transpose(-2,-1) / sqrt(d); scores = scores + bias; softmax(scores) @ v
print(bias.shape)   # torch.Size([4, 8, 8])
print(alibi_slopes(8).round(decimals=4))

ALiBi’s headline virtue is length extrapolation: because the penalty is a smooth linear function of distance defined for any distance, a model trained at length 1024 runs at length 4096+ with little degradation — the bias at distance 3000 is just a bigger version of the bias at distance 800, nothing the model has not effectively seen the shape of. It is also dead simple and cheap. Its limitation is that the strict monotonic decay with distance is an inductive bias toward recency — fine for left-to-right language modeling, but a poor fit for tasks needing strong long-range retrieval, where a far-away token must sometimes dominate. ALiBi powered BLOOM and MPT; RoPE ultimately became more popular for frontier models partly because of this retrieval ceiling, though ALiBi remains an excellent, robust choice.

NoPE: maybe you need nothing¶

A genuinely surprising 2023 result (Kazemnejad et al., The Impact of Positional Encoding on Length Generalization): a decoder-only Transformer with a causal mask and no positional encoding at all (NoPE) can still learn position — and sometimes generalizes to longer lengths better than explicit schemes. How? The causal mask itself breaks permutation symmetry: token \(i\) can attend to \(\{0, \dots, i\}\) but token \(j > i\) can attend to a strictly larger set. The model can count how many tokens are visible (e.g., via attention patterns that effectively measure the size of the attended set), recovering absolute position implicitly. NoPE is mostly a research curiosity and a clarifying conceptual point — it proves position info is latent in causality — rather than a production default, but it is a favorite interview probe and a reminder that the causal mask is itself doing positional work.

Context-Length Extension: Stretching RoPE to 128K¶

You have a model pretrained with RoPE at 4096 tokens. You want it to handle 32K or 128K. Retraining from scratch at long context is enormously expensive (attention is \(\mathcal{O}(n^2)\); see the long-context chapter). The RoPE-scaling family lets you extend context with little or no additional training, by manipulating the frequency ladder. This is one of the most practically important — and interview-hot — topics in the chapter.

The core problem, restated from the worked example: at inference position \(p > L_{\text{train}}\), the slow RoPE dimensions rotate into angles never seen in training. There are three families of fixes, best understood through the metaphor of an instrument with many strings tuned to different frequencies.

Position Interpolation (PI)¶

Position Interpolation (Chen et al., 2023) is the simplest idea: if the model was trained to understand positions \(0\) to \(L\), don’t let positions exceed \(L\) — instead squeeze the longer sequence into the trained range. To extend to length \(L' = s \cdot L\) (scale factor \(s = L'/L\)), replace each position \(m\) with \(m / s\) before computing rotations:

Geometrically, you compress all positions so position 8000 in a 2× extension is treated like position 4000 — an angle the model has seen. Every dimension’s rotation now stays within the trained angular range. The cost: you have reduced the resolution of position (adjacent tokens are now only \(\theta_k/s\) apart in angle, so the model must distinguish finer differences). PI works remarkably well but typically needs a short fine-tune (a few hundred steps) to recover the lost fine-grained resolution.

NTK-aware scaling¶

PI’s flaw is that it scales every frequency uniformly, including the fast (high-frequency, local) dimensions that were not the problem — squeezing them needlessly destroys local resolution. NTK-aware scaling (named for Neural Tangent Kernel intuitions; originally a community/blog contribution, “bloc97”) fixes this by scaling non-uniformly: change the RoPE base instead of the positions, which stretches slow dimensions a lot and fast dimensions barely at all.

The high-frequency pairs (small \(k\)) are nearly unchanged (local precision preserved); the low-frequency pairs (large \(k\)) are stretched the most (so they don’t run off into unseen angles). The headline practical fact: NTK-aware scaling often extends context with zero fine-tuning — you can sometimes 2–4× a model’s context by changing one number, the base, at inference time. This is why “rope_theta” appears as a config knob in every modern serving stack.

import torch

def ntk_aware_inv_freq(head_dim: int, scale: float, base: float = 10000.0):
    """NTK-aware RoPE: stretch the base by scale^(d/(d-2)), then recompute frequencies.

    `scale` = target_len / train_len. High-freq dims barely change;
    low-freq dims stretch the most, keeping their angles in the trained range.
    """
    base_prime = base * (scale ** (head_dim / (head_dim - 2)))
    half = head_dim // 2
    return 1.0 / (base_prime ** (torch.arange(0, half).float() / half))

# Extend a d=128 model's context 4x with ZERO retraining, just by changing the base:
orig = 1.0 / (10000.0 ** (torch.arange(0, 64).float() / 64))
ntk  = ntk_aware_inv_freq(head_dim=128, scale=4.0)
print("fast dim ratio (k=0): ", float(ntk[0] / orig[0]))    # ~1.0  (local: unchanged)
print("slow dim ratio (k=63):", float(ntk[63] / orig[63]))  # <1.0  (long-range: stretched)

YaRN: the production-grade combination¶

YaRN (Yet another RoPE extensioN; Peng et al., 2023) is the method most long-context production models actually ship. It refines NTK with two extra ideas:

1. NTK-by-parts (per-dimension ramp). Instead of one global rule, classify each frequency dimension by its wavelength relative to the training context. Dimensions whose wavelength is short compared to \(L\) (they complete many full rotations within the trained window) are left untouched — they already generalize. Dimensions whose wavelength is long (they never complete a rotation in \(L\)) get interpolated like PI. A smooth ramp blends the two regimes for in-between dimensions. This is strictly better than treating all dimensions identically.
2. Attention temperature scaling. As you stretch context, the softmax distribution over many more keys tends to flatten and entropy rises. YaRN multiplies the attention logits by a temperature factor \(\frac{1}{t}\) defined by \(\sqrt{1/t} = 0.1 \ln(s) + 1\) (a function of the scale \(s\)), so the logits are effectively scaled by \(\big(0.1 \ln s + 1\big)^2\) — this sharpens the flattened distribution back and recovers effective attention. Crucially it costs nothing at runtime: implementations fold the per-vector factor \(m_{\text{scale}} = \sqrt{1/t} = 0.1 \ln(s) + 1\) into the RoPE \(\cos/\sin\) tables and apply it to both \(q\) and \(k\), so the square \(m_{\text{scale}}^2 = 1/t\) arises automatically from the \(q\cdot k\) dot product. (Watch this factor-of-two-in-the-exponent when reading the paper: the paper’s \(\sqrt{1/t}\) is the per-tensor scale, not the logit scale.)

The empirical result is that YaRN extends context further than PI or plain NTK and with less fine-tuning — it is the basis of many “32K → 128K” and “128K → 1M” model releases. The lineage is worth memorizing: PI (uniform squeeze, needs fine-tune) → NTK-aware (non-uniform via base, often training-free) → YaRN (per-dimension ramp + logit temperature, the production default). HuggingFace exposes these through a rope_scaling/rope_parameters config whose type is one of linear/dynamic/yarn/llama3/longrope; the dynamic variant recomputes the scale on the fly so short sequences stay un-degraded, and Llama 3.1’s llama3 type applies a frequency-split rescaling (separate low- and high-frequency factors) that has itself become a common default.

import torch

def rope_scaling_summary():
    """The decision the lineage encodes, in code-as-table form."""
    return {
        "linear / PI":  "pos -> pos/s; uniform squeeze; simplest; usually needs a short fine-tune",
        "ntk-aware":    "base -> base * s^(d/(d-2)); non-uniform; OFTEN training-free for 2-4x",
        "yarn":         "per-dim ramp (NTK-by-parts) + logit temperature; best quality, longest reach",
        "dynamic-ntk":  "recompute scale per actual seq_len; no penalty on short inputs",
    }
for k, v in rope_scaling_summary().items():
    print(f"{k:14s}: {v}")

In practice you rarely write any of this: you edit a config or pass a serving flag. The three places that matter are the HF config (training and reference inference), the inference server, and the GGUF/llama.cpp path.

# 1) HuggingFace config.json — the source of truth a checkpoint carries with it.
#    {"rope_theta": 10000.0,
#     "rope_scaling": {"rope_type": "yarn", "factor": 4.0,
#                      "original_max_position_embeddings": 8192}}
#    Recent transformers releases normalize this into a single `rope_parameters`
#    dict (absorbing rope_theta) and keep `rope_scaling` as a read-compatible alias.

# 2) vLLM — override the checkpoint's rotary config at serve time, no retraining.
vllm serve my-org/my-8k-model \
  --max-model-len 32768 \
  --hf-overrides '{"rope_scaling": {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 8192}}'
# (older vLLM releases expose dedicated --rope-scaling / --rope-theta flags instead)

# 3) llama.cpp — the same knobs on a quantized GGUF model.
llama-cli -m model.gguf -c 32768 --rope-scaling yarn --yarn-orig-ctx 8192

Whichever route you take, factor and original_max_position_embeddings must describe the original pretraining length, and the server’s --max-model-len must not exceed factor × original — the two are a matched pair (recent transformers versions will warn you when the explicit factor disagrees with max_position_embeddings / original_max_position_embeddings), and mismatching them is exactly the double-counting failure described in the warning below. The capstone works this end to end for Stack-100M, rescaling the base from 10,000 to about 42,000 for a 2048 → 8192 extension and then continuing training at the new length (Mid-Training).

Putting It Together: Choosing and Reasoning About Position¶

We have traversed the whole landscape. Step back and the design space is a small, comprehensible table.

The story of the field is a march from Approach A’s absolute encodings (sinusoidal, learned) toward relative ones, with the two survivors being RoPE (rotate \(q,k\) — exact relative, tunable extrapolation, the default for almost all frontier models) and ALiBi (a linear score penalty — supremely simple and robust extrapolation, recency-biased). The context-extension toolkit — PI, NTK-aware, YaRN — exists because RoPE’s frequency ladder is a continuous, manipulable object, turning “train short, serve long” from impossible (learned tables) into a config flag.

From here, RoPE’d attention slots into a full Transformer block and a from-scratch GPT; its interaction with GQA/MLA matters because RoPE is applied per-head and MLA must carefully preserve the rotation through its low-rank compression; and the long-context economics it enables drive inference and scaling-law decisions throughout the stack.

Further reading¶

• Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin — Attention Is All You Need (2017). Introduces the sinusoidal positional encoding and the rotation-by-fixed-offset property.
• Su, Lu, Pan, Murtadha, Wen, Liu — RoFormer: Enhanced Transformer with Rotary Position Embedding (2021). The RoPE paper; the 2-D complex-multiplication derivation and the \(d\)-dimensional block-rotation construction.
• Press, Smith, Lewis — Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (2021). The ALiBi method and its extrapolation story.
• Raffel, Shazeer, Roberts, et al. — Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (T5, 2019). The bucketed relative-position scalar bias on attention scores.
• Shaw, Uszkoreit, Vaswani — Self-Attention with Relative Position Representations (2018), and Dai et al. — Transformer-XL (2019). The origins of relative position representations.
• Chen, Wong, Chen, Tian — Extending Context Window of Large Language Models via Positional Interpolation (2023). Position Interpolation.
• Peng, Quesnelle, Fan, Shippole — YaRN: Efficient Context Window Extension of Large Language Models (2023). NTK-by-parts plus attention temperature; the dominant production extension method.
• Kazemnejad, Padhi, Natesan Ramamurthy, Das, Reddy — The Impact of Positional Encoding on Length Generalization in Transformers (2023). The NoPE result for decoder-only models.

Exercises¶

1. A colleague claims that a Transformer with no positional scheme can still tell “dog bites man” from “man bites dog” because the words are in different places in the input tensor. Using the permutation-equivariance argument from this chapter, explain precisely why they are wrong for a bidirectional (no causal mask) encoder. Then explain why the situation is subtly different for a decoder-only model with a causal mask (the NoPE result).

2. Using the sinusoidal formula from the chapter with d_model = 8 and base = 10000, compute the four per-pair divisors div_term[k] (for \(k=0,1,2,3\)) and the corresponding wavelengths of each sine/cosine. Which dimension pair has the finest positional resolution, and which has the longest range?

3. A model has head dimension \(d = 128\) and was trained with RoPE at base = 10000. You want to serve it at \(4\times\) the training length (\(s = 4\)) using NTK-aware scaling.

(a) Compute the new base \(\text{base}'\). (b) By what factor does the fastest pair’s frequency \(\theta_0\) change, and by what factor does the slowest pair’s frequency \(\theta_{63}\) change? Interpret both. © Contrast this with what Position Interpolation would do to those same two frequencies.

4. The chapter implements the Llama rotate_half convention (pairing dim \(k\) with dim \(k+d/2\)) and verifies the relative-position property. Implement the interleaved convention instead (pairs are adjacent dims \((0,1),(2,3),\dots\)), and numerically verify that it too satisfies \(\langle f(q,m), f(k,n)\rangle = g(q,k,\,m-n)\).

import torch

def rope_interleaved(x, pos, base=10000.0):
    """Apply RoPE (interleaved convention) to a single vector x at position `pos`.
    Pairs are adjacent dims (0,1),(2,3),...; pair k rotates by pos * base^(-2k/d)."""
    d = x.shape[-1]
    half = d // 2
    k = torch.arange(half).float()             # pair indices 0..d/2-1
    theta = base ** (-2.0 * k / d)             # angular frequency per pair
    ang = pos * theta                          # (d/2,)
    cos, sin = ang.cos(), ang.sin()
    x0 = x[0::2]                                # even dims -> first  element of each pair
    x1 = x[1::2]                                # odd  dims -> second element of each pair
    out = torch.empty_like(x)
    out[0::2] = x0 * cos - x1 * sin
    out[1::2] = x0 * sin + x1 * cos
    return out

torch.manual_seed(7)
d = 64
q = torch.randn(d); k = torch.randn(d)
m, n = 25, 10                                   # offset m - n = 15
lhs = rope_interleaved(q, m) @ rope_interleaved(k, n)     # absolute positions 25, 10
rhs = rope_interleaved(q, m - n) @ rope_interleaved(k, 0) # relative: offset 15, 0
print(float(lhs), float(rhs))
assert torch.allclose(lhs, rhs, atol=1e-4), "interleaved RoPE relative property failed!"
print("Interleaved RoPE relative-position property verified.")

5. RoPE rotates \(Q\) and \(K\) but leaves \(V\) untouched, and its rotation of key \(n\) depends only on \(n\).

(a) Explain why rotating \(V\) would be a mistake, in terms of what the score versus the value represent. (b) Explain why the “rotation of key \(n\) depends only on \(n\)” fact is exactly what makes RoPE KV-cache-safe, and why a learned absolute table is also cache-safe yet still fails the long-context test for a different reason.

6. You download a long-context model whose config already advertises rope_scaling with a YaRN factor of 8 (extended from a 4K pretraining length to 32K). You want 64K, so you set an additional linear scale of 2 at serving time and observe that (a) perplexity at 64K looks fine but (b) needle-in-a-haystack retrieval collapses even at 16K, well within the advertised 32K. Diagnose both observations using the chapter’s warnings.