mojolearn.Lion#

class mojolearn.Lion(params, lr=0.0001, betas=(0.9, 0.99), weight_decay=0.0, numeric_mode=None)[source]#

Bases: _SeqOptimizer

Lion optimizer (Chen et al., 2023, “Symbolic Discovery of Optimization Algorithms”).

Follows the lion-pytorch statement of the update: p *= 1 - lr * weight_decay; p -= lr * sign(beta1 * m + (1 - beta1) * g); m = beta2 * m + (1 - beta2) * g. Defaults are the paper’s.

Parameters:

  • params (ndarray or list of ndarray) – Parameters to update in place.

  • lr (float, default 1e-4) – Learning rate.

  • betas (tuple of float, default (0.9, 0.99)) – Interpolation factor of the update direction and decay of the momentum.

  • weight_decay (float, default 0.0) – Decoupled weight decay.

  • numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode of the update. None uses the library default.

Variables:

  • t (int) – Number of steps taken.

  • exp_avg (ndarray) – Momentum buffer, flat over all parameters.

Notes

params must be C-contiguous float32 NumPy arrays, updated in place; anything else raises TypeError. step(grads) takes one gradient per parameter, in the same order and shapes; float64 gradients raise TypeError. lr is read at every step, so you can change opt.lr between steps, or set opt.lr_schedule to a schedule (StepLR, ExponentialLR, OneCycleLR, or any object with a one-based lr_at(t)), which sets lr before each step. state_dict() and load_state_dict() save and restore the step count, learning rate and state buffers.

Not supported: parameter groups, maximize, foreach/fused/capturable/ differentiable, sparse gradients and float64 parameters.

Examples

>>> import numpy as np, mojolearn
>>> w = np.zeros(128, dtype=np.float32)
>>> opt = mojolearn.Lion([w], lr=1e-4).step([np.ones_like(w)])
__init__(params, lr=0.0001, betas=(0.9, 0.99), weight_decay=0.0, numeric_mode=None)[source]#

Methods

load_state_dict(sd)

state_dict()

step(grads)

One update of every param from grads (same shapes, same order).

Attributes

exp_avg

Momentum buffer, flat over all parameters.

resident_

True once the state lives on the device.

state

property exp_avg#

Momentum buffer, flat over all parameters.