The Transformer, Explained

This is the hands-on post, and it is intentionally code-heavy. The previous posts built understanding from concept through mechanics; here you assemble a complete, runnable Transformer block — scaled dot-product attention, multi-head attention, a full encoder layer with residuals — and verify it against PyTorch's own implementation so you trust every line.

Each snippet runs in order and builds on the last. The most valuable thing you can do is run it and then print the attention weights to see, concretely, which tokens the model is looking at. That inspection is how the abstract mechanism from the previous post turns into intuition you can debug with.

This is the atom of the whole architecture: scaled dot-product attention in a handful of lines. The function reads d_k from the query, computes scores as Q times K-transpose divided by the square root of d_k, optionally fills masked positions with negative infinity so softmax assigns them zero weight, applies softmax, and returns both the context vectors and the attention weights.

Returning the weights alongside the output is a deliberate choice — it is what lets you inspect attention later. The masked_fill with negative infinity is the single trick that powers both causal masking for decoders and padding masks for batches; everything else in the function is the plain four-step formula from the mechanics post, now expressed in real PyTorch.

The multi-head module wraps that atom into a reusable layer. In the constructor it stores the number of heads and the per-head dimension d_k, and creates two linear layers: one that projects the input into queries, keys, and values all at once (3 times d_model wide), and one output projection. Combining the three projections into a single qkv linear is a common efficiency that does the same work as three separate matrices.

In forward, the qkv output is reshaped to expose the head dimension and permuted so each head can attend independently, attention runs across all heads at once, and the per-head outputs are transposed back together and reshaped to the full width before the output projection mixes them. This is the exact shape choreography from the mechanics post, now as a module you can drop into any model.

The encoder block assembles the full layer from its parts, and its forward method is where the architecture's defining pattern becomes code. It holds the multi-head attention, a two-layer feed-forward network with a ReLU in the middle, and two LayerNorms. The forward pass applies attention and adds the result back to the input before normalizing — the first residual — then applies the feed-forward network and adds it back before normalizing again — the second residual.

Reading x = self.n1(x + self.attn(x)) literally spells out 'add the sub-layer output to its input, then normalize.' That one line is the residual-and-LayerNorm wrapper from the mechanics post made concrete. Note this is the post-norm arrangement; the mistakes post will explain why deep models often move the norm before the sub-layer instead.

This snippet builds the causal mask that converts an encoder-style block into a GPT-style autoregressive decoder. torch.tril produces a lower-triangular matrix of ones, and reading the printed result row by row shows the rule precisely: token 0 sees only itself, token 1 sees tokens 0 and 1, token 2 sees 0 through 2, and so on, with every future position blocked.

When this mask is passed into the attention function, the upper triangle becomes negative infinity before softmax, so each token assigns exactly zero weight to anything after it. This is the entire difference between an encoder and a decoder — not a different architecture, just a triangular mask. Forgetting it is the most damaging Transformer bug there is, which is why the final post leads with it.

The flow diagram traces a tensor through the block and is the picture to hold in your head while reading the code. A batch of shape (B, n, 512) enters multi-head attention, which mixes information across the n tokens; the result is added back and normalized; the feed-forward network then processes each token independently; and a final add-and-normalize produces the output, which has the identical shape it started with.

That shape preservation — same in, same out — is what allows blocks to stack arbitrarily, since the output of one is a valid input to the next. Most wiring bugs in a Transformer are shape bugs, so being able to recite this flow lets you locate a mismatch by reasoning about which step changed the wrong dimension.

Sanity-checking against the library is how you earn confidence in a from-scratch build. The snippet runs a random batch of shape (2, 6, 512) through your hand-built EncoderBlock and confirms the output shape, then runs the same batch through PyTorch's own nn.TransformerEncoderLayer and confirms it produces an identical shape.

Matching shapes does not prove the numerics are bit-identical — the two implementations initialize weights differently and may order operations slightly differently — but it proves your block is a valid, drop-in Transformer layer with the correct interface. This habit of checking a custom implementation against a trusted reference is exactly how you catch the silent shape and broadcasting mistakes that the final post catalogs.

This snippet shows how to look inside the model and see what it is actually attending to. It runs the qkv projection, reshapes to expose the heads, and calls the attention function while keeping the returned weights. The weights have shape (batch, heads, n, n): for each example and each head, an n-by-n matrix where entry (i, j) is how much token i attends to token j.

Printing one head's matrix, rounded, makes attention tangible — you can read off which tokens each query weights most heavily. This is not just a curiosity; inspecting attention weights is a standard debugging and interpretability technique. If a head's weights look uniform or stuck, something upstream may be wrong, and seeing the matrix is often faster than reasoning about it.

These shape facts are where most practitioners lose time wiring a Transformer. The input is consistently (batch, sequence, d_model); the attention weights are (batch, heads, n, n); d_model must divide evenly by the number of heads or the per-head split is undefined; the mask must have a shape that broadcasts against the score tensor; and LayerNorm normalizes the last (feature) dimension, not the batch — a key difference from BatchNorm.

Keeping these straight turns most debugging from trial and error into reasoning. When a Transformer throws a shape error, it is almost always one of these: a head count that does not divide d_model, or a mask whose dimensions will not broadcast over the (batch, heads, n, n) scores. Knowing the canonical shapes lets you spot the culprit immediately.

This slide collects the gotchas that do not crash but quietly ruin results — the through-line to the next post. Forgetting the causal mask lets a decoder peek at future tokens, producing a training loss that looks suspiciously good while generation fails, because the model learned to rely on information it will not have at inference. Forgetting the 1/sqrt(d_k) scale lets scores blow up and gradients vanish. And dropping the residual connection — writing self.attn(x) instead of x + self.attn(x) — starves the early blocks of gradient.

The common thread is that none of these raise an exception; they hand you a model that looks like it is working but is broken in a specific, diagnosable way. Recognizing the symptom-to-cause mapping is exactly what separates a working Transformer from a mysteriously bad one.

This recap consolidates the build: scaled dot-product attention is scale, optional mask, softmax, then weighted values; multi-head attention reshapes into heads, attends, concatenates, and projects; a block is attention then Add+Norm then feed-forward then Add+Norm; a triangular mask turns the block autoregressive; and you should verify your shapes against PyTorch's own Transformer layers.

With a working, verified block in hand, the only thing left is to learn the traps — the quiet mistakes that make Transformers train wrong or run out of memory. That is the focus of the final post, several of which you have already glimpsed in the gotchas slide.

The teaser points to the common-mistakes post. You now have a Transformer block that runs and matches the library; the next step is hardening your instincts against the quiet failure modes — a missing causal mask, dropped positional information, unbudgeted quadratic cost, mishandled padding, and unscaled scores — each with the one-line fix that resolves it.