burn_optim/optim/
adafactor.rs

1use burn_core as burn;
2
3use burn::{
4    config::Config,
5    tensor::{DType, Device, FloatDType, Tensor},
6};
7
8use crate::{LearningRate, RecordState, grad_clipping::GradientClippingConfig};
9
10use super::{ModuleOptimizer, Optimizer};
11
12#[cfg(not(feature = "std"))]
13#[allow(unused_imports)]
14use num_traits::Float as _;
15
16/// [`Adafactor`] configuration.
17///
18/// By default, the supplied learning rate caps the relative step size at
19/// `min(lr, 1 / sqrt(step))`, which is then multiplied by the parameter's root mean square
20/// (RMS). A learning rate of `0.01` gives the schedule proposed in the paper.
21/// Disable both `relative_step` and `scale_parameter` to use an absolute learning rate
22/// controlled entirely by an external scheduler.
23#[derive(Config, Debug)]
24pub struct AdafactorConfig {
25    /// Positive constant added to squared gradients for numerical stability.
26    #[config(default = 1e-30)]
27    epsilon_1: f32,
28    /// Positive lower bound on the parameter RMS when scaling the step size.
29    #[config(default = 1e-3)]
30    epsilon_2: f32,
31    /// Positive upper bound on the RMS of the normalized update before learning rate scaling.
32    /// This is separate from gradient clipping.
33    #[config(default = 1.0)]
34    clip_threshold: f32,
35    /// Negative exponent in the second moment decay schedule `beta_2 = 1 - step^decay_rate`.
36    #[config(default = -0.8)]
37    decay_rate: f64,
38    /// Whether to cap the supplied learning rate by `1 / sqrt(step)`.
39    #[config(default = true)]
40    relative_step: bool,
41    /// Whether to multiply the step size by `max(epsilon_2, RMS(parameter))`.
42    #[config(default = true)]
43    scale_parameter: bool,
44    /// Decoupled weight decay coefficient, multiplied by the supplied learning rate.
45    #[config(default = 0.0)]
46    weight_decay: f32,
47    /// Optional gradient clipping applied before the Adafactor update.
48    grad_clipping: Option<GradientClippingConfig>,
49}
50
51/// Adafactor optimizer with factored second moments and update clipping.
52///
53/// Matrices store only row and column second moments. Higher-rank tensors factor their
54/// last two dimensions independently for each leading index; vectors store a full second
55/// moment. No first moment is stored. Half-precision parameters use `f32` updates and state.
56///
57/// See [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://arxiv.org/abs/1804.04235).
58/// Configured by [`AdafactorConfig`].
59#[derive(Clone)]
60pub struct Adafactor {
61    epsilon_1: f32,
62    epsilon_2: f32,
63    clip_threshold: f32,
64    decay_rate: f64,
65    relative_step: bool,
66    scale_parameter: bool,
67    weight_decay: f32,
68}
69
70/// Adafactor state for a single parameter tensor.
71#[derive(RecordState, Clone)]
72pub struct AdafactorState<const D: usize> {
73    /// Number of updates applied to this parameter.
74    pub time: usize,
75    /// Moving average of squared gradients, reduced over the last dimension for matrices
76    /// and higher-rank tensors (shape `[..., rows, 1]`). Vectors retain their full shape.
77    pub second_moment: Tensor<D>,
78    /// Column second moment for matrices and higher-rank tensors (shape `[..., 1, columns]`).
79    /// Absent for vectors.
80    pub column_second_moment: Option<Tensor<D>>,
81}
82
83impl AdafactorConfig {
84    /// Build the per-parameter Adafactor optimizer.
85    ///
86    /// Use [`Self::init`] to construct a whole-module optimizer with the configured gradient
87    /// clipping, or pass this optimizer to [`ModuleOptimizer::with_group`].
88    ///
89    /// # Panics
90    ///
91    /// Panics if either epsilon or the clipping threshold is not positive and finite,
92    /// the decay exponent is not negative and finite, or weight decay is not nonnegative
93    /// and finite.
94    pub fn build(&self) -> Adafactor {
95        assert!(
96            self.epsilon_1.is_finite() && self.epsilon_1 > 0.0,
97            "Adafactor epsilon_1 must be positive and finite"
98        );
99        assert!(
100            self.epsilon_2.is_finite() && self.epsilon_2 > 0.0,
101            "Adafactor epsilon_2 must be positive and finite"
102        );
103        assert!(
104            self.clip_threshold.is_finite() && self.clip_threshold > 0.0,
105            "Adafactor clip_threshold must be positive and finite"
106        );
107        assert!(
108            self.decay_rate.is_finite() && self.decay_rate < 0.0,
109            "Adafactor decay_rate must be negative and finite"
110        );
111        assert!(
112            self.weight_decay.is_finite() && self.weight_decay >= 0.0,
113            "Adafactor weight_decay must be nonnegative and finite"
114        );
115
116        Adafactor {
117            epsilon_1: self.epsilon_1,
118            epsilon_2: self.epsilon_2,
119            clip_threshold: self.clip_threshold,
120            decay_rate: self.decay_rate,
121            relative_step: self.relative_step,
122            scale_parameter: self.scale_parameter,
123            weight_decay: self.weight_decay,
124        }
125    }
126
127    /// Initialize a whole-module Adafactor optimizer.
128    pub fn init(&self) -> ModuleOptimizer {
129        let mut optimizer = ModuleOptimizer::from(self.build());
130        if let Some(config) = &self.grad_clipping {
131            optimizer = optimizer.with_grad_clipping(config.init());
132        }
133        optimizer
134    }
135}
136
137impl Optimizer for Adafactor {
138    type State<const D: usize> = AdafactorState<D>;
139
140    fn step<const D: usize>(
141        &self,
142        lr: LearningRate,
143        tensor: Tensor<D>,
144        grad: Tensor<D>,
145        state: Option<Self::State<D>>,
146    ) -> (Tensor<D>, Option<Self::State<D>>) {
147        let dtype = tensor.dtype();
148        // Squaring and adding epsilon in half precision can underflow, even for ordinary
149        // gradients. Keep the accumulated statistics and the whole update in f32.
150        let (tensor, grad) = match dtype {
151            DType::F16 | DType::BF16 => (tensor.cast(FloatDType::F32), grad.cast(FloatDType::F32)),
152            _ => (tensor, grad),
153        };
154
155        let squared_grad = grad.clone().square().add_scalar(self.epsilon_1);
156        let (second_moment, column_second_moment) = if D >= 2 {
157            (
158                squared_grad.clone().mean_dim(D - 1),
159                Some(squared_grad.mean_dim(D - 2)),
160            )
161        } else {
162            (squared_grad, None)
163        };
164
165        let state = match state {
166            Some(mut state) => {
167                state.time += 1;
168                let new_weight = (state.time as f64).powf(self.decay_rate);
169                let old_weight = 1.0 - new_weight;
170                state.second_moment = state
171                    .second_moment
172                    .mul_scalar(old_weight)
173                    .add(second_moment.mul_scalar(new_weight));
174                state.column_second_moment = column_second_moment.map(|column| {
175                    state
176                        .column_second_moment
177                        .take()
178                        .expect("Factored Adafactor state must contain a column second moment")
179                        .mul_scalar(old_weight)
180                        .add(column.mul_scalar(new_weight))
181                });
182                state
183            }
184            None => AdafactorState {
185                time: 1,
186                second_moment,
187                column_second_moment,
188            },
189        };
190
191        let mut update = match &state.column_second_moment {
192            Some(column) => {
193                let row = state.second_moment.clone();
194                // Normalize the row factor first to avoid underflow from the product of
195                // small moments. Broadcasting avoids storing a full second moment.
196                let row_scale = row.clone().div(row.mean_dim(D - 2)).sqrt();
197                grad.div(row_scale).div(column.clone().sqrt())
198            }
199            None => grad.div(state.second_moment.clone().sqrt()),
200        };
201
202        let clipping = rms(update.clone())
203            .div_scalar(self.clip_threshold)
204            .clamp_min(1.0);
205        update = update.div(clipping.unsqueeze());
206
207        let step_size = if self.relative_step {
208            lr.min(1.0 / (state.time as f64).sqrt())
209        } else {
210            lr
211        };
212        if self.scale_parameter {
213            let scale = rms(tensor.clone()).clamp_min(self.epsilon_2);
214            update = update.mul(scale.unsqueeze());
215        }
216        update = update.mul_scalar(step_size);
217
218        let tensor = if self.weight_decay == 0.0 {
219            tensor
220        } else {
221            tensor.mul_scalar(1.0 - lr * self.weight_decay as f64)
222        };
223        let tensor = tensor - update;
224        let tensor = match dtype {
225            DType::F16 | DType::BF16 => tensor.cast(FloatDType::from(dtype)),
226            _ => tensor,
227        };
228
229        (tensor, Some(state))
230    }
231
232    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device) -> Self::State<D> {
233        state.second_moment = state.second_moment.to_device(device);
234        state.column_second_moment = state
235            .column_second_moment
236            .map(|column| column.to_device(device));
237        state
238    }
239}
240
241fn rms<const D: usize>(tensor: Tensor<D>) -> Tensor<1> {
242    tensor.square().mean().sqrt()
243}
244
245#[cfg(test)]
246mod tests;
burn_optim/optim/adafactor.rs

burn_optim/optim/
adafactor.rs
