Lion Sign Update

Problem

Perform one Lion optimizer update on a scalar parameter.

With the interpolation c = beta1 * m + (1 - beta1) * g, take update = sign(c) and set x_new = x - lr * (update + weight_decay * x). Then update the momentum m_new = beta2 * m + (1 - beta2) * g. Return [x_new, m_new].

Hint

Only the sign of the update is used; momentum is updated after x.

Test cases

case 1

in: [1,0,0.5,0.1,0.9,0.99,0]

expected: [0.9,0.0050000000000000044]

case 2

in: [1,0.1,0.5,0.01,0.9,0.99,0.01]

expected: [0.9899,0.10400000000000001]

case 3

in: [-2,-0.3,-1.2,0.05,0.9,0.99,0]

expected: [-1.95,-0.309]

case 4

in: [0.5,0,0,0.1,0.9,0.99,0]

expected: [0.5,0]

Starter codePython

def lion_update(x, m, g, lr, beta1, beta2, weight_decay):
    # Returns [x_new, m_new]
    # Your code here
    pass

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