Lecture 4 — Attention

Project connection. Project Step 4 implements one causal self-attention head, turns the simplex, scaling, and causality results below into numerical tests, and visualizes attention matrices on real text and controlled positional kernels.

Chapter overview. The Step 3 MLP reads a fixed number of preceding tokens by concatenating their embeddings. Attention instead gives every position content-dependent access to all permitted positions. Queries and keys decide where to read, values determine what is returned, and softmax makes the result a convex average. We will calculate one head completely before studying its geometry. Causal masking restricts every average to the past, while positional information breaks an otherwise exact permutation symmetry. The final sections reinterpret the same operation as kernel smoothing and account for its quadratic cost.

0. Embeddings and Transformers

Recall that we work with a vocabulary $\mathcal{V}$ of size $V$, which we will enumerate for convenience. An embedding is a map from tokens to vectors:

$$e: \mathcal{V} \rightarrow \mathbb{R}^d$$

The dimension $d$ is called the model width. The embedding matrix for the embedding $e$ is the matrix $E \in M_{V \times d}(\mathbb{R})$ such that the embedding of the $t$-th token is the $t$-th row.

In a neural network, as we saw, we concatenated the embeddings of the tokens in the context. Then this longer vector was modified as it passed through each layer of the network, and finally results in an output that is interpreted as a probability distribution on the next token.

In a transformer architecture (the model we will study next, namely a decoder-only causal single-head transformer, to be defined formally shortly), the token positions in the context each have a residual stream, which is a vector $x_n$ initially computed as a function of two things: (1) the embedding of the token in that position, and (2) the position itself. At the $n$-th position, this is called the residual stream at position $n$. Each stream is modified at each layer of the network, until at the end, each stream is unembedded (we’ll define this later) and fed through softmax. These two operations form the language modelling head (although sometimes this terminology only refers to the first operation). The output of the language modelling head on the last stream is interpreted as the probability distribution on the next token. The picture is like this:

We need to examine what is happening in the puffy cloud. Inside the cloud there are several layers, or steps, to this process. These layers are grouped into transformer blocks that have a standardized layer structure, and are stacked up, each with its own parameters. So now our picture looks like this:

Inside one common transformer block style, we have, in order of execution: a layer norm, an attention layer, another layer norm and then a feedforward layer. The layer norms are layers that normalize vectors. The feedforward layer is a small feedforward neural network.

The most important layer is the attention layer. An attention layer can be multi-headed, but for now we will assume we have just one attention head (which we will describe formally in the next section). In the single-headed case, the attention layer looks like the following. A row vector $\mathbf x_n$ in the $n$-th stream is holding some type of “meaning” for the LLM. As mentioned above, the initial meaning of that vector is just a function of the embedding of the $n$-th token and the integer $n$ (its position in the context). In this layer, the weights of the model dictate that the new value $\mathbf{x}_n'$ of the $n$-th stream $\mathbf{x}_n$ should be a function of itself and all the streams preceding it (i.e. corresponding to the $j$-th token, where $j < n$). The form of this function is this:

$$\mathbf{x}_n' = \mathbf{x}_n + \underbrace{\sum_{j \le n} \alpha_{j,n} \mathbf{x}_j T}_{\mathbf{a}_n},\quad (\alpha_{j,n})_{j=1}^{n} = \operatorname{softmax}\left( (\mathbf{x}_n G \mathbf{x}_j^\top)_{j=1}^n \right),$$

for some matrices $T$ and $G$. The matrices $T$ and $G$ are part of the parameters $\theta$ of the system, trained in some way. For now, observe that $G$ controls how much $\mathbf{x}_j$ affects $\mathbf{x}_n$, and $T$ controls how we transform the resulting linear combination. The way that $G$ controls the interaction of one token with previous tokens is referred to as attention. When the pairing of two streams through $G$ is large, the earlier token has a large relative effect on the later one, and we say that the later token is attending strongly to the earlier one.

We call $\mathbf{a}_n$ the attention output at position $n$. It is added to $\mathbf{x}_n$ to update the stream. This style of update (computing a vector to add to the existing stream) is called residual connection, and the stream is sometimes called the residual stream for this reason. The feedforward network’s output is also added into the stream; we will discuss this later.

1. Presenting the context to the transformer

We now describe how a context, as a raw sequence of tokens, is formatted for the use of a transformer block.

As in the last section, let $\mathcal{V}$ be a vocabulary of size $V$, enumerated in some order. We have an embedding map

$$e: \mathcal{V} \rightarrow \mathbb{R}^d$$

We typically take the output of $e$ to be a column vector. The embedding matrix for the embedding $e$ is the matrix $E \in M_{V \times d}(\mathbb{R})$ such that the embedding of the $t$-th token is the $t$-th row transposed.

The map above allows us to embed information about which token we are considering as a vector in $\mathbb{R}^d$. We also wish to carry positional information: which position is the token occupying in the context. A system of additive positional vectors is a map

$$p : \{1, \ldots, T \} \rightarrow \mathbb{R}^d, \quad t \mapsto p_t$$

We should think of $p_t$ as carrying the information “we are at position $t$”. We treat $p_t$ as a column vector.

These two maps allow us to feed our context into a transformer block as a sequence of streams, as follows. Let $T$ be the length of the context, as a sequence of tokens $x_1, \ldots, x_T$. For each position $t$, we create a vector of $\mathbb{R}^d$ which contains information about the token $x_t$ and its position $t$. This is

$$r : \mathcal{V}^T \rightarrow \mathbb{R}^{T\times d}, (x_t)_{t=1}^T \mapsto \begin{pmatrix} X_1^\top \\ \vdots \\ X_T^\top \end{pmatrix}, \quad X_t := e(x_t) + p_t.$$

This represents the context as a $T \times d$ matrix $X$ whose rows are $X_t^\top$, $t=1, \ldots, T$.

We refer to an element of $\mathbb{R}^{T \times d}$ as a residual stream, and the first step of the transformer is to initialize the residual stream to $X := r(x_1, \ldots, x_T)$.

---

2. The transformer head

In the introductory material, we conflated the transformer layer with the attention head, in a single-head model. We will now define an attention head more precisely. In general, an attention head and transformer layer are not exactly the same, even in a single-head model.

Fix positive integers $d_k$ and $d_v$ (to be used as dimensions of vector spaces). Let

$$W_Q,W_K\in\mathbb R^{d\times d_k}, \qquad W_V\in\mathbb R^{d\times d_v}$$

be matrices called the query, key, and value projections or projection matrices. These matrices are part of the parameters of the system, to be adjusted during training.

Here the matrix $V$ of values is distinguished by context from the scalar $V=|\mathcal V|$.

When $Q,K,V$ are all obtained from the same stream $X$, the operation is called self-attention. In cross-attention, queries come from one stream and keys and values from another. This lecture studies self-attention.

For a matrix $S\in\mathbb R^{T\times T}$, row-wise softmax is defined by

$$(\operatorname{softmax}_{\mathrm{row}}S)_{ij} = \frac{e^{S_{ij}}}{\sum_{\ell=1}^Te^{S_{i\ell}}}.$$

Entrywise,

$$S_{ij}=\frac{\langle q_i,k_j\rangle}{\sqrt{d_k}}, \qquad A_{ij}=\frac{e^{S_{ij}}}{\sum_{\ell=1}^Te^{S_{i\ell}}}, \qquad \operatorname{Attn}(X)_i=\sum_{j=1}^TA_{ij}v_j.$$

Here $\operatorname{Attn}(X)_i$ denotes the column-vector form of output row $i$. The index order matters:

$$A_{ij} = \text{weight assigned by query position }i \text{ to key position }j.$$

Now the attention head can be described by the following diagram, illustrating the algorithm taking as input the stream $X$, computing intermediate matrices $Q, K, V, S, A$ and producing attention output $AV$.

With reference to our introduction, the key and query matrices determine $G$, and $T$ is determined from the value matrix and a final operation that comes later.

The intuition is as follows. The product $QK^\top$ is comparing every pair $(\text{query (row)}, \text{key (column)}) = (q_i, k_j)$. When the resulting row-column product (normalized, and called the score) is large, the residual vector at the query position will attend strongly to the residual vector at the key position, meaning the vector in the key position will strongly influence the vector in the query position during the transformer block algorithm. In other words, the score matrix tells the query vectors which key vectors to attend to.

The score matrix is row-wise softmaxed, as a type of normalization. Then the result $A$, the attention matrix, is multiplied by the value matrix $V$. Now we need to actually exert the influence the scores are indicating we should. You might think we simply add a multiple of the stream vector in a well-attended-to key position directly to the stream vector in the query position. Instead, we multiply by $V$. The role of $V$ is to interpret a stream vector in some way (by a linear transformation) before it is added in to the stream vector in the query position.

In summary:

• $q_i$ is what position $i$ is looking for.
• $k_j$ is how position $j$ advertises itself.
• $v_j$ is the information position $j$ offers if selected.

The score matrix has a useful rank restriction.

Proof

The first identity follows from $K^\top=(XW_K)^\top=W_K^\top X^\top$. Since $S=QK^\top/\sqrt{d_k}$ factors through $\mathbb R^{d_k}$, the rank of the product is at most $d_k$. $\square$

Remark 2.5 (nonlinearity). The map $X\mapsto\operatorname{Attn}(X)$ is generally nonlinear, thanks to the softmax. Applying a nonlinear function row by row can increase matrix rank. Low-dimensional query and key spaces therefore do not imply that $A$ is a low-rank matrix.

---

3. Attention as averaging

Attention scores may be positive or negative and have no fixed scale. After softmax, however, every row belongs to a familiar geometric object. Let

$$\Delta^{T-1} = \left\{ a\in\mathbb R^T: a_j\geq0,\ \sum_{j=1}^Ta_j=1 \right\}.$$

This is the standard $(T-1)$-dimensional probability simplex. Its relative interior consists of the vectors with every coordinate strictly positive.

For a finite collection of vectors $v_1,\ldots,v_T$, their convex hull is

$$\operatorname{conv}\{v_1,\ldots,v_T\} = \left\{ \sum_{j=1}^Ta_jv_j: a_j\geq0,\ \sum_{j=1}^Ta_j=1 \right\}.$$

Proof

Every exponential $e^{s_j}$ is positive. Dividing by their positive sum gives

$$\operatorname{softmax}(s)_j>0, \qquad \sum_{j=1}^T\operatorname{softmax}(s)_j=1.$$

Thus the softmax vector lies in the relative interior of the simplex. By Definition 2.2,

$$\operatorname{Attn}(X)_i = \sum_{j=1}^TA_{ij}v_j,$$

and the coefficients are nonnegative and sum to one. This is a convex combination by definition. $\square$

The same statement has a probabilistic form. For each query position $i$, let $J_i$ be a random source index with

$$\mathbb P(J_i=j)=A_{ij}.$$

Then

$$\operatorname{Attn}(X)_i = \mathbb E[v_{J_i}].$$

Attention is therefore literally an expectation of the value vector under a distribution selected by the query and keys.

There are several useful scalar consequences. If $u\in\mathbb R^{d_v}$, then

$$\min_j\langle u,v_j\rangle \leq \left\langle u,\operatorname{Attn}(X)_i\right\rangle \leq \max_j\langle u,v_j\rangle.$$

Taking the Euclidean norm and using the triangle inequality gives

$$\left\|\operatorname{Attn}(X)_i\right\|_2 \leq \sum_jA_{ij}\|v_j\|_2 \leq \max_j\|v_j\|_2.$$

More generally, Jensen’s inequality says that every convex $\phi:\mathbb R^{d_v}\to\mathbb R$ satisfies

$$\phi\!\left(\operatorname{Attn}(X)_i\right) \leq \sum_jA_{ij}\phi(v_j).$$

Remark 3.4 (scope of the corollary). The value projection can create new feature directions before averaging, and an output projection and residual addition can transform the result afterward. The corollary is the precise statement that the averaging step $AV$ itself does not extrapolate beyond its current values. It is not a claim that a complete transformer remains inside the convex hull of its input token embeddings.

The qualification “unmasked” matters. Causal masking will set forbidden coordinates exactly to zero, so a causal row lies in a lower-dimensional face of $\Delta^{T-1}$. Among its permitted coordinates, finite scores still give strictly positive weights.

Geometric aside 3.6. The attention row gives barycentric coordinates in the convex hull of the values. When the values are affinely independent, those coordinates are unique. When they are affinely dependent, different attention rows can produce the same output. Thus an attention pattern need not be recoverable from the output vector alone. The simplex picture separates the selection rule, encoded by $Q$ and $K$, from the information being averaged, encoded by $V$.

---

4. Causal masking

There’s one important aspect we have so far ignored. For token prediction, we wish to train the model to predict later tokens from earlier ones. When the representation at position $i$ is used to predict the next token, it should depend on the tokens up through position $i$ but not on the later positions $i+1,\ldots,T$.

Right now, our score matrix $S$ can have non-zero entries in position $(i,j)$ where $j > i$. In other words, stream vectors associated to later tokens in the context can contribute to earlier ones. The solution is causal masking, which happens inside the attention operator. Everything preceding this section described an unmasked transformer.

Thus $A$ is lower triangular, meaning $A_{ij}=0$ whenever $j>i$, and row-stochastic, meaning its entries are nonnegative and every row sums to one. Row $i$ lies in the face of the simplex supported on its first $i$ coordinates.

Proof

Row $i$ uses $q_i$, the keys $k_\ell$ with $\ell\leq i$, the values $v_\ell$ with $\ell\leq i$, and a normalizing sum over $\ell\leq i$. These quantities depend only on $X_1,\ldots,X_i$. None depends on $X_j$ for $j>i$, so the output itself does not depend on $X_j$. Its derivative with respect to that variable is therefore zero. $\square$

The convex-hull result also acquires a prefix form:

$$\operatorname{Attn}(X)_i \in \operatorname{conv}\{v_1,\ldots,v_i\}.$$

Implementation aside 4.4 (mask before softmax). Setting forbidden weights to zero after an unmasked softmax, without renormalizing, leaves the row sums below one and changes the derivatives; masking the scores to $-\infty$ before softmax normalizes over exactly the allowed positions. Numerical implementations use negative infinity or, when required by a numerical format, a sufficiently large negative number.

Remark 4.5 (parallel training). In teacher-forced training, the full token sequence is available to the computer, and losses for all $T$ positions can be computed in one forward pass. Inputs and targets are shifted: the representation after token $x_i$ predicts $x_{i+1}$. Allowing position $i$ to attend to itself therefore reveals $x_i$, not the target $x_{i+1}$. The causal mask proves that no position uses a future target, so parallel training is consistent with left-to-right generation.

---

5. Log-sum-exp, temperature, and score scaling

Every row of attention applies softmax to a vector of scores. We now study that map as a mathematical object. Define the log-sum-exp function by

$$\operatorname{lse}(s) = \log\sum_{j=1}^Te^{s_j}, \qquad s\in\mathbb R^T.$$

For a probability vector $p\in\Delta^{T-1}$, let $\operatorname{diag}(p)$ denote the diagonal matrix with diagonal $p$, and let $\mathbf1=(1,\ldots,1)^\top$. A symmetric matrix $M$ is positive semidefinite if

$$u^\top Mu\geq0$$

for every vector $u$. The Hessian of a twice differentiable scalar function is the matrix of its second partial derivatives.

Proof

Direct differentiation gives

$$\frac{\partial}{\partial s_i}\operatorname{lse}(s) = \frac{e^{s_i}}{\sum_je^{s_j}} = p_i.$$

Differentiating once more,

$$\frac{\partial p_i}{\partial s_j} = p_i\bigl(\mathbf1_{\{i=j\}}-p_j\bigr),$$

which is the $(i,j)$ entry of $\operatorname{diag}(p)-pp^\top$. For any $u\in\mathbb R^T$,

$$\begin{aligned} u^\top(\operatorname{diag}(p)-pp^\top)u &= \sum_ip_iu_i^2-\left(\sum_ip_iu_i\right)^2\\ &= \operatorname{Var}_{I\sim p}(u_I) \geq0. \end{aligned}$$

Because $p_i>0$ for every $i$, this variance is zero exactly when all $u_i$ are equal. The nullspace is therefore $\mathbb R\mathbf1$. Positive semidefiniteness of the Hessian gives convexity, and its positive definiteness on $\mathbf1^\perp$ gives strict convexity there. $\square$

The calculation simultaneously gives the Jacobian of softmax:

$$D\operatorname{softmax}(s) = \operatorname{diag}(p)-pp^\top.$$

In coordinates, a perturbation $h\in\mathbb R^T$ changes probability $p_i$ to first order by

$$\bigl(D\operatorname{softmax}(s)h\bigr)_i = p_i\left(h_i-\sum_jp_jh_j\right).$$

Only the centered part of $h$ matters. Indeed,

$$\operatorname{lse}(s+c\mathbf1) = c+\operatorname{lse}(s), \qquad \operatorname{softmax}(s+c\mathbf1) = \operatorname{softmax}(s).$$

The common-shift direction is exactly the null direction found in Proposition 5.1. This is the same softmax geometry developed in Lecture 2; here each attention row supplies its own logit vector.

5.1 Temperature as entropy regularization

For $\tau>0$, the vector $\operatorname{softmax}(s/\tau)$ has temperature $\tau$. The entropy of $a\in\Delta^{T-1}$ is

$$H(a) = -\sum_ja_j\log a_j,$$

with $0\log0=0$. For probability vectors $a$ and $p$ with $p_j>0$, their Kullback—Leibler divergence is

$$D_{\mathrm{KL}}(a\|p) = \sum_{j:a_j>0}a_j\log\frac{a_j}{p_j}.$$

Gibbs’ inequality states that this quantity is nonnegative and is zero exactly when $a=p$.

Proof

Put

$$Z=\sum_je^{s_j/\tau}, \qquad p=\operatorname{softmax}(s/\tau).$$

Since $\log p_j=s_j/\tau-\log Z$, every $a\in\Delta^{T-1}$ satisfies

$$\begin{aligned} \langle a,s\rangle+\tau H(a) &= \tau\sum_ja_j\log p_j +\tau\log Z -\tau\sum_ja_j\log a_j\\ &= \tau\log Z-\tau D_{\mathrm{KL}}(a\|p). \end{aligned}$$

Gibbs’ inequality makes the last expression at most $\tau\log Z$, with equality exactly when $a=p$. Hence the displayed maximizer exists and is unique. $\square$

The formula says that softmax balances two desires. The linear term $\langle a,s\rangle$ rewards mass on large scores. The entropy term $\tau H(a)$ rewards a spread-out distribution. Temperature is the relative strength of that entropy reward.

Remark 5.4 (temperature limits). Let $M=\{j:s_j=\max_\ell s_\ell\}$. Dividing numerator and denominator by $\exp(\max_\ell s_\ell/\tau)$ shows that, as $\tau\to0^+$, $\operatorname{softmax}(s/\tau)$ tends to the uniform distribution on $M$. As $\tau\to\infty$, every $e^{s_j/\tau}$ tends to $1$, so the distribution tends to the uniform distribution on all $T$ positions. The same variational principle governs sampling temperature in Lecture 7.

5.2 Why the score is divided by $\sqrt{d_k}$

The score normalization follows from a variance calculation.

Proof

Write

$$\langle q,k\rangle = \sum_{m=1}^{d_k}q_mk_m.$$

Each product has mean

$$\mathbb E[q_mk_m] = \mathbb E[q_m]\mathbb E[k_m] = 0$$

and variance

$$\mathbb E[q_m^2k_m^2] = \mathbb E[q_m^2]\mathbb E[k_m^2] = \sigma^4.$$

Independence makes the products $q_mk_m$ independent, and in particular uncorrelated, across $m$. The mean of their sum is zero and its variance is the sum $d_k\sigma^4$. $\square$

Consequently,

$$\operatorname{Var} \left( \frac{\langle q,k\rangle}{\sqrt{d_k}} \right) = \sigma^4.$$

With unit coordinate variance, the expected scales are therefore:

$d_k$	standard deviation of $q^\top k$	standard deviation of $q^\top k/\sqrt{d_k}$
$16$	$4$	$1$
$64$	$8$	$1$
$256$	$16$	$1$

Remark 5.6 (why divide by $\sqrt{d_k}$). Without normalization, score gaps typically grow with the head dimension. Large gaps move a softmax row toward a simplex vertex, a regime called softmax saturation. If $p$ tends to a one-hot vector, then $\operatorname{diag}(p)-pp^\top$ tends entrywise to zero. Gradients passing from the attention weights back to the scores can therefore become very small. Dividing by $\sqrt{d_k}$ keeps the initial score scale approximately independent of $d_k$.

The word “initial” is important. The lemma is an idealized calculation, not an invariant of training. Queries and keys in a trained model need not have independent coordinates, equal variances, or mean zero. Even at initialization, a self-score compares two projections of the same stream row, so exact independence depends on the initialization model. The calculation motivates the architectural scale; it does not assert that every trained score matrix has variance one.

Probabilistic aside 5.7. Under suitable moment hypotheses, the central limit theorem makes the normalized sum $\langle q,k\rangle/\sqrt{d_k}$ approximately Gaussian for large $d_k$. The variance calculation, rather than the normal approximation, is the essential justification. Project Step 4 measures the scale directly instead of assuming that a finite random sample exactly matches the asymptotic model.

---

6. Permutation equivariance and positional information

Attention compares content at different rows. Does the operator itself know which row is “first” or “third”? The answer for unmasked attention is no: if its input rows are reordered, its output rows undergo exactly the same reordering.

Let $S_T$ be the group of permutations of $\{1,\ldots,T\}$. For $\pi\in S_T$, let $P_\pi\in\{0,1\}^{T\times T}$ be its permutation matrix, with the convention

$$(P_\pi X)_t = X_{\pi^{-1}(t)}.$$

Define the corresponding token action by

$$(\pi\mathbin{\cdot}x)_t = x_{\pi^{-1}(t)}.$$

Thus left multiplication by $P_\pi$ and the action $\pi\mathbin{\cdot}x$ use the same reordering.

A map $F:\mathbb R^{T\times d}\to\mathbb R^{T\times d'}$ is permutation-equivariant if

$$F(P_\pi X) = P_\pi F(X)$$

for every $\pi$ and $X$. Equivariance is not invariance. An invariant map would satisfy $F(P_\pi X)=F(X)$; an equivariant map changes its output, but in a completely prescribed way.

Proof

Row-wise linear projection commutes with row permutation:

$$Q(P_\pi X) = (P_\pi X)W_Q = P_\pi Q(X).$$

The same identity holds for $K$ and $V$. Therefore

$$\begin{aligned} S(P_\pi X) &= \frac{(P_\pi Q)(P_\pi K)^\top}{\sqrt{d_k}}\\ &= P_\pi S(X)P_\pi^\top. \end{aligned}$$

The matrix on the right simultaneously permutes the rows and columns of $S$. Row-wise softmax commutes with this simultaneous permutation, so

$$A(P_\pi X) = P_\pi A(X)P_\pi^\top.$$

Finally, using $P_\pi^\top P_\pi=I$,

$$\begin{aligned} \operatorname{Attn}(P_\pi X) &= \bigl(P_\pi AP_\pi^\top\bigr)(P_\pi V)\\ &= P_\pi AV\\ &= P_\pi\operatorname{Attn}(X). \end{aligned}$$

$\square$

The simultaneous row—column action on $A$ has a concrete meaning. If source position $j$ and query position $i$ are both renamed by the same permutation, the weight coupling those two renamed positions is unchanged. The attention pattern follows the objects being permuted rather than remaining attached to absolute row numbers.

Position-wise maps are equivariant for the same reason. If $\phi:\mathbb R^d\to\mathbb R^{d'}$ is applied separately with shared parameters to every row, then

$$\phi_{\mathrm{row}}(P_\pi X) = P_\pi\phi_{\mathrm{row}}(X).$$

Sums and compositions of equivariant maps remain equivariant. Thus stacking unmasked attention with shared row-wise MLPs does not remove the symmetry.

6.1 Where positional embeddings break the symmetry

Let $E_x\in\mathbb R^{T\times d}$ have row $E_{x_t,:}$ and let $p\in\mathbb R^{T\times d}$ have row $p_t^\top$. Define the token-sequence embedding maps

$$\iota_0(x)_t = E_{x_t}^\top$$

and

$$\iota_p(x)_t = E_{x_t}^\top+p_t.$$

Without positions, embedding commutes with token permutation:

$$\iota_0(\pi\mathbin{\cdot}x) = P_\pi\iota_0(x).$$

With fixed, nonconstant position vectors, in general

$$\begin{aligned} \iota_p(\pi\mathbin{\cdot}x) &= P_\pi E_x+p,\\ P_\pi\iota_p(x) &= P_\pi E_x+P_\pi p. \end{aligned}$$

These are unequal because permuting the tokens leaves the positional matrix fixed in the first expression but permutes it in the second.

Proof

Without positional information, the embedding map is equivariant. Theorem 6.1 and the observation above show that each subsequent attention or position-wise map is equivariant. Equivariance is preserved under composition, so the entire composite is equivariant. With fixed position-dependent information, the embedding identity already fails in general, as the two displayed formulas show. $\square$

Precision remark 6.3. Adding positional vectors to a particular numerical matrix $X$ does not make the function $X\mapsto\operatorname{Attn}(X)$ cease to be equivariant. Theorem 6.1 holds for every input matrix, including one whose rows happen to contain positional information. Symmetry is broken by the preceding map from token sequences to matrices, because positions stay fixed while tokens are permuted. Step 4 tests both statements separately.

This distinction is easiest to remember as a comparison of two experiments.

1. Permute completed stream rows. Starting from a numerical matrix $X$, compare $\operatorname{Attn}(P_\pi X)$ with $P_\pi\operatorname{Attn}(X)$. They are equal.
2. Permute tokens before fixed-position embedding. Compare $\operatorname{Attn}(\iota_p(\pi\cdot x))$ with $P_\pi\operatorname{Attn}(\iota_p(x))$. They are generally unequal.

The first experiment tests the attention map. The second tests the composite from discrete tokens through positional embedding and attention.

6.2 Additive, sinusoidal, and rotary position information

Two common positional constructions illustrate absolute and relative geometry.

1. Additive positional embeddings. Use learned vectors $p_1,\ldots,p_{T_{\max}}$, as in the course project, or fixed sinusoidal coordinates. At angular frequency $\omega$, the two-coordinate feature

   $$r_t(\omega) = \bigl(\sin(\omega t),\cos(\omega t)\bigr)$$

   satisfies

   $$\langle r_i(\omega),r_j(\omega)\rangle = \cos\bigl(\omega(i-j)\bigr).$$

   Thus absolute vectors also make relative offsets accessible through inner products. The standard fixed construction uses a range of frequencies, for example

   $$p_{t,2m} = \sin\!\left(t/10000^{2m/d}\right), \qquad p_{t,2m+1} = \cos\!\left(t/10000^{2m/d}\right).$$

2. Rotary positional embeddings (RoPE). In one two-dimensional query/key coordinate plane, let

   $$R_\theta = \begin{pmatrix} \cos\theta&-\sin\theta\\ \sin\theta&\cos\theta \end{pmatrix} \in SO(2).$$

   Rotate the query at position $i$ by $R_{i\omega}$ and the key at position $j$ by $R_{j\omega}$. Then

   $$\begin{aligned} \langle R_{i\omega}q_i,R_{j\omega}k_j\rangle &= q_i^\top R_{i\omega}^\top R_{j\omega}k_j\\ &= q_i^\top R_{(j-i)\omega}k_j. \end{aligned}$$

   The positional contribution to the query—key comparison depends on the relative offset $j-i$. RoPE uses several such planes with different frequencies.

The model still has to learn how to use these features. A position encoding makes order available; it does not prescribe a particular linguistic algorithm.

6.3 What the causal mask contributes

Remark 6.3 (the causal mask). A causal mask is not invariant under arbitrary permutations and therefore already breaks the full $S_T$ symmetry. In fact, the finite total order has no nontrivial permutation preserving every relation $j\leq i$. The mask supplies ordered prefix structure, while explicit positional information supplies richer access to absolute positions, offsets, and distances.

A two-position example shows the failure of equivariance. Set all query—key scores equal and use distinct values $v_1,v_2$. Causal attention gives

$$F(v_1,v_2) = \left( v_1,\frac{v_1+v_2}{2} \right).$$

If the two input rows are swapped first, then

$$F(v_2,v_1) = \left( v_2,\frac{v_1+v_2}{2} \right).$$

Swapping the original output instead gives

$$P F(v_1,v_2) = \left( \frac{v_1+v_2}{2},v_1 \right),$$

which is generally different. The first row is special because the causal order permits it to read only itself.

---

7. Attention as kernel smoothing

The formula for an attention row is not unique to neural networks. It has the same normalized-weight form as a classical nonparametric regression estimator.

The word kernel is often used for $\kappa$. Here it means a nonnegative similarity or weighting function. It need not be a symmetric positive-definite kernel in the sense of reproducing-kernel Hilbert spaces.

For attention, set

$$\kappa(q_i,k_j) = \exp\!\left( \frac{\langle q_i,k_j\rangle}{\sqrt{d_k}} \right).$$

This quantity is strictly positive for finite query and key vectors.

Proof

By the definition of row-wise softmax,

$$A_{ij} = \frac{\kappa(q_i,k_j)} {\sum_\ell\kappa(q_i,k_\ell)}$$

in the unmasked case. Substituting this expression into $\sum_jA_{ij}v_j$ gives the first formula. Under a causal mask, $e^{-\infty}=0$ removes the terms with $j>i$ from numerator and denominator. $\square$

This is more flexible than a classical smoother in three ways.

1. The query locations $q_i$ and key locations $k_j$ are learned projections of the current features.
2. The kernel can be asymmetric because $W_Q$ and $W_K$ need not agree.
3. The targets $v_j$ are themselves learned projections rather than fixed observed responses.

The normalization still performs the same mathematical job: it turns nonnegative similarities into barycentric weights.

7.1 Three instructive special cases

Interpretability remark 7.3. A trained attention head may be studied through its learned weighting rule. A row may emphasize the previous token, a recent delimiter, or a position whose token matches the current token. This suggests hypotheses about the head’s algorithm.

---

8. Computational cost and multi-head attention

Attention replaces the fixed context width of an $n$-gram MLP by access to a length-$T$ sequence. That flexibility has a computational price.

The notation $O(g(T,d))$ suppresses fixed multiplicative constants and lower-order terms in an operation count. For one head, the main forward operations are:

operation	output shape	arithmetic cost
$XW_Q$ and $XW_K$	two $T\times d_k$ matrices	$O(Tdd_k)$
$XW_V$	$T\times d_v$	$O(Tdd_v)$
$QK^\top$	$T\times T$	$O(T^2d_k)$
row-wise softmax	$T\times T$	$O(T^2)$
$AV$	$T\times d_v$	$O(T^2d_v)$

For fixed widths, the two products involving the attention matrix are quadratic in sequence length. A causal mask removes roughly half the query—key pairs but does not change the $O(T^2)$ asymptotic order.

Materializing either the score matrix or the attention matrix requires $O(T^2)$ memory, in addition to the $O(T(d_k+d_v))$ storage for projected queries, keys, and values. The quadratic matrix is often the main activation-memory bottleneck for long sequences.

The three projection matrices of one bias-free head contain

$$d(2d_k+d_v)$$

parameters. This count does not depend on $T$. Attention therefore has a useful separation: its learned parameter count can remain fixed as the context grows, even though the arithmetic and temporary memory grow with the number of position pairs.

Aside 8.1 (streaming softmax). A softmax-weighted value sum can be evaluated without storing a complete row of weights. For scores $s_j$, choose $m=\max_js_j$ and write

$$\frac{\sum_je^{s_j}v_j}{\sum_je^{s_j}} = \frac{\sum_je^{s_j-m}v_j}{\sum_je^{s_j-m}}.$$

When keys and values arrive in blocks, maintain a running maximum, a rescaled denominator, and a rescaled weighted-value numerator. If a new block has a larger maximum, multiply the old accumulators by the appropriate exponential factor before adding the new block. This is an exact algebraic reorganization, apart from floating-point rounding. FlashAttention combines this streaming identity with hardware-aware blocking so that the full $T\times T$ matrix need not be stored. It reduces memory traffic and activation storage without changing the general quadratic arithmetic count of dense attention.

8.1 Autoregressive inference and the key—value cache

Training evaluates all rows in parallel. Generation is sequential: after producing a new token, the model appends one new stream row and predicts again. At step $t$, the keys and values of positions $1,\ldots,t-1$ have already been computed and do not change within a fixed layer. A key—value cache, or KV cache, stores them.

For one new query, comparison with $t$ cached keys costs $O(td_k)$ and the weighted value sum costs $O(td_v)$. The cache for one head through length $T$ uses

$$O\bigl(T(d_k+d_v)\bigr)$$

memory. Summed over $T$ generated positions, dense cached attention still takes quadratic arithmetic, but it avoids recomputing the complete old prefix at every step. A naive full-prefix evaluation at every generation step would repeatedly pay a quadratic prefix cost and therefore accumulate a cubic attention cost.

8.2 Preview of multi-head attention

A single head produces one attention distribution per query position. It must use the same distribution for every coordinate of its value output. In Lecture 5, multi-head attention runs $H$ heads with independent query, key, and value projections, concatenates their outputs, and applies an output projection.

When each head has width $d/H$, all heads together have $3d^2$ query, key, and value weights, and the output projection contributes another $d^2$ weights, ignoring biases. Different heads can therefore implement different weighting rules at the same position while the total projection parameter count remains $O(d^2)$. Position-wise MLPs, residual connections, and LayerNorm then complete the transformer block.

---

9. Connection to Project Step 4

Project Step 4 isolates one causal head so that its structural properties can be tested before training a full transformer. This is a useful scientific order: first establish what the implementation must satisfy for every parameter value, then examine the patterns produced by particular parameters.

9.1 From proofs to executable assertions

A minibatch of size $B$ adds a leading coordinate to every shape:

object	batched shape
input $X$	$B\times T\times d$
$Q,K$	$B\times T\times d_k$
$V$	$B\times T\times d_v$
scores and weights	$B\times T\times T$
output	$B\times T\times d_v$

Matrix multiplication acts on the final two coordinates and preserves the batch coordinate. Before running the implementation, writing these shapes beside every line catches most transposition mistakes.

Each main test in Step 4 is a theorem specialized to finite arrays.

assertion	mathematical source
all weights are nonnegative and rows sum to one	Proposition 3.1
the strict upper triangle is zero	Definition 4.1
earlier outputs survive a future perturbation	Proposition 4.3
future input-gradient blocks are zero	Proposition 4.3
raw score scale grows like $\sqrt{d_k}$	Lemma 5.5
row-permuted unmasked inputs give row-permuted outputs	Theorem 6.1

Some equalities are structural and some are numerical. A stored causal mask can make the upper triangle exactly zero. Floating-point row sums are usually only close to one. A perturbation test should compare outputs that execute the same deterministic arithmetic on the same earlier data; gradient tests should allow for the numerical tolerance appropriate to the chosen data type.

9.2 Reading an attention heatmap

A heatmap displays $A_{ij}$ with query position $i$ on one axis and source position $j$ on the other. Several visible features occur before any learning.

• The strict upper triangle is blank because it is masked.
• The first row is forced to be $(1,0,\ldots,0)$; it has only one permitted source.
• Row $i$ distributes unit mass across $i$ permitted columns. Even exchangeable random scores therefore look different near the top and bottom of the triangle.
• Different random seeds produce different columns and diagonals by chance. An untrained pattern is not evidence of a learned linguistic algorithm.

The Gaussian positional experiment supplies a controlled comparison. Its bandwidth is known, so a narrow diagonal or a broad running average has a mathematical explanation. Trained heatmaps can later be compared with these baselines.

9.3 What this head changes from Step 3

The Step 3 neural $n$-gram compresses a fixed concatenated context with one MLP. Its first-layer parameter count grows with the chosen context length. The Step 4 head uses shared projections at every position, and its parameter count is independent of the current $T$. Its attention matrix then selects a different convex combination for every input and query position.

This does not make one head a language model. It has no multi-head output projection, residual connection, normalization, position-wise MLP, or unembedding to vocabulary logits. Lecture 5 assembles those pieces. Step 4’s goal is narrower and foundational: make the data-dependent communication operator correct before multiplying and stacking it.

---

Summary

• Scaled dot-product self-attention projects each stream row into a query, key, and value. Query—key inner products determine a probability distribution; the output is the corresponding value average.

• The unmasked score matrix factors through $\mathbb R^{d_k}$ and has rank at most $d_k$. The attention matrix after row-wise softmax need not have low rank.

• Every attention row lies in a simplex, so each output lies in the convex hull of the available values and can be written as an expectation. The value projection chooses the points being averaged; the query and key projections choose their barycentric weights.

• Causal masking normalizes each row over its prefix. The result is lower triangular and row-stochastic, and output row $i$ depends only on input rows $1,\ldots,i$. This makes parallel teacher-forced training consistent with left-to-right generation.

• Softmax is the gradient of log-sum-exp. Its Jacobian is the categorical covariance matrix $\operatorname{diag}(p)-pp^\top$. Temperature trades score maximization against entropy, and the common-shift direction is invisible.

• Under an idealized independent-coordinate model, an unscaled dot product has variance $d_k\sigma^4$. Dividing by $\sqrt{d_k}$ keeps its variance independent of head dimension and reduces initial softmax saturation.

• Unmasked self-attention is permutation-equivariant as a function of its stream rows. Fixed positional information breaks symmetry in the map from token sequences to stream matrices; the attention formula itself remains equivariant. A causal mask separately supplies ordered prefix structure.

• Attention is a learned Nadaraya—Watson smoother. Uniform scores give a mean, Gaussian positional scores give ordinary sequence smoothing, and low temperature approaches retrieval from the highest-scoring key.

• Dense attention has quadratic arithmetic and, naively, quadratic memory in sequence length. Streaming softmax reduces stored intermediates, and a KV cache avoids recomputing old keys and values during generation.

---

Exercises (paired with Step 4)

A star marks a Project Step 4 task.

1. ★ Reproduce Examples 2.3 and 4.2 without looking at the displayed answers. Label every matrix dimension, compute all three causal output rows, and locate each one in the convex hull of its permitted values.

2. ★ Implement one causal head. Test simplex membership, a zero upper triangle, invariance of earlier outputs under a future perturbation, and zero future derivatives. Match each assertion to a precise result above. Explain why permitted gradients are not guaranteed to be nonzero.

3. ★ Measure raw and scaled score variance for $d_k\in\{16,64,256\}$. Record a typical largest weight in each row and relate the observed saturation to Proposition 5.1 and Lemma 5.5. State why the empirical standard deviations need not equal their theoretical values exactly.

4. Prove Proposition 5.2, including existence and uniqueness of the maximizer. Prove both temperature limits in Remark 5.4, treating multiple score maximizers explicitly.

5. Starting from the $2\times2$ zero score matrix, compare the two procedures in Implementation aside 4.4: mask the scores before softmax, and zero the forbidden weight after an unmasked softmax. Compute the resulting attention matrices and every derivative $\partial A_{1j}/\partial S_{1\ell}$ under each procedure. What changes if post-softmax masking is followed by exact renormalization?

6. Prove Theorem 6.1. Then verify that attention remains equivariant as a function of stream rows even when those rows already contain positional vectors. Locate exactly where equivariance of the composite from tokens fails.

7. ★ Implement the masked Gaussian positional kernel for several bandwidths. Verify both limits in Example 7.4 numerically. Then use the displayed two-dimensional queries and keys to reproduce the same weights by scaled dot-product attention.

8. Prove $\operatorname{rank}(QK^\top)\leq d_k$. Construct a $T>d_k$ example in which the row-wise softmax of $QK^\top$ has rank larger than $d_k$, showing that the rank bound does not pass through softmax.

9. Show that $X\mapsto\operatorname{Attn}(X)$ is Lipschitz on every bounded subset of $\mathbb R^{T\times d}$. Explain why the bilinear score map prevents this argument from giving one global Lipschitz constant on the whole space.

10. Fix keys and values and regard one unmasked output as a function $y(q)$ of its query. Prove that for a query perturbation $h$,

    $$Dy(q)[h] = \frac1{\sqrt{d_k}} \sum_jA_j\langle h,k_j\rangle(v_j-y).$$

    Interpret the centered value $v_j-y$ and determine when this differential is zero for every $h$.

11. Apply Jensen’s inequality to prove the norm bound in Section 3. Give conditions for equality when the Euclidean norm is strictly convex along the relevant value segments.

12. Prove directly, without Jacobians, that a composition of two causal sequence maps is causal. Then derive the corresponding block lower-triangular Jacobian statement by the chain rule.

13. For one head, compare the arithmetic required to generate $T$ tokens (a) by reevaluating the full prefix at every step and (b) with a KV cache. State the cache memory, and distinguish parameter memory from activation memory.

14. Prove the rotary identity

    $$(R_{i\omega}q)^\top(R_{j\omega}k) = q^\top R_{(j-i)\omega}k.$$

    What information about $i$ and $j$ is absent from this expression? Compare that relative encoding with a learned additive positional matrix.

---

Pointers

Bahdanau, Cho, and Bengio, Neural Machine Translation by Jointly Learning to Align and Translate (2015); Vaswani et al., Attention Is All You Need (2017); Nadaraya (1964) and Watson (1964) on kernel regression; Phuong and Hutter, Formal Algorithms for Transformers (2022); Elhage et al., A Mathematical Framework for Transformer Circuits (2021); Su et al., RoFormer (2021); and Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022). See the resources page, Project Step 4, and Lecture 5.