rlx_optim/
adam.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Adam — Adaptive Moment Estimation (Kingma & Ba, 2014).
6//!
7//! # Update rule
8//!
9//! For each parameter, with `t` the 1-based iteration index:
10//!
11//! ```text
12//! g_t   = ∇L(θ_{t-1}) + λ·θ_{t-1}            // L2 decay folded in
13//! m_t   = β₁·m_{t-1} + (1 − β₁)·g_t
14//! v_t   = β₂·v_{t-1} + (1 − β₂)·g_t²
15//! m̂_t   = m_t / (1 − β₁ᵗ)                    // bias correction
16//! v̂_t   = v_t / (1 − β₂ᵗ)
17//! θ_t   = θ_{t-1} − lr · m̂_t / (√v̂_t + ε)
18//! ```
19//!
20//! # When to use
21//!
22//! Reliable default for transformer pre-training and most non-vision
23//! workloads. Per-parameter memory is **2×** the parameter size (one
24//! `m`, one `v`). If you need decoupled weight decay (recommended for
25//! transformers) use [`crate::AdamW`] instead.
26
27use std::collections::HashMap;
28
29use crate::Optimizer;
30use crate::common::{zeros_entry, zip4_for_each};
31
32/// Bias-corrected first/second moment optimizer.
33///
34/// Per-tensor state: two `f32` buffers (`m`, `v`) of the same shape as
35/// the parameter.
36#[derive(Debug, Clone)]
37pub struct Adam {
38    /// Learning rate. Typical: `1e-3` for from-scratch CNNs, `1e-4`
39    /// for transformer fine-tuning.
40    pub lr: f32,
41    /// First-moment EMA decay β₁ ∈ \[0, 1). Default `0.9`.
42    pub beta1: f32,
43    /// Second-moment EMA decay β₂ ∈ \[0, 1). Default `0.999`.
44    pub beta2: f32,
45    /// Stability constant in the denominator. Default `1e-8`.
46    pub eps: f32,
47    /// L2 weight decay coefficient. **Folded into the gradient**
48    /// (the "classic Adam" rule); use [`crate::AdamW`] for decoupled
49    /// decay. Default `0.0`.
50    pub weight_decay: f32,
51    step: u64,
52    m: HashMap<String, Vec<f32>>,
53    v: HashMap<String, Vec<f32>>,
54}
55
56impl Adam {
57    /// Construct with the given learning rate and the standard
58    /// (β₁, β₂, ε) = (0.9, 0.999, 1e-8) defaults.
59    pub fn new(lr: f32) -> Self {
60        Self {
61            lr,
62            beta1: 0.9,
63            beta2: 0.999,
64            eps: 1e-8,
65            weight_decay: 0.0,
66            step: 0,
67            m: HashMap::new(),
68            v: HashMap::new(),
69        }
70    }
71
72    /// Override (β₁, β₂).
73    pub fn with_betas(mut self, b1: f32, b2: f32) -> Self {
74        self.beta1 = b1;
75        self.beta2 = b2;
76        self
77    }
78
79    /// Override the denominator stability constant ε.
80    pub fn with_eps(mut self, eps: f32) -> Self {
81        self.eps = eps;
82        self
83    }
84
85    /// Override the L2 weight-decay coefficient.
86    pub fn with_weight_decay(mut self, wd: f32) -> Self {
87        self.weight_decay = wd;
88        self
89    }
90
91    /// 1-based iteration counter. Starts at 1 (so the first call to
92    /// `step()` sees `t=1`), advances on [`Optimizer::end_iteration`].
93    pub fn current_step(&self) -> u64 {
94        self.step + 1
95    }
96}
97
98impl Optimizer for Adam {
99    fn set_lr(&mut self, lr: f32) {
100        self.lr = lr;
101    }
102
103    fn step(&mut self, name: &str, _shape: &[usize], param: &mut [f32], grad: &[f32]) {
104        debug_assert_eq!(param.len(), grad.len());
105        let t = (self.step + 1) as f64;
106        let b1 = self.beta1 as f64;
107        let b2 = self.beta2 as f64;
108        let bc1 = 1.0 - b1.powf(t);
109        let bc2 = 1.0 - b2.powf(t);
110        let eps = self.eps as f64;
111        let lr = self.lr as f64;
112        let wd = self.weight_decay;
113        let n = param.len();
114        // `self.m` / `self.v` are distinct fields, so the two
115        // `zeros_entry` calls borrow disjoint regions of `self` and
116        // their results can coexist.
117        let m = zeros_entry(&mut self.m, name, n);
118        let v = zeros_entry(&mut self.v, name, n);
119        zip4_for_each(param, m, v, grad, |p, mi, vi, gi| {
120            let g = (gi + wd * *p) as f64;
121            let new_m = b1 * *mi as f64 + (1.0 - b1) * g;
122            let new_v = b2 * *vi as f64 + (1.0 - b2) * g * g;
123            *mi = new_m as f32;
124            *vi = new_v as f32;
125            let m_hat = new_m / bc1;
126            let v_hat = new_v / bc2;
127            *p -= (lr * m_hat / (v_hat.sqrt() + eps)) as f32;
128        });
129    }
130
131    fn end_iteration(&mut self) {
132        self.step += 1;
133    }
134}
rlx_optim/adam.rs

rlx_optim/
adam.rs
