mlx_rs/optimizers/
adafactor.rs

1use std::{borrow::Cow, collections::HashMap, rc::Rc};
2
3use mlx_internal_macros::{generate_builder, Buildable};
4
5use crate::{
6    array,
7    error::AdafactorBuildError,
8    ops::{
9        matmul, maximum, mean, mean_axes, minimum, rsqrt, sqrt, square, zeros_dtype, zeros_like,
10    },
11    utils::Updatable,
12    Array,
13};
14
15use super::*;
16
17fn rms(inputs: &Array) -> crate::error::Result<Array> {
18    sqrt(&mean(&square(inputs)?, None)?)
19}
20
21fn approvate_exp_moving_avg(
22    exp_avg_sq_row: &Array,
23    exp_avg_sq_col: &Array,
24) -> crate::error::Result<Array> {
25    let rfactor = rsqrt(&exp_avg_sq_row.divide(&mean_axes(exp_avg_sq_row, &[-1], true)?)?)?;
26    let cfactor = rsqrt(exp_avg_sq_col)?;
27    matmul(&rfactor.expand_dims(-1)?, &cfactor.expand_dims(0)?)
28}
29
30/// Type alias for the epsilon values used in Adafactor builder
31pub type AdafactorEps = (f32, f32);
32
33/// State of the Adafactor optimizer.
34#[derive(Debug, Clone)]
35pub struct AdafactorState {
36    pub(crate) step: Array,
37    pub(crate) exp_avg_sq_row: Option<Array>,
38    pub(crate) exp_avg_sq_col: Option<Array>,
39    pub(crate) exp_avg_sq: Option<Array>,
40    pub(crate) exp_avg: Option<Array>,
41}
42
43impl OptimizerState for State<AdafactorState> {
44    type UnflattenError = UnflattenError;
45
46    fn state_projection(&mut self) -> Result<StateProjection<'_>, StateProjectionError> {
47        let mut projection = StateProjection::new();
48        for (key, state) in self {
49            projection.required(format!("{key}.step"), &mut state.step)?;
50            projection.optional(format!("{key}.exp_avg_sq_row"), &mut state.exp_avg_sq_row)?;
51            projection.optional(format!("{key}.exp_avg_sq_col"), &mut state.exp_avg_sq_col)?;
52            projection.optional(format!("{key}.exp_avg_sq"), &mut state.exp_avg_sq)?;
53            projection.optional(format!("{key}.exp_avg"), &mut state.exp_avg)?;
54        }
55        Ok(projection)
56    }
57
58    fn unflatten<I, K>(input: I) -> Result<Self, Self::UnflattenError>
59    where
60        Self: Sized,
61        I: IntoIterator<Item = (K, Array)>,
62        K: Ord + AsRef<str> + Into<Rc<str>>,
63    {
64        let mut state = State::new();
65        let iter = input
66            .into_iter()
67            .sorted_by(|a, b| a.0.as_ref().cmp(b.0.as_ref()));
68
69        for (k, v) in iter {
70            let key = k.into();
71            let mut parts = key.rsplit('.');
72            let suffix = parts.next().ok_or(UnflattenError::InvalidKey)?;
73            let prefix = parts.next().ok_or(UnflattenError::InvalidKey)?;
74
75            let prefix = Rc::from(prefix);
76            let state = state.entry(prefix).or_insert_with(|| AdafactorState {
77                step: array!(AdafactorState::DEFAULT_STEP),
78                exp_avg_sq_row: None,
79                exp_avg_sq_col: None,
80                exp_avg_sq: None,
81                exp_avg: None,
82            });
83
84            match suffix {
85                "step" => state.step = v,
86                "exp_avg_sq_row" => state.exp_avg_sq_row = Some(v),
87                "exp_avg_sq_col" => state.exp_avg_sq_col = Some(v),
88                "exp_avg_sq" => state.exp_avg_sq = Some(v),
89                "exp_avg" => state.exp_avg = Some(v),
90                _ => return Err(UnflattenError::InvalidKey),
91            }
92        }
93
94        Ok(state)
95    }
96}
97
98impl AdafactorState {
99    /// Default value for `step`
100    pub const DEFAULT_STEP: i32 = 0;
101
102    fn new(parameter: &Array, beta1_is_some: bool) -> crate::error::Result<Self> {
103        let step = array!(Self::DEFAULT_STEP);
104        let mut exp_avg_sq_row = None;
105        let mut exp_avg_sq_col = None;
106        let mut exp_avg_sq = None;
107        let mut exp_avg = None;
108
109        if parameter.ndim() >= 2 {
110            let shape = parameter.shape();
111            let dtype = parameter.dtype();
112
113            let row_shape = &shape[..shape.len() - 1];
114            exp_avg_sq_row = Some(zeros_dtype(row_shape, dtype)?);
115
116            let mut col_shape = shape[..shape.len() - 2].to_vec();
117            col_shape.push(*shape.last().unwrap());
118            exp_avg_sq_col = Some(zeros_dtype(&col_shape, dtype)?);
119        } else {
120            exp_avg_sq = Some(zeros_like(parameter)?);
121        };
122
123        if beta1_is_some {
124            exp_avg = Some(zeros_like(parameter)?);
125        }
126
127        Ok(Self {
128            step,
129            exp_avg_sq_row,
130            exp_avg_sq_col,
131            exp_avg_sq,
132            exp_avg,
133        })
134    }
135}
136
137/// `Option<Array>`. Type alias for the learning rate used in Adafactor builder due to limitation in
138/// the `generate_builder` macro
139pub type AdafactorBuilderLr = Option<f32>;
140
141/// Type alias for the learning rate used in Adafactor
142pub type AdafactorLr = Option<Array>;
143
144/// `Option<f32>` Type alias for the beta1 used in Adafactor builder due to limitation in the
145/// `generate_builder` macro
146pub type AdafactorBuilderBeta1 = Option<f32>;
147
148/// Type alias for the beta1 used in Adafactor
149pub type AdafactorBeta1 = Option<Array>;
150
151generate_builder! {
152    /// The Adafactor optimizer.
153    ///
154    /// Our Adafactor implementation follows the original paper: `Adafactor:
155    /// Adaptive Learning Rates with Sublinear Memory Cost
156    /// <https://arxiv.org/abs/1804.04235>
157    #[derive(Debug, Clone, Buildable)]
158    #[buildable(root = crate)]
159    #[builder(
160        build_with = build_adafactor,
161        err = AdafactorBuildError,
162        root = crate
163    )]
164    pub struct Adafactor {
165        /// The learning rate.
166        #[builder(optional, default = Adafactor::DEFAULT_LR)]
167        pub lr: Option<f32>,
168
169        /// The first term is added to the square of the gradients to improve numerical stability.
170        /// Default to [`Adafactor::DEFAULT_EPS`].
171        #[builder(optional, ty_override = AdafactorEps, default = Adafactor::DEFAULT_EPS)]
172        pub eps: (Array, Array),
173
174        /// Clips the unscaled update. Default to [`Adafactor::DEFAULT_CLIP_THRESHOLD`].
175        #[builder(optional, ty_override = f32, default = Adafactor::DEFAULT_CLIP_THRESHOLD)]
176        pub clip_threshold: Array,
177
178        /// Coefficient for the running average of the squared gradient. Default to
179        /// [`Adafactor::DEFAULT_DECAY_RATE`].
180        #[builder(optional, ty_override = f32, default = Adafactor::DEFAULT_DECAY_RATE)]
181        pub decay_rate: Array,
182
183        /// If set then the first moment will be used.
184        #[builder(optional, ty_override = AdafactorBuilderBeta1, default = Adafactor::DEFAULT_BETA1)]
185        pub beta1: AdafactorBeta1,
186
187        /// The weight decay. Default to [`Adafactor::DEFAULT_WEIGHT_DECAY`].
188        #[builder(optional, default = Adafactor::DEFAULT_WEIGHT_DECAY)]
189        pub weight_decay: f32,
190
191        /// If `true` the `learningRate` will be scaled by `max(eps.0, RMS(parameter))`. Default to
192        /// [`Adafactor::DEFAULT_SCALE_PARAMETER`].
193        #[builder(optional, default = Adafactor::DEFAULT_SCALE_PARAMETER)]
194        pub scale_parameter: bool,
195
196        /// If `true` the `learningRate` will be ignored and the relative step size will be
197        /// computed. Default to [`Adafactor::DEFAULT_RELATIVE_STEP`].
198        #[builder(optional, ty_override = bool, default = Adafactor::DEFAULT_RELATIVE_STEP)]
199        pub relative_step: bool,
200
201        /// If `true` the relative step size will be calculated by the current step. Default to
202        /// [`Adafactor::DEFAULT_WARMUP_INIT`].
203        #[builder(optional, default = Adafactor::DEFAULT_WARMUP_INIT)]
204        pub warmup_init: bool,
205
206        /// Inner state.
207        #[builder(ignore)]
208        pub state: State<AdafactorState>,
209    }
210}
211
212/// Builds a new [`Adafactor`] optimizer.
213fn build_adafactor(builder: AdafactorBuilder) -> Result<Adafactor, AdafactorBuildError> {
214    let eps = builder.eps;
215    let clip_threshold = builder.clip_threshold;
216    let decay_rate = builder.decay_rate;
217    let weight_decay = builder.weight_decay;
218    let scale_parameter = builder.scale_parameter;
219    let relative_step = builder.relative_step;
220    let warmup_init = builder.warmup_init;
221
222    if builder.lr.is_none() && !relative_step {
223        return Err(AdafactorBuildError::LrIsNoneAndRelativeStepIsFalse);
224    }
225
226    Ok(Adafactor {
227        lr: builder.lr,
228        eps: (array!(eps.0), array!(eps.1)),
229        clip_threshold: array!(clip_threshold),
230        decay_rate: array!(decay_rate),
231        beta1: builder.beta1.map(Array::from),
232        weight_decay,
233        scale_parameter,
234        relative_step,
235        warmup_init,
236        state: State::new(),
237    })
238}
239
240impl Adafactor {
241    /// Default value for `lr`
242    pub const DEFAULT_LR: Option<f32> = None;
243
244    /// Default values for `eps`
245    pub const DEFAULT_EPS: (f32, f32) = (1e-30, 1e-3);
246
247    /// Default value for `clip_threshold`
248    pub const DEFAULT_CLIP_THRESHOLD: f32 = 1.0;
249
250    /// Default value for `decay_rate`
251    pub const DEFAULT_DECAY_RATE: f32 = -0.8;
252
253    /// Default value for `weight_decay`
254    pub const DEFAULT_WEIGHT_DECAY: f32 = 0.0;
255
256    /// Default value for `scale_parameter`
257    pub const DEFAULT_SCALE_PARAMETER: bool = true;
258
259    /// Default value for `relative_step`
260    pub const DEFAULT_RELATIVE_STEP: bool = true;
261
262    /// Default value for `warmup_init`
263    pub const DEFAULT_WARMUP_INIT: bool = false;
264
265    /// Default value for `beta1`
266    pub const DEFAULT_BETA1: Option<f32> = None;
267}
268
269fn get_mut_or_insert_with<'a, T, E>(
270    map: &'a mut HashMap<Rc<str>, T>,
271    key: &Rc<str>,
272    f: impl FnOnce() -> Result<T, E>,
273) -> Result<&'a mut T, E> {
274    if !map.contains_key(key) {
275        map.insert(key.clone(), f()?);
276    }
277
278    Ok(map.get_mut(key).unwrap())
279}
280
281fn compute_lr(
282    relative_step: bool,
283    warmup_init: bool,
284    lr: Option<f32>,
285    scale_parameter: bool,
286    eps: &(Array, Array),
287    step: &Array,
288    parameter_rms: &Array,
289) -> crate::error::Result<Array> {
290    let relative_step_size = if relative_step {
291        let min_step = if warmup_init {
292            // SAFETY: `step` is a single-element array and won't panic.
293            array!(1e-6) * step
294        } else {
295            array!(1e-2)
296        };
297        // SAFETY: `step` is a single-element array and won't panic.
298        minimum(min_step, array!(1.0) / sqrt(step)?)?
299    } else {
300        // SAFETY: This is already checked in the `build` stage.
301        array!(lr.expect("The learning rate should be set if the relative step is not enabled"))
302    };
303
304    let mut parameter_scale = array!(1.0);
305    if scale_parameter {
306        parameter_scale = maximum(&eps.1, parameter_rms)?;
307    }
308
309    parameter_scale.multiply(relative_step_size)
310}
311
312impl Optimizer for Adafactor {
313    type State = State<AdafactorState>;
314
315    fn state(&self) -> &Self::State {
316        &self.state
317    }
318
319    fn state_mut(&mut self) -> &mut Self::State {
320        &mut self.state
321    }
322
323    fn update_single(
324        &mut self,
325        key: &std::rc::Rc<str>,
326        gradient: &Array,
327        parameter: &mut Array,
328    ) -> crate::error::Result<()> {
329        let beta1_is_some = self.beta1.is_some();
330        let state = get_mut_or_insert_with(&mut self.state, key, || {
331            AdafactorState::new(parameter, beta1_is_some)
332        })?;
333
334        state.step = state.step.add(array!(1))?;
335
336        let gradient_shape = gradient.shape();
337        let factored = gradient_shape.len() >= 2;
338        let step = &state.step;
339
340        let parameter_rms = rms(parameter)?;
341        let lr = compute_lr(
342            self.relative_step,
343            self.warmup_init,
344            self.lr,
345            self.scale_parameter,
346            &self.eps,
347            step,
348            &parameter_rms,
349        )?;
350        let beta2 = array!(1.0).subtract(&step.power(&self.decay_rate)?)?;
351
352        let mut update: Cow<Array> = Cow::Owned(gradient.square()?.add(&self.eps.0)?);
353
354        let one_minus_beta2 = array!(1.0).subtract(&beta2)?;
355        if factored {
356            // SAFETY: These fields are created in the `new` when ndim >= 2 and won't panic.
357            let exp_avg_sq_row = state.exp_avg_sq_row.as_mut().unwrap();
358            let exp_avg_sq_col = state.exp_avg_sq_col.as_mut().unwrap();
359
360            *exp_avg_sq_row = beta2
361                .multiply(&*exp_avg_sq_row)?
362                .add(&one_minus_beta2.multiply(&update.mean_axes(&[-1], None)?)?)?;
363            *exp_avg_sq_col = beta2
364                .multiply(&*exp_avg_sq_col)?
365                .add(&one_minus_beta2.multiply(&update.mean_axes(&[-2], None)?)?)?;
366
367            update = Cow::Owned(approvate_exp_moving_avg(
368                &*exp_avg_sq_row,
369                &*exp_avg_sq_col,
370            )?);
371            update = Cow::Owned(update.multiply(gradient)?);
372        } else {
373            // SAFETY: This field is created in the `new` when ndim < 2 and won't panic.
374            let exp_avg_sq = state.exp_avg_sq.as_mut().unwrap();
375
376            *exp_avg_sq = beta2
377                .multiply(&*exp_avg_sq)?
378                .add(&one_minus_beta2.multiply(&update)?)?;
379            update = Cow::Owned(rsqrt(&*exp_avg_sq)?.multiply(gradient)?);
380        }
381
382        let update_rms = rms(&update)?;
383        let max = maximum(array!(1.0), update_rms.divide(&self.clip_threshold)?)?;
384        update = Cow::Owned(update.divide(max)?);
385        update = Cow::Owned(lr.multiply(update)?);
386
387        if let Some(beta1) = &self.beta1 {
388            // SAFETY: This field is created in the `new` when beta1 is set and won't panic.
389            let exp_avg = state.exp_avg.as_mut().unwrap();
390            let one_minus_beta1 = array!(1.0).subtract(beta1)?;
391            *exp_avg = beta1
392                .multiply(&*exp_avg)?
393                .add(&one_minus_beta1.multiply(&update)?)?;
394            update = Cow::Borrowed(&*exp_avg);
395        }
396
397        if self.weight_decay != 0.0 {
398            let rhs = parameter.multiply(array!(-self.weight_decay).multiply(lr)?)?;
399            *parameter = parameter.add(rhs)?;
400        }
401
402        *parameter = parameter.subtract(&update)?;
403
404        Ok(())
405    }
406}
407
408impl Updatable for Adafactor {
409    optimizer_updatable_state_methods!();
410}
411
412impl_updatable_for_mut_optimizer!(Adafactor);
mlx_rs/optimizers/adafactor.rs

mlx_rs/optimizers/
adafactor.rs
