Sophia Clipped Ratio Update

Problem

Perform one Sophia optimizer update with a clipped ratio and a Hessian EMA.

Update m = beta1 * m + (1 - beta1) * g and h = beta2 * h + (1 - beta2) * g^2, then compute ratio = m / max(rho * h, eps), clip it to [-1, 1], and set x = x - lr * ratio. Return [x, m, h].

Hint

The Hessian EMA in the denominator makes the clipped ratio curvature-aware.

Test cases

case 1

in: [1,0,0,0.5,0.1,0.9,0.99,0.05,1e-8]

expected: [0.9,0.04999999999999999,0.0025000000000000022]

case 2

in: [1,0.1,2,0.5,0.01,0.9,0.99,0.1,1e-8]

expected: [0.992938209331652,0.14,1.9825]

case 3

in: [-2,-0.3,0.4,-1.2,0.05,0.8,0.95,0.2,1e-8]

expected: [-1.95,-0.4799999999999999,0.45200000000000007]

case 4

in: [0,0,0,0,0.1,0.9,0.99,0.1,1e-8]

expected: [0,0,0]

Starter codePython

def sophia_update(x, m, h, g, lr, beta1, beta2, rho, eps):
    # Returns [x, m, h]
    # Your code here
    pass

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