Adam

You will learn

How Adam combines momentum with a running scale, and why its bias correction makes step 1 a fixed size.

Try it

Give the three gradients very different sizes and read the weights after step 1; then find where the path tops out.

specs/ml/optimizer/adam.t27

// SPDX-License-Identifier: Apache-2.0
// t27/specs/
// Adam adapts learning rate per parameter | φ² + 1/φ² = 3 | TRINITY

// #8980: step only checked its gradients with an undeclared `math::is_finite`
// and updated nothing; the tests asserted `result != undefined`. step is now
// Adam (Kingma and Ba, 2015) with bias correction. It refuses (returns false,
// changes nothing) when a gradient is NaN or infinite, which is what the old body
// validated.

module Adam;
    use base::types;
    use math::constants;

    // ═══════════════════════════════════════════════════════════
    // 1. Constants
    // ═══════════════════════════════════════════════════════════

    const DEFAULT_BETA1 : f32 = 0.9;
    const DEFAULT_BETA2 : f32 = 0.999;
    const DEFAULT_EPSILON : f32 = 1e-8;

    // ═══════════════════════════════════════════════════════════
    // 2. Types
    // ═══════════════════════════════════════════════════════════

    pub const AdamConfig = struct {
        learning_rate : f32,
        beta1 : f32,
        beta2 : f32,
        epsilon : f32,
        weight_decay : f32,
    };

    // ═══════════════════════════════════════════════════════════
    // 3. Core Functions
    // ═══════════════════════════════════════════════════════════

    // b^t for a whole t >= 0
    fn powi(b: f32, t: u32) -> f32 {
        var r : f32 = 1.0;
        var n : u32 = 0;
        while (n < t) {
            r = r * b;
            n = n + 1;
        }
        return r;
    }

    // NaN is the one value not equal to itself; inf - inf is NaN
    fn is_finite(x: f32) -> bool {
        return x == x and (x - x) == (x - x);
    }

    // One Adam step number t (1-based) over params, with moment buffers m and v:
    //   g = grad + wd p
    //   m = b1 m + (1 - b1) g          v = b2 v + (1 - b2) g^2
    //   m_hat = m / (1 - b1^t)          v_hat = v / (1 - b2^t)
    //   p = p - lr m_hat / (sqrt(v_hat) + eps)
    fn step(params: []f32, grads: []const f32, m: []f32, v: []f32, t: u32, config: AdamConfig) -> bool {
        var i : usize = 0;
        while (i < grads.len) {
            if (!is_finite(grads[i])) {
                return false;
            }
            i = i + 1;
        }
        const c1 = 1.0 - powi(config.beta1, t);
        const c2 = 1.0 - powi(config.beta2, t);
        i = 0;
        while (i < params.len) {
            const g = grads[i] + config.weight_decay * params[i];
            m[i] = config.beta1 * m[i] + (1.0 - config.beta1) * g;
            v[i] = config.beta2 * v[i] + (1.0 - config.beta2) * g * g;
            const m_hat = m[i] / c1;
            const v_hat = v[i] / c2;
            params[i] = params[i] - config.learning_rate * m_hat / (sqrt(v_hat) + config.epsilon);
            i = i + 1;
        }
        return true;
    }

    fn cfg(lr: f32) -> AdamConfig {
        return AdamConfig{ .learning_rate = lr, .beta1 = DEFAULT_BETA1, .beta2 = DEFAULT_BETA2, .epsilon = DEFAULT_EPSILON, .weight_decay = 0.0 };
    }

    // which: 0 = param, 1 = m, 2 = v after one step on params (1, 2, 3) with grads (g0, g1, g2)
    fn first_step(g0: f32, g1: f32, g2: f32, which: u8, index: usize) -> f32 {
        var params : [3]f32 = [1.0, 2.0, 3.0];
        const grads : [3]f32 = [g0, g1, g2];
        var m : [3]f32 = [0.0] ** 3;
        var v : [3]f32 = [0.0] ** 3;
        _ = step(&params, &grads, &m, &v, 1, cfg(0.001));
        if (which == 1) {
            return m[index];
        }
        if (which == 2) {
            return v[index];
        }
        return params[index];
    }

    // a / b at run time, where IEEE gives NaN for 0 / 0 and inf for 1 / 0
    fn ratio(a: f32, b: f32) -> f32 {
        return a / b;
    }

    fn accepts(g: f32) -> bool {
        var params : [1]f32 = [1.0];
        const grads : [1]f32 = [g];
        var m : [1]f32 = [0.0];
        var v : [1]f32 = [0.0];
        return step(&params, &grads, &m, &v, 1, cfg(0.001)) and params[0] != 1.0;
    }

    // Minimise (x - 3)^2 from 0
    fn descend(steps: u32) -> f32 {
        var params : [1]f32 = [0.0];
        var m : [1]f32 = [0.0];
        var v : [1]f32 = [0.0];
        var t : u32 = 1;
        while (t <= steps) {
            const grads : [1]f32 = [2.0 * (params[0] - 3.0)];
            _ = step(&params, &grads, &m, &v, t, cfg(0.1));
            t = t + 1;
        }
        return params[0];
    }

    // ═══════════════════════════════════════════════════════════
    // TDD: Tests (from .tri behaviors)
    // ═══════════════════════════════════════════════════════════

    // Bias correction makes the first step lr * sign(g), whatever |g| is
    test step_basic_case
        given p0 = first_step(0.1, -5.0, 0.0, 0, 0)
        when p1 = first_step(0.1, -5.0, 0.0, 0, 1)
        then @abs(p0 - 0.999) < 0.000001 and @abs(p1 - 2.001) < 0.000001 and first_step(0.1, -5.0, 0.0, 0, 2) == 3.0

    // m = 0.1 g and v = 0.001 g^2 after the first step
    test step_with_gradients
        given m2 = first_step(0.1, 0.2, 0.3, 1, 2)
        when v2 = first_step(0.1, 0.2, 0.3, 2, 2)
        then @abs(m2 - 0.03) < 0.0000001 and @abs(v2 - 0.00009) < 0.0000001

    // Finite gradients are applied; NaN and inf are refused
    test step_validates_gradients
        given ok = accepts(1.5) and accepts(-2.3) and accepts(4.7)
        when refused = !accepts(ratio(0.0, 0.0)) and !accepts(ratio(1.0, 0.0))
        then ok and refused

    test step_descends_a_quadratic
        given result = descend(500)
        then @abs(result - 3.0) < 0.01

    invariant adam_betas_are_below_one
        assert DEFAULT_BETA1 < 1.0 and DEFAULT_BETA2 < 1.0 and DEFAULT_BETA1 < DEFAULT_BETA2 and DEFAULT_EPSILON > 0.0

Open the lesson's spec in the player ↗

All lessons

• Module 1 · One neuron

  A neuron sums its inputs and bends the sum: ReLU, the sigmoid and tanh, and where each one stops passing a gradient back.

  1. 01 ReLU: keep the positive part
  2. 02 The sigmoid: a number into a probability
  3. 03 tanh and saturation

• Module 2 · New bends

  The smooth curves of current models: GELU with its dip, Swish with its knob, and ELU with a floor instead of a cliff.

  1. 04 GELU: a dip below zero
  2. 05 Swish and SiLU
  3. 06 ELU: a floor, not a cliff

• Module 3 · A layer

  What sits around the neurons: softmax turns scores into a choice, an embedding turns a word into numbers, dropout switches units off.

  1. 07 Softmax and temperature
  2. 08 Embeddings
  3. 09 Dropout

• Module 4 · Pictures

  Reading a picture in patches: a window averaged over its patch, max pooling, and batch normalisation of features on different scales.

  1. 10 Patches of a picture
  2. 11 Max pooling
  3. 12 Batch normalisation

• Module 5 · The error

  How wrong a guess is, as one number: mean squared error, cross-entropy for probabilities, and the Huber loss for outliers.

  1. 13 Mean squared error
  2. 14 Cross-entropy
  3. 15 The Huber loss

• Module 6 · Downhill

  Walking down the loss: gradient descent with momentum, RMSProp that ignores the gradient's size, and Adam that combines them.

  1. 16 Gradient descent
  2. 17 RMSProp
  3. 18 Adam

• Module 7 · Learning

  Where the gradients come from and what they achieve: backpropagation through a layer, a network that learns XOR, and LAMB for large batches.

  1. 19 Backpropagation
  2. 20 Learning XOR
  3. 21 LAMB and large batches

• Module 8 · Memory and attention

  Networks that read a sequence: a recurrent cell and its fading echo, the gates of an LSTM, and attention over every position at once.

  1. 22 A recurrent cell
  2. 23 LSTM gates
  3. 24 Attention

• Module 9 · The transformer and the reward

  The parts of a transformer and of an agent: positions as clocks, one encoder block, and PPO, which learns from reward without lurching.

  1. 25 Positional encoding
  2. 26 An encoder block
  3. 27 PPO and the reward