Adam: момент и масштаб

Вы узнаете

Как Adam соединяет момент со скользящим масштабом и почему поправка смещения делает шаг 1 фиксированного размера.

Попробовать

Дайте трём градиентам очень разные размеры и прочитайте веса после шага 1; затем найдите, где путь достигает вершины.

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

Открыть спеку урока в плеере ↗

Все уроки

• Модуль 1 · Один нейрон

  Нейрон складывает входы и изгибает сумму: ReLU, сигмоида и tanh, и где каждая перестаёт пропускать градиент назад.

  1. 01 ReLU: оставить положительную часть
  2. 02 Сигмоида: число в вероятность
  3. 03 tanh и насыщение

• Модуль 2 · Новые изгибы

  Гладкие кривые нынешних моделей: GELU с провалом, Swish с ручкой и ELU с полом вместо обрыва.

  1. 04 GELU: провал ниже нуля
  2. 05 Swish и SiLU
  3. 06 ELU: пол, а не обрыв

• Модуль 3 · Слой

  Что окружает нейроны: softmax превращает оценки в выбор, эмбеддинг превращает слово в числа, dropout выключает нейроны.

  1. 07 Softmax и температура
  2. 08 Эмбеддинги
  3. 09 Dropout: выключенные нейроны

• Модуль 4 · Картинки

  Чтение картинки по участкам: окно, усреднённое по участку, max pooling и пакетная нормализация признаков разного масштаба.

  1. 10 Участки картинки
  2. 11 Max pooling: максимум окна
  3. 12 Пакетная нормализация

• Модуль 5 · Ошибка

  Насколько неверна догадка, одним числом: среднеквадратичная ошибка, перекрёстная энтропия для вероятностей и функция Хьюбера для выбросов.

  1. 13 Среднеквадратичная ошибка
  2. 14 Перекрёстная энтропия
  3. 15 Функция Хьюбера

• Модуль 6 · Под гору

  Спуск по функции потерь: градиентный спуск с моментом, RMSProp без оглядки на величину градиента и Adam, который соединяет оба.

  1. 16 Градиентный спуск
  2. 17 RMSProp: шаг без масштаба
  3. 18 Adam: момент и масштаб

• Модуль 7 · Обучение

  Откуда берутся градиенты и чего они добиваются: обратное распространение через слой, сеть, которая учит XOR, и LAMB для больших пакетов.

  1. 19 Обратное распространение
  2. 20 Учим XOR
  3. 21 LAMB и большие пакеты

• Модуль 8 · Память и внимание

  Сети, которые читают последовательность: рекуррентная ячейка и её затухающее эхо, вентили LSTM и внимание ко всем позициям сразу.

  1. 22 Рекуррентная ячейка
  2. 23 Вентили LSTM
  3. 24 Внимание

• Модуль 9 · Трансформер и награда

  Части трансформера и агента: позиции как часы, один блок энкодера и PPO, который учится на награде без рывков.

  1. 25 Позиционное кодирование
  2. 26 Блок энкодера
  3. 27 PPO и награда