Inside the Transformer

A Transformer first turns each token of its input into a vector. In this sentence each of the five tokens gets four numbers, drawn as a column of shaded cells.

At this point a vector only describes its own word. The vector for “chased” is the same in every sentence, no matter who did the chasing.

To mean something here, “chased” needs information from the other tokens: who chased, and what was chased.

How can one token gather useful information from the others?

Attention gives every other token a weight, then mixes their information according to those weights. Here “chased” takes most from “dog” and “ball”.

The result is a new vector for “chased” that carries its context. The rest of this essay builds the operation that produces these weights.

Here are three word vectors with two numbers each. “dog” and “puppy” point in nearly the same direction. “car” points elsewhere. Stacked as rows, they form a matrix $X$ of shape (3, 2).

Pick “dog”. To measure how related another vector is to it, take the dot product. Multiply matching entries and add them up.

Do the same for every word. X @ x computes all three dot products at once, one per row of $X$.

A dot product is large when two vectors point the same way, near zero when they are at right angles, and negative when they point apart. “puppy” scores 0.90, “car” only 0.15.

These scores are not weights yet. They could be negative, and they do not add up to anything in particular. Softmax fixes both. First it exponentiates every score, which makes it positive.

Then it divides each by their total, 7.11. The results are positive and sum to 1, and the highest score still gets the most weight.

Now use the weights. The new vector for “dog” is the weighted sum of all three vectors. In the figure the three scaled vectors are laid end to end.

The weights decide how much to take from each vector. The output is the mix.

The row of scores we computed belongs to “dog”. In the full picture, each row belongs to one token, the one that asks. Each column is a token that is looked at.

Filling in every row gives the score matrix $S = XX^{\top}$, of shape (T, T) for $T$ tokens. Entry $(i, j)$ is the dot product of token $i$ with token $j$.

This matrix is symmetric. “dog” scores “car” exactly as “car” scores “dog”.

Softmax is applied to each row separately. That is what axis=-1 says: normalize along the last axis, across the columns of a row. Each row of $A$is one token's weights and sums to 1. The columns do not.

One matrix product now does the mixing for every token. Row $i$ of $AX$ is token $i$'s weighted sum of all rows of $X$, so the output has one new vector per token, shape (T, D) with $D$ = 2 numbers per token.

Its first row is the vector for “dog” from before.

Back to the sentence. $X$ has one row per token and $D$ = 4 numbers per row, shape (T, D).

In the simple version, each token vector did three jobs at once. Attention gives each job its own vector.

Each role is a learned linear projection of $X$. The queries are $Q = XW_Q$.

The row of “chased” in $X$ is $[0.0,\ 1.0,\ 0.0,\ 0.0]$, so its query is simply the second row of $W_Q$: $[2.0,\ 0.0,\ 0.0,\ 0.0]$.

Keys and values come from two more weight matrices. All three have shape (D, D), so $Q$, $K$ and $V$ have the same shape as $X$.

These weights were picked by hand so the numbers stay readable. In a trained model they are learned, and much less tidy.

Row $i$ of $Q$, $K$ and $V$ all belong to token $i$. “chased” now has a query $[2.0,\ 0.0,\ 0.0,\ 0.0]$ that matches keys with a large first entry, which in this example are the nouns. Its key $[0.0,\ 1.0,\ 0.0,\ 0.0]$ is what other queries will find.

Because the query and key come from different matrices, what a token looks for no longer has to resemble what it is.

The steps are the same as before. Only the inputs change: scores now compare one token's query with every token's key.

Turning $K$ on its side gives $K^{\top}$, shape (D, T). The product $QK^{\top}$ holds every query–key dot product, shape (T, T). Rows are queries, columns are keys. In the figure, each score sits where its row of $Q$ meets its column of $K^{\top}$.

Unlike $XX^{\top}$, this matrix is not symmetric. “chased” scores “dog” at 2.0, but “dog” scores “chased” at 1.5.

Every score is then divided by $\sqrt{d_k}$, where $d_k$ is the length of the query and key vectors. Here $d_k = 4$, so we divide by 2.

A dot product adds up $d_k$ products. If the entries of $q$ and $k$ are independent with mean 0 and variance 1, the sum has variance $d_k$. At $d_k = 512$ scores would routinely reach the tens, softmax would put nearly all weight on one token, and its gradients would all but vanish. Dividing by $\sqrt{d_k}$ brings the variance back to 1.

Softmax along each row turns scores into weights. “chased” gives 0.32 to “dog” and to “ball”, and 0.12 to each of the others.

The weights mix the values, not the keys. $A$ becomes the left operand and multiplies $V$: (T, T) times (T, D) gives (T, D), one output per token. The figure uses the same layout as before: an output cell sits where its row of weights meets its column of $V$.

The output row of “chased” is its weighted sum of value rows.

Its last entry, 0.48, comes entirely from “dog”: 0.32 × 1.5. What a token hands over is its value, which can carry information its key never showed.

Every part of the formula now has a place in the figure.

$QK^{\top}$ compares queries with keys. $\sqrt{d_k}$keeps the scores in a range where softmax still spreads its weight. Softmax runs along each row, so every token's weights sum to 1. Multiplying by $V$ mixes the values.

In NumPy, single-head attention is six lines, with the same names as the figures. softmax is a small helper that exponentiates and divides by the sum along the given axis.

PyTorch is nearly identical. K.transpose(-2, -1) swaps the last two dimensions, and the softmax axis is called dim. dim=-1 still means along each row.

Models process a batch of $B$ sequences at once, so $X$ has shape (B, T, D). The code does not change. @ multiplies over the last two dimensions and treats the leading one as a batch, and transpose(-2, -1) leaves $B$ alone. Each sequence gets its own (T, T) weights.

As a module, the three weight matrices become nn.Linear layers without bias, and self.W_Q(X) computes X @ W_Q. (nn.Linear stores its matrix transposed and multiplies by weight.T. The result is the same.)

One attention operation gives each token one set of weights. But a token may need several kinds of information at once. “chased” wants its subject and its object. A noun may want its article.

Multi-head attention runs several attention operations side by side. Each head has its own query, key and value projections, so each can learn its own way of routing information. Every head reads the same $X$.

With our weights, head 1 links “chased” and the nouns. Head 2 links the nouns with “the”. Those are readings of a hand-built example. Some heads in trained models show patterns this clear. Many do not.

Each head produces its own output. The outputs are concatenated and multiplied by one more learned matrix, $W_O$, which mixes the heads back together.

In code, the heads are not separate modules. One projection computes the queries for all heads at once, and the result is split. Follow the shapes for a batch of $B$ sequences.

$W_Q$ is a full (D, D) matrix, so every column of $Q$ is computed from all $D$ features of a token.

view regroups the last dimension into $H$ heads of $D_{\text{head}} = D / H$ numbers each. Here $D = 4$, $H = 2$ and $D_{\text{head}} = 2$. No number changes.

The split happens after the projection. Head 1 gets the first two columns of $Q$, which is $X$ times the first two columns of $W_Q$. Each head still sees all of $X$.

transpose(1, 2) moves the head dimension in front of $T$. Each pair of batch index and head now holds an ordinary (T, D_head) matrix. Everything after this treats $B$ and $H$ alike, as batch dimensions.

Scores and weights are computed exactly as before, once per head, giving (B, H, T, T). The scale is $\sqrt{D_{\text{head}}}$, the length of each head's queries and keys.

weights @ V gives each head its own output, (B, H, T, D_head).

transpose(1, 2) moves the heads back behind $T$, and reshape lays them side by side: (B, T, H, D_head) becomes (B, T, D). That is the concatenation.

It has to be reshape rather than view. After a transpose the tensor is no longer laid out contiguously in memory, and reshape copies it when needed.

Finally $W_O$, another (D, D) matrix, mixes information across heads. The output has the shape of $X$.

Because each head works in $D / H$ dimensions, $H$ heads cost about as much as one head of full width.

Swap “dog” and “ball”. The sentence means something else, but the output for “dog” is exactly the same. Every step so far treats the tokens as a set. Scores depend on the vectors, never on where they sit.

For the same reason, the two copies of “the” get identical outputs.

The fix is to add a position vector to every token vector before the first block. It has the same width $D$, and it is added entry by entry.

Position is not a new axis. The vector is still $D$ numbers wide. Each number now mixes what the token is with where it is.

The same word at two positions now starts from two different vectors, so attention can tell them apart.

Position vectors can be fixed sine and cosine waves, as in the original Transformer, or learned like the token embeddings. The model at the end of this essay learns them. The values here are made up.

A Transformer block wraps multi-head attention with a few standard parts. Attention is the only place in the block where tokens exchange information.

A residual connection adds the block's input to the attention output, so attention only has to learn a change to each vector. LayerNorm then rescales each token's vector to mean 0 and variance 1, followed by a learned scale and shift.

The MLP is two linear layers with a nonlinearity between them, usually $4D$ wide inside. It is applied to every token separately, with the same weights. Attention moved information between tokens. The MLP transforms what each token now holds.

A second residual connection and LayerNorm finish the block. The output has shape (T, D), like the input.

This is the arrangement of the original Transformer. Many recent models normalize before each part instead, X + self.attn(self.norm1(X)), which trains more reliably in deep stacks.

A model starts with two lookup tables, one row per vocabulary token and one row per position. Their sum is $X$, shape (B, T, D). The position part has shape (T, D) and is broadcast over the batch.

$X$ passes through a Transformer block. Out comes a tensor of the same shape, every vector now informed by the others.

Because the shape never changes, blocks stack. Each block works on what the blocks before it gathered, so information can travel in several hops and be combined along the way.

A final linear layer turns every vector into one score per word in the vocabulary. A language model reads the scores at each position as a guess for the token that comes next.

A model that predicts the next token is trained on every position of a sentence at once. At “chased” it must guess “the”, so it may not see the words that follow. Plain attention lets every token see every other token.

The fix happens between the scaled scores and the softmax. Every score where a token looks at a later position, everything above the diagonal, is replaced by $-\infty$.

Because $e^{-\infty} = 0$, those positions get weight 0, and softmax spreads each row over the positions that remain. The row of “chased” now covers three tokens instead of five.

The first token can only attend to itself. The last one sees the whole sentence, so its row is the same as without the mask. Nothing else in the model changes.