EasyDeepLearn

Why compute softmax + cross-entropy jointly via log-sum-exp?

medium

Answer

  • Naive softmax  =  exp(zi)  /  sum(exp(zj))\operatorname{softmax}\; = \;\operatorname{exp}(z_{i})\; / \;\mathrm{sum}(\operatorname{exp}(z_{j})) overflows for large logits.
  • Compute logsoftmax(zi)  =  zi    logsumexp\operatorname{log} - \operatorname{softmax}(z_{i})\; = \;z_{i}\; - \;\operatorname{log}_{\mathrm{sum}}\operatorname{exp}(z) where logsumexp\operatorname{log}_{\mathrm{sum}}\operatorname{exp}(z) = max(z)  +  log\operatorname{max}(z)\; + \;\operatorname{log}(sum(exp(z    max(z))\operatorname{exp}(z\; - \;\operatorname{max}(z)))).
  • Numerically stable.
  • Cross-entropy then becomes -sum(target * log-softmax) — computed in one op (nn.CrossEntropyLoss in PyTorch), avoiding a separate softmax that could overflow before the log.
Check yourself — multiple choice
  • Only for aesthetic reasons
  • Numerical stability — compute log-softmax directly with logsumexp\operatorname{log}_{\mathrm{sum}}\operatorname{exp} trick to avoid overflow
  • It's faster on CPU only
  • Because softmax is always non-differentiable

logsumexp\operatorname{log}_{\mathrm{sum}}\operatorname{exp} trick: log(sum  exp(zi))  =  max(z)  +  log\operatorname{log}(\mathrm{sum}\;\operatorname{exp}(z_{i}))\; = \;\operatorname{max}(z)\; + \;\operatorname{log}(sum(exp(zi    max(z))\operatorname{exp}(z_{i}\; - \;\operatorname{max}(z)))) — no overflow.

#fundamentals#losses

Practise Deep Learning

214 interview questions in this topic.

Related questions