Token embeddings and positional encoding

Step 1 Understand it

The problem it solves. Attention compares every token with every other token, but it treats the input as an unordered set: shuffle the words and each output just moves with its word. "dog bites man" and "man bites dog" would look the same. Position information fixes that, and how you encode positions decides how well a model handles long contexts.

Watch it happen

The idea

Each token id becomes a vector (its embedding). Attention on its own does not know word order, so we add position information: a learned position vector, sine waves, or rotations (RoPE).

How it works, step by step

1. Look up each token id in the embedding table: a vector of size d_model (768 in GPT-2 small).
2. Make a position vector for each position 0, 1, 2, … of the same size.
3. Learned positions: a second table with one row per position, trained like any weight (GPT-2).
4. Sinusoidal positions: fixed sine and cosine waves of different frequencies, so every position gets a unique pattern (the original transformer).
5. Add the two: x = token_embedding + position_embedding. The sum goes into the first transformer block.
6. RoPE (rotary embeddings, used by Llama and many others) instead rotates each query and key by an angle that depends on its position, so attention scores depend on the relative distance between tokens.

Picture it

flowchart LR
  T["token ids<br/>'dog', 'bites', 'man'"] --> TE["token embeddings"]
  P["positions<br/>0, 1, 2"] --> PE["position embeddings"]
  TE --> ADD(("+"))
  PE --> ADD
  ADD --> X["input to the transformer"]

Analogy

Seat numbers in a theatre. Everyone is a person (the embedding), but the seat number (the position) tells you who sits next to whom.

Where the analogy breaks

Seats are fixed. A model trained on 1,000 positions often struggles at 10,000, so long-context models need special tricks.

Worked example

'dog bites man' and 'man bites dog' use the same tokens but mean different things. Without positions, the model sees the same bag of words.

x = tok_emb[ids] + pos_emb[np.arange(len(ids))]   # shape (T, d)

Must-know points

• Without positions, attention is blind to order.
• Learned positions can't go beyond the longest position seen in training.
• Sinusoidal and rotary schemes generalise further, and can be stretched for longer contexts.
• RoPE encodes relative position directly in the query·key dot product.
• Embeddings and positions are both d_model wide, so they can simply be added.
• Many models tie the input embedding table to the output layer (weight tying) to save parameters.

Common mistakes

• Forgetting positions entirely in a hand-built transformer.
• Running a model past its trained context length and expecting good results.
• Adding positions with the wrong shape, so they broadcast across the batch instead of the sequence.

Step 2 Go further (optional)

Optional. The lesson above is written to be enough on its own. If you want a second explanation, start with the one marked Worth it.

• Read

  The Illustrated Transformer (positional encoding section)

  Jay Alammar

• Read

  Rotary Embeddings: A Relative Revolution

  EleutherAI blog

• Code

  Build a Large Language Model (From Scratch), Ch 2 (token and positional embeddings)

  Sebastian Raschka · code free

Step 3 Practise

Try each exercise yourself first: on paper, out loud or in code. Then open the answer and compare it with yours.

1. ExplainSay it in plain words, with no notes.

   Explain positional encoding with the theatre-seats picture.

   Show answer

   In a theatre, every seat holds a person, and each person is who they are wherever they sit: that is the token embedding. But the seat number tells you who is next to whom, who is at the front, who is at the back: that is the position.

   Attention on its own sees the audience as a crowd with no seat numbers, so "dog bites man" and "man bites dog" look like the same group of people. Positional encoding pins a seat number onto everyone before the show starts, by adding a position vector to each token's vector, so the model can tell the order.

2. DeriveWrite the maths, or the steps, from memory.

   Show that self-attention without positions gives the same outputs, just reordered, if you shuffle the input tokens. Why does that mean we need positions?

   Show answer

   Self-attention for inputs X (one token per row) is Attn(X) = softmax(Q·Kᵀ/√d)·V with Q = XW_Q, K = XW_K, V = XW_V.

   Shuffle the tokens: X becomes P·X, where P is a permutation matrix (it reorders rows; PᵀP = I).
   Then Q, K and V become PQ, PK and PV, and the scores become (PQ)(PK)ᵀ = P·(QKᵀ)·Pᵀ: the same scores, with rows and columns reordered the same way.
   Softmax works row by row, so softmax(P·S·Pᵀ) = P·softmax(S)·Pᵀ.
   The output: P·softmax(S)·Pᵀ·P·V = P·softmax(S)·V = P·Attn(X).

   So shuffling the input just shuffles the output in the same way: every token gets exactly the same vector as before, wherever it sits. The layer computes a function of the set of tokens, not their order, so "dog bites man" and "man bites dog" would look identical. The MLP works on each token separately, so it can't add order either. That is why positions must be added to the input, as position embeddings or rotary encodings.

3. BuildCode it from scratch in NumPy, with no ML library.

   Implement sinusoidal position encodings and plot them as a heatmap. Then add learned position embeddings to your model.

   Show answer

   Sinusoidal encoding gives position p and dimension pair i the values sin(p / 10000^(2i/d)) and cos(p / 10000^(2i/d)). Low dimensions wave quickly, high ones slowly, so every position gets a unique pattern, like the hands of a clock running at different speeds, and nearby positions get similar patterns. The heatmap shows fast stripes on the left fading into slow ones on the right.

   Learned position embeddings are simpler: a table with one trainable vector per position, added to the token embedding. The shared gpt.py already does this: self.pos = nn.Embedding(block_size, n_embd), and the forward pass adds self.pos(torch.arange(T)) to the token embeddings. The model learns whatever position pattern helps, but can't handle positions beyond block_size.

   pythonCopy

   import numpy as np
   import matplotlib.pyplot as plt
   
   def sinusoidal(n_positions, d):
       pos = np.arange(n_positions)[:, None]
       i = np.arange(0, d, 2)[None, :]
       angles = pos / 10000 ** (i / d)
       pe = np.zeros((n_positions, d))
       pe[:, 0::2], pe[:, 1::2] = np.sin(angles), np.cos(angles)
       return pe
   
   pe = sinusoidal(100, 64)
   plt.imshow(pe, aspect="auto", cmap="RdBu"); plt.colorbar()
   plt.xlabel("dimension"); plt.ylabel("position"); plt.title("sinusoidal position encoding")
   plt.show()
   cos = lambda a, b: a @ b / (np.linalg.norm(a) * np.linalg.norm(b))
   print("similarity of position 10 with 11:", round(cos(pe[10], pe[11]), 3), "| with 60:", round(cos(pe[10], pe[60]), 3))

4. BreakShow when it fails, and why.

   Train your small GPT with no position information at all. Compare the generated text with and without it.

   Show answer

   Training the same GPT with no position information gives a clearly worse validation loss than with learned positions, and the text gets worse: words fall apart into letter salad more often, and line structure (names followed by a colon, lines of similar length) is weaker.

   It is not a total failure, and that is interesting. A model with a causal mask can partly work out position by itself: the first token can only attend to itself, the second to two tokens, and so on, so the "amount of context" leaks some order information. But it is a weak, indirect signal, and an encoder with no mask would get none at all. That is why every practical transformer adds positions.

   pythonCopy

   import torch
   from gpt import GPT, Config, train, estimate_loss, decode, encode
   
   for positions in [True, False]:
       torch.manual_seed(0)
       model = train(GPT(Config(positions=positions)), steps=1500, log_every=0)
       torch.manual_seed(1)
       sample = decode(model.generate(torch.tensor([encode("\n")]), 200)[0].tolist())
       print(f"positions={positions}: val loss {estimate_loss(model):.3f}\n{sample}\n")

Step 4 Prove it

Check you understood

Tick each check when you can do it without notes