Lion Vector Sign Step

Problem

Perform one Lion update on a weight vector with decoupled weight decay.

For each coordinate let c_i = beta1 * m_i + (1 - beta1) * g_i and take u_i = sign(c_i). Update w_i = w_i - lr * (u_i + wd * w_i), then set m_i = beta2 * m_i + (1 - beta2) * g_i. Return [w_new, m_new] as two lists.

Hint

Only the sign of the interpolated momentum enters the update; decay is decoupled.

Test cases

case 1

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

expected: [[0.9,2.1],[0.0050000000000000044,-0.0050000000000000044]]

case 2

in: [[0.5,-0.5],[0.1,-0.2],[0,0],0.01,0.9,0.99,0.1]

expected: [[0.4895,-0.4895],[0.099,-0.198]]

case 3

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

expected: [[0.995,0.995],[0,0]]

case 4

in: [[-1],[-0.5],[-2],0.2,0.8,0.9,0]

expected: [[-0.8],[-0.6499999999999999]]

Starter codePython

def lion_vector_step(w, m, g, lr, beta1, beta2, wd):
    # Returns [w_new, m_new]
    # Your code here
    pass

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