From scores to a distribution.

A dot product can give scores such as [-2, 0, 3]. Those numbers cannot directly serve as our previous lesson’s weighted-average weights: one is negative, and the total is not one. We need a transformation that preserves which scores are larger while producing a distribution over the available choices.

Softmax does this in two stages. First turn scores into positive evidence. Then divide each evidence value by the total evidence.

1. A probability distribution needs a stated set of choices

For a finite set of mutually exclusive choices, probabilities are nonnegative and sum to one. A vector like [0.2, 0.3, 0.5] is incomplete without a meaning for its entries. Does the third entry mean a source position, a vocabulary symbol, or a class label?

Attention normalizes over source positions. A next-token prediction normalizes over vocabulary symbols. Both can use softmax while answering different questions. A 0.5 attention weight is not a 50% claim that a particular word should be generated next.

2. What does exp mean?

exp(s) means the mathematical constant e raised to the score s; e is approximately 2.718. You do not need to memorize many powers. Three properties matter here: exp is always positive, larger inputs give larger outputs, and adding a common amount to inputs multiplies all outputs by the same factor.

log is the natural logarithm, the inverse of exp. For example, exp(log(2)) = 2. We choose scores [0, log(2), log(3)], approximately [0, 0.6931, 1.0986]. Their exponentials are exactly [1, 2, 3]. Those sum to 6, so normalization gives [1/6, 2/6, 3/6].

3. The denominator connects the choices

Increase one score while keeping the others fixed. Its evidence increases, but the denominator also increases. Its probability rises while the other probabilities fall. Softmax therefore does not transform each coordinate independently: choices compete through a shared denominator.

If all scores are equal, all evidence values are equal, and the result is a uniform distribution. With four choices, each receives 0.25. It does not matter whether those equal scores are all 0, all 100, or all −100.

4. Temperature changes concentration

At temperature 0.5, our scores are doubled before exponentiation. The evidence becomes [1, 4, 9], and the probabilities become [1/14, 4/14, 9/14]. The largest choice now gets about 0.643 instead of 0.5.

As temperature increases, differences shrink and the distribution moves toward uniformity. As it approaches zero from above, mass concentrates on the maximum-scoring choices. We do not set temperature to zero: the formula would divide by zero. Tied maxima remain tied.

Temperature is a control in our experiment. The later attention formula also divides dot products by the square root of key width for a separate scale-related reason. Do not confuse that architectural scaling with a freely adjusted decoding temperature.

5. Stable arithmetic gives the same distribution

Very large positive scores can overflow when exponentiated by a computer. Before exp, subtract the largest score in each row. That makes the largest shifted score zero, so its exponential is one, and all others are at most one.

This does not change probabilities: subtracting the same amount from all scores multiplies every exponential by the same positive factor. That factor appears in both numerator and denominator and cancels. We change the arithmetic range, not the intended distribution.

For multiple rows of scores, subtract each row’s own maximum and divide by each row’s own total. Normalizing the entire matrix at once would make unrelated receivers compete with each other.

def softmax(scores: FloatArray, temperature: float = 1.0) -> FloatArray:
    """Normalize the final axis; temperature must be positive."""
    if temperature <= 0:
        raise ValueError("Temperature must be positive")
    scaled = scores / temperature
    shifted = scaled - scaled.max(axis=-1, keepdims=True)
    evidence = np.exp(shifted)
    probabilities = evidence / evidence.sum(axis=-1, keepdims=True)
    return probabilities

def softmax(scores: Tensor, temperature: float = 1.0) -> Tensor:
    if temperature <= 0:
        raise ValueError("Temperature must be positive")
    scaled = scores / temperature
    probabilities = scaled.softmax(dim=-1)
    return probabilities

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

The final axis is selected with axis=-1 or dim=-1. A negative axis index counts from the end. keepdims=True keeps the reduced dimension at length one so subtraction and division broadcast back across the correct row. PyTorch’s softmax performs the stable normalization internally.

6. A mask removes choices before normalization

When a choice must be forbidden, attention typically replaces its score with negative infinity before softmax. Its exponential is zero, so it receives zero weight. Setting a score to zero does not remove it: exp(0) is one.

There must still be at least one allowed choice. If every choice is removed, no valid normalized distribution remains. The attention implementation checks this condition. The intermediate course adds causal and padding masks after the core operation is familiar.

Carry this forward

We can now turn scores into weights and use weights to mix values. There is one more missing foundation: where do the projection weights come from? We will make a single number learn before asking a full Transformer to do it.