ml/optimizers

Lion

Lion is an learned optimizer that only uses the gradient sign for the update vector. The paper Symbolic Discovery of Optimization Algorithms presents how they used symbolic AI.

def train(weight, gradient, momentum, lr):
    # default beta1 = 0.9, beta2 = 0.99
    # only tracks momentum and uses sign to compute update
    update = interp(gradient, momentum, beta1)
    update = sign(update)
    momentum = interp(gradient, momentum, beta2)

    # regular weight decay
    weight_decay = weight * lambda
    update = update + weight_decay

    update = update * lr
    return update, momentum