AdaBelief Belief Update

Problem

Perform one AdaBelief update on a scalar parameter.

Update m = beta1 * m + (1 - beta1) * g and s = beta2 * s + (1 - beta2) * (g - m) ** 2 + eps, bias-correct both, then set x = x - lr * m_hat / (sqrt(s_hat) + eps). Return [x, m, s].

Hint

The belief variance tracks the surprise (g - m) squared, not g squared.

Test cases

case 1

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

expected: [0.8888916347405214,0.04999999999999999,0.0002025100000000002]

case 2

in: [1,0.1,0.01,0.5,5,0.01,0.9,0.999,1e-8]

expected: [0.9975993306863485,0.14,0.010119610000000001]

case 3

in: [-2,-0.3,0.4,1.2,10,0.05,0.8,0.99,1e-8]

expected: [-2,-5.551115123125783e-17,0.41040001000000004]

Starter codePython

def adabelief_update(x, m, s, g, t, lr, beta1, beta2, eps):
    # Returns [x, m, s]
    # Your code here
    pass

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