10.3 Scaled Dot-Product and Multi-Head Attention: First, Get the Mask Semantics Right

RNNs compress history into an updated state. Attention instead forms a query at each position, compares it with a set of keys, and combines the corresponding values using the resulting weights.

Attention is a differentiable retrieval or weighted aggregation. What's actually prone to errors isn't the Q/K/V analogy, but the shapes, scaling, softmax dimensions, and mask semantics.

This lesson's objectives

• Derive the shape of scaled dot-product attention and the scaling factor;
• Distinguish self-attention, cross-attention, causal attention, and padding attention;
• Correctly construct and test Boolean/additive masks;
• Understand multi-head projection, concatenation, and output mapping;
• Identify the complexity and interpretive limitations of attention.

1. Q, K, V are Learned Projections

The input is represented as $X\in {\mathbb{R}}^{B\times T\times d_{model}}$. Self-attention typically computes:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V.$$

The query/key determines the match score, and the value provides the content to be aggregated. These roles are learned through training and are not guaranteed to correspond to database keys/values or human-readable semantic labels.

In cross-attention, Q comes from the target/decoder states, and K/V comes from the source/encoder states, so the query length $L$ can differ from the source length $S$.

2. Scaled Dot-Product Attention

Single-headed:

$$A=\mathrm{softmax}(\frac{Q{K}^T}{\sqrt{d_k}}+M),$$$$O=AV.$$

If:

Q: [B, L, d_k]
K: [B, S, d_k]
V: [B, S, d_v]
scores/A: [B, L, S]
O: [B, L, d_v]

Softmax is applied along the key/source dimension $S$. The sum of allowed weights for each query row is 1.

If the Q/K components have approximately zero mean, unit variance, and are independent, the dot-product variance grows with $d_k$; dividing by $\sqrt{d_k}$ brings the typical scale back, reducing premature softmax saturation. This is an approximate justification and does not guarantee that actual trained representations satisfy the independence assumption.

3. The mask must be applied before Softmax

Additive mask uses $-\mathrm{\infty }$ at prohibited positions to make their softmax weights zero. A finite negative constant only suppresses weights relative to allowed scores; it is not an unconditional zero guarantee. Multiplying the softmax weights by zero after softmax would leave the remaining weights unnormalized unless re-normalized.

Common masks:

• key padding mask: all queries ignore source padding;
• causal mask: position $t$ doesn't look ahead to future $>t$;
• local/block/sparse mask: allows only specified connections;
• cross-attention source mask: ignore encoder padding.

Query padding must also be masked in the loss/output; masking only the padded keys won't automatically remove the padded query outputs.

4. Boolean Mask Semantics Are Not Uniform

This is an engineering high-risk point. Currently in the PyTorch API:

• nn.MultiheadAttention's Boolean key_padding_mask/attn_mask: True means disallowed/ignored;
• torch.nn.functional.scaled_dot_product_attention's Boolean attn_mask: True indicates participation is allowed.

Migrating from one API to another might require taking the opposite. Don't guess based on variable names; write a $3\times 3$ manual example to verify that forbidden weights are zero.

If all keys in a query are masked, the softmax will have no valid distribution and may produce NaN or undefined behavior. Batch construction must ensure at least one valid key, or explicitly define empty-row results.

5. Build a single-head version from scratch

The runnable Python 3 / PyTorch 2.14 CPU example below supports finite floating-point tensors with the stated three-dimensional shapes, matched batches, and nonempty key/feature dimensions. Its Boolean mask must be on the same device and broadcast to [B,L,S].

import math
import torch

def scaled_dot_product_attention(query, key, value, allowed=None):
    # query: [B, L, D], key: [B, S, D], value: [B, S, Dv]
    if query.ndim != 3 or key.ndim != 3 or value.ndim != 3:
        raise ValueError("query, key and value must be three-dimensional")
    if not query.shape[0] == key.shape[0] == value.shape[0]:
        raise ValueError("batch dimensions must match")
    if query.shape[-1] != key.shape[-1] or key.shape[1] != value.shape[1]:
        raise ValueError("key dimensions and source lengths must match")
    if key.shape[1] == 0 or query.shape[-1] == 0:
        raise ValueError("source length and key dimension must be positive")
    scores = query @ key.transpose(-2, -1) / math.sqrt(query.shape[-1])
    if not torch.isfinite(scores).all() or not torch.isfinite(value).all():
        raise ValueError("attention scores and values must be finite")

    if allowed is not None:
        if allowed.dtype != torch.bool:
            raise TypeError("allowed must be a Boolean tensor")
        # Broadcast allowed to [B, L, S]; True means participation.
        allowed = torch.broadcast_to(allowed, scores.shape)
        if not torch.all(allowed.any(dim=-1)):
            raise ValueError("Every query must have at least one allowed key")
        scores = scores.masked_fill(~allowed, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ value
    return output, weights

q = torch.randn(2, 4, 8)
k = torch.randn(2, 6, 8)
v = torch.randn(2, 6, 5)
out, weight = scaled_dot_product_attention(q, k, v)
assert out.shape == (2, 4, 5)
assert torch.allclose(weight.sum(dim=-1), torch.ones(2, 4))

q = k = torch.zeros(1, 3, 2)
v = torch.tensor([[[2.], [4.], [8.]]])
allowed = torch.ones(3, 3, dtype=torch.bool).tril()
out, weight = scaled_dot_product_attention(q, k, v, allowed)
assert torch.all(weight[0][~allowed] == 0)
assert torch.allclose(out[0, :, 0], torch.tensor([2., 3., 14. / 3.]))

With zero Q/K scores, each causal row averages just its allowed values: 2, then 3, then $14/3$. A reversed mask changes these results immediately, which makes this a more useful check than output shape alone.

Prioritize using framework-optimized kernels for production code; this version is used to establish shape/mask tests and does not handle fused precision, dropout, GQA, or high-performance layouts.

6. Causal Mask

A decoder self-attention layer of length $T$, allowed matrix:

q0: k0
q1: k0 k1
q2: k0 k1 k2
...

Lower triangular (including diagonal). A causal mask prevents tokens $t$ from accessing future targets during training. If labels are shifted by one position without applying the mask, the model would inadvertently "peek" at the answers through hidden representations, leading to abnormally low training loss.

The masks for Prefix-LM, bidirectional encoder, and sequence-to-sequence decoder are different. Don't treat "Transformer mask" as a single template.

7. Multi-head Attention

Project $d_{model}$ into $H$ groups of head subspaces:

$$hea{d}_h=Attention(X{W}_Q^{(h)},X{W}_K^{(h)},X{W}_V^{(h)}),$$$$MHA(X)=Concat(hea{d}_1,\dots ,hea{d}_H)W_O.$$

A common implementation requires that $d_{model}$ be divisible by $H$, so that $d_h=d_{model}/H$. Increasing the number of heads doesn't necessarily increase the total projection parameters, but it does alter the dimension per head, kernel efficiency, and representation decomposition.

“Each head learns grammar, coreference, and position” is just one observational analysis, not a training constraint. Heads can be redundant, mixed, or unstable, and human-readable roles should not be assumed.

8. Framework API Performance and Dropout

nn.MultiheadAttention(batch_first=True) accepts [B,T,D] input. When only the output is needed, set need_weights=False so the framework is more likely to use an optimized scaled-dot-product path.

When directly calling the current F.scaled_dot_product_attention, dropout_p>0 applies dropout based on parameters and does not automatically read the module's train/eval state; the module must choose the probability explicitly.

Inside a module method, assuming self.dropout is the configured probability, the relevant fragment is:

dropout_p = self.dropout if self.training else 0.0

Fused Flash/memory-efficient/math kernels may switch based on dtype, device, and mask shape, leading to floating-point differences. Performance tests must record the backend and not rely solely on latency estimates from formulas.

9. Attention Has No Positional Information

Without positional signals or a fixed position-dependent mask, permuting the tokens of self-attention permutes its output in the same way. A causal or local mask already imposes positional structure; the permutation statement holds only if that mask is permuted along with both token axes.

Positions can be injected via mechanisms such as absolute embedding, sinusoidal, relative bias, and rotary, each affecting length extrapolation and cache implementation, detailed in Section 10.4.

10. Complexity Goes Beyond one line of $O(T^2)$

The self-attention score matrix is approximately [B,H,T,T]; the score computation and storage for standard dense attention scale with $T^2$, while projection and FFN still depend on $T{d}_{model}^2$.

In short sequences with large hidden dimensions, the projection/FFN layers may dominate; in long sequences, the attention matrix often becomes the bottleneck. FlashAttention-style algorithms reduce materialized memory/I/O and preserve exact attention semantics (allowing floating-point implementation differences), but they don't automatically turn all theoretical complexity into linear scaling.

Sparse/local/linear attention changes connections or kernel approximations, and the task quality and actual acceleration must be verified.

11. Attention Weights Are Not Causal Explanations

A high weight indicates the coefficient of the value mixture in the current head/layer, and it does not mean that the token makes a unique contribution to the final output:

• Value vectors and output projection affect the outcome;
• residual/FFN/subsequent layers have other paths;
• Multiple sets of weights can produce similar outputs;
• Intervention tokens will simultaneously modify Q/K/V.

Weights can be used for diagnosis, but should be combined with ablation, gradient/perturbation, and counterfactual analysis; causality cannot be claimed directly.

Common Misconceptions

• Q/K/V are fixed semantic fields: They are learned projections.
• Mask then multiply by 0 is sufficient: Mask and test normalization before softmax.
• In PyTorch, all Boolean masks have True = keep: The semantics can vary by API.
• Each head specializes in a human concept: No such constraint exists.
• Attention weight is just the explanation: The final output still includes value, projection, and other paths.

Practice

1. Write Q/K/V/scores/output shapes for cross-attention.
2. Construct a $3\times 3$ causal mask and verify that future weights are zero.
3. Compare MultiheadAttention with SDPA Boolean mask semantics.
4. Mask all of a query and observe and fix the failure.
5. Estimate the memory and computation cost of attention scores and FFN under different values of $T$, $D$, and $H$.

Summary

Scaled dot-product attention uses Q/K similarity to produce a normalized mixture of values for each query. Scaling controls the softmax input scale, masks define the information flow, and multi-head attention computes in parallel across different learned subspaces. It excels at directly connecting positional information and introduces quadratic matrices, API semantics, and explanation risks.

In the next lesson, we'll focus on residual connections, normalization, and the FFN, add positional and causal objectives, and put together the encoder, decoder, and a complete Transformer.