tensorlogic_train/optimizers/
lamb.rs

1//! LAMB optimizer (Layer-wise Adaptive Moments optimizer for Batch training).
2//!
3//! LAMB is designed for large batch training and uses layer-wise adaptation
4//! of the learning rate based on the ratio of parameter and update norms.
5//!
6//! Reference: You et al., "Large Batch Optimization for Deep Learning:
7//! Training BERT in 76 minutes", ICLR 2020
8
9use super::common::{compute_gradient_norm, GradClipMode, Optimizer, OptimizerConfig};
10use crate::{TrainError, TrainResult};
11use scirs2_core::ndarray::{Array, Ix2};
12use std::collections::HashMap;
13
14/// LAMB optimizer (Layer-wise Adaptive Moments optimizer for Batch training).
15/// Designed for large batch training, uses layer-wise adaptation.
16#[derive(Debug)]
17pub struct LambOptimizer {
18    config: OptimizerConfig,
19    /// First moment estimates.
20    m: HashMap<String, Array<f64, Ix2>>,
21    /// Second moment estimates.
22    v: HashMap<String, Array<f64, Ix2>>,
23    /// Timestep counter.
24    t: usize,
25}
26
27impl LambOptimizer {
28    /// Create a new LAMB optimizer.
29    pub fn new(config: OptimizerConfig) -> Self {
30        Self {
31            config,
32            m: HashMap::new(),
33            v: HashMap::new(),
34            t: 0,
35        }
36    }
37
38    /// Apply gradient clipping if configured.
39    fn clip_gradients(&self, gradients: &mut HashMap<String, Array<f64, Ix2>>) {
40        if let Some(clip_value) = self.config.grad_clip {
41            match self.config.grad_clip_mode {
42                GradClipMode::Value => {
43                    for grad in gradients.values_mut() {
44                        grad.mapv_inplace(|g| g.max(-clip_value).min(clip_value));
45                    }
46                }
47                GradClipMode::Norm => {
48                    let total_norm = compute_gradient_norm(gradients);
49                    if total_norm > clip_value {
50                        let scale = clip_value / total_norm;
51                        for grad in gradients.values_mut() {
52                            grad.mapv_inplace(|g| g * scale);
53                        }
54                    }
55                }
56            }
57        }
58    }
59
60    /// Compute L2 norm of an array.
61    fn compute_norm(arr: &Array<f64, Ix2>) -> f64 {
62        arr.iter().map(|&x| x * x).sum::<f64>().sqrt()
63    }
64}
65
66impl Optimizer for LambOptimizer {
67    fn step(
68        &mut self,
69        parameters: &mut HashMap<String, Array<f64, Ix2>>,
70        gradients: &HashMap<String, Array<f64, Ix2>>,
71    ) -> TrainResult<()> {
72        let mut clipped_gradients = gradients.clone();
73        self.clip_gradients(&mut clipped_gradients);
74        self.t += 1;
75        let lr = self.config.learning_rate;
76        let beta1 = self.config.beta1;
77        let beta2 = self.config.beta2;
78        let eps = self.config.epsilon;
79        let weight_decay = self.config.weight_decay;
80        for (name, param) in parameters.iter_mut() {
81            let grad = clipped_gradients.get(name).ok_or_else(|| {
82                TrainError::OptimizerError(format!("Missing gradient for parameter: {}", name))
83            })?;
84            if !self.m.contains_key(name) {
85                self.m.insert(name.clone(), Array::zeros(param.raw_dim()));
86                self.v.insert(name.clone(), Array::zeros(param.raw_dim()));
87            }
88            let m = self
89                .m
90                .get_mut(name)
91                .expect("m initialized for all parameters");
92            let v = self
93                .v
94                .get_mut(name)
95                .expect("v initialized for all parameters");
96            *m = &*m * beta1 + &(grad * (1.0 - beta1));
97            let grad_squared = grad.mapv(|g| g * g);
98            *v = &*v * beta2 + &(grad_squared * (1.0 - beta2));
99            let m_hat = &*m / (1.0 - beta1.powi(self.t as i32));
100            let v_hat = &*v / (1.0 - beta2.powi(self.t as i32));
101            let adam_step = &m_hat / &v_hat.mapv(|v_val| v_val.sqrt() + eps);
102            let update = &adam_step + &param.mapv(|p| p * weight_decay);
103            let param_norm = Self::compute_norm(param);
104            let update_norm = Self::compute_norm(&update);
105            let trust_ratio = if param_norm > 0.0 && update_norm > 0.0 {
106                param_norm / update_norm
107            } else {
108                1.0
109            };
110            *param = &*param - &(update * (lr * trust_ratio));
111        }
112        Ok(())
113    }
114
115    fn zero_grad(&mut self) {}
116
117    fn get_lr(&self) -> f64 {
118        self.config.learning_rate
119    }
120
121    fn set_lr(&mut self, lr: f64) {
122        self.config.learning_rate = lr;
123    }
124
125    fn state_dict(&self) -> HashMap<String, Vec<f64>> {
126        let mut state = HashMap::new();
127        state.insert("t".to_string(), vec![self.t as f64]);
128        for (name, m_val) in &self.m {
129            state.insert(format!("m_{}", name), m_val.iter().copied().collect());
130        }
131        for (name, v_val) in &self.v {
132            state.insert(format!("v_{}", name), v_val.iter().copied().collect());
133        }
134        state
135    }
136
137    fn load_state_dict(&mut self, state: HashMap<String, Vec<f64>>) {
138        if let Some(t_vals) = state.get("t") {
139            self.t = t_vals[0] as usize;
140        }
141        for (key, values) in state {
142            if let Some(name) = key.strip_prefix("m_") {
143                if let Some(m) = self.m.get(name) {
144                    let shape = m.raw_dim();
145                    if let Ok(arr) = Array::from_shape_vec(shape, values) {
146                        self.m.insert(name.to_string(), arr);
147                    }
148                }
149            } else if let Some(name) = key.strip_prefix("v_") {
150                if let Some(v) = self.v.get(name) {
151                    let shape = v.raw_dim();
152                    if let Ok(arr) = Array::from_shape_vec(shape, values) {
153                        self.v.insert(name.to_string(), arr);
154                    }
155                }
156            }
157        }
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use scirs2_core::ndarray::array;
165
166    #[test]
167    fn test_lamb_optimizer() {
168        let config = OptimizerConfig {
169            learning_rate: 0.001,
170            weight_decay: 0.01,
171            ..Default::default()
172        };
173        let mut optimizer = LambOptimizer::new(config);
174        let mut params = HashMap::new();
175        params.insert("w".to_string(), array![[1.0, 2.0], [3.0, 4.0]]);
176        let mut grads = HashMap::new();
177        grads.insert("w".to_string(), array![[0.1, 0.1], [0.1, 0.1]]);
178        optimizer.step(&mut params, &grads).expect("unwrap");
179        let w = params.get("w").expect("unwrap");
180        assert!(w[[0, 0]] < 1.0);
181    }
182}
tensorlogic_train/optimizers/lamb.rs

tensorlogic_train/optimizers/
lamb.rs
