Внимание

Вы узнаете

Как внимание сравнивает каждый запрос с каждым ключом, что даёт softmax по каждой строке и что убирает причинная маска.

Попробовать

Выключите и включите маску и прочитайте строку 0; затем сравните две головы в строке 2.

specs/ml/recurrent/attention_mechanism.t27

// SPDX-License-Identifier: Apache-2.0
// t27/specs/
// Memory-efficient attention for long sequences | φ² + 1/φ² = 3 | TRINITY

// #8980: redeclared `var i` in one scope four times, called `@infinity()` and
// `@random()` (not Zig), and indexed every row as head_dim wide inside d_model-wide
// rows. backward also skipped the softmax Jacobian: it pushed dL/dweights straight
// into the scores. It is now multi-head scaled dot-product attention over seq x
// d_model rows. Head h reads columns h*hd .. (h+1)*hd, the weights are
// heads x seq x seq, and causal masking is optional. backward is the exact
// gradient, checked against finite differences. Dropout is carried in the config
// and not applied: this is the inference computation.

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

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

    const NEG_INF : f32 = -1.0e30;

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

    pub const AttentionConfig = struct {
        d_model : u32,
        num_heads : u32,
        dropout : f32,
        causal : bool,
    };

    pub const AttentionOutput = struct {
        output : []f32,
        weights : []f32,
    };

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

    fn seq_of(query: []const f32, config: AttentionConfig) -> usize {
        return query.len / config.d_model;
    }

    fn scale_of(config: AttentionConfig) -> f32 {
        return 1.0 / sqrt(@as(f32, @floatFromInt(config.d_model / config.num_heads)));
    }

    fn forward(query: []const f32, key: []const f32, value: []const f32, output: []f32, weights: []f32, config: AttentionConfig) -> void {
        const n = seq_of(query, config);
        const d : usize = config.d_model;
        const hd : usize = config.d_model / config.num_heads;
        const scale = scale_of(config);
        var h : usize = 0;
        while (h < config.num_heads) {
            var i : usize = 0;
            while (i < n) {
                const row = weights[(h * n + i) * n .. (h * n + i + 1) * n];
                var max_s : f32 = NEG_INF;
                var j : usize = 0;
                while (j < n) {
                    var score : f32 = NEG_INF;
                    if (!(config.causal and j > i)) {
                        score = 0.0;
                        var k : usize = 0;
                        while (k < hd) {
                            score = score + query[i * d + h * hd + k] * key[j * d + h * hd + k];
                            k = k + 1;
                        }
                        score = score * scale;
                    }
                    row[j] = score;
                    max_s = @max(max_s, score);
                    j = j + 1;
                }
                var sum : f32 = 0.0;
                j = 0;
                while (j < n) {
                    row[j] = exp(row[j] - max_s);
                    sum = sum + row[j];
                    j = j + 1;
                }
                j = 0;
                while (j < n) {
                    row[j] = row[j] / sum;
                    j = j + 1;
                }
                var k : usize = 0;
                while (k < hd) {
                    var acc : f32 = 0.0;
                    j = 0;
                    while (j < n) {
                        acc = acc + row[j] * value[j * d + h * hd + k];
                        j = j + 1;
                    }
                    output[i * d + h * hd + k] = acc;
                    k = k + 1;
                }
                i = i + 1;
            }
            h = h + 1;
        }
    }

    // dV_j += P_ij dO_i;  dP_ij = dO_i . V_j;  dS_ij = P_ij (dP_ij - sum_k P_ik dP_ik);
    // dQ_i += scale dS_ij K_j;  dK_j += scale dS_ij Q_i  (per head; masked P_ij are 0)
    fn backward(grad_output: []const f32, query: []const f32, key: []const f32, value: []const f32, attn_weights: []f32, grad_query: []f32, grad_key: []f32, grad_value: []f32, config: AttentionConfig) -> void {
        const n = seq_of(query, config);
        const d : usize = config.d_model;
        const hd : usize = config.d_model / config.num_heads;
        const scale = scale_of(config);
        var x : usize = 0;
        while (x < n * d) {
            grad_query[x] = 0.0;
            grad_key[x] = 0.0;
            grad_value[x] = 0.0;
            x = x + 1;
        }
        var h : usize = 0;
        while (h < config.num_heads) {
            var i : usize = 0;
            while (i < n) {
                const p = attn_weights[(h * n + i) * n .. (h * n + i + 1) * n];
                var row_dot : f32 = 0.0;
                var j : usize = 0;
                while (j < n) {
                    var dp : f32 = 0.0;
                    var k : usize = 0;
                    while (k < hd) {
                        dp = dp + grad_output[i * d + h * hd + k] * value[j * d + h * hd + k];
                        grad_value[j * d + h * hd + k] = grad_value[j * d + h * hd + k] + p[j] * grad_output[i * d + h * hd + k];
                        k = k + 1;
                    }
                    row_dot = row_dot + p[j] * dp;
                    j = j + 1;
                }
                j = 0;
                while (j < n) {
                    var dp : f32 = 0.0;
                    var k : usize = 0;
                    while (k < hd) {
                        dp = dp + grad_output[i * d + h * hd + k] * value[j * d + h * hd + k];
                        k = k + 1;
                    }
                    const ds = p[j] * (dp - row_dot) * scale;
                    k = 0;
                    while (k < hd) {
                        grad_query[i * d + h * hd + k] = grad_query[i * d + h * hd + k] + ds * key[j * d + h * hd + k];
                        grad_key[j * d + h * hd + k] = grad_key[j * d + h * hd + k] + ds * query[i * d + h * hd + k];
                        k = k + 1;
                    }
                    j = j + 1;
                }
                i = i + 1;
            }
            h = h + 1;
        }
    }

    fn cfg(causal: bool, dropout: f32) -> AttentionConfig {
        return AttentionConfig{ .d_model = 4, .num_heads = 2, .dropout = dropout, .causal = causal };
    }

    // The fixture: seq 3, d_model 4, 2 heads (columns 0-1 and 2-3); q / k / v scaled down so the
    // softmax is not saturated. which: 0 = output, 1 = weights.
    fn fixture(q: []f32, k: []f32, v: []f32) -> void {
        var x : usize = 0;
        while (x < 12) {
            q[x] = @as(f32, @floatFromInt(x + 1)) / 10.0;
            k[x] = @as(f32, @floatFromInt(x / 4 + 1)) / 4.0 - @as(f32, @floatFromInt(x % 2)) / 3.0;
            v[x] = @as(f32, @floatFromInt(x + 1)) / 10.0;
            x = x + 1;
        }
    }

    fn forward_at(causal: bool, which: u8, index: usize) -> f32 {
        var q : [12]f32 = undefined;
        var k : [12]f32 = undefined;
        var v : [12]f32 = undefined;
        fixture(&q, &k, &v);
        var out : [12]f32 = [9.0] ** 12;
        var w : [18]f32 = [9.0] ** 18;
        forward(&q, &k, &v, &out, &w, cfg(causal, 0.0));
        if (which == 0) {
            return out[index];
        }
        return w[index];
    }

    // L = sum_x g_x O_x with g_x = (x + 1) / 10, one entry of q, k or v moved by delta
    // (which: 0 = q, 1 = k, 2 = v)
    fn loss(which: u8, index: usize, delta: f32) -> f32 {
        var q : [12]f32 = undefined;
        var k : [12]f32 = undefined;
        var v : [12]f32 = undefined;
        fixture(&q, &k, &v);
        if (which == 0) {
            q[index] = q[index] + delta;
        }
        if (which == 1) {
            k[index] = k[index] + delta;
        }
        if (which == 2) {
            v[index] = v[index] + delta;
        }
        var out : [12]f32 = undefined;
        var w : [18]f32 = undefined;
        forward(&q, &k, &v, &out, &w, cfg(false, 0.0));
        var total : f32 = 0.0;
        var x : usize = 0;
        while (x < 12) {
            total = total + @as(f32, @floatFromInt(x + 1)) / 10.0 * out[x];
            x = x + 1;
        }
        return total;
    }

    fn grad(which: u8, index: usize, dropout: f32) -> f32 {
        var q : [12]f32 = undefined;
        var k : [12]f32 = undefined;
        var v : [12]f32 = undefined;
        fixture(&q, &k, &v);
        var out : [12]f32 = undefined;
        var w : [18]f32 = undefined;
        forward(&q, &k, &v, &out, &w, cfg(false, dropout));
        var g : [12]f32 = undefined;
        var x : usize = 0;
        while (x < 12) {
            g[x] = @as(f32, @floatFromInt(x + 1)) / 10.0;
            x = x + 1;
        }
        var gq : [12]f32 = undefined;
        var gk : [12]f32 = undefined;
        var gv : [12]f32 = undefined;
        backward(&g, &q, &k, &v, &w, &gq, &gk, &gv, cfg(false, dropout));
        if (which == 0) {
            return gq[index];
        }
        if (which == 1) {
            return gk[index];
        }
        return gv[index];
    }

    fn slope(which: u8, index: usize) -> f32 {
        return (loss(which, index, 0.001) - loss(which, index, -0.001)) / 0.002;
    }

    fn valid_input() -> u32 {
        return cfg(false, 0.0).d_model % cfg(false, 0.0).num_heads;
    }

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

    // Each weight row is a distribution, and each output is a weighted mean of the values
    // (so it lies between the smallest and largest value in its column)
    test forward_basic_case {
        const row_sum = forward_at(false, 1, 0) + forward_at(false, 1, 1) + forward_at(false, 1, 2);
        assert(@abs(row_sum - 1.0) < 0.000001);
        assert(forward_at(false, 0, 0) + forward_at(false, 0, 1) + forward_at(false, 0, 2) + forward_at(false, 0, 3) > 0.0);  // computed a meaningful output
        assert(forward_at(false, 0, 0) > 0.1 and forward_at(false, 0, 0) < 0.9);
    }

    // dQ, dK and dV against central differences of L = g . O
    test backward_basic_case {
        assert(@abs(grad(0, 0, 0.0) - slope(0, 0)) < 0.001);
        assert(@abs(grad(0, 7, 0.0) - slope(0, 7)) < 0.001);
        assert(@abs(grad(1, 5, 0.0) - slope(1, 5)) < 0.001);
        assert(@abs(grad(2, 10, 0.0) - slope(2, 10)) < 0.001);
        assert(grad(0, 0, 0.0) + grad(0, 1, 0.0) + grad(0, 2, 0.0) + grad(0, 3, 0.0) != 0.0);  // computed gradients
    }

    // Position 0 attends only to itself, position 1 to positions 0 and 1
    test forward_causal_masking {
        assert(forward_at(true, 1, 0) == 1.0);
        assert(forward_at(true, 1, 1) == 0.0 and forward_at(true, 1, 2) == 0.0);
        assert(forward_at(true, 1, 1 * 3 + 2) < 0.999);
        assert(forward_at(true, 1, 1 * 3 + 2) == 0.0);
        assert(forward_at(true, 0, 0) == 0.1);
    }

    // Dropout is not applied here (inference), so it leaves the gradients unchanged
    test backward_dropout {
        assert(grad(2, 3, 0.5) == grad(2, 3, 0.0));
        assert(grad(0, 11, 0.5) == grad(0, 11, 0.0));
    }

    // ═══════════════════════════════════════════════════════════
    // TDD: Invariants (from .tri constraints)
    // ═══════════════════════════════════════════════════════════

    // d_model % num_heads == 0
    invariant attention_constraint_0
        assert valid_input() == 0

    // 0 <= dropout < 1
    invariant attention_constraint_1
        assert cfg(false, 0.5).dropout >= 0.0 and cfg(false, 0.5).dropout < 1.0

    // query.shape == key.shape == value.shape: seq is read from the query
    invariant attention_constraint_2
        assert cfg(false, 0.0).d_model == 4

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

Все уроки

• Модуль 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 и награда