Matrix multiplication and map composition

Section 1.4 Matrix multiplication and map composition

Subsection Matrix product by columns

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.1. Computing a matrix product using columns.

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{.}\) Thus

\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*}

🔗

Subsection Entry formula by row-column dot products

Theorem 1.4.2.

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{equation*} (AB)_{ij}=\sum_{\ell=1}^{n}a_{i\ell}b_{\ell j}. \end{equation*}

🔗

Activity 1.4.2. Computing a matrix entry.

\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{equation*} (2,4,0)\cdot(-4,-1,1)=2(-4)+4(-1)+0(1)=-12. \end{equation*}

🔗

Theorem 1.4.3. 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{.}\)
🔗

🔗

🔗

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.

🔗

Warning 1.4.4. Order matters.

Let

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

Then

\begin{equation*} RS=\begin{bmatrix}0\amp-1\\2\amp0\end{bmatrix}, \qquad SR=\begin{bmatrix}0\amp-2\\1\amp0\end{bmatrix}. \end{equation*}

Rotating after a horizontal stretch is not the same as stretching after a rotation.

🔗

Activity 1.4.3. Composition shape check.

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.4. Applying one matrix to many points.

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

np.set_printoptions(precision=3, suppress=True)

# 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],
], dtype=float)

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

Y = A @ X
Y

Output:

🔗

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

What are the columns of X?
🔗

🔗
What are the columns of Y?
🔗

🔗
Which transformation from the gallery is this?
🔗

🔗
Why does the bottom edge stay fixed?
🔗

🔗
Why does the top edge shift right?
🔗

🔗

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, and points with \(y=1\) shift right by \(1\text{.}\)

🔗

Warning 1.4.5.

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.

🔗

Note 1.4.6. Later square-grid visualizations.

The matrix computation \(Y=AX\) applies one matrix to every column of \(X\text{.}\) For a \(2\times 2\) matrix, the images of the coordinate directions determine the whole grid.

🔗

Later, in Activity 3.1.3, the same columns-of-points convention is used for rules that are not matrix maps. Then four corners may not tell the whole story. A grid of input points gives more information.

🔗

Subsection Many token queries and keys

Example 1.4.7. 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 token \(i\text{,}\) the vector \(\mathbf{q}_i\) is its query: what token \(i\) uses to decide which tokens matter for its update.

🔗

For token \(j\text{,}\) the vector \(\mathbf{k}_j\) is its key: what token \(j\) uses to be compared against a query.

🔗

For token \(j\text{,}\) the vector \(\mathbf{v}_j\) is its value: what token \(j\) 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 updating token \(i\text{.}\)

🔗

Activity 1.4.5. Many-token shape and entry check.

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{equation*} S_{24} = \begin{bmatrix} q_{21} \amp q_{22} \amp q_{23} \end{bmatrix} \begin{bmatrix} k_{41}\\ k_{42}\\ k_{43} \end{bmatrix} = q_{21}k_{41}+q_{22}k_{42}+q_{23}k_{43}. \end{equation*}

\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 token \(i\)’s query with token \(j\)’s key. A larger score means token \(j\) is more relevant when updating token \(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*}

So each updated token vector is a weighted average of value vectors.

🔗

Activity 1.4.6. Many-token attention output.

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 updated 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

\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.8.

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.7. The same shape check in code.

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