Add Adam optimizer (#590)

Author: R055A

src/machine_learning/mod.rs

@@ -8,3 +8,4 @@ pub use self::loss_function::mae_loss;
 pub use self::loss_function::mse_loss;
 
 pub use self::optimization::gradient_descent;
+pub use self::optimization::Adam;

src/machine_learning/optimization/adam.rs

@@ -0,0 +1,288 @@
+//! # Adam (Adaptive Moment Estimation) optimizer
+//!
+//! The `Adam (Adaptive Moment Estimation)` optimizer is an adaptive learning rate algorithm used
+//! in gradient descent and machine learning, such as for training neural networks to solve deep
+//! learning problems. Boasting memory-efficient fast convergence rates, it sets and iteratively
+//! updates learning rates individually for each model parameter based on the gradient history.
+//!
+//! ## Algorithm:
+//!
+//! Given:
+//!   - α is the learning rate
+//!   - (β_1, β_2) are the exponential decay rates for moment estimates
+//!   - ϵ is any small value to prevent division by zero
+//!   - g_t are the gradients at time step t
+//!   - m_t are the biased first moment estimates of the gradient at time step t
+//!   - v_t are the biased second raw moment estimates of the gradient at time step t
+//!   - θ_t are the model parameters at time step t
+//!   - t is the time step
+//!
+//! Required:
+//!   θ_0
+//!
+//! Initialize:
+//!   m_0 <- 0
+//!   v_0 <- 0
+//!   t <- 0
+//!
+//! while θ_t not converged do
+//!   m_t = β_1 * m_{t−1} + (1 − β_1) * g_t
+//!   v_t = β_2 * v_{t−1} + (1 − β_2) * g_t^2
+//!   m_hat_t = m_t / 1 - β_1^t
+//!   v_hat_t = v_t / 1 - β_2^t
+//!   θ_t = θ_{t-1} − α * m_hat_t / (sqrt(v_hat_t) + ϵ)
+//!
+//! ## Resources:
+//!   - Adam: A Method for Stochastic Optimization (by Diederik P. Kingma and Jimmy Ba):
+//!       - [https://arxiv.org/abs/1412.6980]
+//!   - PyTorch Adam optimizer:
+//!       - [https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam]
+//!
+pub struct Adam {
+    learning_rate: f64, // alpha: initial step size for iterative optimization
+    betas: (f64, f64),  // betas: exponential decay rates for moment estimates
+    epsilon: f64,       // epsilon: prevent division by zero
+    m: Vec<f64>,        // m: biased first moment estimate of the gradient vector
+    v: Vec<f64>,        // v: biased second raw moment estimate of the gradient vector
+    t: usize,           // t: time step
+}
+
+impl Adam {
+    pub fn new(
+        learning_rate: Option<f64>,
+        betas: Option<(f64, f64)>,
+        epsilon: Option<f64>,
+        params_len: usize,
+    ) -> Self {
+        Adam {
+            learning_rate: learning_rate.unwrap_or(1e-3), // typical good default lr
+            betas: betas.unwrap_or((0.9, 0.999)),         // typical good default decay rates
+            epsilon: epsilon.unwrap_or(1e-8),             // typical good default epsilon
+            m: vec![0.0; params_len], // first moment vector elements all initialized to zero
+            v: vec![0.0; params_len], // second moment vector elements all initialized to zero
+            t: 0,                     // time step initialized to zero
+        }
+    }
+
+    pub fn step(&mut self, gradients: &Vec<f64>) -> Vec<f64> {
+        let mut model_params = vec![0.0; gradients.len()];
+        self.t += 1;
+
+        for i in 0..gradients.len() {
+            // update biased first moment estimate and second raw moment estimate
+            self.m[i] = self.betas.0 * self.m[i] + (1.0 - self.betas.0) * gradients[i];
+            self.v[i] = self.betas.1 * self.v[i] + (1.0 - self.betas.1) * gradients[i].powf(2f64);
+
+            // compute bias-corrected first moment estimate and second raw moment estimate
+            let m_hat = self.m[i] / (1.0 - self.betas.0.powi(self.t as i32));
+            let v_hat = self.v[i] / (1.0 - self.betas.1.powi(self.t as i32));
+
+            // update model parameters
+            model_params[i] -= self.learning_rate * m_hat / (v_hat.sqrt() + self.epsilon);
+        }
+        model_params // return updated model parameters
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_adam_init_default_values() {
+        let optimizer = Adam::new(None, None, None, 1);
+
+        assert_eq!(optimizer.learning_rate, 0.001);
+        assert_eq!(optimizer.betas, (0.9, 0.999));
+        assert_eq!(optimizer.epsilon, 1e-8);
+        assert_eq!(optimizer.m, vec![0.0; 1]);
+        assert_eq!(optimizer.v, vec![0.0; 1]);
+        assert_eq!(optimizer.t, 0);
+    }
+
+    #[test]
+    fn test_adam_init_custom_lr_value() {
+        let optimizer = Adam::new(Some(0.9), None, None, 2);
+
+        assert_eq!(optimizer.learning_rate, 0.9);
+        assert_eq!(optimizer.betas, (0.9, 0.999));
+        assert_eq!(optimizer.epsilon, 1e-8);
+        assert_eq!(optimizer.m, vec![0.0; 2]);
+        assert_eq!(optimizer.v, vec![0.0; 2]);
+        assert_eq!(optimizer.t, 0);
+    }
+
+    #[test]
+    fn test_adam_init_custom_betas_value() {
+        let optimizer = Adam::new(None, Some((0.8, 0.899)), None, 3);
+
+        assert_eq!(optimizer.learning_rate, 0.001);
+        assert_eq!(optimizer.betas, (0.8, 0.899));
+        assert_eq!(optimizer.epsilon, 1e-8);
+        assert_eq!(optimizer.m, vec![0.0; 3]);
+        assert_eq!(optimizer.v, vec![0.0; 3]);
+        assert_eq!(optimizer.t, 0);
+    }
+
+    #[test]
+    fn test_adam_init_custom_epsilon_value() {
+        let optimizer = Adam::new(None, None, Some(1e-10), 4);
+
+        assert_eq!(optimizer.learning_rate, 0.001);
+        assert_eq!(optimizer.betas, (0.9, 0.999));
+        assert_eq!(optimizer.epsilon, 1e-10);
+        assert_eq!(optimizer.m, vec![0.0; 4]);
+        assert_eq!(optimizer.v, vec![0.0; 4]);
+        assert_eq!(optimizer.t, 0);
+    }
+
+    #[test]
+    fn test_adam_init_all_custom_values() {
+        let optimizer = Adam::new(Some(1.0), Some((0.001, 0.099)), Some(1e-1), 5);
+
+        assert_eq!(optimizer.learning_rate, 1.0);
+        assert_eq!(optimizer.betas, (0.001, 0.099));
+        assert_eq!(optimizer.epsilon, 1e-1);
+        assert_eq!(optimizer.m, vec![0.0; 5]);
+        assert_eq!(optimizer.v, vec![0.0; 5]);
+        assert_eq!(optimizer.t, 0);
+    }
+
+    #[test]
+    fn test_adam_step_default_params() {
+        let gradients = vec![-1.0, 2.0, -3.0, 4.0, -5.0, 6.0, -7.0, 8.0];
+
+        let mut optimizer = Adam::new(None, None, None, 8);
+        let updated_params = optimizer.step(&gradients);
+
+        assert_eq!(
+            updated_params,
+            vec![
+                0.0009999999900000003,
+                -0.000999999995,
+                0.0009999999966666666,
+                -0.0009999999975,
+                0.000999999998,
+                -0.0009999999983333334,
+                0.0009999999985714286,
+                -0.00099999999875
+            ]
+        );
+    }
+
+    #[test]
+    fn test_adam_step_custom_params() {
+        let gradients = vec![9.0, -8.0, 7.0, -6.0, 5.0, -4.0, 3.0, -2.0, 1.0];
+
+        let mut optimizer = Adam::new(Some(0.005), Some((0.5, 0.599)), Some(1e-5), 9);
+        let updated_params = optimizer.step(&gradients);
+
+        assert_eq!(
+            updated_params,
+            vec![
+                -0.004999994444450618,
+                0.004999993750007813,
+                -0.004999992857153062,
+                0.004999991666680556,
+                -0.004999990000020001,
+                0.004999987500031251,
+                -0.004999983333388888,
+                0.004999975000124999,
+                -0.0049999500004999945
+            ]
+        );
+    }
+
+    #[test]
+    fn test_adam_step_empty_gradients_array() {
+        let gradients = vec![];
+
+        let mut optimizer = Adam::new(None, None, None, 0);
+        let updated_params = optimizer.step(&gradients);
+
+        assert_eq!(updated_params, vec![]);
+    }
+
+    #[ignore]
+    #[test]
+    fn test_adam_step_iteratively_until_convergence_with_default_params() {
+        const CONVERGENCE_THRESHOLD: f64 = 1e-5;
+        let gradients = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+
+        let mut optimizer = Adam::new(None, None, None, 6);
+
+        let mut model_params = vec![0.0; 6];
+        let mut updated_params = optimizer.step(&gradients);
+
+        while (updated_params
+            .iter()
+            .zip(model_params.iter())
+            .map(|(x, y)| x - y)
+            .collect::<Vec<f64>>())
+        .iter()
+        .map(|&x| x.powi(2))
+        .sum::<f64>()
+        .sqrt()
+            > CONVERGENCE_THRESHOLD
+        {
+            model_params = updated_params;
+            updated_params = optimizer.step(&gradients);
+        }
+
+        assert!(updated_params < vec![CONVERGENCE_THRESHOLD; 6]);
+        assert_ne!(updated_params, model_params);
+        assert_eq!(
+            updated_params,
+            vec![
+                -0.0009999999899999931,
+                -0.0009999999949999929,
+                -0.0009999999966666597,
+                -0.0009999999974999929,
+                -0.0009999999979999927,
+                -0.0009999999983333263
+            ]
+        );
+    }
+
+    #[ignore]
+    #[test]
+    fn test_adam_step_iteratively_until_convergence_with_custom_params() {
+        const CONVERGENCE_THRESHOLD: f64 = 1e-7;
+        let gradients = vec![7.0, -8.0, 9.0, -10.0, 11.0, -12.0, 13.0];
+
+        let mut optimizer = Adam::new(Some(0.005), Some((0.8, 0.899)), Some(1e-5), 7);
+
+        let mut model_params = vec![0.0; 7];
+        let mut updated_params = optimizer.step(&gradients);
+
+        while (updated_params
+            .iter()
+            .zip(model_params.iter())
+            .map(|(x, y)| x - y)
+            .collect::<Vec<f64>>())
+        .iter()
+        .map(|&x| x.powi(2))
+        .sum::<f64>()
+        .sqrt()
+            > CONVERGENCE_THRESHOLD
+        {
+            model_params = updated_params;
+            updated_params = optimizer.step(&gradients);
+        }
+
+        assert!(updated_params < vec![CONVERGENCE_THRESHOLD; 7]);
+        assert_ne!(updated_params, model_params);
+        assert_eq!(
+            updated_params,
+            vec![
+                -0.004999992857153061,
+                0.004999993750007814,
+                -0.0049999944444506185,
+                0.004999995000005001,
+                -0.004999995454549587,
+                0.004999995833336807,
+                -0.004999996153849113
+            ]
+        );
+    }
+}

src/machine_learning/optimization/mod.rs

@@ -1,3 +1,6 @@
 mod gradient_descent;
 
+mod adam;
+
+pub use self::adam::Adam;
 pub use self::gradient_descent::gradient_descent;