simd-anim
The course

Beyond GEMM · 2/3

4 scores · stable softmax

02 Chapter 2 · softmax

Softmax

Subtract the max, then exp, then divide. The largest exp is 1. A naive exp can overflow while this one stays finite.

Instruction
softmax (x)ᵢ = exp(xᵢ − max x) / Σⱼ exp(xⱼ − max x)

Four scores. Mild and sharp stay in a range where exp(x) is finite either way. Overflow uses values near 1000, where exp(x) is Inf and exp(x − max) is not. The weights of the stable path still sum to 1. Same ratios, no overflow — the subtraction is the whole trick.

softmax · max, subtract, exp, normalize

Raw scores. Play to mark the max.

Arrow keys step, space plays or pauses.

Reading the picture

  1. The max is the whole trick

    Subtracting the max shifts the peak to 0. exp(0) is 1, and every other exp is smaller. The ratios do not change, so the weights do not change.

  2. Naive exp overflows

    exp(1000) is Inf. The overflow preset shows those values going Inf, while the shifted ones stay between −3 and 0 and the weights still sum to 1.

  3. Where it sits in attention

    Attention is QKᵀ, this softmax, then a mix with V. The GEMM pages never do this step. Open the attention page

Mild scores: 1.20, 0.30, 2.00, 0.70.