Gradient descent — Rosenbrock, μ = 0.92

//! 2-D gradient descent on the Rosenbrock landscape.
//!
//!   f(x, y) = (a − x)² + b (y − x²)²,  a = 1, b = 100
//!   ∇f     = ( −2(a − x) − 4 b x (y − x²),  2 b (y − x²) )
//!
//! Momentum (Polyak heavy-ball):
//!   v ← μ v − η ∇f
//!   x ← x + v
//!
//! Pure, deterministic. V-class L0-closed.

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct GdParams {
    pub a: f32,
    pub b: f32,
    pub lr: f32,
    pub momentum: f32,
}

impl Default for GdParams {
    fn default() -> Self {
        Self { a: 1.0, b: 100.0, lr: 1.0e-3, momentum: 0.92 }
    }
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct GdState {
    pub x: f32,
    pub y: f32,
    pub vx: f32,
    pub vy: f32,
    pub step: u64,
}

impl Default for GdState {
    fn default() -> Self {
        Self { x: -1.5, y: 2.5, vx: 0.0, vy: 0.0, step: 0 }
    }
}

pub fn rosenbrock(x: f32, y: f32, a: f32, b: f32) -> f32 {
    let dx = a - x;
    let dy = y - x * x;
    dx * dx + b * dy * dy
}

pub fn grad_rosenbrock(x: f32, y: f32, a: f32, b: f32) -> (f32, f32) {
    let gx = -2.0 * (a - x) - 4.0 * b * x * (y - x * x);
    let gy = 2.0 * b * (y - x * x);
    (gx, gy)
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GdStepReport {
    pub bit_ops: u64,
    pub loss: f32,
    pub grad_norm: f32,
}

pub fn gd_step(state: &mut GdState, p: &GdParams) -> GdStepReport {
    let (gx, gy) = grad_rosenbrock(state.x, state.y, p.a, p.b);
    state.vx = p.momentum * state.vx - p.lr * gx;
    state.vy = p.momentum * state.vy - p.lr * gy;
    state.x += state.vx;
    state.y += state.vy;
    state.step += 1;
    let loss = rosenbrock(state.x, state.y, p.a, p.b);
    let gn = (gx * gx + gy * gy).sqrt();
    // ~20 f32 ops/step.
    GdStepReport { bit_ops: 20 * 32, loss, grad_norm: gn }
}

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

    #[test]
    fn minimum_is_at_one_one() {
        let f = rosenbrock(1.0, 1.0, 1.0, 100.0);
        assert!(f.abs() < 1e-6, "f(1,1) = {f}");
        let (gx, gy) = grad_rosenbrock(1.0, 1.0, 1.0, 100.0);
        assert!(gx.abs() < 1e-6 && gy.abs() < 1e-6);
    }

    #[test]
    fn descent_reduces_loss() {
        let p = GdParams::default();
        let mut s = GdState::default();
        let l0 = rosenbrock(s.x, s.y, p.a, p.b);
        for _ in 0..5_000 {
            gd_step(&mut s, &p);
        }
        let l1 = rosenbrock(s.x, s.y, p.a, p.b);
        assert!(l1 < l0 * 0.05, "loss {l1} should be ≪ {l0}");
    }

    #[test]
    fn converges_near_minimum() {
        let p = GdParams::default();
        let mut s = GdState::default();
        for _ in 0..20_000 {
            gd_step(&mut s, &p);
        }
        let dist = ((s.x - 1.0).powi(2) + (s.y - 1.0).powi(2)).sqrt();
        assert!(dist < 0.5, "ended at ({}, {}), distance {dist}", s.x, s.y);
    }
}

This is the exact Rust file compiled into the WASM module the page just loaded. Every line that runs on your device is here. The receipt is a function of this code, not a bespoke benchmark.