Attention Mechanism

This is the engine-room post. The previous two established what attention is and why it matters; here we replace intuition with the explicit formula, tracing exactly what scaled dot-product attention does to its inputs and why each step is there.

The payoff is that the mechanism stops being a black box. By the end you will have seen the four operations — score, scale, softmax, blend — written out, understood the otherwise-mysterious division by the square root of d_k, and seen how multi-head attention runs the trick several times in parallel. Doing this trace once is worth more than reading ten high-level descriptions.

The whole mechanism fits in one line: Attention(Q, K, V) equals softmax of (Q times K-transpose, divided by the square root of d_k), all multiplied by V. It looks dense, but it reads naturally left to right. First you score each query against each key with a dot product, producing the matrix Q times K-transpose. Then you scale those scores by dividing by the square root of the key dimension. Then softmax turns each row of scores into a set of weights that sum to one. Finally you multiply those weights by the values to get the output.

Every later subtlety — masking, multi-head, the efficiency literature — is a variation on this single formula. Memorizing it pays off, but the goal of this post is to make you understand each piece well enough that you could reconstruct it.

Breaking the formula into four named steps makes it easy to hold in your head and to debug. Step one, score: compute S equals Q times K-transpose, which is every query-key dot product at once. Step two, scale: divide S by the square root of d_k. Step three, weights: apply softmax across each row of S to get the attention matrix A. Step four, blend: multiply A by V to produce the output.

Laying it out this way clarifies which parts are pure linear algebra (the two matrix multiplies in steps one and four) and which parts shape the distribution (the scaling and softmax in steps two and three). When attention misbehaves, it is almost always one of these four steps — most often a missing scale or a missing mask — so knowing them by number makes diagnosis faster.

The pipeline diagram traces a batch of inputs through the four stages: queries dotted against keys to produce raw scores, those scores divided by the square root of d_k to control their scale, softmax to convert them into normalized weights, and a final multiplication by values to blend. Each stage does exactly one job.

Visualizing it as a pipeline reinforces that scoring and blending are separate concerns from the normalization in between. The two ends are matrix multiplies that move information; the middle two stages shape how sharply the model focuses. Holding this four-stage picture makes the from-scratch code on the next slides read as a direct transcription rather than something new.

The scaling factor is the step most people skip and most regret skipping. A dot product of two vectors with d_k components is a sum of d_k terms, so its variance grows with the dimension. In a high-dimensional space the raw scores can become large, and large inputs push softmax into a saturated regime where one weight is essentially one and all the others are essentially zero.

In that saturated regime the gradient of softmax is nearly flat, so almost no learning signal flows back — training stalls or converges poorly. Dividing the scores by the square root of d_k counteracts the dimension-driven growth, keeping the score variance near one so softmax stays in its responsive, well-graded region. It is a small constant with an outsized effect on trainability, which is why the mistakes post treats omitting it as a serious bug.

This from-scratch implementation makes the formula undeniable. The softmax helper subtracts the row maximum before exponentiating — a standard numerical-stability trick that prevents overflow without changing the result — then divides by the row sum. The attention function reads off the formula directly: it grabs d_k from the last dimension, computes the scaled scores with a matrix multiply and the square-root divisor, applies softmax to get weights, and returns the weighted blend of values along with the weights themselves.

Returning the weights as well as the output is deliberate: those weights are the focus map discussed in the previous post, and inspecting them is how you debug and interpret attention. Seeing the entire mechanism in a handful of lines confirms that the conceptual four steps map one-to-one onto real, runnable code.

Softmax is the step that turns arbitrary scores into usable weights, and understanding it explains attention's 'soft' character. It exponentiates each score and divides by the sum of the exponentials in that row, producing a set of non-negative numbers that sum to one — a probability distribution over the input positions. Larger scores become larger weights, but every position still receives a non-zero share.

That last property is what makes attention a soft, differentiable selection rather than a hard pick of a single position. Because every value contributes at least a little and the operation is smooth, gradients can flow to all positions and the model can gradually adjust where it focuses. A hard argmax would be non-differentiable and untrainable; softmax is the differentiable relaxation that makes end-to-end learning possible.

The causal mask is what lets a single attention mechanism serve as a left-to-right language model. During training you feed the model an entire sequence at once for efficiency, but a token at position t must not be allowed to attend to tokens at positions after t — those are the future words it is supposed to predict. The fix is to set the scores for all future positions to negative infinity before applying softmax.

Because the exponential of negative infinity is zero, those future positions receive exactly zero weight, so each position's output depends only on itself and earlier positions. This is what reconciles parallel training with autoregressive generation: the model sees the whole sequence during training for speed, but the mask guarantees it never uses information it would not have at inference time. Forgetting this mask is one of the most damaging silent bugs, as the final post details.

This network diagram depicts multi-head attention as a fan-out and fan-in: a single input is projected into several parallel heads, each of which performs its own attention, and their outputs are concatenated and projected back to the model dimension. The labels mark the three stages — input, the h parallel heads, and the concatenate-and-project step that merges them.

The key visual idea is parallelism within a layer. Rather than one attention computation, multi-head attention runs several at once on different learned projections of the same input. Each head is a full scaled dot-product attention; the diagram simply shows that they operate side by side and are recombined, which is the structure the next code slide makes concrete with tensor shapes.

This snippet tells the shape story behind multi-head attention, which is where most implementation confusion lives. You start with a tensor of shape (batch, time, d_model). To create h heads, you reshape the model dimension into h groups of size d_model divided by h, giving (batch, time, h, d_k), then transpose so the head dimension comes before the time dimension, yielding (batch, h, time, d_k).

With heads in front of time, every head's attention is just a batched matrix multiplication over the last two dimensions, computed for all heads at once. The single most common multi-head bug is getting this reshape-and-transpose wrong, so internalizing that d_model splits into h times d_k, and that the head axis must precede the sequence axis for the batched matmul, saves a great deal of debugging.

Multi-head attention exists because a single attention pass collapses everything into one weighting, which is limiting. With multiple heads, the model runs several attentions in parallel, each on its own learned linear projection of the input, so different heads can specialize in different kinds of relationships. Empirically, one head may track local syntax, another long-range coreference, another positional patterns.

After each head computes its own weighted blend, the outputs are concatenated and passed through a final linear projection that mixes them back into the model dimension. The result is a representation that integrates several distinct 'views' of how the tokens relate, which is richer than any single attention could produce. This is why real architectures use multi-head attention rather than a single head, at essentially the same total compute.

This recap pins down the mechanics: the order is score, scale, softmax, blend; you divide by the square root of d_k to keep softmax smooth and gradients alive; softmax makes a soft, differentiable selection so every value contributes and learning can flow; the causal mask blocks attention to future positions; and multi-head attention runs several attentions in parallel on different projections to capture multiple relationships.

With the math traced by hand, you are ready to build the real thing. The next post assembles a working multi-head attention layer in PyTorch — including the masking and the head reshaping — so the formula becomes a runnable artifact.

The teaser points to the hands-on build. Having traced the four steps, understood the scaling and the mask, and seen how heads split and merge, the next post wires it all into a real PyTorch module: scaled dot-product attention as a function, a full multi-head wrapper, a causal mask, and a forward pass whose shapes you can check.