Matrix multiplication and map composition

Section 1.4 Matrix multiplication and map composition

Subsection Matrix product formulas

Definition 1.4.1. Matrix-matrix product.

Let \(A\) be an \(m\times n\) matrix, let \(B\) be an \(n\times k\) matrix, and write

\begin{equation*} B=\begin{bmatrix}\mathbf{b}_1 \amp \mathbf{b}_2 \amp \cdots \amp \mathbf{b}_k\end{bmatrix} \end{equation*}

where \(\mathbf{b}_j\) is column \(j\) of \(B\text{.}\) The product \(AB\) is the \(m\times k\) matrix

\begin{equation*} AB=\begin{bmatrix}A\mathbf{b}_1 \amp A\mathbf{b}_2 \amp \cdots \amp A\mathbf{b}_k\end{bmatrix}. \end{equation*}

🔗

Activity 1.4.2. Computing a matrix product using columns (U1-LO6).

Let

\begin{equation*} A=\begin{bmatrix}1 \amp -3 \amp 5\\3 \amp 1 \amp 4\end{bmatrix}, \qquad B=\begin{bmatrix}5 \amp 1\\-2 \amp -6\\-1 \amp 0\end{bmatrix}. \end{equation*}

Compute \(AB\text{.}\)

🔗

Solution.

The columns of \(B\) are \(\mathbf{b}_1=(5,-2,-1)\) and \(\mathbf{b}_2=(1,-6,0)\text{.}\) Using the row dot-product rule in Theorem 1.3.3,

\begin{align*} A\mathbf{b}_1 \amp=\begin{bmatrix}1(5)+(-3)(-2)+5(-1)\\3(5)+1(-2)+4(-1)\end{bmatrix}=\begin{bmatrix}6\\9\end{bmatrix},\\ A\mathbf{b}_2 \amp=\begin{bmatrix}1(1)+(-3)(-6)+5(0)\\3(1)+1(-6)+4(0)\end{bmatrix}=\begin{bmatrix}19\\-3\end{bmatrix}. \end{align*}

Placing these columns side by side gives

\begin{equation*} AB=\begin{bmatrix}A\mathbf{b}_1 \amp A\mathbf{b}_2\end{bmatrix} =\begin{bmatrix}6 \amp 19\\9 \amp -3\end{bmatrix}. \end{equation*}

🔗

Theorem 1.4.3.

If \(A\) is \(m\times n\) and \(B\) is \(n\times k\text{,}\) then the \((i,j)\)-entry of \(AB\) is the dot product of row \(i\) of \(A\) with column \(j\) of \(B\text{:}\)

\begin{align*} (AB)_{ij}\amp=\sum_{\ell=1}^{n}a_{i\ell}b_{\ell j}\\ \amp=a_{i1}b_{1j}+a_{i2}b_{2j}+\cdots+a_{in}b_{nj}. \end{align*}

🔗

Why is this true?.

By the definition of matrix multiplication, column \(j\) of \(AB\) is \(A\mathbf{b}_j\text{,}\) where \(\mathbf{b}_j=(b_{1j},\ldots,b_{nj})\) is column \(j\) of \(B\text{.}\) The row dot-product rule for matrix-vector multiplication in Theorem 1.3.3 therefore gives

\begin{align*} (AB)_{ij} \amp=(A\mathbf{b}_j)_i\\ \amp=\operatorname{row}_i(A)\cdot\mathbf{b}_j\\ \amp=\sum_{\ell=1}^{n}a_{i\ell}b_{\ell j}. \end{align*}

🔗

Note 1.4.4. Row-column interpretation.

The sum above is the dot product of row \(i\) of \(A\) and column \(j\) of \(B\text{.}\) Viewing both as \(n\)-vectors,

\begin{equation*} (AB)_{ij}=(a_{i1},\ldots,a_{in})\cdot(b_{1j},\ldots,b_{nj}). \end{equation*}

Pair entries in corresponding positions, multiply, and add.

🔗

Activity 1.4.5. Computing a matrix entry (U1-LO6).

\begin{equation*} A=\begin{bmatrix}-3 \amp -4 \amp 1\\2 \amp 4 \amp 0\\1 \amp -4 \amp -5\end{bmatrix}, \qquad B=\begin{bmatrix}-1 \amp 2 \amp -4\\-4 \amp -3 \amp -1\\4 \amp 3 \amp 1\end{bmatrix}, \end{equation*}

what is the \((2,3)\)-entry of \(AB\text{?}\)

🔗

Solution.

Use row \(2\) of \(A\) and column \(3\) of \(B\text{:}\)

\begin{align*} (2,4,0)\cdot(-4,-1,1) \amp= 2(-4)+4(-1)+0(1)\\ \amp= -12. \end{align*}

🔗

Remark 1.4.6. Matrix-vector products as matrix products.

A vector \(\mathbf{x}\in\mathbb R^n\) can be viewed as an \(n\times 1\) matrix. With this convention, if \(A\) is an \(m\times n\) matrix, the matrix-vector product \(A\mathbf{x}\) is a special case of matrix multiplication: its \(m\times 1\) result represents a vector in \(\mathbb R^m\text{,}\) with entries given by the same row dot products.

🔗

Theorem 1.4.7. Properties of matrix multiplication.

Assume that \(A,B,C\) are matrices of appropriate sizes and \(r\) is a scalar. Then

\(A(BC)=(AB)C\text{.}\)
🔗

🔗
\((A+B)C=AC+BC\text{.}\)
🔗

🔗
\(C(A+B)=CA+CB\text{.}\)
🔗

🔗
\(A(rB)=r(AB)=(rA)B\text{.}\)
🔗

🔗
\((AB)^T=B^TA^T\text{.}\)
🔗

🔗

🔗

Why is this true?.

We prove property 5. Let \(A\) be \(m\times n\) and \(B\) be \(n\times k\text{.}\) Then \(A^T\) is \(n\times m\) and \(B^T\) is \(k\times n\text{.}\) Both \((AB)^T\) and \(B^TA^T\) have shape \(k\times m\text{.}\) Transposition exchanges row and column indices, so the entry formula in Theorem 1.4.3 gives

\begin{align*} ((AB)^T)_{ij} \amp=(AB)_{ji}\\ \amp=\sum_{\ell=1}^{n}a_{j\ell}b_{\ell i}\\ \amp=\sum_{\ell=1}^{n}b_{\ell i}a_{j\ell}\\ \amp=\sum_{\ell=1}^{n}b_{i\ell}^T a_{\ell j}^T\\ \amp=\sum_{\ell=1}^{n}(B^T)_{i\ell}(A^T)_{\ell j}\\ \amp=(B^TA^T)_{ij}. \end{align*}

Here \(b_{i\ell}^T\) and \(a_{\ell j}^T\) denote the indicated entries of \(B^T\) and \(A^T\text{,}\) as in Definition 1.2.13. We can reverse the two scalar factors in each summand because real-number multiplication is commutative. Every corresponding entry agrees, so \((AB)^T=B^TA^T\text{.}\)

🔗

Note 1.4.8. Shape habit for matrix products.

Before forming \(AB\text{,}\) check the inner dimensions. If \(A\) is \(m\times n\) and \(B\) is \(n\times k\text{,}\) then the shared inner dimension \(n\) makes the product possible, and the outer dimensions give \(AB\) shape \(m\times k\text{:}\) inner dimensions match, outer dimensions survive.

🔗

Practice shape habit in MatriX Match.

MatriX Match

🔗

Activity 1.4.9. Applying one matrix to many points (U1-LO5, U1-LO8).

Put the vertices of a shape in the columns of a matrix \(X\text{.}\) Then \(Y=AX\) contains the transformed vertices.

🔗

import numpy as np

# Columns are vertices of the unit square.
# The final column repeats the first vertex to close the shape.
X = np.array([
    [0., 1., 1., 0., 0.],
    [0., 0., 1., 1., 0.],
])

A = np.array([
    [1., 1.],
    [0., 1.],
])

Y = A @ X
Y

Output:

🔗

array([[0., 1., 2., 1., 0.],
       [0., 0., 1., 1., 0.]])

What are the columns of Y?
🔗

🔗
Which transformation from the gallery is this?
🔗

🔗
Why does the bottom edge stay fixed?
🔗

🔗

Solution.

The columns of X are the vertices of the unit square:

\begin{equation*} \begin{bmatrix}0\\0\end{bmatrix},\quad \begin{bmatrix}1\\0\end{bmatrix},\quad \begin{bmatrix}1\\1\end{bmatrix},\quad \begin{bmatrix}0\\1\end{bmatrix},\quad \begin{bmatrix}0\\0\end{bmatrix}. \end{equation*}

The final column repeats the first vertex to close the shape. The columns of Y are their images under the map \(\mathbf{x}\mapsto A\mathbf{x}\text{:}\)

\begin{equation*} \begin{bmatrix}0\\0\end{bmatrix},\quad \begin{bmatrix}1\\0\end{bmatrix},\quad \begin{bmatrix}2\\1\end{bmatrix},\quad \begin{bmatrix}1\\1\end{bmatrix},\quad \begin{bmatrix}0\\0\end{bmatrix}. \end{equation*}

This matrix is a horizontal shear. It sends \((x,y)\) to \((x+y,y)\text{.}\) Points with \(y=0\) stay fixed.

🔗

Warning 1.4.10.

In this geometric example, the columns of \(X\) are points, so \(Y=AX\) applies the same matrix to every point. In many data tables, observations are stored as rows. For example, a document matrix may have one document per row. Always check what rows and columns represent before interpreting a product.

🔗

Subsection A missing rule: \(AB\) need not equal \(BA\)

Multiplication of real numbers is commutative: \(ab=ba\text{.}\) Matrix multiplication is not commutative in general. For two square matrices \(A\) and \(B\) of the same size, we say that \(A\) and \(B\) commute when

\begin{equation*} AB=BA. \end{equation*}

Some pairs commute and some do not.

🔗

Example 1.4.11. A product defined in only one order.

Let

\begin{equation*} A=\begin{bmatrix}1\amp2\\0\amp1\end{bmatrix}, \qquad B=\begin{bmatrix}1\amp0\amp2\\0\amp1\amp3\end{bmatrix}. \end{equation*}

The matrix \(A\) is \(2\times2\) and \(B\) is \(2\times3\text{.}\) The product \(AB\) is defined because the two columns of \(A\) match the two rows of \(B\text{,}\) and

\begin{equation*} AB=\begin{bmatrix}1\amp2\amp8\\0\amp1\amp3\end{bmatrix}. \end{equation*}

🔗

In the reverse order, \(B\) has three columns but \(A\) has only two rows, so \(BA\) is not defined. Thus knowing that \(AB\) exists does not guarantee that \(BA\) exists, let alone that the two products are equal.

🔗

Activity 1.4.12. Do these matrices commute? (U1-LO6).

Let

\begin{equation*} A=\begin{bmatrix}1\amp2\\3\amp1\end{bmatrix}, \qquad B=\begin{bmatrix}2\amp1\\1\amp1\end{bmatrix}. \end{equation*}

Compute \(AB\) and \(BA\text{.}\) Do \(A\) and \(B\) commute?

🔗

Solution.

Both matrices are \(2\times2\text{,}\) so both products are defined and are \(2\times2\) matrices. Each entry is the dot product of a row of the left matrix with a column of the right matrix. For \(AB\text{,}\) this gives

\begin{align*} AB \amp=\begin{bmatrix}1\amp2\\3\amp1\end{bmatrix}\begin{bmatrix}2\amp1\\1\amp1\end{bmatrix}\\ \amp=\begin{bmatrix}1\cdot2+2\cdot1\amp1\cdot1+2\cdot1\\3\cdot2+1\cdot1\amp3\cdot1+1\cdot1\end{bmatrix}\\ \amp=\begin{bmatrix}4\amp3\\7\amp4\end{bmatrix}. \end{align*}

🔗

For \(BA\text{,}\) we instead use the rows of \(B\) and the columns of \(A\text{:}\)

\begin{align*} BA \amp=\begin{bmatrix}2\amp1\\1\amp1\end{bmatrix}\begin{bmatrix}1\amp2\\3\amp1\end{bmatrix}\\ \amp=\begin{bmatrix}2\cdot1+1\cdot3\amp2\cdot2+1\cdot1\\1\cdot1+1\cdot3\amp1\cdot2+1\cdot1\end{bmatrix}\\ \amp=\begin{bmatrix}5\amp5\\4\amp3\end{bmatrix}. \end{align*}

🔗

The products differ: for example, the entry in row 1, column 1 is \(4\) in \(AB\) and \(5\) in \(BA\text{.}\) Thus \(AB\ne BA\text{,}\) so \(A\) and \(B\) do not commute.

🔗

Subsection Matrix multiplication as composition

If \(T_A(\mathbf{x})=A\mathbf{x}\) and \(T_B(\mathbf{x})=B\mathbf{x}\text{,}\) then applying \(B\) first and then \(A\) gives

\begin{equation*} T_A(T_B(\mathbf{x}))=A(B\mathbf{x})=(AB)\mathbf{x}. \end{equation*}

Matrix multiplication represents composition of matrix maps.

🔗

Activity 1.4.13. Composition shape check (U1-LO3, U1-LO6).

Suppose \(A\) is \(4\times 3\) and \(B\) is \(3\times 2\text{.}\) What is the shape of \(AB\text{,}\) and what is the domain and codomain of the composed map \(\mathbf{x}\mapsto A(B\mathbf{x})\text{?}\)

🔗

Solution.

The product \(AB\) is \(4\times 2\text{.}\) The composed map takes inputs from \(\mathbb R^2\) and outputs vectors in \(\mathbb R^4\text{.}\)

🔗

Activity 1.4.14. Commute or not? (U1-LO6).

For each pair below:

🔗

compute both products;
🔗

🔗
determine whether \(AB=BA\text{;}\)
🔗

🔗

Pair A: two distinct diagonal transformations. Use the horizontal stretch and reflection across the \(y\)-axis:

\begin{equation*} S=\begin{bmatrix}2\amp0\\0\amp1\end{bmatrix}, \qquad F=\begin{bmatrix}-1\amp0\\0\amp1\end{bmatrix}. \end{equation*}

🔗

Pair B: a shared axis label is not enough. Use the same horizontal stretch and the horizontal shear:

\begin{equation*} S=\begin{bmatrix}2\amp0\\0\amp1\end{bmatrix}, \qquad H=\begin{bmatrix}1\amp1\\0\amp1\end{bmatrix}. \end{equation*}

🔗

Pair C: reflection across the line \(y=x\text{.}\) Use the same horizontal stretch and

\begin{equation*} M=\begin{bmatrix}0\amp1\\1\amp0\end{bmatrix}. \end{equation*}

🔗

Solution.

For Pair A, row-column dot products give

\begin{align*} SF \amp= \begin{bmatrix}(2,0)\cdot(-1,0)\amp(2,0)\cdot(0,1)\\(0,1)\cdot(-1,0)\amp(0,1)\cdot(0,1)\end{bmatrix}=\begin{bmatrix}-2\amp0\\0\amp1\end{bmatrix},\\ FS \amp= \begin{bmatrix}(-1,0)\cdot(2,0)\amp(-1,0)\cdot(0,1)\\(0,1)\cdot(2,0)\amp(0,1)\cdot(0,1)\end{bmatrix}=\begin{bmatrix}-2\amp0\\0\amp1\end{bmatrix}. \end{align*}

The matrices commute. Flipping the horizontal coordinate and doubling it give the same result in either order.

🔗

For Pair B, row-column dot products give

\begin{align*} SH \amp= \begin{bmatrix}(2,0)\cdot(1,0)\amp(2,0)\cdot(1,1)\\(0,1)\cdot(1,0)\amp(0,1)\cdot(1,1)\end{bmatrix}=\begin{bmatrix}2\amp2\\0\amp1\end{bmatrix},\\ HS \amp= \begin{bmatrix}(1,1)\cdot(2,0)\amp(1,1)\cdot(0,1)\\(0,1)\cdot(2,0)\amp(0,1)\cdot(0,1)\end{bmatrix}=\begin{bmatrix}2\amp1\\0\amp1\end{bmatrix}. \end{align*}

The matrices do not commute. Both transformations involve the horizontal coordinate, but this does not suffice for the matrices to commute.

🔗

For Pair C, row-column dot products give

\begin{align*} SM \amp= \begin{bmatrix}(2,0)\cdot(0,1)\amp(2,0)\cdot(1,0)\\(0,1)\cdot(0,1)\amp(0,1)\cdot(1,0)\end{bmatrix}=\begin{bmatrix}0\amp2\\1\amp0\end{bmatrix},\\ MS \amp= \begin{bmatrix}(0,1)\cdot(2,0)\amp(0,1)\cdot(0,1)\\(1,0)\cdot(2,0)\amp(1,0)\cdot(0,1)\end{bmatrix}=\begin{bmatrix}0\amp1\\2\amp0\end{bmatrix}. \end{align*}

The matrices do not commute. The reflection exchanges the coordinate directions. Stretching before that exchange is not the same as stretching afterward.

🔗

Activity 1.4.15. Two horizontal shears (U1-LO6).

Let

\begin{equation*} H=\begin{bmatrix}1\amp1\\0\amp1\end{bmatrix}, \qquad G=\begin{bmatrix}1\amp2\\0\amp1\end{bmatrix}. \end{equation*}

The matrix \(G\) is a stronger horizontal shear.

🔗

Justify the claim that applying \(H\) twice has the same effect as applying \(G\) once.
🔗

🔗
Compute \(HG\) and \(GH\text{.}\) Do the matrices commute?
🔗

🔗

Solution.

Applying \(H\) twice is represented by \(H^2\text{,}\) and

\begin{equation*} H^2=\begin{bmatrix}1\cdot1+1\cdot0\amp1\cdot1+1\cdot1\\0\cdot1+1\cdot0\amp0\cdot1+1\cdot1\end{bmatrix}=\begin{bmatrix}1\amp2\\0\amp1\end{bmatrix}=G. \end{equation*}

Thus \(H(H\mathbf{x})=G\mathbf{x}\) for every input \(\mathbf{x}\text{.}\)

🔗

🔗
Computing both products directly, row-column dot products give

\begin{align*} HG \amp= \begin{bmatrix}(1,1)\cdot(1,0)\amp(1,1)\cdot(2,1)\\(0,1)\cdot(1,0)\amp(0,1)\cdot(2,1)\end{bmatrix}=\begin{bmatrix}1\amp3\\0\amp1\end{bmatrix},\\ GH \amp= \begin{bmatrix}(1,2)\cdot(1,0)\amp(1,2)\cdot(1,1)\\(0,1)\cdot(1,0)\amp(0,1)\cdot(1,1)\end{bmatrix}=\begin{bmatrix}1\amp3\\0\amp1\end{bmatrix}. \end{align*}

These matrices are distinct and non-diagonal, but they commute. The two horizontal shear amounts add in either order.

🔗

Part 1 also gives a shorter explanation: since \(G=H^2\text{,}\) part 1 of Theorem 1.4.7 gives

\begin{equation*} HG=H(HH)=(HH)H=GH. \end{equation*}

Both orders apply \(H\) three times, so the result of part 1 lets us conclude that the matrices commute without computing both products.

🔗

🔗

🔗

Note 1.4.16.

Some structural patterns settle the question quickly. Two diagonal matrices of the same size commute. For example, if

\begin{equation*} A=\begin{bmatrix}a\amp0\\0\amp b\end{bmatrix}, \qquad B=\begin{bmatrix}c\amp0\\0\amp d\end{bmatrix}, \end{equation*}

then

\begin{align*} AB \amp=\begin{bmatrix}ac\amp0\\0\amp bd\end{bmatrix},\\ BA \amp=\begin{bmatrix}ca\amp0\\0\amp db\end{bmatrix}=\begin{bmatrix}ac\amp0\\0\amp bd\end{bmatrix}=AB. \end{align*}

The equality follows because multiplication of real numbers is commutative.

🔗

Every square matrix commutes with itself. The two horizontal shears in Activity 1.4.15 are closely related to this fact: as the solution shows, \(G=H^2\text{,}\) so both \(HG\) and \(GH\) apply the same matrix \(H\) three times. Part 1 of Theorem 1.4.7 gives

\begin{equation*} HG=H(HH)=(HH)H=GH. \end{equation*}

🔗

In general, compute \(AB\) and \(BA\) and compare.

🔗

Subsection Many token queries and keys

Example 1.4.17. Many token queries at once.

Suppose a sequence has \(L\) tokens. Store the token vectors as rows of a matrix

\begin{equation*} X= \begin{bmatrix} \mathbf{x}_1^T\\ \mathbf{x}_2^T\\ \vdots\\ \mathbf{x}_L^T \end{bmatrix} \in \mathbb R^{L\times d}. \end{equation*}

🔗

A transformer layer forms three views of the same token vectors:

\begin{equation*} Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V. \end{equation*}

The matrices \(W_Q,W_K,W_V\) are weight matrices. Their entries are adjusted during training. For now, treat them as ordinary matrices that assign each token three roles.

🔗

For destination position \(i\text{,}\) the vector \(\mathbf{q}_i\) is its query: what that position uses to decide which source positions are relevant to its attention output.

🔗

For source position \(j\text{,}\) the vector \(\mathbf{k}_j\) is its key: what is compared with the query from a destination position.

🔗

For source position \(j\text{,}\) the vector \(\mathbf{v}_j\) is its value: the information that position can contribute after weights are chosen.

🔗

Write

\begin{equation*} Q= \begin{bmatrix} \mathbf{q}_1^T\\ \mathbf{q}_2^T\\ \vdots\\ \mathbf{q}_L^T \end{bmatrix}, \qquad K= \begin{bmatrix} \mathbf{k}_1^T\\ \mathbf{k}_2^T\\ \vdots\\ \mathbf{k}_L^T \end{bmatrix}. \end{equation*}

Then

\begin{equation*} K^T= \begin{bmatrix} \mathbf{k}_1 \amp \mathbf{k}_2 \amp \cdots \amp \mathbf{k}_L \end{bmatrix}. \end{equation*}

\begin{equation*} QK^T = \begin{bmatrix} \mathbf{q}_1^T\\ \mathbf{q}_2^T\\ \vdots\\ \mathbf{q}_L^T \end{bmatrix} \begin{bmatrix} \mathbf{k}_1 \amp \mathbf{k}_2 \amp \cdots \amp \mathbf{k}_L \end{bmatrix} = \begin{bmatrix} \mathbf{q}_1\cdot \mathbf{k}_1 \amp \mathbf{q}_1\cdot \mathbf{k}_2 \amp \cdots \amp \mathbf{q}_1\cdot \mathbf{k}_L\\ \mathbf{q}_2\cdot \mathbf{k}_1 \amp \mathbf{q}_2\cdot \mathbf{k}_2 \amp \cdots \amp \mathbf{q}_2\cdot \mathbf{k}_L\\ \vdots \amp \vdots \amp \ddots \amp \vdots\\ \mathbf{q}_L\cdot \mathbf{k}_1 \amp \mathbf{q}_L\cdot \mathbf{k}_2 \amp \cdots \amp \mathbf{q}_L\cdot \mathbf{k}_L \end{bmatrix}. \end{equation*}

Thus the score matrix

\begin{equation*} S=QK^T \end{equation*}

has entries

\begin{equation*} S_{ij}=\mathbf{q}_i\cdot \mathbf{k}_j. \end{equation*}

Row \(i\) of \(S\) contains the scores used when computing the attention output at position \(i\text{.}\)

🔗

Activity 1.4.18. Many-token shape and entry check (U1-LO3, U1-LO6).

Suppose

\begin{equation*} X\in\mathbb R^{5\times 4} \end{equation*}

and

\begin{equation*} W_Q,W_K,W_V\in\mathbb R^{4\times 3}. \end{equation*}

Let

\begin{equation*} Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,\qquad S=QK^T. \end{equation*}

🔗

What are the shapes of \(Q\text{,}\) \(K\text{,}\) and \(V\text{?}\)
🔗

🔗
What is the shape of \(K^T\text{?}\)
🔗

🔗
What is the shape of \(S=QK^T\text{?}\)
🔗

🔗
Suppose

\begin{equation*} \mathbf{q}_2= \begin{bmatrix} q_{21}\\ q_{22}\\ q_{23} \end{bmatrix}, \qquad \mathbf{k}_4= \begin{bmatrix} k_{41}\\ k_{42}\\ k_{43} \end{bmatrix}. \end{equation*}

Expand the entry \(S_{24}\text{.}\)

🔗

🔗
What does \(S_{24}\) compare?
🔗

🔗
What does row \(2\) of \(S\) contain?
🔗

🔗

Solution.

Since

\begin{equation*} X\in\mathbb R^{5\times 4} \qquad \text{and} \qquad W_Q,W_K,W_V\in\mathbb R^{4\times 3}, \end{equation*}

the matrices \(Q\text{,}\) \(K\text{,}\) and \(V\) all have shape

\begin{equation*} 5\times 3. \end{equation*}

Thus

\begin{equation*} K^T\in\mathbb R^{3\times 5}. \end{equation*}

Therefore

\begin{equation*} S=QK^T \end{equation*}

has shape

\begin{equation*} 5\times 5. \end{equation*}

The \((2,4)\)-entry comes from row \(2\) of \(Q\) and column \(4\) of \(K^T\text{.}\) But column \(4\) of \(K^T\) is the transpose of row \(4\) of \(K\text{.}\) Hence

\begin{align*} S_{24} \amp= \begin{bmatrix} q_{21} \amp q_{22} \amp q_{23} \end{bmatrix} \begin{bmatrix} k_{41}\\ k_{42}\\ k_{43} \end{bmatrix}\\ \amp= q_{21}k_{41}+q_{22}k_{42}+q_{23}k_{43}. \end{align*}

\begin{equation*} S_{24}=\mathbf{q}_2\cdot \mathbf{k}_4. \end{equation*}

This entry compares the query for token \(2\) with the key for token \(4\text{.}\) Row \(2\) of \(S\) contains the scores comparing token \(2\)’s query with all five token keys.

🔗

A later weighting step converts each row of \(S\) into weights that add to \(1\text{.}\) Those weights are then used to average the value vectors.

🔗

From attention scores to attention output.

For this section, an attention calculation has three steps:

\begin{equation*} \text{dot-product scores} \quad \longrightarrow \quad \text{weights} \quad \longrightarrow \quad \text{weighted averages}. \end{equation*}

🔗

The score \(S_{ij}\) compares the query at destination position \(i\) with the key at source position \(j\text{.}\) A larger score means source position \(j\) is more relevant to the attention output at destination position \(i\text{.}\)

🔗

The scores are not yet the output. For each token \(i\text{,}\) a weighting rule converts row \(i\) of \(S\) into numbers

\begin{equation*} \alpha_{i1},\alpha_{i2},\ldots,\alpha_{iL} \end{equation*}

that are nonnegative and add to \(1\text{.}\) These are the attention weights for token \(i\text{.}\) They say how much token \(i\) uses each value vector.

🔗

Let

\begin{equation*} A_{\mathrm{att}}=[\alpha_{ij}] \end{equation*}

be the attention-weight matrix. If the value vectors are stored as rows,

\begin{equation*} V= \begin{bmatrix} \mathbf{v}_1^T\\ \mathbf{v}_2^T\\ \vdots\\ \mathbf{v}_L^T \end{bmatrix}, \end{equation*}

then the attention output is

\begin{equation*} H=A_{\mathrm{att}}V. \end{equation*}

Row \(i\) of \(H\) is

\begin{equation*} \mathbf{h}_i^T = \alpha_{i1}\mathbf{v}_1^T + \alpha_{i2}\mathbf{v}_2^T + \cdots + \alpha_{iL}\mathbf{v}_L^T. \end{equation*}

Equivalently,

\begin{equation*} \mathbf{h}_i = \alpha_{i1}\mathbf{v}_1 + \alpha_{i2}\mathbf{v}_2 + \cdots + \alpha_{iL}\mathbf{v}_L. \end{equation*}

Thus \(H\) is the attention-output matrix, and each row of \(H\) is a weighted average of value vectors. These rows contribute to later token representations; a full transformer layer combines them with the current representations and processes them further.

🔗

Activity 1.4.19. Many-token attention output (U1-LO3, U1-LO6, U1-LO8).

Suppose

\begin{equation*} X\in\mathbb R^{4\times 3}, \qquad W_Q,W_K,W_V\in\mathbb R^{3\times 2}. \end{equation*}

Define

\begin{equation*} Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,\qquad S=QK^T. \end{equation*}

A weighting step converts \(S\) into an attention-weight matrix

\begin{equation*} A_{\mathrm{att}}\in\mathbb R^{4\times 4}. \end{equation*}

The attention output is

\begin{equation*} H=A_{\mathrm{att}}V. \end{equation*}

🔗

What are the shapes of \(Q\text{,}\) \(K\text{,}\) and \(V\text{?}\)
🔗

🔗
What is the shape of \(S=QK^T\text{?}\)
🔗

🔗
What does \(S_{ij}\) compare?
🔗

🔗
What does row \(i\) of \(A_{\mathrm{att}}\) tell you?
🔗

🔗
What is the shape of \(H\text{?}\)
🔗

🔗
Why is each row of \(H\) a weighted average of value vectors?
🔗

🔗
Now suppose row \(2\) of \(A_{\mathrm{att}}\) is

\begin{equation*} \boldsymbol{\alpha}_2^T = \begin{bmatrix} 0 \amp 1/2 \amp 1/4 \amp 1/4 \end{bmatrix} \end{equation*}

and

\begin{equation*} V= \begin{bmatrix} 1 \amp 0\\ 2 \amp 0\\ 0 \amp 4\\ 4 \amp 4 \end{bmatrix}. \end{equation*}

Compute the attention-output vector \(\mathbf{h}_2\text{.}\)

🔗

🔗

Solution.

Since

\begin{equation*} X\in\mathbb R^{4\times 3} \qquad \text{and} \qquad W_Q,W_K,W_V\in\mathbb R^{3\times 2}, \end{equation*}

the matrices \(Q\text{,}\) \(K\text{,}\) and \(V\) are all \(4\times 2\text{.}\)

🔗

The score matrix is

\begin{equation*} S=QK^T. \end{equation*}

Since \(Q\) is \(4\times 2\) and \(K^T\) is \(2\times 4\text{,}\) the matrix \(S\) is \(4\times 4\text{.}\) The entry \(S_{ij}\) compares the query for token \(i\) with the key for token \(j\text{:}\)

\begin{equation*} S_{ij}=\mathbf{q}_i\cdot \mathbf{k}_j. \end{equation*}

🔗

Row \(i\) of \(A_{\mathrm{att}}\) gives the weights token \(i\) uses when averaging the value vectors. Since \(A_{\mathrm{att}}\) is \(4\times 4\) and \(V\) is \(4\times 2\text{,}\)

\begin{equation*} H=A_{\mathrm{att}}V \end{equation*}

has shape \(4\times 2\text{.}\)

🔗

Row \(i\) of \(H\) is a weighted average because

\begin{equation*} \mathbf{h}_i = \alpha_{i1}\mathbf{v}_1 + \alpha_{i2}\mathbf{v}_2 + \alpha_{i3}\mathbf{v}_3 + \alpha_{i4}\mathbf{v}_4, \end{equation*}

where the weights in row \(i\) of \(A_{\mathrm{att}}\) are nonnegative and add to \(1\text{.}\)

🔗

For row \(2\text{,}\)

\begin{equation*} \mathbf{h}_2 = 0 \begin{bmatrix} 1\\ 0 \end{bmatrix} + \frac12 \begin{bmatrix} 2\\ 0 \end{bmatrix} + \frac14 \begin{bmatrix} 0\\ 4 \end{bmatrix} + \frac14 \begin{bmatrix} 4\\ 4 \end{bmatrix}. \end{equation*}

Thus the attention-output vector is

\begin{equation*} \mathbf{h}_2 = \begin{bmatrix} 0\\ 0 \end{bmatrix} + \begin{bmatrix} 1\\ 0 \end{bmatrix} + \begin{bmatrix} 0\\ 1 \end{bmatrix} + \begin{bmatrix} 1\\ 1 \end{bmatrix} = \begin{bmatrix} 2\\ 2 \end{bmatrix}. \end{equation*}

The second token ignores value vector \(1\text{,}\) uses value vector \(2\) with weight \(1/2\text{,}\) and uses value vectors \(3\) and \(4\) with weight \(1/4\) each.

🔗

Warning 1.4.20.

In real transformer attention, the weighting rule usually uses scaled and masked scores before forming weights. For Unit 1, the important facts are: rows of \(A_{\mathrm{att}}\) are weights, the weights add to \(1\text{,}\) and \(A_{\mathrm{att}}V\) forms weighted averages of the rows of \(V\text{.}\)

🔗

Activity 1.4.21. The same shape check in code (U1-LO3, U1-LO8).

The following code performs a many-token shape check similar to the previous activity.

🔗

import numpy as np

X_tokens = np.ones((5, 4))
WQ = np.ones((4, 3))
WK = np.ones((4, 3))
WV = np.ones((4, 3))

Q = X_tokens @ WQ
K_tokens = X_tokens @ WK
V_tokens = X_tokens @ WV
scores_many = Q @ K_tokens.T

Q.shape, K_tokens.shape, V_tokens.shape, scores_many.shape

Output:

🔗

((5, 3), (5, 3), (5, 3), (5, 5))

Which line creates all query-key scores at once?
🔗

🔗
Why is scores_many.shape equal to (5, 5)?
🔗

🔗

Solution.

The line scores_many = Q @ K_tokens.T creates all query-key scores at once. Since Q has shape (5, 3) and K_tokens.T has shape (3, 5), the score matrix has shape (5, 5).

🔗

Prev Top Next