Scaled Dot-Product Attention › Summary

Scaled dot-product attention maps QKV inputs into context-enriched representations through four sequential operations:

• Dot product ($QK^T$): measures pairwise similarity between all queries and all keys, producing an $n \times n$ score matrix with $O(n^2 d_k)$ cost
• Scaling ($/ \sqrt{d_k}$): normalizes scores to have unit variance, preventing softmax saturation and enabling stable gradient flow throughout training
• Softmax: converts scores row-wise into probability distributions over positions, producing attention weights $A_{ij} \geq 0$ with $\sum_j A_{ij} = 1$
• Weighted sum ($AV$): aggregates value vectors according to attention weights, producing a $d_v$-dimensional context-aware representation for each position

The causal mask extends this mechanism to autoregressive generation by zeroing out attention to future positions before softmax. Combined with the padding mask for batched computation, the masking mechanism handles all common production use cases. Both types of masking work by adding $-\infty$ to score matrix entries before softmax, which causes those positions to receive exactly zero weight in the output.

The quadratic complexity $O(n^2 d)$ in both computation and memory is the central scalability challenge. Efficient attention variants address this bottleneck, which we'll examine in the Attention Complexity chapter. Before that, the next chapter covers multi-head attention, which runs multiple instances of scaled dot-product attention in parallel to allow the model to attend to different aspects of the input simultaneously. Multi-head attention is the direct generalization of everything we've covered here: it applies scaled dot-product attention not once but $h$ times, with different learned projections for each head, and concatenates the results. Understanding scaled dot-product attention fully, as we've done in this chapter, is the necessary foundation for understanding why multi-head attention works so well in practice.