The Transformer, Explained

This is the engine-room post. The first two established what a Transformer is and why it matters; here we replace intuition with explicit equations, tracing exactly what attention does to a set of vectors and how the surrounding machinery keeps the deep stack trainable.

The payoff is that words like 'query,' 'key,' 'multi-head,' and 'positional encoding' stop being vocabulary and become operations you have actually computed. Attention is far simpler than its reputation once you write down the arithmetic, and doing the trace once is worth more than reading ten metaphors about it.

The query, key, value framing is the heart of attention, and the search analogy makes it click. Each token's vector is projected by three separate learned matrices into a query, a key, and a value. The query represents what this token is looking for, each key represents what a token has to offer, and each value is the actual content retrieved when a query and key match strongly.

Because all three projections come from the same input sequence, this is called self-attention — every token simultaneously issues a query and offers a key and value. The mechanism is essentially a soft, differentiable lookup: rather than retrieving one matching entry, it retrieves a weighted blend of all values, with the weights determined by how well each key answers the query.

Scaled dot-product attention is four clean steps. First, multiply the query matrix by the transpose of the key matrix to get a score for every token pair — how well each query matches each key. Second, divide those scores by the square root of d_k to keep them in a healthy range. Third, apply softmax across each row so the scores for one query become a set of weights that sum to one. Fourth, multiply those weights by the value matrix to get each token's new, context-mixed representation.

Laying it out as four steps clarifies what each part does: the dot product measures relevance, the scaling stabilizes it, the softmax turns relevance into a probability-like weighting, and the final multiply performs the actual information retrieval. Every attention mechanism you will meet is a variation on these four operations.

This from-scratch implementation makes the formula undeniable. It reads the key dimension d_k from the query shape, computes the scaled score matrix as Q times K-transpose divided by the square root of d_k, applies softmax along the last axis to turn each row into weights, and returns the weights times V.

Reading this, notice how compact the core of a Transformer really is — four lines capture the mechanism behind every large language model. There is nothing hidden: the learned parts are the projections that produced Q, K, and V upstream, while the attention operation itself has no parameters. Seeing it stripped to this minimum is what makes the framework version, and the multi-head wrapper, easy to trust.

The scaling by the square root of d_k is a small detail with outsized importance. When you take the dot product of two vectors of dimension d_k, the magnitude of the result tends to grow with d_k. In a model where d_k might be 64 or more, raw scores can become large, and large scores push softmax toward a near one-hot distribution where almost all weight lands on a single token.

That saturation is poison for training: in the flat regions of a saturated softmax, gradients are tiny, so learning stalls. Dividing by the square root of d_k keeps the variance of the scores around one regardless of dimension, so softmax stays in its sensitive, well-behaved range and gradients flow. It is a one-line fix for a problem that would otherwise scale with model width.

This flow diagram walks one attention step end to end for a single token. The input vector is projected into a query, a key, and a value; the query dots against every key to produce a row of scores; softmax turns that row into weights; and those weights mix the values into a single context vector that becomes the token's new representation.

Reading it left to right ties the equation back to the search metaphor: the token asks a question (query), compares it against what every token offers (keys), and retrieves a weighted blend of content (values). Every token in the sequence runs this same path in parallel, which is what the matrix form computes all at once.

Multi-head attention is the realization that one attention operation is a bottleneck — it can only learn one notion of relevance at a time. So instead of attending once in the full d_model space, the model splits that space into h smaller heads, each with its own query, key, and value projections, and runs attention h times in parallel.

The value is specialization. One head might learn to track subject-verb agreement, another to resolve pronoun references across long distances, another to attend to adjacent tokens. Each head computes its own context vectors, and their outputs are concatenated back to the full width and passed through a final projection that lets the model mix what the heads found. More heads means more simultaneous relationships the model can attend to.

This snippet tells the shape story of multi-head attention, which is where most confusion lives. With d_model of 512 and 8 heads, each head operates in a 64-dimensional space — d_model divided by h. The Q, K, and V tensors are reshaped to expose the head dimension and transposed so each head's slice can be attended independently.

After running attention per head, the outputs are concatenated back to the full 512 dimensions and passed through a learned output projection W_o that mixes information across heads. The key arithmetic to internalize is that the heads partition d_model rather than duplicating it: eight 64-dimensional heads cost roughly the same as one 512-dimensional attention, which is what makes multi-head essentially free.

Positional encoding fixes a fundamental blind spot. Attention is permutation-invariant — if you shuffle the input tokens, it computes the same set of outputs in shuffled order, because nothing in the dot-product mechanism knows where a token sits. Left alone, the model would treat a sentence as an unordered bag of words.

To restore order, a positional signal is added to each token's embedding before the first block. The original Transformer used fixed sinusoids of different frequencies, which let the model infer relative distances; modern models often use learned position embeddings or rotary encodings that inject position directly into the attention computation. Whatever the scheme, the purpose is identical: without it, 'dog bites man' and 'man bites dog' are indistinguishable to the model.

This stack diagram shows the full ordering inside one block, which is exactly what the code post will implement. Multi-head attention runs first to mix information across tokens, then an Add & LayerNorm wraps it in the first residual connection. Next a position-wise feed-forward network refines each token independently, followed by a second Add & LayerNorm and residual.

Visualizing the block as this four-level stack reinforces the alternation that defines a Transformer: a token-mixing step followed by a per-token processing step, each protected by a residual-and-normalize wrapper. Every block in the network repeats this exact pattern, so once you can draw this stack from memory, you can draw the whole model.

The residual-and-LayerNorm wrapper is what makes stacking dozens of these blocks possible. Each sub-layer's output is added back to its own input — that is the residual connection — and the sum is then normalized by LayerNorm. The residual gives gradients a clean, near-identity highway straight through the network, so even very deep stacks can be trained without the signal vanishing.

LayerNorm plays the complementary role of keeping activations well-scaled, normalizing across the feature dimension for each token independently — which, unlike BatchNorm, does not depend on batch size or sequence length, making it ideal for variable-length sequences. Together the residual and the normalization are why a 96-layer Transformer trains at all; remove them and the deep stack collapses, a point the mistakes post drives home.

This recap pins down the mechanics: attention scores are Q times K-transpose, scaled by the square root of d_k, softmaxed into weights, and multiplied by V; that scaling is what keeps softmax in a trainable range; multiple heads attend in parallel and specialize before being concatenated; positional encoding restores the word order that attention ignores; and residual connections plus LayerNorm are what make deep stacks trainable.

With the math traced by hand, you are ready to build the real thing. The next post assembles a working multi-head attention block and a full encoder layer in PyTorch, and verifies it against the library implementation.

The teaser points to the hands-on build. Having traced the forward pass of attention and understood multi-head, positional encoding, and the residual wrapper, the next post turns all of it into runnable PyTorch — scaled dot-product attention, a multi-head module, a complete block, and the causal mask — so the equations become a working artifact you can run and inspect.