Home / Fundamentals / Transformer Architecture

Transformer ArchitectureMidQuestion 2 of 9

Explain scaled dot-product attention, implement it, and justify dividing by √d_k.

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Confirm whether they want derivation, implementation, or both. Framing: "I'll write the computation with shapes, implement causal multi-head attention, justify √d_k with a variance argument, then cover numerical stability, cost and the backward pass."

5–15 min Core explanation

For X ∈ ℝ^{T×d} and head size d_k:

Q = X W_Q,  K = X W_K,  V = X W_V          (T × d_k)
S = Q Kᵀ / √d_k                            (T × T)
P = softmax(S + M)                          row-wise; M = 0 or −∞ (causal mask)
O = P V                                     (T × d_k)

Each row of P is a distribution over positions; each output row is a weighted average of value vectors. Multi-head attention runs H heads of size d/H in parallel, concatenates, and applies W_O.

import torch, math

def causal_attention(x, Wq, Wk, Wv, Wo, n_heads):
    B, T, d = x.shape
    dk = d // n_heads
    def heads(t):                                   # (B, T, d) -> (B, H, T, dk)
        return t.view(B, T, n_heads, dk).transpose(1, 2)
    q, k, v = heads(x @ Wq), heads(x @ Wk), heads(x @ Wv)
    s = (q @ k.transpose(-2, -1)) / math.sqrt(dk)   # (B, H, T, T)
    mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
    s = s.masked_fill(mask, float("-inf"))
    p = s.softmax(dim=-1)
    o = (p @ v).transpose(1, 2).reshape(B, T, d)
    return o @ Wo

15–25 min Deep dive

Why √d_k. Let components of q and k be independent, mean 0, variance 1. Then Var(q·k) = Σ Var(q_i k_i) = d_k, so the standard deviation is √d_k. With d_k = 128, scores have std ≈ 11; softmax of such logits is nearly one-hot, and its Jacobian diag(p) − ppᵀ is nearly zero, so gradients vanish. Dividing by √d_k keeps logits at unit scale at initialization.

Numerical stability. Softmax is computed as exp(s − max(s)). In fp16, use a large negative constant instead of −∞ where kernels require it, and compute softmax in fp32. A fully masked row (e.g. padding queries) produces NaN with −∞ masks, so handle it explicitly.

Backward pass, briefly. With dO given:

dV = Pᵀ dO
dP = dO Vᵀ
dS = P ⊙ (dP − rowsum(dP ⊙ P))      # softmax backward
dQ = dS K / √d_k,   dK = dSᵀ Q / √d_k

Storing P (T × T per head) for the backward pass is what makes naive attention memory-hungry; FlashAttention recomputes it block by block instead.

Cost per token per layer: projections 8d² FLOPs, scores and weighted sum 4·T·d. The quadratic term matches the parameter matmuls (including the MLP's 16d²) at T ≈ 6d.

25–32 min Trade-offs and variants

• Additive attention (a small MLP scoring each pair) predates dot-product attention; dot-product wins on hardware because it's a matmul.
• Separate Q and K projections allow asymmetric relevance; tying them forces score(i, j) = score(j, i).
• QK-norm (normalizing queries and keys) bounds logit growth and stabilizes large-scale training.
• Softmax alternatives (sigmoid attention, linear kernels) trade the normalization for efficiency or other properties.

32–37 min Production example

A team saw NaNs in fp16 fine-tuning only on padded batches: some query rows were fully masked, producing softmax(−inf…). Fixing the mask to leave at least one valid key (or zeroing those rows after softmax) and computing softmax in fp32 resolved it. Using a fused attention kernel (FlashAttention-style) also cut activation memory enough to double the batch size.

37–40 min Likely follow-ups

• What does W_O do? Mixes information across heads and maps back to the model dimension.
• What changes in cross-attention? Queries come from one sequence, keys and values from another; there's no causal mask over the source.
• Is the attention map an explanation of the model's decision? Not reliably; it shows routing in one layer, not causal importance.

What interviewers listen for

• Correct shapes, masking before softmax, and per-row softmax.
• The variance argument and the vanishing-gradient consequence.
• Awareness of memory cost from storing P and how fused kernels avoid it.

Red flags

• Applying the mask after softmax.
• "We divide by √d_k to normalize" without the saturation argument.