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, and the target space $\mathbb{R}^d$ is called the embedding space (or, especially outside the language-modelling literature, the latent space). 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 of }Q\text{)}, \text{key (column of }K^\top\text{)}) = (q_i, k_j)$. When the resulting row-column product (normalized, and called the score) is large, the residual vector $\mathbf x_i$ at the query position $i$ will attend strongly to the residual vector $\mathbf{x}_j$ at the key position $j$, meaning the $\mathbf{x}_j$ will strongly influence $\mathbf{x}_i$ 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 $\mathbf{x}_j$ directly to $\mathbf{x}_i$. Instead, we add $\mathbf{x}_j V$. The role of $V$ is to interpret a stream vector in some way (by a linear transformation) before it is added in to $\mathbf{x}_i$.

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.

Remark 2.6 (why divide by $\sqrt{d_k}$). Suppose the coordinates of $q,k\in\mathbb R^{d_k}$ are independent across coordinates and between the two vectors, with mean zero and variance $\sigma^2$. Then $\langle q,k\rangle=\sum_{m=1}^{d_k}q_mk_m$ is a sum of $d_k$ independent terms, each with mean $\mathbb E[q_m]\mathbb E[k_m]=0$ and variance $\mathbb E[q_m^2]\mathbb E[k_m^2]=\sigma^4$, so $\langle q,k\rangle$ has expectation $0$ and variance $d_k\sigma^4$. The scaling by $1/\sqrt{d_k}$ in the definition of $S$ reduces the variance to $\sigma^4$, independent of the dimension. Without this scaling, the variance can become quite high, and as a result softmax saturates: the result of softmax tends to be near a corner of the simplex. Saturation makes training by gradient descent more difficult, since gradients are small (Proposition 6.1 below computes the Jacobian of softmax; it tends to zero as the softmax output approaches a corner). This is an idealized calculation about initialization, not an invariant of training: queries and keys in a trained model need not have independent, mean-zero coordinates.

---

3. Attention as averaging

Attention scores may be positive or negative and have no fixed scale. After softmax, however, every row belongs to the set

$$\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 last display in the proof can be interpreted probabilistically. For each query position $i$, let $J_i$ be a random variable taking values in $1, \ldots, T$, given by

$$\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 geometry of the queries 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).$$

---

4. Causal masking

There’s one important aspect we have so far ignored, which requires a modification. 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. To avoid this, we use a method called causal masking. 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$ of $\operatorname{Attn}(X)$ depends on $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 of Proposition 3.1 can be strengthened to:

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

Remark 4.4 (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. Temperature

The entropy of $a\in\Delta^{T-1}$ is

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

with $0\log0=0$. The entropy is large when the distribution is spread-out and small when it is concentrated. The following proposition makes this precise.

Proof

Each term $-a_j\log a_j$ is nonnegative, since $0\leq a_j\leq1$, and it vanishes exactly when $a_j\in\{0,1\}$. Hence $H(a)\geq0$, with equality exactly when every coordinate is $0$ or $1$; as the coordinates sum to $1$, that means $a$ is a vertex.

For the upper bound, let $J=\{j:a_j>0\}$. Since $\log$ is strictly concave, Jensen’s inequality gives

$$H(a) = \sum_{j\in J}a_j\log\frac1{a_j} \leq \log\sum_{j\in J}a_j\cdot\frac1{a_j} = \log|J| \leq \log T,$$

with equality in the first inequality exactly when the numbers $1/a_j$, $j\in J$, are all equal, and in the second exactly when $J$ is all of $\{1,\ldots,T\}$. Both hold exactly when $a$ is uniform. $\square$

Let $\tau>0$. The function $\operatorname{softmax}(s/\tau)$, as a function of $s$, is said to have temperature $\tau$. The next proposition says that $\operatorname{softmax}(s/\tau)$, acting on scores $s$, produces a probability distribution which has mass on tokens with large scores, but is also spread-out. The amount of emphasis on being spread out is determined by the temperature. With temperature $1$, we are using standard softmax; with a high temperature, we get more random variation in the final probability distribution.

Proof

We will use the notion of KL divergence. 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 (Lecture 1, Theorem 7.2) states that this quantity is nonnegative and is zero exactly when $a=p$.

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$

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.

---

6. Log-sum-exp and computation remarks

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 6.1. This is the same softmax geometry developed in Lecture 2; here each attention row supplies its own logit vector.

---

7. Positional embeddings

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. In this section we discuss how we embed positional information into the input stream.

7.1 Permutation equivariance

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, defined by the convention

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

i.e. $P_\pi$ acts from the left to permute the rows of $X$ by $\pi^{-1}$.

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$.

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$

A position-wise map applies one function $\phi:\mathbb R^d\to\mathbb R^{d'}$, with shared parameters, separately to every row. Such a map is equivariant as well, since $\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.

Proof

Without positional information, the map $\iota_0$ from token sequences to stream matrices, defined below, is equivariant. Theorem 7.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. $\square$

To see exactly where positional information enters, compare the embedding maps with and without positions.

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.

Thus, supplying position-dependent information in the embedding or attention scores breaks this symmetry in the composite map from token sequences to outputs.

7.2 Types of positional embeddings: learned, sinusoidal, and rotary

Here are some of the standard ways to define positional vectors. Recall from Section 1 that a system of additive positional vectors is a map $t\mapsto p_t\in\mathbb R^d$, defined for $t=1,\ldots,T_{\max}$, where $T_{\max}$ is the longest context the model supports. The column vector $p_t$ is called the $t$-th positional embedding, and the stream is initialized with $X_t=e(x_t)+p_t$. In matrix form, $p\in\mathbb R^{T\times d}$ has rows $p_t^\top$.

1. Learned positional embeddings. It is possible to consider the $p_t$ to be part of the parameters and learn these vectors, exactly as the token embeddings are learned. This is what the course project does.

2. Sinusoidal positional embeddings. Use fixed sinusoidal coordinates. At angular frequency $\omega$, the two-coordinate feature vector defined by

   $$r_t(\omega) := \bigl(\sin(\omega t),\cos(\omega t)\bigr), \quad t=1, \ldots, T_{\max},$$

   satisfies

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

   Observe that using such vectors as positional embeddings, $p_t := r_t(\omega)$, makes relative offsets accessible by taking inner products.

   In higher dimensions, the standard fixed construction stacks such pairs at a range of frequencies. Indexing the $d$ coordinates of $p_t$ from $0$, it sets, for $m=0,1,\ldots,d/2-1$,

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

   so that $(p_{t,2m},p_{t,2m+1})=r_t(\omega_m)$ with $\omega_m=10000^{-2m/d}$: the frequencies range from $1$ down to nearly $1/10000$.

3. 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.

These designs make relative position available to the model, but the model still has to learn how to use these features.

Remark 7.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.

---

8. Expressivity of attention

What weighting rules can an attention head express? To describe them uniformly, 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$

Remark 8.2 (Nadaraya—Watson estimation). The formula in Proposition 8.1 is reminiscent of, but more flexible than, a classical nonparametric regression estimator. Given input—value pairs $(z_j,v_j)$ and a nonnegative weighting function $\kappa(z,z_j)$, the Nadaraya—Watson regression estimate at $z$ is

$$\widehat f(z) = \frac{\sum_j\kappa(z,z_j)v_j} {\sum_j\kappa(z,z_j)},$$

provided the denominator is positive. The word kernel is often used for $\kappa$; here it means only a nonnegative weighting function and has nothing to do with Hilbert spaces. Attention is more flexible in that the locations $q_i,k_j$ and the targets $v_j$ are all learned projections of the current features, and the kernel can be asymmetric because $W_Q$ and $W_K$ need not agree.

8.1 Three instructive special cases

8.2 Heatmaps

A heatmap displays $A_{ij}$ with query position $i$ on one axis and source position $j$ on the other, with darker cells for larger weights. Here are the three examples above, drawn for $T=12$.

Several visible features occur in every causal heatmap, even before any learning.

• The strict upper triangle is blank (grey above) 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.

---

9. Computational cost

How expensive is attention? To measure the complexity of an algorithm, we count the arithmetic operations it performs: each scalar addition, subtraction, multiplication, division, or exponential counts as one operation. For example, multiplying an $m\times n$ matrix by an $n\times p$ matrix by the usual formula takes $mnp$ multiplications and $m(n-1)p$ additions, so about $2mnp$ operations. We are interested in how this count grows with $T$, the length of the context, and $d$, the dimension of the embedding space, along with the head widths $d_k$ and $d_v$. Exact counts depend on implementation details, so we record only the growth rate, using big-O notation.

Thus $O(g)$ suppresses fixed multiplicative constants and lower-order terms: for instance $2T^2d+5Td+T=O(T^2d)$. 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.

Storing 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 dependence on $T$, namely $O(T^2)$, 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.

9.1 Autoregressive generation and the key—value cache

Training evaluates all rows in parallel: the whole training sequence is known in advance, and one forward pass produces the loss at every position (Remark 4.4). Generation is sequential.

Consider one attention head in one layer during this process, and suppose the sequence has grown to length $t$. Its input rows $X_1,\ldots,X_{t-1}$ are the same as they were at the previous step: by causality (Proposition 4.3, applied to every layer below; see Exercise 12), the stream at a position never depends on later tokens, so appending $x_t$ changes nothing at positions $1,\ldots,t-1$. Consequently the keys $k_j=W_K^\top X_j$ and values $v_j=W_V^\top X_j$ for $j<t$ are unchanged as well, and so are the earlier output rows. Only the new row $X_t$, and hence $q_t$, $k_t$, and $v_t$, must be computed, and the only output row needed is

$$\operatorname{Attn}(X)_t = \sum_{j=1}^tA_{tj}v_j, \qquad A_{tj} = \frac{e^{\langle q_t,k_j\rangle/\sqrt{d_k}}} {\sum_{\ell=1}^te^{\langle q_t,k_\ell\rangle/\sqrt{d_k}}}.$$

Proof

With the cache, step $t$ computes three projections of one row, $O(d(2d_k+d_v))$ operations; compares $q_t$ with $t$ keys, $O(td_k)$; normalizes, $O(t)$; and forms the weighted value sum, $O(td_v)$. Summing over $t=1,\ldots,T$ and using $\sum_{t\leq T}t=O(T^2)$ gives the first count. Without the cache, step $t$ recomputes the full prefix, whose score and value products cost $O(t^2(d_k+d_v))$ by the table above, and $\sum_{t\leq T}t^2=O(T^3)$. The cache holds $t-1\leq T$ vectors of length $d_k$ and as many of length $d_v$. $\square$

Summed over $T$ generated positions, dense cached attention still takes quadratic arithmetic, but it avoids recomputing the complete old prefix at every step. Note also that the cache is activation memory, which grows with the context length, rather than parameter memory, which does not: the cache for a long conversation can be far larger than the head’s $d(2d_k+d_v)$ parameters. Multi-head attention, which runs several heads side by side, is the subject of Lecture 5.

---

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 6.1 and Remark 2.6. 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 two procedures: mask the scores before softmax, as in Definition 4.1, 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 7.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 8.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.