Home / Question bank / LLM Fundamentals

LLM FundamentalsMidQuestion 2 of 13

Explain scaled dot-product attention. Why divide 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

Check the depth they want: intuition, the math, or implementation details (masking, shapes, numerical stability). Framing: "I'll write the formula, explain each term with shapes, justify the √d_k scaling with a short variance argument, then cover masking, multi-head attention and cost."

5–15 min Core explanation

Given token representations X of shape n × d:

1. Project to queries, keys and values: Q = XW_Q, K = XW_K, V = XW_V, each n × d_k.
2. Compute similarity scores: S = QKᵀ / √d_k (shape n × n). Entry S[i, j] says how relevant token j is to token i.
3. Apply softmax row by row so each row is a probability distribution.
4. Output softmax(S) · V: each token's output is a weighted average of all value vectors.

Intuition: query = what I'm looking for; key = what I offer for matching; value = the content I hand over if matched. Separating key and value lets a token be found for one reason and contribute different information.

15–25 min Deep dive: why divide by √d_k

Assume query and key components are independent with mean 0 and variance 1. Their dot product is a sum of d_k such products, so its variance is d_k and its standard deviation is √d_k.

With d_k = 128, raw scores have standard deviation ≈ 11. Softmax of numbers that spread out pushes almost all weight onto the maximum. In that saturated regime, gradients of softmax are nearly zero, so learning stalls. Dividing by √d_k restores unit variance and keeps softmax in a useful range.

Implementation details worth mentioning:

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (n, n)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = scores.softmax(dim=-1)                      # rows sum to 1
    return weights @ V                                    # (n, d_k)

• Causal mask: in decoder-only models, position i can't attend to j > i. Set those scores to −∞ before softmax.
• Numerical stability: softmax implementations subtract the row max before exponentiating. Mixed precision makes this essential.
• Multi-head: run h attention heads with d_k = d/h in parallel, concatenate, then project with W_O. Heads specialize (local syntax, coreference, delimiters) at the same total cost as one wide head.

25–32 min Trade-offs and alternatives

• Cost: QKᵀ is n × n, so compute and naive memory are O(n²). FlashAttention computes the exact result in tiles without materializing the full matrix, cutting memory traffic dramatically.
• Additive attention (Bahdanau) uses a small feed-forward network to score pairs; dot-product is faster on GPUs because it's a matrix multiply.
• Alternatives for long sequences: sliding-window, sparse patterns, linear attention, state-space models. They trade some modelling power for sub-quadratic cost.
• Other normalizations: some architectures normalize Q and K (QK-norm) to stabilize training at scale.

32–37 min Production example

When serving a model with long prompts (RAG with 20k tokens of context), prefill time is dominated by attention over the prompt. Practical consequences: use a serving stack with FlashAttention-style kernels, cap retrieved context to what's needed, and use prompt caching for stable prefixes so attention over them isn't recomputed per request.

37–40 min Likely follow-ups

• Why separate W_Q and W_K instead of one matrix? Relevance needn't be symmetric; token i may need j without j needing i.
• What does the attention matrix tell you? Weights per head, useful for intuition but not a faithful explanation of model decisions.
• What changes in cross-attention? Queries come from one sequence (decoder) and keys/values from another (encoder output).
• Where does position come in? Not in this formula; via positional embeddings or RoPE applied to Q and K.

What interviewers listen for

• Correct shapes and the role of each of Q, K, V.
• The variance argument for √d_k, not just "it normalizes".
• Masking, multi-head, and the O(n²) cost.

Red flags

• Forgetting softmax is applied per row.
• Claiming scaling is about "making values small" without the saturation/gradient reason.
• Not knowing the causal mask.