Stable softmax, cross-entropy and loss reductions

1. Losses from maximum likelihood

   Understanding Deep Learning - Chapters 3 to 7

   Chapter 5, Sections 5.1 to 5.2 (Eq 5.1 to 5.6, printed pp. 56 to 60) for likelihood to negative log-likelihood; Section 5.3 (Eq 5.7 to 5.15, pp. 61 to 64) for least squares; Section 5.4 (Eq 5.16 to 5.20, pp. 64 to 66) for sigmoid plus binary cross-entropy; Section 5.5 (Eq 5.21 to 5.24, pp. 67 to 69) and Section 5.7 (Eq 5.27 to 5.31, pp. 71 to 72) for softmax plus multiclass cross-entropy

   Done
2. Naive softmax and the spelled-out cross-entropy

   D2L - tensors, autograd, MLP and backpropagation

   The Softmax: the three-step naive code plus the warning that it is not robust and not for serious use; The Cross-Entropy Loss: y_hat[[0, 1], y] indexing and the reduce_mean of negative logs

   Depth · Do not copy the naive softmax or the log-of-probability loss into your own code. Use them only to see what the stable versions below replace.

   Done
3. The max-subtraction trick

   CS231n Linear Classification - Softmax classifier and numeric stability

   Softmax classifier: \(L_i = -f_{y_i} + \log \sum_j e^{f_j}\); Practical issues: Numeric stability: the f -= np.max(f) code turning [123, 456, 789] into [-666, -333, 0]

   Done
4. Masked reductions in `cross_entropy`

   `torch.nn.functional.cross_entropy` - `reduction` and `ignore_index`

   ignore_index and reduction parameters: ignored targets contribute no gradient, and mean averages over the non-ignored targets only

   Done

Focus on

• Max subtraction is exact, not approximate: shifting every logit by \(c\) multiplies numerator and denominator by \(e^c\), which cancels. The largest shifted logit is 0, so every \(\exp\) input is at most 0 and no term overflows.
• The fused per-example loss is \(L_i = -f_{y_i} + \log \sum_j \exp(f_j)\), and its logit gradient is \(\operatorname{softmax}(z)\) minus one-hot(\(y\)). This matches centered differences to about 1e-9 in float64.
• Binary cross-entropy with logits must use the log-sigmoid form, never sigmoid followed by log: a confident sigmoid rounds to exactly 0 or 1 and the log is then infinite. Its logit gradient is \(\operatorname{sigmoid}(z) - y\).
• Mean divides every gradient by \(N\); a masked mean divides by the valid count and gives masked positions exactly zero gradient. D2L cross_entropy uses reduce_mean; UDL least squares (Eq 5.11) sums.

Reproduce closed-book

Traps

• Taking \(\log(\operatorname{softmax}(z))\) in two steps. Even with max subtraction, a very small probability can underflow to 0 and its log becomes -inf. log_softmax(z) = \((z - m) - \log \sum \exp(z - m)\), with \(m\) the max, never forms the probability.
• Clipping probabilities after a naive softmax to fix overflow. exp(1000) already overflowed to inf and the division already produced nan; clipping cannot recover the true distribution.
• Expecting the stable answer on [1000, 1001, 1002] to be [1, 0, 0]. Only logit gaps matter, so the result equals softmax([0, 1, 2]): about [0.09, 0.245, 0.665].

Practice

Self-test

1. Why does subtracting the maximum logit prevent overflow without changing softmax?

   Show answer Got it Missed

2. What is the gradient of the per-example cross-entropy with respect to the logits?

   Show answer Got it Missed

3. A batch of 5 has 2 masked positions under a masked mean. What is the divisor, and what gradient do masked positions get?

   Show answer Got it Missed

4. Why must binary cross-entropy with logits avoid sigmoid-then-log?

   Show answer Got it Missed

Leads to

• 1.3 Computation graphs and reverse accumulation Push the loss gradient back through the network.
• 1.5 Gradient checks with finite differences Check the fused loss gradient against finite differences.

Comes after

• 2.4 Gradients of logistic, softmax cross-entropy and regularized lossesThat unit fuses log_softmax with max subtraction and uses softmax minus one-hot; this unit derived both.
• 5.7 Stable log-sum-exp, online reductions and mixed precisionApply the log-partition identity to stable probability losses.
• 1.1 Affine maps, nonlinearities and logitsFeed logits into a stable softmax and cross-entropy.
• 3.7 Objective checks, tiny-set overfit and shuffled-label controlsGround the loss in stable softmax and reduction semantics.
• 5.1 Embedding geometry and contrastive logitsImplement the log-sum-exp and reduction the denominator needs.

Resources for this unit

• Understanding Deep Learning - Chapters 3 to 7
  Simon J. D. Princeprimarytheory
  Read Chapters 3-7, emphasizing Chapter 5 for likelihood-derived losses and Chapter 7 for backpropagation and initialization. Use notebooks 3.1, 4.1, 5.1, 6.1 and 7.1 only as equation and plot companions.
• D2L - tensors, autograd, MLP and backpropagation
  D2L authorsrequiredimplementation
  Use chapter_preliminaries ndarray and autograd, then MLP, backpropagation, and numerical stability and initialization. Recreate scalar operations, affine maps, activation, reduction, stable cross-entropy, topological backward and unbroadcast in a blank file.
• CS231n Linear Classification - Softmax classifier and numeric stability
  Stanford CS231nreferencereference
  Read the Softmax classifier section for the cross-entropy form \(L_i = -f_{y_i} + \log \sum_j e^{f_j}\), then the Practical issues: Numeric stability part for the max-subtraction code f -= np.max(f).
• `torch.nn.functional.cross_entropy` - `reduction` and `ignore_index`
  PyTorchreferencereference
  Read the ignore_index and reduction parameters: ignored targets contribute no gradient, and mean reduction averages over the non-ignored targets only.