Softmax Cross-Entropy Hessian

Problem

For softmax probabilities p = softmax(logits) and the cross-entropy loss L = -log(p_y), the Hessian with respect to the logits is diag(p) - p p^T.

Given the logits, return this n x n Hessian as a nested list. Subtract the maximum logit before exponentiating for numerical stability.

Hint

The Hessian is the covariance of the one-hot label under the softmax distribution.

Test cases

case 1

in: [[0,0]]

expected: [[0.25,-0.25],[-0.25,0.25]]

case 2

in: [[0,1]]

expected: [[0.19661193324148185,-0.19661193324148185],[-0.19661193324148185,0.19661193324148185]]

case 3

in: [[1,1,1]]

expected: [[0.2222222222222222,-0.1111111111111111,-0.1111111111111111],[-0.1111111111111111,0.2222222222222222,-0.1111111111111111],[-0.1111111111111111,-0.1111111111111…

case 4

in: [[2,-1]]

expected: [[0.045176659730912005,-0.045176659730912144],[-0.045176659730912144,0.04517665973091214]]

Starter codePython

def softmax_ce_hessian(logits):
    # Your code here
    pass

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