rlx_optim/
radam.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//! RAdam — Rectified Adam (Liu et al., 2019, "On the Variance of the
6//! Adaptive Learning Rate and Beyond").
7//!
8//! # Motivation
9//!
10//! Early in training, `v_t` is built from very few samples and its
11//! variance is huge — which makes Adam's effective learning rate
12//! noisy at the same iterations where stability matters most.
13//! Practitioners "fix" this with an LR warm-up; RAdam derives a
14//! *closed-form* warm-up from the variance of the inverse-square-root
15//! of `v̂_t`.
16//!
17//! # Update rule
18//!
19//! Let `ρ_∞ = 2/(1−β₂) − 1` and the "SMA length"
20//! `ρ_t = ρ_∞ − 2t·β₂ᵗ / (1 − β₂ᵗ)`. Define the rectification term
21//!
22//! ```text
23//! r_t = √( ((ρ_t − 4)(ρ_t − 2)·ρ_∞) / ((ρ_∞ − 4)(ρ_∞ − 2)·ρ_t) )
24//! ```
25//!
26//! When `ρ_t > 4` the second moment is "stable enough" — use the
27//! corrected Adam step with `r_t` scaling:
28//!
29//! ```text
30//! θ_t = θ_{t-1} − lr · r_t · m̂_t / (√v̂_t + ε)
31//! ```
32//!
33//! Otherwise (`ρ_t ≤ 4`, early steps), fall back to SGD-with-momentum
34//! using the first moment alone:
35//!
36//! ```text
37//! θ_t = θ_{t-1} − lr · m̂_t
38//! ```
39//!
40//! # When to use
41//!
42//! Drop-in replacement for Adam when you don't want to hand-tune a
43//! warm-up schedule. Same memory cost as Adam.
44
45use std::collections::HashMap;
46
47use crate::Optimizer;
48use crate::common::zeros_entry;
49
50/// Rectified Adam. Per-tensor state: two `f32` buffers.
51#[derive(Debug, Clone)]
52pub struct RAdam {
53    /// Learning rate.
54    pub lr: f32,
55    /// First-moment EMA decay β₁. Default `0.9`.
56    pub beta1: f32,
57    /// Second-moment EMA decay β₂. Default `0.999`.
58    pub beta2: f32,
59    /// Denominator stability constant. Default `1e-8`.
60    pub eps: f32,
61    /// L2 weight-decay coefficient (folded into the gradient — like
62    /// classical Adam, **not** decoupled). Default `0.0`.
63    pub weight_decay: f32,
64    step: u64,
65    m: HashMap<String, Vec<f32>>,
66    v: HashMap<String, Vec<f32>>,
67}
68
69impl RAdam {
70    /// Construct with `(β₁, β₂, ε, λ) = (0.9, 0.999, 1e-8, 0.0)`.
71    pub fn new(lr: f32) -> Self {
72        Self {
73            lr,
74            beta1: 0.9,
75            beta2: 0.999,
76            eps: 1e-8,
77            weight_decay: 0.0,
78            step: 0,
79            m: HashMap::new(),
80            v: HashMap::new(),
81        }
82    }
83
84    /// Override (β₁, β₂).
85    pub fn with_betas(mut self, b1: f32, b2: f32) -> Self {
86        self.beta1 = b1;
87        self.beta2 = b2;
88        self
89    }
90
91    /// Override the L2 weight-decay coefficient.
92    pub fn with_weight_decay(mut self, wd: f32) -> Self {
93        self.weight_decay = wd;
94        self
95    }
96}
97
98impl Optimizer for RAdam {
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 rho_inf = 2.0 / (1.0 - b2) - 1.0;
111        let rho_t = rho_inf - 2.0 * t * b2.powf(t) / bc2;
112        let eps = self.eps as f64;
113        let lr = self.lr as f64;
114        let wd = self.weight_decay;
115        // Variance-rectification term `r_t` (Liu et al. eq. 14).
116        let r_t = if rho_t > 4.0 {
117            (((rho_t - 4.0) * (rho_t - 2.0) * rho_inf)
118                / ((rho_inf - 4.0) * (rho_inf - 2.0) * rho_t))
119                .sqrt()
120        } else {
121            0.0
122        };
123        let m = zeros_entry(&mut self.m, name, param.len());
124        let v = zeros_entry(&mut self.v, name, param.len());
125        for i in 0..param.len() {
126            let g = (grad[i] + wd * param[i]) as f64;
127            let mi = b1 * m[i] as f64 + (1.0 - b1) * g;
128            let vi = b2 * v[i] as f64 + (1.0 - b2) * g * g;
129            m[i] = mi as f32;
130            v[i] = vi as f32;
131            let m_hat = mi / bc1;
132            let update = if rho_t > 4.0 {
133                let v_hat = (vi / bc2).sqrt();
134                r_t * m_hat / (v_hat + eps)
135            } else {
136                m_hat
137            };
138            param[i] -= (lr * update) as f32;
139        }
140    }
141
142    fn end_iteration(&mut self) {
143        self.step += 1;
144    }
145}
rlx_optim/radam.rs

rlx_optim/
radam.rs
