Attention Mechanism

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 multi-head self-attention layer, mask it, and run a real forward pass to verify the shapes.

Each snippet is real code that fits together in order — the scaled dot-product function, the module's constructor, its forward pass, a causal mask, and a test call. The most valuable thing you can do is run it, then poke at the returned attention weights to watch where the model focuses. That experimentation is how the formula from the previous post turns into intuition you actually own.

This is scaled dot-product attention as a standalone function, and it is a near-direct transcription of the formula. It reads the head dimension from the last axis of q, computes the scores as q times k-transpose divided by the square root of that dimension, optionally fills masked positions with negative infinity, applies softmax along the last axis to get weights, and returns the weighted blend of v together with the weights.

The one subtlety worth noting is masked_fill with negative infinity before softmax: because the exponential of negative infinity is zero, masked positions receive exactly zero weight. Returning the weights alongside the output is deliberate, since those weights are what you inspect to interpret and debug attention. This single function is the engine that the multi-head module wraps.

The module's constructor sets up everything the forward pass will need. It stores the number of heads and computes d_k, the per-head dimension, as d_model divided by the number of heads — which is why d_model must be evenly divisible by the head count. It then creates two linear layers: one that projects the input to three times d_model so a single matmul can produce queries, keys, and values together, and one output projection that mixes the merged heads back to d_model.

Fusing the Q, K, and V projections into a single Linear of width three-times-d_model is a common efficiency trick — one matrix multiply instead of three, split afterward. Keeping the projections as module attributes means their weights are learned parameters, updated by gradient descent like any other layer.

The forward pass is where splitting, attending, and merging come together. It reads the batch and time dimensions, projects the input and chunks the result into q, k, and v. The split lambda reshapes each from (batch, time, d_model) into (batch, time, heads, d_k) and transposes so the head axis comes before time, giving (batch, heads, time, d_k) — the layout that lets a single batched matmul handle all heads at once.

It then calls the scaled dot-product function to get the per-head context and the attention weights, transposes the context back so the head axis follows time again, and reshapes to merge the heads into a single d_model vector per position. A final output projection mixes the heads, and the module returns both the result and the attention weights. This split-attend-merge pattern is the canonical multi-head implementation.

This snippet builds the causal mask that a decoder needs. torch.tril produces a lower-triangular matrix of ones — ones on and below the diagonal, zeros above — which, read as a boolean, means position t may attend to positions zero through t but not to anything later. Printing it shows the staircase pattern directly.

The mask is consumed by the scaled dot-product function, which fills the zero positions with negative infinity before softmax so future positions get exactly zero weight. Building the mask as a separate, inspectable tensor makes it easy to verify the pattern is correct before trusting it in training — a worthwhile habit, since a wrong mask is a silent bug that inflates training scores while ruining generation.

This is the moment of truth: instantiate the module, feed it a real batch, and check the output shapes. With d_model 64 and 8 heads, an input of shape (2, 5, 64) — two sequences of five tokens — should produce an output of the same (2, 5, 64) shape, since attention preserves the sequence layout and only re-mixes information across positions. The attention weights come back as (2, 8, 5, 5): for each of the two batches and each of the eight heads, a five-by-five map of how much every token attended to every other.

Verifying these shapes is the fastest sanity check that the split-and-merge logic is correct. If the output is not (2, 5, 64), the reshape is wrong; if the weights are not (2, 8, 5, 5), the head dimension is misplaced. Reading the printed shapes is how you confirm the whole module wired up correctly before you ever start training.

The data-flow diagram traces tensor shapes through the module, which is the single most useful thing to hold in your head when wiring attention. The input x of shape (batch, time, 64) is projected to (batch, time, 192) by the fused QKV layer, split into eight heads of (batch, 8, time, 8), passed through attention preserving that shape, then merged and projected back to (batch, time, 64).

Most attention wiring bugs are shape bugs, so being able to recite this chain lets you catch a mismatch by reasoning rather than trial and error. When a tensor operation throws, mentally walking these shapes usually reveals whether the split, the transpose, or the merge is the culprit — far faster than scattering print statements through the forward pass.

These shapes-and-gotchas are where most practitioners lose time with attention. d_model must divide evenly by the number of heads, or the per-head dimension is not an integer and the reshape fails. The transpose that puts heads before time is what enables the batched matmul over the last two axes; get it wrong and heads bleed into each other. The mask must be applied as negative infinity before softmax, not after. The attention weights are shaped (batch, heads, time, time) — one map per head, which matters when you visualize them.

The final point is a subtle PyTorch detail: after a transpose, a tensor's memory may be non-contiguous, so a following reshape can require calling .contiguous() first. None of these raise conceptual errors in your understanding of attention, but each is a concrete trap that turns a correct mental model into broken code if missed.

This snippet shows that, in practice, you rarely write attention by hand. PyTorch 2.x provides scaled_dot_product_attention, a single fused call that runs the score-scale-mask-softmax-blend sequence with optimized, memory-efficient kernels — including a convenient is_causal flag that builds the causal mask for you. For a full layer, torch.nn.MultiheadAttention wraps the projections and head logic in one module.

The reason to learn the from-scratch version first is that the built-in is a black box until you know what it computes. Once you have traced the steps yourself, the fused call is simply a faster, better-tested implementation of code you understand. In real projects you should prefer these built-ins for speed and correctness, reserving the hand-written version for learning and debugging.

This comparison frames the choice between writing attention yourself and calling the built-in. The from-scratch version exposes every step, is ideal for learning, and makes it trivial to inspect the weights — at the cost of being slower and requiring more code that you could get subtly wrong. The built-in is a single fused call backed by optimized kernels, with far less surface area for mistakes, which is what you want in production.

The sensible workflow uses both: implement it once by hand to understand and to debug specific behaviors, then switch to the fused built-in for any real training run. Understanding the manual version is precisely what lets you trust the built-in and diagnose it when something looks off, so the two are complementary rather than competing.

This recap consolidates the build: scaled dot-product attention is score, scale, mask, softmax, blend; you project the input to Q, K, and V and split them into heads; you apply the causal mask before softmax; you merge the heads and project back to d_model; and in production you reach for the fused built-in rather than the hand-written loop.

With a working, correctly-masked attention layer in hand — one whose shapes you have verified — the only thing left is to learn the traps. The final post catalogs the quiet mistakes that make attention produce wrong results without ever raising an error.

The teaser points to the common-mistakes post. You now have a multi-head attention module that runs and whose shapes check out; the next step is hardening your instincts against the silent failures — a dropped scaling factor, a missing causal or padding mask, swapped self- and cross-attention inputs, and over-trusting the weight maps — each with its clean one-line fix.