Attention, Queries, Keys, and Values

Definition Let the rows of a matrix $X$ in $\mathbb{R}^{n \times d}$ be the vector representations of the $n$ positions of a sequence, one row per position. Self-attention computes three matrices from $X$ with learned projection matrices,

where $W_Q$ and $W_K$ have shape $d \times d_k$ and $W_V$ has shape $d \times d_v$. Row $i$ of $Q$ is the query vector $\mathbf{q}_i$ of position $i$, row $j$ of $K$ is the key vector $\mathbf{k}_j$ of position $j$, and row $j$ of $V$ is its value vector $\mathbf{v}_j$. A key and a value with the same index belong to the same position, so $K$ and $V$ always have the same number of rows. Queries and keys must have the same width $d_k$, because they are compared by dot products, while the value width $d_v$ may differ. In cross-attention the queries come from one sequence and the keys and values from another, so the number of queries $n_q$ and the number of keys $n_k$ can differ. In self-attention $n_q = n_k = n$.

Definition Let $Q$ in $\mathbb{R}^{n_q \times d_k}$ hold $n_q$ query vectors as rows, $K$ in $\mathbb{R}^{n_k \times d_k}$ hold $n_k \ge 1$ key vectors, and $V$ in $\mathbb{R}^{n_k \times d_v}$ hold the matching value vectors. Scaled dot-product attention first forms the score matrix

of shape $n_q \times n_k$, whose entry $S_{ij} = \mathbf{q}_i^{\top}\mathbf{k}_j / \sqrt{d_k}$ compares query $i$ with key $j$. It applies the softmax to each row of $S$ separately, which gives the attention weight matrix $A$ of the same shape, and returns

an $n_q \times d_v$ matrix whose row $i$ is the output for query $i$. In one formula the output is $\operatorname{softmax}(QK^{\top}/\sqrt{d_k})V$ with the softmax taken along each row, over the keys. The products are defined exactly when $Q$ and $K$ have the same number of columns and $K$ and $V$ have the same number of rows. The factor $1/\sqrt{d_k}$ has a specific reason. If the $2d_k$ entries of a query $\mathbf{q}$ and a key $\mathbf{k}$ are independent random variables with mean $0$ and variance $1$, then $\mathbf{q}^{\top}\mathbf{k}$ has mean $0$ and variance $d_k$, so the scaled score has variance $1$. Learned vectors need not satisfy these assumptions, so the scaling keeps typical scores moderate without guaranteeing any range.

Definition The attention weights are the entries of $A$. Without a mask, for each query $i$ and key $j$,

With finite scores every weight is positive and every row adds to $1$, so the output $\mathbf{o}_i = \sum_j A_{ij}\mathbf{v}_j$ is a convex combination of the value vectors. Each coordinate of $\mathbf{o}_i$ therefore lies between the smallest and the largest value of that coordinate among the value vectors. Equal scores in a row give equal weights. Adding the same constant to every score of a row changes nothing, which is why implementations subtract the row maximum before exponentiating. A weight $A_{ij}$ says how much value vector $j$ contributes to the output of query $i$ in one head of one layer. It is not a measure of how much the original token at position $j$ influences the model’s final prediction.

Definition An attention mask is a matrix $M$ of shape $n_q \times n_k$ with $M_{ij} = 1$ when query $i$ may use key $j$ and $M_{ij} = 0$ when it may not. Masked attention sets each masked score to $-\infty$ before the row softmax, with $\exp(-\infty) = 0$. This is the same as taking the softmax over the allowed keys only,

Masked weights are exactly $0$, the allowed weights of each row add to $1$, and output $i$ is a convex combination of the allowed values only. The formula needs at least one allowed key in every row. In a fully masked row the denominator is $0$, so no probability distribution exists. The convention in this course is to reject such a row with an error. Another system might return a zero vector with a flag marking the row as invalid, but whichever convention is used must be stated and its downstream use checked. A common use is the padding mask. When sequences of different lengths are padded to one length to form a batch, the mask removes the padding keys, and the outputs at padded query positions are left out of the loss.

Definition For self-attention over a sequence read from left to right, with positions $1, \dots, n$, the causal mask allows key $j$ for query $i$ exactly when $j \le i$. Its matrix is lower triangular, with ones on and below the diagonal, so $A_{ij} = 0$ for every $j > i$. The diagonal is allowed, because the token at position $i$ is already known when position $i$ is used to predict the next token. Every row allows at least its own position, so a causal mask alone never produces a fully masked row. When every attention layer uses a causal mask and all other operations act on each position separately, the output at position $i$ depends only on the inputs at positions $1$ to $i$.

Definition Multi-head attention with $h$ heads runs $h$ scaled dot-product attentions side by side, each with its own projections. In self-attention, head $r = 1, \dots, h$ forms $Q^{(r)} = XW_Q^{(r)}$, $K^{(r)} = XW_K^{(r)}$ and $V^{(r)} = XW_V^{(r)}$ and computes

an $n_q \times d_v$ matrix. The head outputs are concatenated along the feature axis into an $n_q \times hd_v$ matrix, which a learned output matrix $W_O$ of shape $hd_v \times d_{\text{out}}$ maps back to one vector per query,

Usually every head uses the same mask, and a common choice is $d_k = d_v = d/h$ and $d_{\text{out}} = d$, so the result has the model’s width again. Because each head has its own projections, different heads can give the same query different weights.