optax_lion

phasic.optax_wrapper.optax_lion(
    learning_rate=0.0001,
    b1=0.9,
    b2=0.99,
    weight_decay=0.0,
)

Create Optax Lion optimizer wrapped for phasic.

Lion (Evolved Sign Momentum) is a memory-efficient optimizer discovered through program search. It uses sign-based updates and typically requires a lower learning rate than Adam.

Parameters

learning_rate : float or optax.Schedule = 1e-4

Learning rate. Typically 3-10x smaller than Adam.

b1 : float = 0.9

Exponential decay rate for first moment.

b2 : float = 0.99

Exponential decay rate for the gradient moving average.

weight_decay : float = 0.0

Weight decay coefficient.

Returns

: OptaxOptimizer

Wrapped Lion optimizer compatible with phasic SVGD.

References

Chen et al. (2023). Symbolic Discovery of Optimization Algorithms. arXiv:2302.06675. https://arxiv.org/abs/2302.06675