Log Softmax

The logarithm of Softmax computed in one stable pass via logsumexp; the input to NLL loss, KL divergence, and any log-probability computation.

Equation

lo g softmax (x)_{i} = x_{i} - lo g j \sum e^{x_{j}} = (x_{i} - m) - lo g j \sum e^{x_{j} - m}, m = j max x_{j}

Gradient (with $y = softmax (x)$ and upstream $g$ ):

\frac{\partial L}{\partial x _{i}} = g_{i} - y_{i} j \sum g_{j}

From scratch

import torch
 
def log_softmax(x, dim=-1):
    m = x.max(dim, keepdim=True).values
    z = x - m
    return z - torch.log(torch.exp(z).sum(dim, keepdim=True))
 
def log_softmax_grad(out, g, dim=-1):
    # out = log_softmax(x) from forward, g = dL/dout
    return g - out.exp() * g.sum(dim, keepdim=True)

Notes

Strictly better than softmax().log(): the latter underflows to $- \infty$ whenever a probability rounds to 0.
Gradients stay bounded even for very confident predictions, which is why Cross-Entropy Loss fuses the two.
$logsumexp$ is the $max$ smoothed by the temperature-1 Softmax; its gradient is the softmax.
Handle $- \infty$ masks by setting masked logits to a large negative finite value, not -inf, or an all-masked row yields nan.