Softmax in Attention — Converting Scores to Probabilities

Why it matters

Softmax is a bottleneck in attention. Understanding its cost and numerical behavior explains why FlashAttention exists and how it optimizes.

The architecture

Formula: softmax(x)_i = exp(x_i) / sum_j exp(x_j).

Naive implementation overflows: exp(x) for large x is infinity. Numerical stability: subtract max first — softmax(x) = softmax(x - max(x)).

How it works end to end

Log-sum-exp: log(sum(exp(x_i))) = max(x) + log(sum(exp(x_i - max(x)))). Same stability trick, different formulation.

Softmax gradient: d softmax(x)_i / d x_j = softmax(x)_i (δ_ij - softmax(x)_j). O(n²) full Jacobian; usually computed implicitly via chain rule.

Compute cost: for attention with sequence length N, softmax over N values happens N times = O(N²). Same as attention matrix.