All problems
Temperature-Scaled Softmax Gradient
Problem
With temperature scaling the probabilities are s = softmax(x / tau), and the Jacobian is (1/tau) * (diag(s) - s s^T).
Return row i of that Jacobian: (1/tau) * s_i * (delta_ij - s_j) for each j. Given the logits, the row index i, and tau > 0, return the row.
Hint
Low temperature sharpens the distribution and inflates the gradient by 1/tau.
Test cases
case 1
in: [[0,0],0,2]
expected: [0.125,-0.125]
case 2
in: [[1,1,1],2,1]
expected: [-0.1111111111111111,-0.1111111111111111,0.22222222222222224]
case 3
in: [[2,-1],0,0.5]
expected: [0.004933018582719862,-0.004933018582720097]
case 4
in: [[0,1],0,1]
expected: [0.19661193324148185,-0.19661193324148185]
Starter codePython
def temperature_softmax_grad(logits, i, tau):
# Your code here
pass4 test cases · run your code against them in the editor