Softmax Cross-Entropy Hessian Matrix

Problem

For softmax probabilities p, the Hessian of the cross-entropy loss with respect to the logits is diag(p) - p p^T: diagonal entries p_i (1 - p_i) and off-diagonal entries -p_i p_j.

Given the logits, return this Hessian as a nested list. Subtract the max logit before exponentiating.

Hint

Start from diag(p) and subtract the outer product p p^T.

Test cases

case 1

in: [[0,0]]

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

case 2

in: [[1,2,3]]

expected: [[0.08192506906499324,-0.022033044520174298,-0.05989202454481893],[-0.022033044520174298,0.1848364465099787,-0.1628034019898044],[-0.05989202454481893,-0.162803…

case 3

in: [[0,0,0,0]]

expected: [[0.1875,-0.0625,-0.0625,-0.0625],[-0.0625,0.1875,-0.0625,-0.0625],[-0.0625,-0.0625,0.1875,-0.0625],[-0.0625,-0.0625,-0.0625,0.1875]]

Starter codePython

def softmax_ce_hessian_matrix(logits):
    # Your code here
    pass

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