Positional Information

Dive into Deep Learning · §10.4

Positional information
permutation blindness · sinusoids and tables · RoPE · train short, test long

Attention ignores order

A one-line theorem

Proposition. Unmasked self-attention is permutation equivariant: f(\boldsymbol{\Pi}\mathbf{X}) = \boldsymbol{\Pi} f(\mathbf{X}) for every permutation matrix \boldsymbol{\Pi}.

Proof sketch. Projections act row-wise; scores become \boldsymbol{\Pi}\mathbf{S}\boldsymbol{\Pi}^\top; row-softmax commutes with permutations; \boldsymbol{\Pi}^\top\boldsymbol{\Pi} = \mathbf{I}. \blacksquare

  • Stacking equivariant layers keeps the model equivariant — depth is no cure.
  • “dog bites man” = “man bites dog”, as far as attention can tell.

The shuffle check

max |Y[perm] - Y_perm|: 6.33e-08

Caveat for later: a causal mask breaks the symmetry — position i attends over i+1 tokens. That loophole becomes NoPE.

Sinusoidal encodings

Add a designed position vector to each token (Vaswani et al., 2017):

p_{i, 2j} = \sin\big(i/10000^{2j/d}\big), \qquad p_{i, 2j+1} = \cos\big(i/10000^{2j/d}\big)

def sinusoidal_encoding(max_len, num_hiddens):
    theta = torch.arange(max_len)[:, None] / 10000 ** (
        torch.arange(0, num_hiddens, 2) / num_hiddens)
    return torch.stack([torch.sin(theta), torch.cos(theta)],
                       -1).reshape(max_len, num_hiddens)

P = sinusoidal_encoding(60, 32)
d2l.plot(torch.arange(60), P[:, 6:10].T, xlabel='position',
         figsize=(6, 2.5), legend=[f'column {d}' for d in range(6, 10)])

A continuous binary counter

Low bits flip fast, high bits flip slow — same pattern, smooth values:

Learned tables (BERT, GPT-2) instead train one free vector per position — and say nothing about positions beyond the trained length.

The rotation hidden in the sinusoids

Moving from position i to i + \delta is a rotation of each two-column pair, by an angle depending only on \delta:

\begin{bmatrix} \cos(\delta\omega_j) & \sin(\delta\omega_j) \\ -\sin(\delta\omega_j) & \cos(\delta\omega_j) \end{bmatrix}\begin{bmatrix} p_{i, 2j} \\ p_{i, 2j+1} \end{bmatrix} = \begin{bmatrix} p_{i+\delta, 2j} \\ p_{i+\delta, 2j+1} \end{bmatrix}

But additively encoded, the model must learn to exploit it — a query–key product of sums has four terms. Why not put the rotation into the score directly?

Rotary position embeddings (RoPE)

The scheme of essentially every open-weights model

Demand scores that see only the offset: (\mathbf{R}_i \mathbf{q})^\top(\mathbf{R}_j \mathbf{k}) = \mathbf{q}^\top \mathbf{R}_{j-i}\, \mathbf{k} ⇒ the \mathbf{R}_i form a rotation group: rotate feature pair m of queries and keys by i\,\omega_m.

image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

RoPE in code

def rope(x, offset=0):
    """Rotate feature pairs of x (..., num_steps, d) by position angles."""
    d = x.shape[-1]
    pos = torch.arange(x.shape[-2], dtype=torch.float32,
                       device=x.device) + offset
    inv_freq = 10000.0 ** (-torch.arange(0, d, 2, device=x.device) / d)
    theta = pos[:, None] * inv_freq[None, :]
    cos, sin = torch.cos(theta), torch.sin(theta)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin,
                        x1 * sin + x2 * cos], -1).flatten(-2)

Invariance, measured

Shift all positions jointly — RoPE scores must not move:

RoPE, all positions shifted by   1: max score change 1.7e-06
RoPE, all positions shifted by  17: max score change 1.5e-06
RoPE, all positions shifted by 480: max score change 1.4e-05
additive sinusoidal, shifted by  17: max score change 2.8e+00

Relative position is now a property of the architecture, not something training must discover.

Train short, test long

Contexts grow after deployment. Two schemes designed for extrapolation:

  • ALiBi (Press et al., 2022): no encoding; subtract a per-head linear distance penalty \mathrm{score}_{ij} = \mathbf{q}_i^\top\mathbf{k}_j/\sqrt{d} - m_h\,(i-j), slopes m_h = 2^{-8h/H} (power-of-two H). Distance 400 is just “further”, never “unseen”.
  • NoPE (Kazemnejad et al., 2023): nothing at all — the causal mask leaks position; how much is usable?

An attention-only language model

TinyCharLM: token embedding + stacked residual causal attention + tied head. No FFN, no LayerNorm, no projection biases — attention is the only machinery mixing information across positions. Positional scheme is a constructor choice: 'learned' · 'sinusoidal' · 'rope' · 'alibi' · 'none'.

class TinyCharLM(nn.Module):
    """Attention-only character-level language model."""
    def __init__(self, vocab_size, num_hiddens=128, num_heads=4, num_blks=2,
                 pos='rope', max_len=512, bias=False):
        super().__init__()
        self.num_heads, self.pos = num_heads, pos
        self.token_emb = nn.Embedding(vocab_size, num_hiddens)
        nn.init.normal_(self.token_emb.weight, std=0.02)
        if pos == 'learned':
            self.pos_emb = nn.Embedding(max_len, num_hiddens)
            nn.init.normal_(self.pos_emb.weight, std=0.02)
        if pos == 'sinusoidal':
            theta = torch.arange(max_len)[:, None] / 10000 ** (
                torch.arange(0, num_hiddens, 2) / num_hiddens)
            P = torch.stack([torch.sin(theta), torch.cos(theta)], -1)
            self.register_buffer('P', P.reshape(max_len, num_hiddens))
        self.blks = nn.ModuleList([nn.ModuleDict(dict(
            qkv=nn.Linear(num_hiddens, 3 * num_hiddens, bias=bias),
            proj=nn.Linear(num_hiddens, num_hiddens, bias=bias)))
            for _ in range(num_blks)])

    def _rope(self, x):
        d = x.shape[-1]
        pos = torch.arange(x.shape[-2], dtype=torch.float32, device=x.device)
        inv_freq = 10000.0 ** (-torch.arange(0, d, 2, device=x.device) / d)
        theta = pos[:, None] * inv_freq[None, :]
        cos, sin = torch.cos(theta), torch.sin(theta)
        x1, x2 = x[..., 0::2], x[..., 1::2]
        return torch.stack([x1 * cos - x2 * sin,
                            x1 * sin + x2 * cos], -1).flatten(-2)

    def _alibi(self, T, device):
        h = torch.arange(1, self.num_heads + 1, device=device)
        slopes = 2.0 ** (-8.0 * h / self.num_heads)
        pos = torch.arange(T, device=device, dtype=torch.float32)
        return slopes[:, None, None] * (pos[None, :] - pos[:, None])

    def _attend(self, blk, H):
        B, T, D = H.shape
        q, k, v = blk['qkv'](H).chunk(3, -1)
        q, k, v = (u.reshape(B, T, self.num_heads, -1).transpose(1, 2)
                   for u in (q, k, v))
        if self.pos == 'rope':
            q, k = self._rope(q), self._rope(k)
        scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
        if self.pos == 'alibi':
            scores = scores + self._alibi(T, H.device)
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool,
                                     device=H.device), 1)
        scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min)
        weights = F.softmax(scores, dim=-1)
        out = (weights @ v).transpose(1, 2).reshape(B, T, D)
        return blk['proj'](out), weights

    def _embed(self, X):
        H = self.token_emb(X)
        if self.pos == 'learned':
            H = H + self.pos_emb(torch.arange(X.shape[1], device=X.device))
        if self.pos == 'sinusoidal':
            H = H + self.P[:X.shape[1]]
        return H

    def forward(self, X):
        H = self._embed(X)
        for blk in self.blks:
            out, _ = self._attend(blk, H)
            H = H + out
        return F.linear(H, self.token_emb.weight)  # Tied output head

    def attention_weights(self, X):
        """Per-block attention maps, each (batch, num_heads, T, T)."""
        H, maps = self._embed(X), []
        for blk in self.blks:
            out, weights = self._attend(blk, H)
            maps.append(weights)
            H = H + out
        return maps

Five models, one corpus

Character-level Time Machine, context 128, 3000 steps each:

   learned: final training loss 1.81
sinusoidal: final training loss 1.96
      rope: final training loss 1.54
     alibi: final training loss 2.01
      none: final training loss 2.24

The verdict at 4× the training length

   learned:     6.0     15.9     26.1
sinusoidal:     6.9     19.8     35.1
      rope:     5.1     14.9     69.5
     alibi:     7.3      7.3      7.2
      none:     9.2      9.4      9.5
  • Training length: RoPE best (~5), learned close behind, sinusoidal ≈ ALiBi (~7), NoPE worst (~9).
  • Length 512: absolute schemes blow up; RoPE blows up too — relative in form is not relative in practice (offsets > 127 were never trained).
  • ALiBi: flat. NoPE: flat, from a weaker start.

Stretching a trained model

RoPE’s dial is continuous → rescale angles so length 4L maps into the trained range: position interpolation (Chen et al., 2023) took Llama 2k → 32k with a brief fine-tune; YaRN (Peng et al., 2024) rescales fast and slow frequencies differently.

Interpolation beats extrapolation — this is how long-context RoPE models are actually made.

Recap

  • Unmasked attention is permutation equivariant — proven and measured.
  • Absolute encodings (sinusoidal, learned) restore order but stop at the trained length.
  • RoPE rotates queries and keys: scores depend only on offsets, by construction.
  • Train short, test long: absolute and RoPE collapse at 4L; ALiBi and NoPE stay flat; PI/YaRN rescue RoPE by interpolation.
  • TinyCharLM — attention-only, positional scheme as an argument — returns later as our specimen for reading trained attention.