Struct Lion

pub struct Lion {
    pub lr: f32,
    pub betas: (Array, Array),
    pub weight_decay: f32,
    pub state: State,
}

The Lion optimizer [1].

Since updates are computed through the sign operation, they tend to have larger norm than for other optimizers such as SGD and Adam. We recommend a learning rate that is 3-10x smaller than AdamW and a weight decay 3-10x larger than AdamW to maintain the strength `(lr

• wd)`. Our Lion implementation follows the original paper. In detail,

[1]: Chen, X. Symbolic Discovery of Optimization Algorithms. arXiv preprint arXiv:2302.06675.

Fields

lr: f32

The learning rate.

betas: (Array, Array)

The coefficients used for computing running averages of the gradient and its square. Default to Lion::DEFAULT_BETAS.

weight_decay: f32

The weight decay. Default to Lion::DEFAULT_WEIGHT_DECAY.

state: State

Inner state.

Implementations

impl Lion

pub fn new(lr: impl Into<f32>) -> Self

Creates a new instance of Lion.

impl Lion

pub const DEFAULT_BETAS: (f32, f32) = _

Default values for betas

pub const DEFAULT_WEIGHT_DECAY: f32 = 0f32

Default value for weight_decay

Trait Implementations

impl Buildable for Lion

type Builder = LionBuilder

The builder type for this buildable type

impl Builder<Lion> for LionBuilder

type Error = Infallible

Error with building

fn build(self) -> Result<Lion, Self::Error>

Build the type

impl Clone for Lion

fn clone(&self) -> Lion

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Lion

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Optimizer for Lion

type State = HashMap<Rc<str>, Array>

State of the optimizer.

fn state(&self) -> &Self::State

Get the state of the optimizer.

fn state_mut(&mut self) -> &mut Self::State

Get the mutable state of the optimizer.

fn update_single( &mut self, key: &Rc<str>, gradient: &Array, parameter: &mut Array, ) -> Result<(), Exception>

Update a single parameter with the given gradient.

fn update<M>( &mut self, model: &mut M, gradients: impl Borrow<FlattenedModuleParam>, ) -> Result<()>

where M: ModuleParameters,

Apply the gradients to the parameters of the model and update the model with the new parameters.

impl Updatable for &mut Lion

fn updatable_states(&self) -> impl IntoIterator<Item = &Array>

Returns a list of references to the updatable states.

fn updatable_states_mut(&mut self) -> impl IntoIterator<Item = &mut Array>

Returns a list of mutable references to the updatable states.

impl Updatable for Lion

fn updatable_states(&self) -> impl IntoIterator<Item = &Array>

Returns a list of references to the updatable states.

fn updatable_states_mut(&mut self) -> impl IntoIterator<Item = &mut Array>

Returns a list of mutable references to the updatable states.