AdaBelief Corrected Variance

Problem

Update the AdaBelief belief variance and bias-correct it.

Given the running variance s, the current gradient g, the first moment m, step t, decay beta2, and eps, compute s_new = beta2 * s + (1 - beta2) * (g - m)^2 + eps, then return [s_new, s_new / (1 - beta2^t)]. If the correction denominator is zero, return s_new for both entries.

Hint

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

Test cases

case 1

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

expected: [0.0002500100000000002,0.25001]

case 2

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

expected: [0.010150010000000001,2.0340660660335996]

case 3

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

expected: [0.40410001,4.226195141101516]

Starter codePython

def adabelief_corrected_variance(s, g, m, t, beta2, eps):
    # Returns [s_new, s_hat]
    # Your code here
    pass

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