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

RNN pushes history into a continuously updated state. An intelligence officer proposes an alternative querying method: when handling a current task, directly query all relevant locations in the archive and then aggregate the values based on matching strength.

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=XW_Q,qquad K=XW_K,qquad V=XW_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=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}+M\right), $$

$$ 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 adds $-\infty$ or a sufficiently negative value to positions to be prohibited, making softmax weights zero. 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\times3$ 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-header version from scratch

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]
    scores = query @ key.transpose(-2, -1) / math.sqrt(query.shape[-1])

    if allowed is not None:
        # allowed May be broadcast To [B, L, S]; True Indicates allowed.
        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))

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:

$$ head_h=Attention(XW_Q^{(h)},XW_K^{(h)},XW_V^{(h)}), $$

$$ MHA(X)=Concat(head_1,\ldots,head_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 we shouldn’t force-naming on each head.

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 typically needs to be passed:

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

If you permute the token order of X, the self-attention output without position signals is permuted accordingly; it is permutation equivariant with respect to set-like inputs and doesn't distinguish between the first and the tenth position.

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 $Td_{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 automatically divides into human-like concepts: 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\times3$ 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.