Cross-Entropy Loss

Negative log-likelihood of the true class under a Softmax over logits. Used as the default loss function for classification and next-token prediction.

Equation

For logits $z\in {\mathbb{R}}^K$ and target class $y$:

$$L=-\log \mathrm{softmax}(z{)}_y=\mathrm{logsumexp}(z)-z_y=m+\log \sum _j{e}^{z_j-m}-z_y,m=\underset{j}{\max }{z}_j$$

The gradient is the residual between prediction and one-hot target:

$$\frac{\partial L}{\partial z_i}=\mathrm{softmax}(z{)}_i-1[i=y]$$

From scratch

import torch
 
def cross_entropy(logits, target, reduction="mean"):
    # logits: (n, k), target: (n,) int64
    m = logits.max(-1, keepdim=True).values
    z = logits - m
    logZ = torch.log(torch.exp(z).sum(-1))            # log-sum-exp, shifted
    loss = logZ - z.gather(-1, target[:, None]).squeeze(-1)
    return loss.mean() if reduction == "mean" else loss
 
def cross_entropy_grad(logits, target):
    p = torch.softmax(logits, -1)
    return p - torch.nn.functional.one_hot(target, logits.size(-1)).to(p.dtype)

Notes

• Never compute softmax then log: fuse via Log Softmax or the shifted logsumexp above.
• Compute the loss in fp32 even under bf16: the logsumexp reduction is the precision-critical step.
• Equivalent to KL$(p^{\star }\Vert p)$ up to a constant, since the target entropy is fixed; see Knowledge Distillation Loss.
• Encourages overconfidence; Label Smoothing and Focal Loss are the usual counterweights.