Embedding Input Tokens for Attention

Self-attention lets each token in a sequence look at the other tokens and decide how much they matter. To make that possible, the model does not use the raw embedding of a token directly. Instead, it creates three new views of the same token: a query, a key, and a value. Each view is produced by multiplying the embedded token by its own trainable weight matrix.

This article focuses on the step that produces those three views. We start from a single input token, run the projection, and trace where each shape comes from. The attention scores and context vectors that follow are handled in other articles in this module.

Think of the embedded token as a row of numbers. A projection is a fixed rule, stored as a matrix, that converts that row into a different row. Write the input token as a row vector $x$. In this lesson we pay attention to the second input element, $x^{(2)}$, because it is the current token. Its role is to ask questions, so its query is the first one we compute.

The shape of $x^{(2)}$ is $(1 \times d_{in})$. It has one row and as many columns as the token has numbers. Shape is the first thing to check, because the rest of the projection depends on it.

The diagram shows the same token feeding three branches. Each branch multiplies the same input by a different matrix. Because the input is identical and only the matrix changes, the three outputs take on separate roles even though they start from one token.

The same pattern applies to all three vectors. The query is a single product of the input row and the query weight matrix: $q^{(2)} = x^{(2)} W_q$. The key and value follow exactly the same rule with different matrices: $k^{(2)} = x^{(2)} W_k$ and $v^{(2)} = x^{(2)} W_v$.

Each weight matrix has shape $(d_{in} \times d_{out})$. Its rows must match the input length $d_{in}$, because each input number multiplies one row of the matrix. Its columns set the output length $d_{out}$. Putting these together, each projected vector has shape $(1 \times d_{out})$.

The code below picks out the second input token, records its shape, and applies three weight matrices. The values are small hand-picked numbers chosen so the arithmetic can be checked by hand. In the source setup the matrices are filled with random starting values after a fixed random seed, so the same numbers appear on every run. Those values are parameters that training would adjust.

The three vectors come from one token, but each has a different job in attention. The query is the view used when a token asks questions about the other tokens. The key is the view other tokens use when they check how relevant this token is. The value is the view that carries content forward once relevance is decided. The same token produces all three, and other tokens use them when they attend to it.

Notice that the three vectors have the same shape. The shape comes only from d_out, so the roles are separated by the matrices, not by the sizes of the outputs.

Most projection errors show up first in shape, not in values. When a result is not what you expected, follow the chain from the input shape to the weight matrix shape and check each step.