Scaled Dot-Product Attention with Trainable Weights

Each weight matrix is built with shape d_in by d_out, so the rows of W must match the length of the input token, and the columns of W set how long the projected vector is.

Programming

Why One Token Needs Three Views

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 single step that produces those three views. We start from one input token, compute its projections, and trace where each shape comes from. The attention scores and context vectors that follow are handled in other articles in this module.

Tracing One Token Through Three Projections

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. Because the rule is a matrix, the same token can pass through three different rules and come out as three different vectors. The three rules are W_q for queries, W_k for keys, and W_v for values.

The example treats the second input token as the current token. Its job is to ask questions, so its query is the one computed first. The same token also produces a key and a value, which other tokens will use when they attend to it. For this lesson we compute only these three vectors for that one token, which keeps the arithmetic visible.

How does a single input vector split into three different outputs, and which weight matrix produces each one?

The three lines that produce the views share one structure: the input token on the left, a weight matrix on the right, and the result stored under a new name. Only the weight matrix changes from line to line. That is why the three outputs differ even though the input is identical. The diagram shows the same token feeding three branches, each with its own matrix and its own role.

Reading Output Shapes from d_in and d_out

Write the input token as a row vector x and each weight matrix as W with its own subscript. The projection for the query is then one product, and the key and value follow exactly the same rule with different matrices. Shape is where most projection errors show up first, so it is worth tracking carefully.

Rule

The rule is that the inner dimensions must agree and the outer dimensions set the result. The input row has d_in columns, so the weight matrix must have d_in rows to match it. The weight matrix has d_out columns, and those columns set the length of the output. Only the number of columns in each weight matrix sets the output length.

Given a weight matrix of size d_in by d_out, what shape does the projected vector have, and which dimension maps to which?

Key point

Each weight matrix is built with shape d_in by d_out, so the rows of W must match the length of the input token, and the columns of W set how long the projected vector is.

Worked Projection with Small Numbers

Worked example

Computing q, k, and v for one token

Let d_in = 3 and d_out = 2. The current token is x(2) = [1, 2, 3]. W_q = [[1, 0], [0, 1], [1, 1]], W_k = [[0, 1], [1, 0], [0, 2]], and W_v = [[2, 0], [0, 2], [0, 0]]. Compute q, k, and v and check each shape.

1
Predict the shape
The input is 1 x 3 and each weight matrix is 3 x 2, so each projected vector should be 1 x 2. Commit to this before computing.
2
Compute q
Each output element is the input dotted with one column of W_q. Column 1 is [1, 0, 1], giving 1 + 3 = 4. Column 2 is [0, 1, 1], giving 2 + 3 = 5. So q = [4, 5].
3
Compute k
Column 1 of W_k is [0, 1, 0], giving 2. Column 2 is [1, 0, 2], giving 1 + 6 = 7. So k = [2, 7].
4
Compute v
Column 1 of W_v is [2, 0, 0], giving 2. Column 2 is [0, 2, 0], giving 4. So v = [2, 4].
5
Verify the shapes
Each result has 2 entries, matching d_out = 2. The three vectors differ because the three matrices differ, even though the input token is the same.

Result

q = [4, 5], k = [2, 7], v = [2, 4]. All three have shape 1 x 2.

Where Projection Errors Show Up First

When a projection result is not what you expected, check three shapes in order: the input, the weight matrix, and the output. The diagram lays out that check as a sequence of comparisons.

At which step do I check the input shape, the weight matrix shape, and the output shape, and what do I compare at each step?

Common mistake
Treating the three projections as one projection
What it looks like
Computing only x @ W_q and reusing that result as the key and value.
Why it fails
Each projection uses its own matrix, so each view is different.
Do this instead
Check which matrix each line uses and make sure W_q, W_k, and W_v each appear once.
Common mistake
Setting the output length with the rows of the weight matrix
What it looks like
Expecting a 1 x 3 output from a 3 x 2 weight matrix because the input has 3 numbers.
Why it fails
The columns of the weight matrix set the output length, which here is d_out = 2.
Do this instead
Read d_out from the columns of W and expect a 1 x d_out result.
Common mistake
Multiplying a row by a matrix whose rows do not match
What it looks like
A 1 x 3 input times a 2 x 2 matrix.
Why it fails
The input has 3 columns but the matrix has only 2 rows, so the inner dimensions do not agree.
Do this instead
Make the weight matrix d_in rows tall before multiplying.
Common mistake
Judging a projection by its values alone instead of its shape

Predict Before You Compute

EASY

A token x(2) has shape 1 x 3 and the query matrix W_q has shape 3 x 4. What shape is the query vector, and which dimension of W_q decides it?

Hints

The output length comes from the columns of the weight matrix.
Here d_out = 4, so the query should have 4 entries.

Your attempt (not saved)

MEDIUM

Using the worked example, change W_v so its second column is [0, 1, 0] instead of [0, 2, 0]. Recompute v before checking the answer, then state which element of v changed and why.