1.1 Linear algebra for language models

A transformer is easier to understand when every matrix has a job and every axis has a name. This chapter develops the small amount of linear algebra needed to read the attention equations.

Vectors as representations

A token ID is an integer. Its embedding is a vector $\mathbf{x}\in\mathbb{R}^{d_{\mathrm{model}}}$, obtained by selecting a row of a learned embedding matrix. The vector’s coordinates are learned features. Individual coordinates need not correspond to human-readable concepts.

A sequence of $n$ tokens becomes a matrix with one token per row:

$$X = \begin{bmatrix} x_{11} & \cdots & x_{1d} \\ \vdots & \ddots & \vdots \\ x_{n1} & \cdots & x_{nd} \end{bmatrix} \in \mathbb{R}^{n\times d}. \tag{1.1}$$(1.1)

Throughout these notes, sequence positions are rows and feature dimensions are columns. Libraries may add batch and attention-head axes, but the same idea still applies.

Linear maps and shapes

Multiplying $X$ by a learned matrix $W\in\mathbb{R}^{d\times k}$ produces $Y=XW\in\mathbb{R}^{n\times k}$. The same map is applied independently to every token position. It mixes features, not positions.

For one output coordinate,

$$y_{ij}=\sum_{r=1}^{d}x_{ir}w_{rj}. \tag{1.2}$$(1.2)

In attention, three such maps create queries, keys, and values. The maps share an input but learn different parameters. Sharing a source does not make their roles interchangeable.

Python implementation

shapes.py

import numpy as np

rng = np.random.default_rng(7)
X = rng.normal(size=(4, 8))       # 4 positions, 8 features
W_Q = rng.normal(size=(8, 3))    # map features to 3 query features
W_K = rng.normal(size=(8, 3))

Q = X @ W_Q                     # (4, 3)
K = X @ W_K                     # (4, 3)
scores = Q @ K.T                # (4, 4): positions against positions
assert scores.shape == (4, 4)

Inner products compare two vectors

The inner product $\mathbf{q}^{\mathsf T}\mathbf{k}=\sum_{r=1}^{d_k}q_rk_r$ returns a scalar. In attention, it is a learned compatibility score. Its magnitude depends on both direction and vector length, so it is not generally a cosine similarity.

For two sequence matrices, $QK^{\mathsf T}$ computes all pairwise query–key scores at once. Entry $(i,j)$ compares the query at position $i$ with the key at position $j$.

A useful distinction

The product $XW$ mixes features separately at each position. The product $AV$, where $A\in\mathbb{R}^{n\times n}$, mixes positions separately for each value feature. Attention combines these two kinds of operation.

Weighted sums

Suppose $a_j\geq0$ and $\sum_j a_j=1$. Then $\mathbf{z}=\sum_j a_j\mathbf{v}_j$ is a weighted average of the value vectors. With fixed values, the result lies in their convex hull.

This describes a single attention head before its output projection. It does not imply that an entire transformer can only average its inputs: learned projections, residual paths, and nonlinear feed-forward layers change what the full network can compute.

Continue reading

The next step is to see how a model constructs the weights $a_j$. The scaled dot-product attention equation (2.3) turns compatibility scores into normalized weights. The attention chapter derives that computation and implements it in NumPy.