Diagonal Natural Gradient Step

Problem

Take a natural gradient step with a diagonal Fisher approximation.

Return x_i - lr * grad_i / (fisher_i + eps) elementwise. If a corrected Fisher entry is not positive, leave that coordinate unchanged.

Hint

Curvature-heavy coordinates receive proportionally smaller steps.

case 1

in: [[0,0],[1,1],[0.25,4],0.1,1e-8]

expected: [-0.3999999840000007,-0.024999999937500003]

case 2

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

expected: [0.75,-2]

case 3

in: [[2],[0],[0],0.1,0]

expected: [2]

Starter codePython

def diag_natural_gradient_step(x, grad, fisher, lr, eps):
    # Your code here
    pass

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