Scaled Dot-Product Attention

Given queries $Q \in \R^{n \times d_k}$, keys $K \in \R^{m \times d_k}$ and values $V \in \R^{m \times d_v}$, attention returns a convex combination of value rows weighted by query–key similarity.

Why $1/\sqrt{d_k}$

If $q$ and $k$ have i.i.d. components with zero mean and unit variance, their dot product has variance $d_k$. Unscaled logits therefore grow with dimension, pushing the softmax toward one-hot and its Jacobian toward zero.

Gradient

Writing $A = \softmax(S)$ with $S = QK^{\top}/\sqrt{d_k}$, the backward pass through the softmax takes the usual diagonal-minus-outer-product form, applied per row.

The $n^2$ term is materialisation, not arithmetic necessity — see flash-attention for the tiled formulation with $O(n)$ memory.

import math
import torch
from torch import Tensor


def attention(
    q: Tensor, k: Tensor, v: Tensor,   # [B, H, N, D]
    mask: Tensor | None = None,
    dropout_p: float = 0.0,
) -> Tensor:
    scale = 1.0 / math.sqrt(q.size(-1))
    logits = (q * scale) @ k.transpose(-2, -1)

    if mask is not None:
        logits = logits.masked_fill(mask, -6e4)

    attn = logits.softmax(dim=-1)
    if dropout_p > 0.0:
        attn = torch.dropout(attn, dropout_p, train=True)
    return attn @ v

Fold the scale into the query before the matmul when running fp16 — scaling after the product overflows for large $d_k$. A fully masked row yields NaN under $-\infty$, so mask with a large negative constant instead when rows can be empty.