optirs_core/optimizers/
lars.rs

1// Layer-wise Adaptive Rate Scaling (LARS) optimizer
2//
3// LARS is an optimization algorithm specifically designed for large batch training
4// in deep neural networks. It scales the learning rate for each layer based on the
5// ratio of the weight norm to the gradient norm.
6//
7// References:
8// - [Large Batch Training of Convolutional Networks](https://arxiv.org/abs/1708.03888)
9
10use crate::error::{OptimError, Result};
11use crate::optimizers::Optimizer;
12use scirs2_core::ndarray::{Array, Dimension, ScalarOperand};
13use scirs2_core::numeric::Float;
14use std::fmt::Debug;
15
16/// Layer-wise Adaptive Rate Scaling (LARS) optimizer
17///
18/// LARS is an optimization algorithm specifically designed for large batch training,
19/// which allows scaling up the batch size significantly without loss of accuracy.
20/// It works by adapting the learning rate per layer based on the ratio of
21/// weight norm to gradient norm.
22///
23/// # Parameters
24///
25/// * `learning_rate` - Base learning rate
26/// * `momentum` - Momentum factor (default: 0.9)
27/// * `weight_decay` - Weight decay factor (default: 0.0001)
28/// * `trust_coefficient` - Trust coefficient for scaling (default: 0.001)
29/// * `eps` - Small constant for numerical stability (default: 1e-8)
30/// * `exclude_bias_and_norm` - Whether to exclude bias and normalization parameters
31///   from LARS trust-ratio scaling (default: true). Bias and normalization
32///   parameters are identified by their rank: tensors with one dimension or fewer
33///   (`ndim() <= 1`) are treated as biases / normalization scales, matching the
34///   convention used by the reference LARS and LAMB implementations.
35///
36/// # Example
37///
38/// ```no_run
39/// use scirs2_core::ndarray::Array1;
40/// use optirs_core::optimizers::{LARS, Optimizer};
41///
42/// let mut optimizer = LARS::new(0.01)
43///     .with_momentum(0.9)
44///     .with_weight_decay(0.0001)
45///     .with_trust_coefficient(0.001);
46///
47/// let params = Array1::zeros(10);
48/// let gradients = Array1::ones(10);
49///
50/// let updated_params = optimizer.step(&params, &gradients).expect("optimizer.step succeeds");
51/// // Parameters are automatically updated
52/// ```
53#[derive(Debug, Clone)]
54pub struct LARS<A: Float> {
55    learning_rate: A,
56    momentum: A,
57    weight_decay: A,
58    trust_coefficient: A,
59    eps: A,
60    exclude_bias_and_norm: bool,
61    /// Momentum buffers, one flat buffer per parameter-tensor index
62    velocity: Option<Vec<Vec<A>>>,
63}
64
65impl<A: Float + ScalarOperand + Debug + Send + Sync> LARS<A> {
66    /// Create a new LARS optimizer with the given learning rate
67    pub fn new(learning_rate: A) -> Self {
68        Self {
69            learning_rate,
70            momentum: A::from(0.9).expect("LARS: default momentum (0.9) must fit in A"),
71            weight_decay: A::from(0.0001)
72                .expect("LARS: default weight_decay (0.0001) must fit in A"),
73            trust_coefficient: A::from(0.001)
74                .expect("LARS: default trust_coefficient (0.001) must fit in A"),
75            eps: A::from(1e-8).expect("LARS: default eps (1e-8) must fit in A"),
76            exclude_bias_and_norm: true,
77            velocity: None,
78        }
79    }
80
81    /// Set the momentum factor
82    pub fn with_momentum(mut self, momentum: A) -> Self {
83        self.momentum = momentum;
84        self
85    }
86
87    /// Set the weight decay factor
88    pub fn with_weight_decay(mut self, weight_decay: A) -> Self {
89        self.weight_decay = weight_decay;
90        self
91    }
92
93    /// Set the trust coefficient
94    pub fn with_trust_coefficient(mut self, trust_coefficient: A) -> Self {
95        self.trust_coefficient = trust_coefficient;
96        self
97    }
98
99    /// Set the epsilon value for numerical stability
100    pub fn with_eps(mut self, eps: A) -> Self {
101        self.eps = eps;
102        self
103    }
104
105    /// Set whether to exclude bias and normalization layers from LARS adaptation
106    pub fn with_exclude_bias_and_norm(mut self, exclude_bias_and_norm: bool) -> Self {
107        self.exclude_bias_and_norm = exclude_bias_and_norm;
108        self
109    }
110
111    /// Reset the optimizer state
112    pub fn reset(&mut self) {
113        self.velocity = None;
114    }
115
116    /// Ensures a momentum buffer exists for `index` with `len` elements
117    fn ensure_state(&mut self, index: usize, len: usize) {
118        let velocity = self.velocity.get_or_insert_with(Vec::new);
119        while velocity.len() <= index {
120            velocity.push(vec![A::zero(); len]);
121        }
122        if velocity[index].len() != len {
123            velocity[index] = vec![A::zero(); len];
124        }
125    }
126
127    /// Performs a LARS update for the parameter tensor at `index`
128    ///
129    /// LARS is a *layer-wise* algorithm: the trust ratio is computed per tensor, and
130    /// each tensor keeps its own momentum buffer.
131    pub fn step_indexed<D: Dimension>(
132        &mut self,
133        index: usize,
134        params: &Array<A, D>,
135        gradients: &Array<A, D>,
136    ) -> Result<Array<A, D>> {
137        if params.shape() != gradients.shape() {
138            return Err(OptimError::DimensionMismatch(format!(
139                "Incompatible shapes: parameters have shape {:?}, gradients have shape {:?}",
140                params.shape(),
141                gradients.shape()
142            )));
143        }
144
145        // A bias / normalization parameter is a rank <= 1 tensor.
146        let is_bias_or_norm = params.ndim() <= 1;
147        let n_params = gradients.len();
148        self.ensure_state(index, n_params);
149
150        // Calculate weight norm and gradient norm over this tensor only.
151        let weight_norm = params.mapv(|x| x * x).sum().sqrt();
152        let grad_norm = gradients.mapv(|x| x * x).sum().sqrt();
153
154        // Determine if we should apply LARS scaling.
155        // Bias and normalization parameters (rank <= 1) are excluded when configured,
156        // and fall back to plain SGD-with-momentum, exactly as the paper prescribes.
157        let should_apply_lars = !(self.exclude_bias_and_norm && is_bias_or_norm);
158
159        // Calculate local learning rate using the trust ratio
160        let local_lr = if should_apply_lars && weight_norm > A::zero() && grad_norm > A::zero() {
161            self.trust_coefficient * weight_norm
162                / (grad_norm + self.weight_decay * weight_norm + self.eps)
163        } else {
164            A::one()
165        };
166
167        let scaled_lr = self.learning_rate * local_lr;
168        let momentum = self.momentum;
169        let weight_decay = self.weight_decay;
170        let use_weight_decay = weight_decay > A::zero();
171
172        let velocity = self
173            .velocity
174            .as_mut()
175            .ok_or_else(|| OptimError::InvalidConfig("LARS state not initialized".to_string()))?;
176        let buffer = velocity.get_mut(index).ok_or_else(|| {
177            OptimError::InvalidConfig(format!("LARS has no velocity buffer for index {}", index))
178        })?;
179
180        let mut updated_params = params.clone();
181        for (slot, (p, g)) in buffer
182            .iter_mut()
183            .zip(updated_params.iter_mut().zip(gradients.iter()))
184        {
185            let grad = if use_weight_decay {
186                *g + weight_decay * *p
187            } else {
188                *g
189            };
190            *slot = momentum * *slot + grad * scaled_lr;
191            *p = *p - *slot;
192        }
193
194        Ok(updated_params)
195    }
196}
197
198impl<A: Float + ScalarOperand + Debug + Send + Sync, D: Dimension + Send + Sync> Optimizer<A, D>
199    for LARS<A>
200{
201    fn step(&mut self, params: &Array<A, D>, gradients: &Array<A, D>) -> Result<Array<A, D>> {
202        self.step_indexed(0, params, gradients)
203    }
204
205    fn step_list(
206        &mut self,
207        params_list: &[&Array<A, D>],
208        gradients_list: &[&Array<A, D>],
209    ) -> Result<Vec<Array<A, D>>> {
210        if params_list.len() != gradients_list.len() {
211            return Err(OptimError::InvalidConfig(format!(
212                "Number of parameter arrays ({}) does not match number of gradient arrays ({})",
213                params_list.len(),
214                gradients_list.len()
215            )));
216        }
217
218        let mut results = Vec::with_capacity(params_list.len());
219        for (index, (params, grads)) in params_list.iter().zip(gradients_list.iter()).enumerate() {
220            results.push(self.step_indexed(index, params, grads)?);
221        }
222        Ok(results)
223    }
224
225    fn set_learning_rate(&mut self, learning_rate: A) {
226        self.learning_rate = learning_rate;
227    }
228
229    fn get_learning_rate(&self) -> A {
230        self.learning_rate
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237    use approx::assert_abs_diff_eq;
238    use scirs2_core::ndarray::Array1;
239
240    #[test]
241    fn test_lars_creation() {
242        let optimizer = LARS::new(0.01);
243        assert_abs_diff_eq!(optimizer.learning_rate, 0.01);
244        assert_abs_diff_eq!(optimizer.momentum, 0.9);
245        assert_abs_diff_eq!(optimizer.weight_decay, 0.0001);
246        assert_abs_diff_eq!(optimizer.trust_coefficient, 0.001);
247        assert_abs_diff_eq!(optimizer.eps, 1e-8);
248        assert!(optimizer.exclude_bias_and_norm);
249    }
250
251    #[test]
252    fn test_lars_builder() {
253        let optimizer = LARS::new(0.01)
254            .with_momentum(0.95)
255            .with_weight_decay(0.0005)
256            .with_trust_coefficient(0.01)
257            .with_eps(1e-6)
258            .with_exclude_bias_and_norm(false);
259
260        assert_abs_diff_eq!(optimizer.momentum, 0.95);
261        assert_abs_diff_eq!(optimizer.weight_decay, 0.0005);
262        assert_abs_diff_eq!(optimizer.trust_coefficient, 0.01);
263        assert_abs_diff_eq!(optimizer.eps, 1e-6);
264        assert!(!optimizer.exclude_bias_and_norm);
265    }
266
267    #[test]
268    fn test_lars_update() {
269        // 1-D tensors are treated as bias / normalization parameters, which LARS
270        // excludes by default. Opt in explicitly to exercise the trust-ratio path.
271        let mut optimizer = LARS::new(0.1)
272            .with_momentum(0.9)
273            .with_weight_decay(0.0)
274            .with_trust_coefficient(1.0)
275            .with_exclude_bias_and_norm(false);
276
277        let params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
278        let gradients = Array1::from_vec(vec![0.1, 0.2, 0.3]);
279
280        // First update
281        let updated_params = optimizer
282            .step(&params, &gradients)
283            .expect("optimizer.step succeeds in test_lars_update");
284
285        // LARS scaling factor with trust_coefficient=1.0 should be:
286        // weight_norm / grad_norm = sqrt(14) / sqrt(0.14) ≈ 10
287        // So the effective learning rate is 0.1 * 10 = 1.0
288        // Scale is approximately 10, but let's check actual value (more precise)
289        let weight_norm = params.mapv(|x| x * x).sum().sqrt();
290        let grad_norm = gradients.mapv(|x| x * x).sum().sqrt();
291        let scale = weight_norm / grad_norm;
292
293        assert_abs_diff_eq!(updated_params[0], 1.0 - 0.1 * scale * 0.1, epsilon = 1e-5);
294        assert_abs_diff_eq!(updated_params[1], 2.0 - 0.1 * scale * 0.2, epsilon = 1e-5);
295        assert_abs_diff_eq!(updated_params[2], 3.0 - 0.1 * scale * 0.3, epsilon = 1e-5);
296
297        // Second update should include momentum
298        let updated_params2 = optimizer
299            .step(&updated_params, &gradients)
300            .expect("step succeeds in test_lars_update");
301
302        // For the second update, the velocity will be updated with momentum
303        // Just check that parameters continue to change in the expected direction
304        assert!(updated_params2[0] < updated_params[0]);
305        assert!(updated_params2[1] < updated_params[1]);
306        assert!(updated_params2[2] < updated_params[2]);
307    }
308
309    #[test]
310    fn test_lars_weight_decay() {
311        let mut optimizer = LARS::new(0.01)
312            .with_momentum(0.0) // No momentum for clarity
313            .with_weight_decay(0.1)
314            .with_trust_coefficient(1.0)
315            .with_exclude_bias_and_norm(false);
316
317        let params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
318        let gradients = Array1::from_vec(vec![0.1, 0.2, 0.3]);
319
320        let updated_params = optimizer
321            .step(&params, &gradients)
322            .expect("optimizer.step succeeds in test_lars_weight_decay");
323
324        // Gradients with weight decay: [0.1, 0.2, 0.3] + 0.1*[1.0, 2.0, 3.0] = [0.2, 0.4, 0.6]
325        // LARS scaling factor includes weight decay in denominator
326        // weight_norm / (grad_norm + weight_decay * weight_norm)
327        // = sqrt(14) / (sqrt(0.56) + 0.1*sqrt(14))
328        let weight_norm = params.mapv(|x| x * x).sum().sqrt();
329        let grad_norm = gradients.mapv(|x| x * x).sum().sqrt();
330        let expected_scale = weight_norm / (grad_norm + 0.1 * weight_norm);
331
332        // Check calculation is approximately correct (allowing for floating point differences)
333        let expected_p0 = 1.0 - 0.01 * expected_scale * (0.1 + 0.1 * 1.0);
334        let expected_p1 = 2.0 - 0.01 * expected_scale * (0.2 + 0.1 * 2.0);
335        let expected_p2 = 3.0 - 0.01 * expected_scale * (0.3 + 0.1 * 3.0);
336
337        assert_abs_diff_eq!(updated_params[0], expected_p0, epsilon = 1e-5);
338        assert_abs_diff_eq!(updated_params[1], expected_p1, epsilon = 1e-5);
339        assert_abs_diff_eq!(updated_params[2], expected_p2, epsilon = 1e-5);
340    }
341
342    #[test]
343    fn test_zero_gradients() {
344        let mut optimizer = LARS::new(0.01);
345        let params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
346        let zero_gradients = Array1::zeros(3);
347
348        let updated_params = optimizer
349            .step(&params, &zero_gradients)
350            .expect("step succeeds in test_zero_gradients");
351
352        // With zero gradients, only weight decay should contribute to the update
353        // With small weight decay (0.0001), changes should be very small
354        assert_abs_diff_eq!(updated_params[0], params[0], epsilon = 1e-3);
355        assert_abs_diff_eq!(updated_params[1], params[1], epsilon = 1e-3);
356        assert_abs_diff_eq!(updated_params[2], params[2], epsilon = 1e-3);
357    }
358
359    #[test]
360    fn test_exclude_bias_and_norm() {
361        let mut optimizer_excluded = LARS::new(0.01)
362            .with_momentum(0.0)
363            .with_weight_decay(0.0)
364            .with_exclude_bias_and_norm(true);
365
366        let mut optimizer_included = LARS::new(0.01)
367            .with_momentum(0.0)
368            .with_weight_decay(0.0)
369            .with_exclude_bias_and_norm(false);
370
371        // Test with parameters that could be bias (small 1D array)
372        let bias_params = Array1::from_vec(vec![0.1, 0.2]);
373        let bias_grads = Array1::from_vec(vec![0.01, 0.02]);
374
375        let updated_excluded = optimizer_excluded
376            .step(&bias_params, &bias_grads)
377            .expect("step succeeds in test_exclude_bias_and_norm");
378        let updated_included = optimizer_included
379            .step(&bias_params, &bias_grads)
380            .expect("step succeeds in test_exclude_bias_and_norm");
381
382        // When excluded, should use base learning rate (but still include momentum calculation)
383        assert_abs_diff_eq!(updated_excluded[0], 0.1 - 0.01 * 0.01, epsilon = 1e-4);
384
385        // When included, should use LARS scaled learning rate
386        let weight_norm = (0.1f64.powi(2) + 0.2f64.powi(2)).sqrt();
387        let grad_norm = (0.01f64.powi(2) + 0.02f64.powi(2)).sqrt();
388        let expected_factor = 0.001 * weight_norm / grad_norm; // trust_coefficient * weight_norm / grad_norm
389
390        assert_abs_diff_eq!(
391            updated_included[0],
392            0.1 - 0.01 * expected_factor * 0.01,
393            epsilon = 1e-5
394        );
395    }
396
397    /// Regression test for the `exclude_bias_and_norm` tautology.
398    ///
399    /// The flag used to be evaluated as `!exclude || weight_norm > 0`, which is true
400    /// for every parameter with a non-zero norm, so the exclusion never actually
401    /// excluded anything. Bias / normalization parameters are rank <= 1 tensors and
402    /// must fall back to the plain (unscaled) learning rate.
403    #[test]
404    fn test_exclude_bias_and_norm_is_decided_by_rank() {
405        use scirs2_core::ndarray::Array2;
406
407        // Rank-1 tensor => treated as a bias, excluded from the trust ratio.
408        let mut bias_opt = LARS::new(0.01)
409            .with_momentum(0.0)
410            .with_weight_decay(0.0)
411            .with_trust_coefficient(1.0)
412            .with_exclude_bias_and_norm(true);
413
414        let bias = Array1::from_vec(vec![1.0f64, 2.0, 3.0]);
415        let bias_grads = Array1::from_vec(vec![0.1f64, 0.2, 0.3]);
416        let updated_bias = bias_opt.step(&bias, &bias_grads).expect("bias step");
417
418        // Plain SGD: p - lr * g
419        assert_abs_diff_eq!(updated_bias[0], 1.0 - 0.01 * 0.1, epsilon = 1e-12);
420        assert_abs_diff_eq!(updated_bias[2], 3.0 - 0.01 * 0.3, epsilon = 1e-12);
421
422        // Rank-2 tensor => a weight matrix, LARS scaling applies even when the
423        // exclusion flag is enabled.
424        let mut weight_opt = LARS::new(0.01)
425            .with_momentum(0.0)
426            .with_weight_decay(0.0)
427            .with_trust_coefficient(1.0)
428            .with_exclude_bias_and_norm(true);
429
430        let weights =
431            Array2::from_shape_vec((3, 1), vec![1.0f64, 2.0, 3.0]).expect("valid 3x1 matrix");
432        let weight_grads =
433            Array2::from_shape_vec((3, 1), vec![0.1f64, 0.2, 0.3]).expect("valid 3x1 matrix");
434        let updated_weights = weight_opt
435            .step(&weights, &weight_grads)
436            .expect("weight step");
437
438        let weight_norm = weights.mapv(|x: f64| x * x).sum().sqrt();
439        let grad_norm = weight_grads.mapv(|x: f64| x * x).sum().sqrt();
440        let scale = weight_norm / (grad_norm + 1e-8);
441        assert_abs_diff_eq!(
442            updated_weights[[0, 0]],
443            1.0 - 0.01 * scale * 0.1,
444            epsilon = 1e-8
445        );
446
447        // The two paths must genuinely differ.
448        assert!((updated_bias[0] - updated_weights[[0, 0]]).abs() > 1e-6);
449    }
450}
optirs_core/optimizers/lars.rs

optirs_core/optimizers/
lars.rs
