18. Scaled Dot-Product Attention From Scratch

Problem

Every transformer-family sequence model, and every efficient-attention variant an interviewer might ask about next, is built on one primitive: scaled dot-product attention. Before touching multi-head splits, masks, or caching, you need to be able to write attention(Q, K, V) = softmax(QKT / √d_k) V cold, with the scaling and softmax math explicit rather than hidden behind a single fused call.

The 1/√d_k scale factor matters for more than convention: if each component of q and k is roughly unit-variance and independent, their dot product has variance that grows linearly with d_k, so its standard deviation grows like √d_k. Without rescaling, scores blow up in magnitude for large d_k, softmax saturates onto one logit, and gradients through the softmax vanish almost everywhere. Dividing by √d_k keeps pre-softmax logits at roughly unit variance regardless of head dimension.

Source: src/18_scaled_dot_product_attention.py

def scaled_dot_product_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None
) -> torch.Tensor: ...

>>> q = torch.randn(1, 2, 4)
>>> k = torch.randn(1, 3, 4)
>>> v = torch.randn(1, 3, 3)
>>> scaled_dot_product_attention(q, k, v).shape
torch.Size([1, 2, 3])

Step-by-Step Approach

1. Compute raw similarity scores between every query and every key: q @ k.transpose(-1, -2), giving shape (batch, seq_q, seq_k).
2. Scale by 1 / sqrt(d_k), where d_k is the last dimension of q (and k) — this keeps the pre-softmax logit magnitudes roughly constant as head dimension changes.
3. If a mask is given, set the scores at blocked positions to -inf before softmax (never zero the post-softmax weights directly, or the remaining weights won't sum to 1).
4. Apply softmax along the last axis (the key axis) so every query's attention weights form a probability distribution over keys.
5. Multiply the attention weights by v: attn_weights @ v, giving shape (batch, seq_q, d_v) — each output row is a weighted average of value vectors.

The key insight is that attention is just two matrix multiplies with a normalization step sandwiched between them: the first matmul turns "how similar is this query to each key" into a score, softmax turns scores into a distribution, and the second matmul turns that distribution into a weighted blend of values.

Reference solution

import math
import torch


def scaled_dot_product_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None
) -> torch.Tensor:
    d_k = q.shape[-1]
    # (batch, seq_q, d_k) @ (batch, d_k, seq_k) -> (batch, seq_q, seq_k)
    scores = q @ k.transpose(-1, -2) / math.sqrt(d_k)
    if mask is not None:
        # masked positions get -inf pre-softmax so their softmax weight is ~0;
        # never zero the *weights* post-softmax, that would leave the
        # remaining weights not summing to 1
        scores = scores.masked_fill(~mask, float("-inf"))
    # softmax over the key dimension (last axis) -> each query's weights sum to 1
    attn_weights = torch.softmax(scores, dim=-1)
    # (batch, seq_q, seq_k) @ (batch, seq_k, d_v) -> (batch, seq_q, d_v)
    return attn_weights @ v

Key Functions & Tricks

• tensor.transpose(-1, -2) — swaps the last two dims so the key matmul contracts over d_k instead of the sequence axis; safer than .T once a batch dim is present.
• math.sqrt(d_k) as a Python float scalar — cheaper than a tensor op and broadcasts fine against the score tensor.
• tensor.masked_fill(~mask, float("-inf")) — sets blocked score entries to -inf so exp(-inf) == 0 zeroes their softmax weight exactly.
• torch.softmax(scores, dim=-1) — internally uses the log-sum-exp / max-subtraction trick, so it stays numerically stable even with -inf entries in the row.
• Batched @ (torch.matmul) — broadcasts over the leading batch dimension automatically as long as the last two dims align for matrix multiply.

How to Recognize This Pattern