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 log⁡−softmax⁡(zi)  =  zi  −  log⁡sumexp⁡\operatorname{log} - \operatorname{softmax}(z_{i})\; = \;z_{i}\; - \;\operatorname{log}_{\mathrm{sum}}\operatorname{exp}(z) where log⁡sumexp⁡\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 log⁡sumexp⁡\operatorname{log}_{\mathrm{sum}}\operatorname{exp} trick to avoid overflow
  • It's faster on CPU only
  • Because softmax is always non-differentiable

log⁡sumexp⁡\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