Skip to content

Embeddings

The path from text to the next token in seven steps: tokens, embeddings, query/key/value, mask and heads, the block, the output and generation. The current step is Embeddings. ATTENTION TRANSFORMERS Tokenstext → ids Embeddingsids → X ▲ you are here Q · K · Vone head Masktwo heads BlockLN · MLP Outputlogits → p Generationpick text → → next token

Step 2 of 7. In: the ids [0, 2, 4, 5, 1] of step 1. Out: the matrix \(X\), one row of \(d = 4\) numbers per token — the input of every step that follows. ← Tokens · next step: Query, key, value →

An id is a name. Id 4 (␣climb) is not twice id 2 (␣cat), and it is not "closer" to id 5 than to id 11; feeding the raw integer to a network would assert an order and a distance that are not there (the data chapter makes the same point about any categorical column). What the network needs is a vector per token — a list of numbers with room for many properties at once — and it learns those vectors from data.

From tokens to vectors: the embedding layer

Before there is anything to compare, the text has to become vectors. It reaches the model as integer indices, and the first layer is a table \(E \in \mathbb{R}^{V \times d}\) with one row per vocabulary token. The forward pass is a lookup:

\[ y_t = E[i_t] \qquad\text{which is the same as}\qquad y_t = \text{one-hot}(i_t)^\top E \]

\(E\) is learned like any other weight matrix. GPT-2 small has \(V = 50{,}257\) rows of \(d = 768\) numbers: \(38.6\) million parameters, 31% of the whole model.

The lookup, on the running example

The toy's table has \(V = 12\) rows of \(d = 4\) numbers. The rows of the five tokens of "the cat climbed the" are marked:

id token \(E[\text{id}]\)
0 the [−0.5, −0.5, 0, 1]
1 ␣the [−0.5, 0, −0.5, 1]
2 ␣cat [1, −0.5, 0, −0.5]
3 ␣dog [1, −0.5, −0.5, 0]
4 ␣climb [−0.5, 1, 0, −0.5]
5 ed [0, 0.5, −0.5, 0]
6 ␣roof [0, −0.5, 1, −0.5]
7 ␣wall [1, 0, −1, 0]
8 ␣sofa [0.5, −0.5, −0.5, 0.5]
9 ␣floor [0.5, 0, −0.5, 0]
10 . [0, 0, 1, −1]
11 <|endoftext|> [0.5, −0.5, 0.5, −0.5]

The ids [0, 2, 4, 5, 1] pick rows 0, 2, 4, 5 and 1, in that order. No arithmetic happens; the same id gets the same row wherever it appears and whatever its neighbours are — which is exactly the limitation attention will fix. (The toy's numbers were picked by hand. A trained table has no readable columns: what a dimension means is spread over many of them.)

Where position comes from

Nothing in the attention of the next two pages uses the order of the tokens: attention is permutation-equivariant: shuffle the tokens and the outputs shuffle with them, with identical values. On its own it cannot tell "dog bites man" from "man bites dog". Position has to be injected, and how it is injected turned out to matter a great deal for long context.

Scheme Mechanism Verdict
Sinusoidal1 Fixed sin/cos vectors added to the input embeddings Historical. Extrapolates poorly.
Learned absolute A trainable vector per position Historical (BERT, GPT-2). Cannot exceed the trained length at all.
ALiBi3 A linear distance penalty added to the scores Extrapolates well; largely superseded
RoPE2 Rotate the Q and K vectors by an angle proportional to position Universal. Llama, Qwen, Mistral, Gemma, DeepSeek
NoPE / hybrid No positional signal in some layers; causality alone carries order Increasingly common in long-context and hybrid stacks

RoPE is worth understanding rather than memorizing. Picture a clock: each token's hand is turned a little further for every position. The angle between two hands does not depend on what time it is, only on how many positions separate them. In symbols: take Q and K in pairs of dimensions, treat each pair as a complex number, and multiply by \(e^{i\,m\theta}\) where \(m\) is the position. Because rotations compose by adding angles, the dot product between position \(m\) and position \(n\) depends only on \(m - n\):

\[ \langle R_m q,\; R_n k \rangle \;=\; \langle q,\; R_{n-m} k \rangle \]

Absolute rotations in, relative position out — with no extra parameters and no extra tensor to add. Each pair of dimensions turns at a different speed \(\theta\), like the hour, minute and second hands: the fast ones tell close neighbours apart, the slow ones tell long distances apart. That property is also why context windows can be extended after training: rescale the frequencies \(\theta\) and the same weights address a longer sequence (position interpolation, NTK scaling, YaRN4). Every "we extended it to 1M tokens" announcement is, mechanically, this plus fine-tuning.

Position in the running example

GPT-2 uses the second row of that table, learned absolute positions: a second table \(P\), one row per position (1024 × 768 in GPT-2), added to the token's row. The toy does the same with 8 positions:

\[ x_t = E[i_t] + P_t \]
position \(t\) token (id) \(E[i_t]\) \(P_t\) \(x_t\)
1 the (0) [−0.5, −0.5, 0, 1] [0, 0, 0, 0] [−0.5, −0.5, 0, 1]
2 ␣cat (2) [1, −0.5, 0, −0.5] [0.5, 0, −0.5, 0] [1.5, −0.5, −0.5, −0.5]
3 ␣climb (4) [−0.5, 1, 0, −0.5] [0, 0.5, 0, −0.5] [−0.5, 1.5, 0, −1]
4 ed (5) [0, 0.5, −0.5, 0] [−0.5, 0, 0.5, 0] [−0.5, 0.5, 0, 0]
5 ␣the (1) [−0.5, 0, −0.5, 1] [0, −0.5, 0, 0.5] [−0.5, −0.5, −0.5, 1.5]

Stacked, the five rows are the matrix every following page starts from:

\[ X = \begin{bmatrix} -0.5 & -0.5 & 0 & 1 \\ 1.5 & -0.5 & -0.5 & -0.5 \\ -0.5 & 1.5 & 0 & -1 \\ -0.5 & 0.5 & 0 & 0 \\ -0.5 & -0.5 & -0.5 & 1.5 \end{bmatrix} \qquad (n \times d = 5 \times 4) \]

(Every row of the toy's \(E\) and \(P\) adds up to zero, so every row of \(X\) has mean 0; that only keeps the next page's arithmetic short.) The other schemes, for comparison:

  • Sinusoidal (2017): no table, the position is computed. For \(d = 4\) it is \([\sin t,\ \cos t,\ \sin(t/100),\ \cos(t/100)]\), with \(t\) counted from 0: position 0 is \([0, 1, 0, 1]\), position 1 is \([0.841, 0.540, 0.010, 1.000]\), position 2 is \([0.909, -0.416, 0.020, 1.000]\). The first pair turns fast and tells neighbours apart; the second turns slowly and tells long distances apart.
  • RoPE: nothing is added here. Inside attention, each pair of numbers of \(q\) and \(k\) is rotated by an angle proportional to the position. With one pair and \(\theta = 0.5\), the same \(q = k = [1, 0]\) at positions \((3, 1)\) and at \((10, 8)\) give the same score, \(\cos(2 \times 0.5) = 0.540\); at \((4, 1)\) the score is \(\cos(1.5) = 0.071\). Only the distance counts.

The panel builds \(X\) step by step. Switch the position scheme, or swap ␣cat and ␣climb: without a position, the swap only swaps two rows of \(X\).

The way back: a sparse gradient

The two forms of the lookup give the same vector, but the second explains the backward pass. Multiplying by a vector that is zero everywhere except at one position produces a gradient that is likewise zero everywhere except in that row, with repetitions summed:

\[ \frac{\partial L}{\partial E[i]} \;=\; \sum_{t \,:\, i_t = i} \frac{\partial L}{\partial y_t} \]

Three consequences:

  • The update is sparse. In a step over a few thousand tokens, most rows of \(E\) receive exactly zero gradient. That is why frameworks keep a separate path for embeddings (sparse=True, IndexSelect) rather than treating \(E\) as just another dense matrix.
  • A rare token learns slowly. The panel draws tokens from a Zipf law, as real language does: the first few appear several times per sequence, the long tail appears almost never, and its rows barely move.
  • The table is enormous, and often reused. At 128k rows and 4096 columns that is 525 million parameters in the input alone. Tying that matrix to the output layer, which projects back to the vocabulary, stores one such table instead of two — and it works because both speak about the same space.
Check yourself — swap “␣cat” and “␣climb”. What happens to X, with and without positions?

Without positions, \(X\) keeps the same five rows in a different order: rows 2 and 3 trade places and nothing else changes. Attention without a mask would then produce the same outputs, swapped, and could not tell "cat climb" from "climb cat". With learned positions the rows themselves change: ␣cat at position 3 is \(E[2] + P_3 = [1, -0.5, 0, -0.5] + [0, 0.5, 0, -0.5] = [1, 0, 0, -1]\), not the \([1.5, -0.5, -0.5, -0.5]\) it had at position 2.

Check yourself — in a batch of 4,000 tokens, which rows of E receive a gradient?

Only the rows of the ids that appear in the batch — at most 4,000 of GPT-2's 50,257, usually far fewer, because frequent tokens repeat. Every other row gets exactly zero, which is why a rare token's vector learns slowly.

← Tokens · next step: Query, key, value →

  1. Vaswani, A., et al. (2017). Attention Is All You Need — NeurIPS. Scaled dot-product attention, multi-head, sinusoidal positions. ↩

  2. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. Published in Neurocomputing 568 (2024). RoPE. ↩

  3. Press, O., Smith, N., & Lewis, M. (2022). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation — ICLR. ↩

  4. Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2024). YaRN: Efficient Context Window Extension of Large Language Models — ICLR. RoPE frequency rescaling, done carefully. ↩