Softmax from the definition

The definition:

softmax(x)_i = exp(x_i) / sum_j exp(x_j)

Transcribed directly:

e = x.exp()
return e / e.sum(dim=-1, keepdim=True)

On [1000., 1001.] that returns [nan, nan].

What happened

float32 tops out around 3.4e38. exp(1000) is about 1e434, so e is [inf, inf], the sum is inf, and inf / inf is nan. Nothing raised; the result is a tensor of the right shape and dtype full of nothing.

The fix, and why it is exact

Subtract the row maximum before exponentiating:

z = x - x.max(dim=-1, keepdim=True).values
e = z.exp()
return e / e.sum(dim=-1, keepdim=True)

This changes nothing mathematically. Multiply the numerator and denominator of the definition by exp(-c) and it cancels, for any c. Choosing c as the row maximum makes the largest exponent exactly exp(0) = 1, so nothing can overflow, and the smallest underflows to 0 which is the correct answer to within float precision anyway.

This is not an approximation or a stabilising fudge. It is the same function, evaluated somewhere the arithmetic works.

Note the keepdim

x.max(dim=-1) returns a named tuple of values and indices, and keepdim=True keeps the reduced axis so the subtraction broadcasts across the row rather than against the wrong axis. Both details have appeared earlier in the course; this is where they combine.

In practice

Use torch.softmax(x, dim=-1), which does this and more. Write it once by hand because the same trick is logsumexp, is inside every cross-entropy implementation, and is the reason attention subtracts a max you never see.

Your task

def softmax(x: torch.Tensor) -> torch.Tensor

Return the softmax over the last axis, stable for large inputs, without calling torch.softmax, torch.nn.functional.softmax or logsumexp.