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|: 0.00e+00

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 = jnp.arange(max_len)[:, None] / 10000 ** (
        jnp.arange(0, num_hiddens, 2) / num_hiddens)
    return jnp.stack([jnp.sin(theta), jnp.cos(theta)],
                     -1).reshape(max_len, num_hiddens)

P = sinusoidal_encoding(60, 32)
d2l.plot(jnp.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 = jnp.arange(x.shape[-2], dtype=jnp.float32) + offset
    inv_freq = 10000.0 ** (-jnp.arange(0, d, 2) / d)
    theta = pos[:, None] * inv_freq[None, :]
    cos, sin = jnp.cos(theta), jnp.sin(theta)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return jnp.stack([x1 * cos - x2 * sin,
                      x1 * sin + x2 * cos], -1).reshape(x.shape)

Invariance, measured

Shift all positions jointly — RoPE scores must not move:

RoPE, all positions shifted by   1: max score change 1.5e-03
RoPE, all positions shifted by  17: max score change 1.7e-03
RoPE, all positions shifted by 480: max score change 1.8e-03
additive sinusoidal, shifted by  17: max score change 3.3e+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(nnx.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, rngs=None):
        rngs = nnx.Rngs(0) if rngs is None else rngs
        self.num_heads, self.pos = num_heads, pos
        init = nnx.initializers.normal(0.02)
        self.token_emb = nnx.Embed(vocab_size, num_hiddens,
                                   embedding_init=init, rngs=rngs)
        if pos == 'learned':
            self.pos_emb = nnx.Embed(max_len, num_hiddens,
                                     embedding_init=init, rngs=rngs)
        if pos == 'sinusoidal':
            theta = jnp.arange(max_len)[:, None] / 10000 ** (
                jnp.arange(0, num_hiddens, 2) / num_hiddens)
            P = jnp.stack([jnp.sin(theta), jnp.cos(theta)], -1)
            self.P = nnx.Cache(P.reshape(max_len, num_hiddens))
        self.blks = nnx.List([nnx.Dict(
            qkv=nnx.Linear(num_hiddens, 3 * num_hiddens, use_bias=bias,
                           rngs=rngs),
            proj=nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,
                            rngs=rngs))
            for _ in range(num_blks)])

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

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

    def _attend(self, blk, H):
        B, T, D = H.shape
        q, k, v = jnp.split(blk['qkv'](H), 3, axis=-1)
        q, k, v = (u.reshape(B, T, self.num_heads, -1).swapaxes(1, 2)
                   for u in (q, k, v))
        if self.pos == 'rope':
            q, k = self._rope(q), self._rope(k)
        scores = q @ k.swapaxes(-2, -1) / math.sqrt(q.shape[-1])
        if self.pos == 'alibi':
            scores = scores + self._alibi(T)
        mask = jnp.triu(jnp.ones((T, T), dtype=bool), 1)
        scores = jnp.where(mask, jnp.finfo(scores.dtype).min, scores)
        weights = jax.nn.softmax(scores, axis=-1)
        out = (weights @ v).swapaxes(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(jnp.arange(X.shape[1]))
        if self.pos == 'sinusoidal':
            H = H + self.P[:X.shape[1]]
        return H

    def __call__(self, X):
        H = self._embed(X)
        for blk in self.blks:
            out, _ = self._attend(blk, H)
            H = H + out
        return self.token_emb.attend(H)  # 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.84
sinusoidal: final training loss 2.09
      rope: final training loss 1.56
     alibi: final training loss 1.96
      none: final training loss 2.25

The verdict at 4× the training length

   learned:     6.2     18.0     31.1
sinusoidal:     7.8     17.4     38.3
      rope:     5.1     10.3     37.0
     alibi:     7.1      7.1      7.1
      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.