Log-Softmax Cross-Entropy from Logits

Problem

Compute the cross-entropy loss in nats directly from raw logits for one labelled sample using the log-sum-exp trick. The returned value is logsumexp(logits) - logits[label], evaluated stably by subtracting the maximum logit. This is the negative log-probability of the true class.

Hint

Subtract the max logit before exponentiating, then add it back through the log-sum.

Test cases

case 1

in: [[0,0],0]

expected: 0.6931471805599453

case 2

in: [[2,1,0],0]

expected: 0.4076059644443806

case 3

in: [[1,2,3],2]

expected: 0.4076059644443806

case 4

in: [[1],0]

expected: 0

case 5

in: [[-1,-2],1]

expected: 1.3132616875182228

Starter codePython

import math
def log_softmax_cross_entropy(logits, label):
    # Your code here
    pass

5 test cases · run your code against them in the editor