Struct Adafactor 

pub struct Adafactor {
    pub lr: Option<f32>,
    pub beta2_decay: f32,
    pub eps1: f32,
    pub eps2: f32,
    pub clip_threshold: f32,
    pub weight_decay: f32,
    /* private fields */
}

Adafactor — factored-second-moment optimizer.

Per-tensor state: a rows-vector + a cols-vector for 2-D parameters (sublinear in rows·cols), or a full EMA for non-2-D.

Fields

lr: Option<f32>

Optional manual learning rate. None ⇒ use the “relative step” rule min(1/√t, 1e-2) · max(ε₂, RMS(θ)) from the paper. Default None.

beta2_decay: f32

β₂_t decay-rate exponent. β₂_t = 1 − tˣ with x = -0.8 (default) means slow decay early, full decay asymptotically.

eps1: f32

Squared-gradient stability constant added before each row / column average. Default 1e-30.

eps2: f32

RMS-of-parameter floor for the relative-step rule. Default 1e-3.

clip_threshold: f32

Update-RMS clipping threshold (Shazeer & Stern §6). Default 1.0.

weight_decay: f32

Decoupled weight-decay coefficient λ. Default 0.0.

Implementations

impl Adafactor

pub fn new() -> Self

Construct with paper defaults (no manual lr ⇒ relative step, decay_rate = -0.8, ε₁=1e-30, ε₂=1e-3, clip=1.0, λ=0.0).

pub fn with_lr(self, lr: f32) -> Self

Switch from the relative-step rule to a manual learning rate.

pub fn with_weight_decay(self, wd: f32) -> Self

Override the decoupled-decay coefficient.

Trait Implementations

impl Clone for Adafactor

fn clone(&self) -> Adafactor

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Adafactor

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

Formats the value using the given formatter.

impl Default for Adafactor

fn default() -> Self

Returns the “default value” for a type.

impl Optimizer for Adafactor

fn step(&mut self, name: &str, shape: &[usize], param: &mut [f32], grad: &[f32])

fn end_iteration(&mut self)

Advance the global step counter. Most algorithms increment per call to step, so most implementations leave this a no-op.

fn step_batch(&mut self, items: &mut [OptItem<'_>])

Batched step over ALL parameters in one call. Default: sequential step per item — bit-identical to the per-parameter loop. Optimizers whose parameter groups are independent (e.g. Muon on the 2-D weight matrices vs AdamW on the embeddings/biases/norms) can override this to run the groups on separate threads; because the groups touch disjoint parameters and disjoint optimizer state, the result is bit-for-bit the same as the serial loop — only the wall-clock (the CPU-side optimizer bubble) shrinks toward max(group_times) instead of their sum.

fn set_lr(&mut self, _lr: f32)

Set the base learning rate (for LR schedules / warmup). Default is a no-op for algorithms without a scalar lr (e.g. Adafactor’s relative step size); every algorithm in this crate that has an lr field overrides this to update it.

fn lr_scale(&self, _name: &str) -> f32

Per-tensor multiplier on the effective learning rate. Default is 1.0 for every name. Override when wrapping this crate to support per-name LR schedules (e.g. embedding-vs-attention splits, or the Gaussian-splat attribute-typed LR setup). The CPU impls in this crate currently honor this only when the caller passes a pre-scaled lr for the relevant call — backends are encouraged to consult it inside their fused kernel.