Gradient of Cross-Entropy

Problem

For softmax probabilities p = softmax(logits) and the label y, the cross-entropy loss is L = -log(p_y).

Its gradient with respect to the logits is p - onehot(y), where the y-th component is p_y - 1 and all others are p_j. Return the gradient vector. Subtract the maximum logit for stability.

Hint

The softmax makes the loss gradient a probability vector minus a one-hot vector.

Test cases

case 1

in: [[0,0],0]

expected: [-0.5,0.5]

case 2

in: [[0,1],1]

expected: [0.2689414213699951,-0.2689414213699951]

case 3

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

expected: [0.3333333333333333,0.3333333333333333,-0.6666666666666667]

case 4

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

expected: [-0.047425873177566635,0.04742587317756679]

Starter codePython

def cross_entropy_grad(logits, y):
    # Your code here
    pass

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