Last modified: 2026-10-18
Source
Adapted from Logan Sizemore’s Spring 2025 lectures on attention and transformers.
The models in Neural Networks take a fixed-length feature vector and produce an output. Language does not arrive that way: a sentence is a sequence, its length varies, and what a word means depends on the words around it. This chapter builds the mechanism that handles all three — attention — and then assembles it into the transformer block, the building block of modern language models [1].
A language model assigns probabilities to sequences of words
In practice this is almost always posed as next-token prediction: given the tokens so far, give a distribution over what comes next
lamb
A token is the unit the model actually predicts. It is usually a subword fragment rather than a whole word, so that a fixed vocabulary can spell any input
Watch
Sanderson explains how a model predicts the next word, from the start [2, Ch. 5].
The earliest language models estimate next-token probability from counts
An \(n\)-gram model conditions on the previous \(n-1\) tokens only:
\[P(w_t \mid w_1, \dots, w_{t-1}) \approx P(w_t \mid w_{t-n+1}, \dots, w_{t-1})\]
For a bigram (\(n=2\)) model this is a table of transition probabilities, one row per vocabulary item
Two problems follow immediately:
The table grows too large. A vocabulary of \(V\) tokens needs \(V^n\) entries, almost all of which are never observed
The context is too short. With \(n=2\), generation degenerates — “the river flows to the sea and the river flows to the sea and …” — because the model cannot see that it has already produced this text
Enlarging \(n\) trades the second problem for the first. Instead, we want a representation that can look arbitrarily far back without needing \(V^n\)
Read
GeeksforGeeks explains \(n\)-gram language models [3].
\[\text{"lamp"} \;\longmapsto\; \begin{bmatrix} 0.6 & 0.1 & 0.8 & 0.3 & 0.5 \end{bmatrix}^\top\]
Such a vector is a word embedding. The model learns it rather than a person designing it. The geometry it learns carries meaning: tokens used in similar contexts are close to one another
That geometry makes the analogy arithmetic work: the vector offset from Paris to France points in roughly the same direction as the offset from Rome to Italy
A fixed lookup gives each word form exactly one vector, so every sense of a homonym has to share it:
Figurative use has the same problem without the homonymy. In “she is a real shark”, the word names an animal but describes a person, and the animal’s vector says nothing about that
So we want a contextualized embedding: a vector for this token in this sentence. We build it by moving the static embedding towards the other tokens present
Watch
Watch
Sanderson explains word embeddings, from 12:27 [2, Ch. 5].
Read
GeeksforGeeks explains word embeddings [6].
\[c_t = \frac{1}{2k+1}\sum_{j=t-k}^{t+k} x_j\]
Neighbors often carry the needed information. In “the pitcher threw a fastball”, the word immediately after pitcher tells us it is a person rather than a jug
But the window weights tokens by position, when we want to weight them by relevance. In “the crane by the lake”, the token that disambiguates crane is lake, three positions away, while the two tokens in between contribute nothing
Attention replaces the fixed positional weights with weights the model computes from the content of the tokens themselves [7]
Take a sentence of \(T\) tokens with embeddings \(x_1, \dots, x_T\), each in \(\reals^d\), and build the contextualized embedding for position \(i\).
\[s_{ij} = \dprod{x_i}{x_j} \tag{1}\]
The inner product is positive when \(x_j\) has a positive projection onto \(x_i\), zero when the two are orthogonal, and negative when the projection is negative — so a large \(s_{ij}\) means the two tokens point the same way in embedding space
In “the crane by the lake”, \(s_{\texttt{crane},\texttt{lake}}\) should come out largest
Watch
Starmer explains cosine similarity [8], the inner product with the lengths divided out.
\[\attnweight_{ij} = \frac{\expf{s_{ij}}}{\sumto{k}{T} \expf{s_{ik}}}\]
Read
GeeksforGeeks explains the softmax function [9].
\[c_i = \sumto{j}{T} \attnweight_{ij}\, x_j\]
crane moves towards lake
The three roles a token plays in those steps have names:
The token doing the asking is the query
The tokens being compared against are the keys
The vectors that get averaged are the values
Written with all \(T\) tokens stacked into the rows of \(Q\), \(K\) and \(V\), the whole computation is one expression:
\[\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\paren{\frac{Q\tp{K}}{\sqrt{d_k}}} V\]
Steps 1 to 3 have none. Every quantity is a fixed function of the input embeddings, so there is nothing for gradient descent to learn
We fix this by giving the query, key and value roles their own linear layers:
\[Q = X W^Q, \qquad K = X W^K, \qquad V = X W^V\]
Now \(W^Q, W^K, W^V\) are the parameters. The model learns what to compare (the query and key projections) separately from what to pass along (the value projection). That is why we use three matrices rather than one
Together, the projections and the three steps form one self-attention block. The block is “self” attention because the queries, keys and values all come from the same sequence
Watch
Watch
Read
GeeksforGeeks explains self-attention [12].
One block learns one notion of relevance. But crane relates to lake for a different reason than a verb relates to its subject, and a single attention distribution has to average those notions together
So we run \(h\) blocks in parallel, each with its own \(W^Q_m, W^K_m, W^V_m\), and concatenate their outputs:
\[\operatorname{MultiHead}(X) = \begin{bmatrix} \operatorname{head}_1 & \cdots & \operatorname{head}_h \end{bmatrix} W^O\]
Each head learns a different kind of relationship. In a trained model the heads are interpretable often enough to be worth looking at: some track syntactic dependencies, some track coreference, some attend almost entirely to the previous token
In practice we give each head \(d/h\) dimensions rather than \(d\). Then multi-head attention costs about the same as single-head attention on the full width
Watch
Read
GeeksforGeeks explains multi-head attention [13].
\[\operatorname{FFN}(z) = \max(0,\, z W_1 + b_1) W_2 + b_2\]
\[z = \operatorname{LayerNorm}(x + \operatorname{MultiHead}(x))\] \[y = \operatorname{LayerNorm}(z + \operatorname{FFN}(z))\]
The residual connections make depth trainable, for the same reason as in Backpropagation: they give the gradient a path to the input that does not pass through the sublayer’s weights
That whole assembly is one transformer block, the T in Generative Pre-trained Transformer (GPT)
Attention is permutation-equivariant: permute the rows of \(X\) and every \(c_i\) is permuted with them, unchanged. The mechanism has no notion of order
That is a serious problem for language, where “the owl saw the cat” and “the cat saw the owl” use the same tokens. The standard fix adds a positional encoding to each embedding before the first block. Then a token’s vector carries where it sits as well as what it is
Watch
Watch
The first lecture of Stanford’s Transformers and Large Language Models course (CME 295) goes from tokenization and embeddings through attention to the full transformer [15], Lecture 1.
Read
GeeksforGeeks explains:
Read
The University of Washington’s Machine Learning course (CSE 446) slides cover attention and transformers [19], Lectures 22 and 23.
\[X \;\rightarrow\; \text{block} \;\rightarrow\; \text{block} \;\rightarrow\; \cdots \;\rightarrow\; \text{block} \;\rightarrow\; H\]
A stack of them is the main body of a large language model. To predict the next token, we take the final-layer vector at the last position and pass it through one more linear layer with a softmax, producing a distribution over the whole vocabulary
Training is next-token prediction on very large quantities of text, most of it collected from the web. The models are large: GPT-3 has 175 billion trainable parameters [20]. Parameter counts for later commercial models are mostly undisclosed, so read any figure you meet for them as an unverified estimate, and check where it came from before repeating it
The model learns more than language from that corpus. It also absorbs the information in the text, the copyrighted material in the text, and the biases in the text. We take up this point in Ethics and Fairness
Watch
Starmer explains decoder-only transformers [21], the kind behind ChatGPT.
Read
GeeksforGeeks introduces large language models [22].
Nothing in the attention mechanism is specific to words. It needs a set of vectors and a reason to compare them, which many kinds of data supply:
Images, cut into patches, each patch embedded as a token
Audio, as a sequence of spectrogram frames
Proteins, as a sequence of amino-acid residues
As a result, one architecture now does work that used to need a different model family per modality
Read
GeeksforGeeks explains the vision transformer [23], which embeds image patches as tokens.
The code in this section computes scaled dot-product attention for a four-token sentence with a three-dimensional embedding, and prints the attention distribution for the first token. We chose the embeddings so that the first and last tokens point in a similar direction (the crane/lake case). The weight the first token puts on the last should then be the largest of its weights on the other three tokens.
import numpy as np
def softmax_rows(scores):
shifted = scores - scores.max(axis=1, keepdims=True)
weights = np.exp(shifted)
return weights / weights.sum(axis=1, keepdims=True)
def attention(queries, keys, values):
scores = queries @ keys.T / np.sqrt(keys.shape[1])
weights = softmax_rows(scores)
return weights @ values, weights
# "the crane by the lake": token 0 and token 3 point the same way.
# Hand-picked 3-dimensional embeddings keep the arithmetic checkable by eye;
# a trained model learns embeddings with hundreds of dimensions.
x = np.array([[0.8, 0.1, 0.4], # crane
[0.0, 1.0, 0.1], # by
[0.1, 0.9, 0.0], # the
[0.7, 0.2, 0.5]]) # lake
context, weights = attention(x, x, x)
print(np.round(weights[0], 3))[0.298 0.203 0.206 0.293][0.464 0.476 0.286]softmax_rows <- function(scores) {
shifted <- sweep(scores, 1, apply(scores, 1, max))
weights <- exp(shifted)
sweep(weights, 1, rowSums(weights), FUN = "/")
}
attention <- function(queries, keys, values) {
scores <- queries %*% t(keys) / sqrt(ncol(keys))
weights <- softmax_rows(scores)
list(context = weights %*% values, weights = weights)
}
# "the crane by the lake": token 1 and token 4 point the same way.
# Hand-picked 3-dimensional embeddings keep the arithmetic checkable by eye;
# a trained model learns embeddings with hundreds of dimensions.
x <- rbind(
c(0.8, 0.1, 0.4), # crane
c(0.0, 1.0, 0.1), # by
c(0.1, 0.9, 0.0), # the
c(0.7, 0.2, 0.5) # lake
)
result <- attention(x, x, x)
print(round(result$weights[1, ], 3))[1] 0.298 0.203 0.206 0.293[1] 0.464 0.476 0.286function softmax_rows(scores)
shifted = scores .- maximum(scores, dims = 2)
weights = exp.(shifted)
weights ./ sum(weights, dims = 2)
end;
function attention(queries, keys, values)
scores = queries * keys' / sqrt(size(keys, 2))
weights = softmax_rows(scores)
(context = weights * values, weights = weights)
end;
# "the crane by the lake": token 1 and token 4 point the same way.
# Hand-picked 3-dimensional embeddings keep the arithmetic checkable by eye;
# a trained model learns embeddings with hundreds of dimensions.
x = [0.8 0.1 0.4; # crane
0.0 1.0 0.1; # by
0.1 0.9 0.0; # the
0.7 0.2 0.5 # lake
];
result = attention(x, x, x);
println(round.(result.weights[1, :], digits = 3))[0.298, 0.203, 0.206, 0.293][0.464, 0.476, 0.286]The function softmax_rows subtracts each row’s maximum before exponentiating. That subtraction changes nothing mathematically, because the constant cancels between numerator and denominator. Without it, a score of a few hundred overflows to Inf and the whole row becomes NaN. Every softmax used in practice does this.
Vaswani et al. [24] — the paper the transformer block comes from
Lin et al. [25] — a survey of the variants that followed, free on arXiv:2106.04554
Amidi and Amidi [15] — Stanford’s course on transformers and large language models, whose first lecture, Transformers (1 h 44 min), goes from tokenization and embeddings through attention to the full transformer; its syllabus lists what each later lecture covers
Starmer [26], on word embeddings, attention and transformers, illustrated, with PyTorch tutorials