Sinusoidal Position Encoding

Add information about order

In unmasked self-attention with no position-dependent terms, reordering the input rows reorders the output rows in the same way. This is called permutation equivariance. It does not mean the output arrays are identical: they move with the input. A causal mask already introduces an ordering constraint, so the unmasked statement does not apply to every attention layer.

A position encoding gives each sequence position a vector. Adding it to token embeddings lets the attention layer use both token content and position. Sinusoidal encodings compute these vectors from a fixed formula, with no learned parameters or random initialization.

This exercise returns the encoding matrix alone. You can add it to token features later. The formula can be evaluated beyond a training sequence's length, but that does not guarantee good predictions on longer sequences.

Alternate sine and cosine features

Let T be the number of positions and D a positive, even feature count. Pair index i runs from 0 through D // 2 - 1. Define:

angle(pos, i) = pos / 10000^(2i / D)
PE[pos, 2i]   = sin(angle(pos, i))
PE[pos, 2i+1] = cos(angle(pos, i))

Here ^ means exponentiation in the formula; Python uses **. Angles are in radians. Both features in a pair use the same angle. Pair zero changes fastest with position. Later pairs use larger divisors and change more slowly. Keep sine and cosine interleaved; putting all sine features before all cosine features would produce a different layout.

For D = 4, the two divisors are 1 and 100. Row zero is [0, 1, 0, 1]. Row one is [sin(1), cos(1), sin(0.01), cos(0.01)].

Build all positions with broadcasting

import jax.numpy as jnp

T, D = 4, 4
pos = jnp.arange(T, dtype=jnp.float32)[:, None]       # (4, 1)
i = jnp.arange(D // 2, dtype=jnp.float32)[None, :]    # (1, 2)
divisors = 10000.0 ** (2.0 * i / D)                 # (1, 2)
angles = pos / divisors                            # (4, 2)
pe = jnp.zeros((T, D), dtype=jnp.float32)
pe = pe.at[:, 0::2].set(jnp.sin(angles))
pe = pe.at[:, 1::2].set(jnp.cos(angles))

assert jnp.array_equal(pe[0], jnp.array([0., 1., 0., 1.]))
assert jnp.allclose(divisors, jnp.array([[1., 100.]]))
assert pe.shape == (4, 4)

token_features = jnp.zeros((2, T, D))
combined = token_features + pe[None, :, :]
assert combined.shape == (2, T, D)

The (T, 1) positions and (1, D // 2) divisors broadcast to one angle for every position and frequency. The final addition shares the same position matrix across both sequences in the batch.

What the pairs tell us

For one frequency w, write a pair as [sin(w*p), cos(w*p)]. Increasing position p by k transforms the pair by:

sin(w*(p+k)) = cos(w*k)*sin(w*p) + sin(w*k)*cos(w*p)
cos(w*(p+k)) = -sin(w*k)*sin(w*p) + cos(w*k)*cos(w*p)

For a fixed offset k, these coefficients do not depend on p. Also, the dot product of two unmodified pairs at positions p and q is cos(w*(p-q)), which depends on their relative distance. These are properties of the encoding itself. Adding token features and applying learned query/key projections does not guarantee attention scores that depend only on relative distance.

Each component stays between -1 and 1. However, each pair has squared norm 1, so the full row has norm sqrt(D / 2). Bounded components alone do not guarantee an appropriate scale relative to the token embeddings.