Scaled Dot-Product Attention in 20 Lines

May 2, 2026

Code note Deep learning

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q, K, V: (B, H, T, d_k)
    d_k = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / d_k ** 0.5
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ V

Why it matters: Without the √d_k divisor, dot products grow with dimension and push softmax into its saturated tail — gradients vanish and training stalls past d_k ≥ 64. The causal mask zeros out future positions before softmax, so each token’s representation depends only on its past — the property that makes autoregressive generation well-defined.