geographdb_core/algorithms/

adam.rs

//! Adam optimiser for flat `Vec<f32>` parameter vectors.
//!
//! This is a pure-Rust implementation of Kingma & Ba (2015). It is intended for
//! small training loops where pulling in a full autograd framework would be
//! overkill — for example, the geometric MLP heads in GeoGraphDB experiments.

/// Gradient norm clipping threshold used by [`Adam::step_clipped`] when none is
/// supplied.
pub const DEFAULT_CLIP_NORM: f32 = 1.0;

/// Adam optimiser state.
///
/// Keeps first- and second-moment vectors `m` and `v` for every parameter and
/// applies bias correction using the timestep `t`.
#[derive(Debug, Clone)]
pub struct Adam {
    learning_rate: f32,
    beta1: f32,
    beta2: f32,
    eps: f32,
    t: u64,
    m: Vec<f32>,
    v: Vec<f32>,
}

impl Adam {
    /// Create a new Adam optimiser for `param_count` parameters.
    ///
    /// Defaults follow the original paper: `lr = 0.001`, `beta1 = 0.9`,
    /// `beta2 = 0.999`, `eps = 1e-8`.
    pub fn new(param_count: usize) -> Self {
        Self::with_lr(0.001, param_count)
    }

    /// Create an Adam optimiser with a custom learning rate.
    pub fn with_lr(learning_rate: f32, param_count: usize) -> Self {
        Self::with_hyperparams(learning_rate, 0.9, 0.999, 1e-8, param_count)
    }

    /// Create an Adam optimiser with full control over all hyperparameters.
    pub fn with_hyperparams(
        learning_rate: f32,
        beta1: f32,
        beta2: f32,
        eps: f32,
        param_count: usize,
    ) -> Self {
        assert!(learning_rate > 0.0, "Adam: learning_rate must be positive");
        assert!(beta1 > 0.0 && beta1 < 1.0, "Adam: beta1 must be in (0, 1)");
        assert!(beta2 > 0.0 && beta2 < 1.0, "Adam: beta2 must be in (0, 1)");
        assert!(eps > 0.0, "Adam: eps must be positive");

        Self {
            learning_rate,
            beta1,
            beta2,
            eps,
            t: 0,
            m: vec![0.0f32; param_count],
            v: vec![0.0f32; param_count],
        }
    }

    /// Number of parameters tracked by this optimiser.
    pub fn param_count(&self) -> usize {
        self.m.len()
    }

    /// Current timestep.
    pub fn timestep(&self) -> u64 {
        self.t
    }

    /// Perform one Adam update on `params` using the supplied gradient.
    ///
    /// `params` and `grad` must have the same length as the optimiser's
    /// parameter count.
    pub fn step(&mut self, params: &mut [f32], grad: &[f32]) {
        assert_eq!(
            params.len(),
            self.m.len(),
            "Adam::step: parameter count mismatch"
        );
        assert_eq!(
            grad.len(),
            self.m.len(),
            "Adam::step: gradient count mismatch"
        );

        self.t += 1;
        let t = self.t as f32;
        let lr = self.learning_rate;
        let beta1 = self.beta1;
        let beta2 = self.beta2;
        let eps = self.eps;
        let bias_correction1 = 1.0 - beta1.powf(t);
        let bias_correction2 = 1.0 - beta2.powf(t);

        for i in 0..params.len() {
            let g = grad[i];
            self.m[i] = beta1 * self.m[i] + (1.0 - beta1) * g;
            self.v[i] = beta2 * self.v[i] + (1.0 - beta2) * g * g;

            let m_hat = self.m[i] / bias_correction1;
            let v_hat = self.v[i] / bias_correction2;
            params[i] -= lr * m_hat / (v_hat.sqrt() + eps);
        }
    }

    /// Perform one Adam update after clipping the gradient to `max_norm`.
    ///
    /// If the global gradient norm is larger than `max_norm`, the whole
    /// gradient is rescaled before the Adam moment update. This matches the
    /// common "clip by global norm" behaviour.
    pub fn step_clipped(&mut self, params: &mut [f32], grad: &[f32], max_norm: f32) {
        assert!(
            max_norm > 0.0,
            "Adam::step_clipped: max_norm must be positive"
        );
        let mut clipped = grad.to_vec();
        clip_gradients(&mut clipped, max_norm);
        self.step(params, &clipped);
    }

    /// Perform one Adam update on multiple parameter slices.
    ///
    /// `params` and `grads` must have the same number of slices, and each
    /// corresponding slice must have the same length.  This avoids flattening
    /// and reloading parameters when a model is stored in separate tensors.
    pub fn step_slices(&mut self, params: &mut [&mut [f32]], grads: &[&[f32]]) {
        assert_eq!(
            params.len(),
            grads.len(),
            "Adam::step_slices: parameter and gradient slice counts mismatch"
        );

        let total_len: usize = params.iter().map(|p| p.len()).sum();
        assert_eq!(
            total_len,
            self.m.len(),
            "Adam::step_slices: total parameter count mismatch"
        );

        self.t += 1;
        let t = self.t as f32;
        let lr = self.learning_rate;
        let beta1 = self.beta1;
        let beta2 = self.beta2;
        let eps = self.eps;
        let bias_correction1 = 1.0 - beta1.powf(t);
        let bias_correction2 = 1.0 - beta2.powf(t);

        let mut mom_off = 0;
        for (param_slice, grad_slice) in params.iter_mut().zip(grads.iter()) {
            assert_eq!(
                param_slice.len(),
                grad_slice.len(),
                "Adam::step_slices: slice length mismatch"
            );
            for i in 0..param_slice.len() {
                let g = grad_slice[i];
                self.m[mom_off] = beta1 * self.m[mom_off] + (1.0 - beta1) * g;
                self.v[mom_off] = beta2 * self.v[mom_off] + (1.0 - beta2) * g * g;

                let m_hat = self.m[mom_off] / bias_correction1;
                let v_hat = self.v[mom_off] / bias_correction2;
                param_slice[i] -= lr * m_hat / (v_hat.sqrt() + eps);

                mom_off += 1;
            }
        }
    }

    /// Perform one Adam update on multiple slices after global norm clipping.
    pub fn step_slices_clipped(
        &mut self,
        params: &mut [&mut [f32]],
        grads: &[&[f32]],
        max_norm: f32,
    ) {
        assert!(
            max_norm > 0.0,
            "Adam::step_slices_clipped: max_norm must be positive"
        );

        let total_len: usize = params.iter().map(|p| p.len()).sum();
        let mut clipped = Vec::with_capacity(total_len);
        for grad_slice in grads {
            clipped.extend_from_slice(grad_slice);
        }
        clip_gradients(&mut clipped, max_norm);

        let mut grad_slices: Vec<&[f32]> = Vec::with_capacity(grads.len());
        let mut off = 0;
        for grad_slice in grads {
            let len = grad_slice.len();
            grad_slices.push(&clipped[off..off + len]);
            off += len;
        }

        self.step_slices(params, &grad_slices);
    }

    /// Reset all moment vectors to zero without changing hyperparameters.
    pub fn reset(&mut self) {
        self.t = 0;
        self.m.fill(0.0);
        self.v.fill(0.0);
    }
}

/// Clip a flat gradient vector by its global L2 norm.
///
/// If `norm <= max_norm` the vector is unchanged; otherwise every element is
/// multiplied by `max_norm / norm`.
pub fn clip_gradients(grad: &mut [f32], max_norm: f32) {
    assert!(max_norm > 0.0, "clip_gradients: max_norm must be positive");
    let norm_sq: f32 = grad.iter().map(|g| g * g).sum();
    let norm = norm_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / norm;
        for g in grad.iter_mut() {
            *g *= scale;
        }
    }
}

/// Compute the global L2 norm of a flat gradient vector.
pub fn grad_norm(grad: &[f32]) -> f32 {
    grad.iter().map(|g| g * g).sum::<f32>().sqrt()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn adam_decreases_simple_quadratic() {
        // Minimise f(x) = (x - 3)^2 with gradient 2*(x - 3).
        let mut x = [0.0f32];
        let mut opt = Adam::with_lr(0.1, 1);
        for _ in 0..200 {
            let grad = [2.0 * (x[0] - 3.0)];
            opt.step(&mut x, &grad);
        }
        assert!(
            (x[0] - 3.0).abs() < 0.01,
            "x converged to {}, expected 3.0",
            x[0]
        );
    }

    #[test]
    fn clipped_step_respects_norm() {
        let mut params = [0.0f32, 0.0f32];
        let grad = [10.0f32, 0.0f32];
        let mut opt = Adam::with_lr(0.1, 2);
        opt.step_clipped(&mut params, &grad, 1.0);
        // With a clipped gradient of [1.0, 0.0] the first parameter should move
        // in the negative direction but not as far as if the raw [10.0, 0.0]
        // gradient had been used.
        assert!(params[0] < 0.0);
        assert!(
            params[0].abs() < 0.5,
            "clipped step moved too far: {}",
            params[0]
        );
    }

    #[test]
    fn clip_gradients_scales_large_norm() {
        let mut g = vec![3.0f32, 4.0f32];
        clip_gradients(&mut g, 1.0);
        assert!(
            (grad_norm(&g) - 1.0).abs() < 1e-5,
            "norm after clipping should be 1.0"
        );
        assert!((g[0] - 0.6).abs() < 1e-5);
        assert!((g[1] - 0.8).abs() < 1e-5);
    }

    #[test]
    fn reset_clears_momentum() {
        let mut opt = Adam::with_lr(0.1, 2);
        let mut params = [1.0f32, 1.0f32];
        opt.step(&mut params, &[1.0f32, 1.0f32]);
        assert!(opt.m.iter().any(|&v| v != 0.0));
        opt.reset();
        assert!(opt.m.iter().all(|&v| v == 0.0));
        assert!(opt.v.iter().all(|&v| v == 0.0));
        assert_eq!(opt.timestep(), 0);
    }

    #[test]
    fn step_slices_matches_flat_step() {
        let mut opt = Adam::with_lr(0.1, 4);

        let mut p_flat = [1.0f32, 2.0, 3.0, 4.0];
        let g_flat = [0.1f32, 0.2, 0.3, 0.4];
        opt.step(&mut p_flat, &g_flat);

        let mut opt2 = Adam::with_lr(0.1, 4);
        let mut a = [1.0f32, 2.0];
        let mut b = [3.0f32, 4.0];
        let g_a = [0.1f32, 0.2];
        let g_b = [0.3f32, 0.4];
        opt2.step_slices(&mut [&mut a, &mut b], &[&g_a, &g_b]);

        let tol = 1e-6;
        assert!((p_flat[0] - a[0]).abs() < tol);
        assert!((p_flat[1] - a[1]).abs() < tol);
        assert!((p_flat[2] - b[0]).abs() < tol);
        assert!((p_flat[3] - b[1]).abs() < tol);
    }

    #[test]
    fn step_slices_clipped_matches_flat_clipped() {
        let mut opt = Adam::with_lr(0.1, 2);
        let mut p_flat = [1.0f32, 1.0];
        let g_flat = [10.0f32, 0.0];
        opt.step_clipped(&mut p_flat, &g_flat, 1.0);

        let mut opt2 = Adam::with_lr(0.1, 2);
        let mut a = [1.0f32];
        let mut b = [1.0f32];
        let g_a = [10.0f32];
        let g_b = [0.0f32];
        opt2.step_slices_clipped(&mut [&mut a, &mut b], &[&g_a, &g_b], 1.0);

        let tol = 1e-6;
        assert!((p_flat[0] - a[0]).abs() < tol);
        assert!((p_flat[1] - b[0]).abs() < tol);
    }
}