Mixing values — Linear algebra for LLMs

Weights times values

=

$O = AV$ : $(T \times T)(T \times d_v) = T \times d_v$ . One output row per token.

Write

O = AV

, shape

T \times d_v

.

One output row is a weighted average

\displaystyle o_{\text{cat}} = 0.056\,v_{\text{the}} + 0.944\,v_{\text{cat}} + 0\,v_{\text{sat}}

\displaystyle = 0.056\,[1, 1] + 0.944\,[0, 1] = [0.056, 1]

Row view of a matrix product (Part 2): row $i$ of $AV$ mixes the rows of $V$ with the weights in row $i$ of $A$ . The weights are positive and sum to 1, so the output is a weighted average of value vectors, never outside their range.

Write

o_i = \sum_j A_{ij}\, v_j

with

A_{ij} \ge 0

,

\sum_j A_{ij} = 1

.

✏️ Your turn

Mix the last row

A_{\text{sat}} \approx [0.14, 0.58, 0.28]

Write $o_{\text{sat}}$ .

Work it out on paper, then go to the next card to check.

Answer

\displaystyle 0.14\,[1,1] + 0.58\,[0,1] + 0.28\,[1,0]

\displaystyle \approx [0.42, 0.72]

With the unrounded weights, row 3 of $O$ is $[0.42, 0.72]$ .

All of attention in one line

\displaystyle \operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V

Project, score, scale, mask, softmax each row, mix the values. You have now computed every piece of it by hand.

Write

\operatorname{softmax}(QK^\top/\sqrt{d_k} + M)\,V

The shapes, start to finish

Step	Shape
$X$	$T \times d$
$Q, K, V$	$T \times d_k$
scores $QK^\top$ , weights $A$	$T \times T$
output $AV$	$T \times d_v$

Write

T \times d \to T \times d_k \to T \times T \to T \times d_v

What it costs

\displaystyle QK^\top\!: \; 2\,T^2 d_k \qquad AV\!: \; 2\,T^2 d_v

Both grow with $T^2$ . The projections cost $2Td\,d_k$ each, linear in $T$ . For short contexts the projections dominate; for long contexts the $T^2$ terms take over.

out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

📌 Recap

Mixing values in one card

$O = AV$ : each output row is a weighted average of value rows.
$\operatorname{Attention} = \operatorname{softmax}(QK^\top/\sqrt{d_k} + M)\,V$ .
Example: $O \approx [1\ 1;\ 0.06\ 1;\ 0.42\ 0.72]$ .
Cost $\approx 4T^2 d_k$ per head for scores and mixing.

Write

The four lines above.