rlx_optim/
sgd.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//! Stochastic Gradient Descent with optional momentum and decoupled
6//! L2 weight decay.
7//!
8//! # Update rules
9//!
10//! Vanilla SGD (`momentum = 0`):
11//!
12//! ```text
13//! θ_{t+1} = θ_t − lr · (g_t + λ·θ_t)
14//! ```
15//!
16//! Polyak momentum (`momentum = μ`, `nesterov = false`):
17//!
18//! ```text
19//! v_{t+1} = μ·v_t + (g_t + λ·θ_t)
20//! θ_{t+1} = θ_t − lr · v_{t+1}
21//! ```
22//!
23//! Nesterov-accelerated SGD (`nesterov = true`):
24//!
25//! ```text
26//! v_{t+1} = μ·v_t + (g_t + λ·θ_t)
27//! θ_{t+1} = θ_t − lr · (g_t + λ·θ_t + μ·v_{t+1})
28//! ```
29//!
30//! # When to use
31//!
32//! The default choice when training CNNs from scratch; with a
33//! well-tuned `lr` schedule it still beats Adam on many vision
34//! benchmarks. Cheap state (one buffer if `momentum > 0`).
35
36use std::collections::HashMap;
37
38use crate::Optimizer;
39use crate::common::zeros_entry;
40
41/// SGD with momentum / Nesterov / L2 weight decay.
42///
43/// All hyperparameters are public so callers can hot-swap them between
44/// iterations (e.g. for a warm-up schedule). State is keyed by
45/// parameter name; the same `Sgd` instance can drive every tensor in
46/// a model.
47#[derive(Debug, Clone)]
48pub struct Sgd {
49    /// Learning rate. No default — pass it to [`Sgd::new`].
50    pub lr: f32,
51    /// Polyak momentum coefficient ∈ \[0, 1\). `0.0` disables momentum
52    /// entirely (and the per-tensor velocity buffer is still allocated
53    /// but unused — set via [`Sgd::with_momentum`] if you want it on).
54    pub momentum: f32,
55    /// Use Nesterov-accelerated momentum. Only meaningful when
56    /// `momentum > 0`.
57    pub nesterov: bool,
58    /// L2 weight decay coefficient λ. Folded into the gradient
59    /// *before* the momentum EMA (classical, **not** decoupled).
60    /// Use [`crate::AdamW`]-style decoupling if you need that.
61    pub weight_decay: f32,
62    v: HashMap<String, Vec<f32>>,
63}
64
65impl Sgd {
66    /// Construct with `lr` and momentum / decay disabled.
67    pub fn new(lr: f32) -> Self {
68        Self {
69            lr,
70            momentum: 0.0,
71            nesterov: false,
72            weight_decay: 0.0,
73            v: HashMap::new(),
74        }
75    }
76
77    /// Enable Polyak (or Nesterov) momentum.
78    pub fn with_momentum(mut self, momentum: f32, nesterov: bool) -> Self {
79        self.momentum = momentum;
80        self.nesterov = nesterov;
81        self
82    }
83
84    /// Set the L2 weight-decay coefficient.
85    pub fn with_weight_decay(mut self, wd: f32) -> Self {
86        self.weight_decay = wd;
87        self
88    }
89}
90
91impl Optimizer for Sgd {
92    fn set_lr(&mut self, lr: f32) {
93        self.lr = lr;
94    }
95
96    fn step(&mut self, name: &str, _shape: &[usize], param: &mut [f32], grad: &[f32]) {
97        debug_assert_eq!(param.len(), grad.len());
98        let v = zeros_entry(&mut self.v, name, param.len());
99        let mu = self.momentum;
100        let wd = self.weight_decay;
101        let lr = self.lr;
102        for i in 0..param.len() {
103            let g = grad[i] + wd * param[i];
104            if mu == 0.0 {
105                param[i] -= lr * g;
106            } else {
107                v[i] = mu * v[i] + g;
108                let update = if self.nesterov { g + mu * v[i] } else { v[i] };
109                param[i] -= lr * update;
110            }
111        }
112    }
113}
rlx_optim/sgd.rs

rlx_optim/
sgd.rs
