Fall 2026
  • Discord
  • Gradescope
  • Syllabus
  • Spring 2026

On this page

  • Queries, Keys, and Values
  • Score Scaling
  • Permutation Equivariance

Self-Attention

To interpret a word in a sentence, a model may need information from words several positions away. Self-attention computes a weight for each pair of tokens and uses those weights to combine their representations. We will derive the queries, keys, and values used in this calculation.

How does attention choose which tokens contribute to an output?


Queries, Keys, and Values

A tokenizer turns a sequence into token identities, and a learned embedding table turns those identities into vectors \(\mathbf x_1,\ldots,\mathbf x_n\in\mathbb R^d\). Stack them as the rows of \(\mathbf X\in\mathbb R^{n\times d}\). At this point each row describes one token without using the others.

Each token produces three vectors through shared learned projections: \(\mathbf q_i,\mathbf k_i\in\mathbb R^{d_k}\) and \(\mathbf v_i\in\mathbb R^{d_v}\). The query is compared with other tokens’ keys. Their match determines how much of each value enters the output. Queries and keys need the same dimension for comparison, while values may have a different dimension.

For query position \(i\in\{1,\ldots,n\}\), compare it with each key position \(j\in\{1,\ldots,n\}\) by a scaled inner product, \[ s_{ij}=\frac{\langle\mathbf q_i,\mathbf k_j\rangle}{\sqrt{d_k}}. \] Softmax across \(j\) turns arbitrary scores into nonnegative attention weights that sum to one, \[ a_{ij}=\frac{e^{s_{ij}}}{\sum_{\ell=1}^n e^{s_{i\ell}}}, \] and the output is the weighted value sum \[ \mathbf o_i=\sum_{j=1}^n a_{ij}\mathbf v_j. \] The vectors are fixed illustrative representations. Changing their order moves each token’s query, key, and value together.

A query scores every key, softmax normalizes the scores, and those probabilities weight the value vectors. Reordering all token representations together preserves the corresponding output.

Softmax chooses the coefficients. The values determine what those coefficients combine.

For all tokens together, collect the projected vectors in row-stacked matrices: \[ \mathbf Q=\mathbf X\mathbf W_Q,\qquad \mathbf K=\mathbf X\mathbf W_K,\qquad \mathbf V=\mathbf X\mathbf W_V, \] where \(\mathbf W_Q,\mathbf W_K\in\mathbb R^{d\times d_k}\) and \(\mathbf W_V\in\mathbb R^{d\times d_v}\). The projections are shared across positions, but the mixing weights are recomputed from the current input.

All rows together give the standard formula \[ \operatorname{Attention}(\mathbf X) =\operatorname{softmax}\!\left( \frac{\mathbf Q\mathbf K^\top}{\sqrt{d_k}} \right)\mathbf V, \] where softmax acts row by row. Every output can use every value in one layer, regardless of distance.

The score matrix \(\mathbf Q\mathbf K^\top\) is \(n\times n\), with one query-key comparison per entry. Softmax makes each row a distribution, and multiplication by \(\mathbf V\) produces \(n\times d_v\) outputs. Attention is therefore a data-dependent mixing matrix applied to the input’s values.

The output is a contextual vector assembled from the values. A key that aligns more strongly with the query receives more weight, so its value contributes more to the output. Gradient descent learns the matching and transmitted representations together.

The weights specify each value vector’s contribution to this layer’s sum. Explaining their effect on the final output also requires accounting for later layers and residual paths.

Score Scaling

At a typical initialization, model query and key coordinates as independent with mean zero and variance one. Their dot product sums \(d_k\) random products, so its scale grows with head width.

Claim: Under these assumptions, the unscaled dot product has mean zero and variance \(d_k\), while division by \(\sqrt{d_k}\) gives variance one.

Proof

For \(\mathbf q,\mathbf k\in\mathbb R^{d_k}\), write the dot product as \[ \langle\mathbf q,\mathbf k\rangle=\sum_{r=1}^{d_k}q_rk_r. \] For each coordinate \(r=1,\ldots,d_k\), let \(Z_r=q_rk_r\). Independence of \(q_r\) and \(k_r\) gives \[ \mathbb E[Z_r]=\mathbb E[q_r]\mathbb E[k_r]=0 \] and \[ \operatorname{Var}(Z_r) =\mathbb E[q_r^2k_r^2]-\mathbb E[Z_r]^2 =\mathbb E[q_r^2]\mathbb E[k_r^2]=1. \] For \(r\ne s\), the coordinate pairs are independent, so \(\operatorname{Cov}(Z_r,Z_s)=0\). The variance-of-a-sum identity therefore gives \[ \mathbb E[\langle\mathbf q,\mathbf k\rangle]=0, \qquad \operatorname{Var}(\langle\mathbf q,\mathbf k\rangle) =\sum_{r=1}^{d_k}\operatorname{Var}(Z_r)=d_k. \] Finally, variance scales quadratically with a constant, so \[ \operatorname{Var}\!\left( \frac{\langle\mathbf q,\mathbf k\rangle}{\sqrt{d_k}} \right) =\frac1{d_k}\operatorname{Var}(\langle\mathbf q,\mathbf k\rangle)=1. \] This proves the score-scale statements using only the earlier sum rule.

The typical unscaled score magnitude therefore grows like \(\sqrt{d_k}\).

Unscaled dot-product scores grow with head width, while scaled scores remain at a stable scale.

Large score gaps make softmax nearly one-hot, where small changes to losing scores have little effect. Dividing by \(\sqrt{d_k}\) keeps initial score variance fixed across head widths without preventing training from later creating sharp attention.

Permutation Equivariance

Nothing in \(\mathbf q_i\), \(\mathbf k_j\), or their inner product uses the position indices \(i\) and \(j\).

Claim: For every permutation matrix \(\mathbf P\in\mathbb R^{n\times n}\), content-only self-attention is permutation-equivariant: \[ \operatorname{Attention}(\mathbf P\mathbf X) =\mathbf P\operatorname{Attention}(\mathbf X). \]

Proof

Write \(\mathbf X'=\mathbf P\mathbf X\), and let \(\rho\) be the permutation of \(\{1,\ldots,n\}\) such that \(\mathbf x'_i=\mathbf x_{\rho(i)}\). Shared projections permute each token’s query, key, and value together, so \(s'_{ij}=s_{\rho(i),\rho(j)}\) and \(\mathbf v'_j=\mathbf v_{\rho(j)}\). For \(i,j\in\{1,\ldots,n\}\), softmax gives \[ a'_{ij} =\frac{e^{s_{\rho(i),\rho(j)}}} {\sum_{\ell=1}^n e^{s_{\rho(i),\rho(\ell)}}} =a_{\rho(i),\rho(j)}, \] because the denominator contains the same terms in a different order. Reindexing the weighted value sum by \(m=\rho(j)\) now gives \[ \begin{aligned} \mathbf o'_i &=\sum_{j=1}^n a_{\rho(i),\rho(j)}\mathbf v_{\rho(j)}\\ &=\sum_{m=1}^n a_{\rho(i),m}\mathbf v_m =\mathbf o_{\rho(i)}. \end{aligned} \] Thus each token receives the same weighted combination after reordering, and \(\operatorname{Attention}(\mathbf P\mathbf X)=\mathbf P\operatorname{Attention}(\mathbf X)\).

Reversing the token order moves a query’s contributors to new positions, but its output stays the same. This is useful for sets but insufficient for language, where order changes meaning. Training cannot recover an ordering variable that the model never receives, so position must be added to the representation or attention rule.

A token’s query scores the keys, softmax turns those scores into routing weights, and the associated values form a contextual representation. Scaling stabilizes the scores across widths. Positional information supplies the order missing from content alone.