Закройте доступ к будущим ключам каузальной маской

Предскажите видимый треугольник

Добавьте к примеру из главы 27 третью позицию:

Здесь один пакет, три позиции токенов, ширина запросов и ключей равна двум, ширина значений — тоже двум. Следовательно,

До маскирования строки исходных оценок равны $[0,6,3]$, $[6,-4,3]$ и $[3,1,3]$. Сначала отметьте ячейки, доступные каждому запросу. Запросу $0$ доступен только ключ $0$; запросу $1$ — ключи $0$ и $1$; запросу $2$ — все три ключа. Итого получается шесть доступных и три закрытые ячейки. Все три диагональные ячейки входят в число доступных.

Примените маску до softmax

Для строки запроса $i$ и столбца ключа $j$ используйте аддитивную маску

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

Пусть тензор масштабированных оценок равен

До нормировки добавьте ограничение видимости, затем смешайте значения:

Для этого примера

Вероятность внимания к каждому закрытому ключу строго равна нулю. Сумма вероятностей по каждому доступному префиксу равна единице:

Получаются строки

Нельзя сначала применить обычный softmax ко всем ключам, а затем обнулить вероятности будущих ключей: в знаменателе вероятностей доступных ключей по-прежнему будут учтены закрытые ключи. Большое по модулю конечное отрицательное число можно использовать как приближение идеальной маски, но математически оно не равно $-\infty$.

Не путайте ограничение видимости с позиционной информацией

• $S$ — тензор масштабированных оценок соответствия запросов и ключей до маскирования.
• $M$ — аддитивная маска видимости.
• $i$ обозначает позицию запроса, а $j$ — позицию ключа.
• $A$ содержит вероятности внимания после исключения будущих ключей.
• $Q$, $K$ и $V$ сохраняют роли запросов, ключей и значений из главы 27.
• $O$ содержит по одной взвешенной смеси доступных строк значений для каждой позиции запроса.
• $B$ — размер пакета, $T$ — число токенов, $d_k$ — ширина запросов и ключей, а $d_v$ — ширина значений.

Полное правило для форм остаётся таким:

Для доступной ячейки нормировка строки имеет вид

а для закрытой ячейки

Это правило задаёт видимость, но не положение. Чтобы различать порядок, строкам нужен отдельный абсолютный или относительный позиционный сигнал. Маски заполнения, переменная длина последовательности, несколько голов, выходная проекция и кеш ключей и значений — тоже отдельные задачи.

От рекуррентного состояния префикса к явной маске декодера

Грейвс, Generating Sequences With Recurrent Neural Networks обучает модель предсказывать следующий элемент, последовательно обрабатывая элементы последовательности. При генерации каждый выбранный элемент становится следующим рекуррентным входом, поэтому будущих элементов ещё не существует. Граница префикса возникает из последовательной рекуррентной обработки, где состояние приходится обновлять шаг за шагом. Это утверждение относится к генерации текста в статье, а не к отдельной модели синтеза рукописного текста с дополнительным условным входом.

Васвани и соавторы, Attention Is All You Need объединяют запросы в матрицы и маскируют самовнимание декодера так, чтобы строка не могла обратиться к последующим позициям. Недопустимым оценкам до softmax присваивается $-\infty$; вместе с эмбеддингами выходной последовательности, сдвинутыми на одну позицию, это означает, что строка $i$ зависит только от известных выходов до позиции $i$.

Поэтому при обучении строки внимания для известных целевых позиций можно вычислять вместе, не позволяя более ранней строке использовать более позднюю цель. При этом обычная авторегрессионная генерация по-прежнему добавляет по одному токену. Декодерный Transformer применяет эту каузальную границу в каждом слое самовнимания. Маска задаёт границу видимости, но не кодирует положение.

Оставьте маску доступной для проверки, а записанные значения — конечными

causal_additive_mask строит обычный тензор формы $[T,T]$. В доступных ячейках находится $0$, а в будущих — $-\infty$:

/// Builds an additive square mask with zero for `key <= query` and negative
/// infinity for future keys.
pub fn causal_additive_mask(tokens: usize) -> Result<Tensor, CausalMaskingError> {
    let elements = tokens
        .checked_mul(tokens)
        .ok_or(CausalMaskingError::MaskTensor(TensorError::ShapeOverflow))?;
    let mut values = Vec::new();
    values
        .try_reserve_exact(elements)
        .map_err(|_| CausalMaskingError::MaskAllocationFailed { elements })?;
    for query in 0..tokens {
        for key in 0..tokens {
            values.push(if key <= query { 0.0 } else { f64::NEG_INFINITY });
        }
    }
    Tensor::from_vec(vec![tokens, tokens], values).map_err(CausalMaskingError::MaskTensor)
}

Дифференцируемый TensorValue отклоняет неконечные данные листовых тензоров. Операция causal_softmax реализует ту же математическую границу, не записывая $-\infty$: она читает только $j\le i$, вычитает максимум доступного префикса, нормирует доступные ячейки и записывает точное значение с плавающей точкой $+0.0$ во все закрытые ячейки.

fn causal_softmax_forward(input: &Tensor) -> Result<Tensor, TensorAutodiffError> {
    if input.rank() < 2 {
        return Err(ModelOpError::CausalSoftmaxRank { rank: input.rank() }.into());
    }
    let queries = input.shape()[input.rank() - 2];
    let keys = input.shape()[input.rank() - 1];
    if queries != keys {
        return Err(ModelOpError::CausalSoftmaxNonSquare { queries, keys }.into());
    }
    if queries == 0 {
        return Err(ModelOpError::CausalSoftmaxEmptyTokens.into());
    }

    let mut probabilities = zeros(input.shape())?;
    let grids = input.len() / (queries * keys);
    for grid in 0..grids {
        for query in 0..queries {
            let row_start = (grid * queries + query) * keys;
            let allowed = &input.as_slice()[row_start..=row_start + query];
            let maximum = allowed.iter().copied().fold(f64::NEG_INFINITY, f64::max);
            let mut exponential_tail = 0.0;
            let mut skipped_one_maximum = false;
            for &score in allowed {
                let shifted = score - maximum;
                if shifted == 0.0 && !skipped_one_maximum {
                    skipped_one_maximum = true;
                } else {
                    exponential_tail += shifted.exp();
                }
            }
            debug_assert!(skipped_one_maximum);
            let denominator = 1.0 + exponential_tail;
            for (key, &score) in allowed.iter().enumerate() {
                let probability = (score - maximum).exp() / denominator;
                probabilities.as_mut_slice()[row_start + key] =
                    if probability == 0.0 { 0.0 } else { probability };
            }
        }
    }
    Ok(probabilities)
}

Операция принимает тензоры оценок ранга не меньше двух, у которых размеры последних двух осей равны. Поэтому она поддерживает форму $[B,H,T,T]$, которая понадобится далее, не смешивая ведущие оси. Операция отклоняет ранг меньше двух, разные размеры последних осей, пустую ось токенов и уже освобождённый операнд.

causal_scaled_dot_product_self_attention повторно использует построение оценок из главы 27, применяет causal_softmax и умножает результат на $V$:

/// Inspectable evidence from one causally masked attention head.
#[derive(Clone, Debug)]
pub struct CausalSelfAttentionForward {
    raw_scores: TensorValue,
    scaled_scores: TensorValue,
    additive_mask: Tensor,
    weights: TensorValue,
    output: TensorValue,
    scale: f64,
    key_width: usize,
    value_width: usize,
}

impl CausalSelfAttentionForward {
    pub fn raw_scores(&self) -> &TensorValue {
        &self.raw_scores
    }

    pub fn dot_products(&self) -> &TensorValue {
        &self.raw_scores
    }

    pub fn scaled_scores(&self) -> &TensorValue {
        &self.scaled_scores
    }

    /// The plain additive mask. It is intentionally not a tape value because
    /// its blocked cells contain negative infinity.
    pub const fn additive_mask(&self) -> &Tensor {
        &self.additive_mask
    }

    pub fn weights(&self) -> &TensorValue {
        &self.weights
    }

    pub fn probabilities(&self) -> &TensorValue {
        &self.weights
    }

    pub fn output(&self) -> &TensorValue {
        &self.output
    }

    pub const fn scale(&self) -> f64 {
        self.scale
    }

    pub const fn key_width(&self) -> usize {
        self.key_width
    }

    pub const fn value_width(&self) -> usize {
        self.value_width
    }

    pub fn into_output(self) -> TensorValue {
        self.output
    }
}

/// Computes one scaled self-attention head with an inclusive-prefix mask.
pub fn causal_scaled_dot_product_self_attention(
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
) -> Result<CausalSelfAttentionForward, CausalMaskingError> {
    let prepared = scaled_self_attention_scores(query, key, value)?;
    let tokens = query.shape()[1];
    let additive_mask = causal_additive_mask(tokens)?;
    let weights = prepared
        .scaled_scores
        .causal_softmax()
        .map_err(autodiff_error(CausalMaskingStage::MaskedSoftmax))?;
    let output = weights
        .matmul(value)
        .map_err(autodiff_error(CausalMaskingStage::ValueMixture))?;

    Ok(CausalSelfAttentionForward {
        raw_scores: prepared.raw_scores,
        scaled_scores: prepared.scaled_scores,
        additive_mask,
        weights,
        output,
        scale: prepared.scale,
        key_width: prepared.key_width,
        value_width: prepared.value_width,
    })
}

Для доступной оценки обратный проход подчиняется формуле

а закрытая оценка получает

Все шесть координат каждого из тензоров $Q$, $K$ и $V$ согласуются с центральными разностями при шаге $10^{-6}$ и допуске $4\times10^{-6}$. Среди проверенных граничных случаев также есть один токен, пустые пакеты, тензоры оценок ранга два и четыре, независимые ведущие оси, очень большие закрытые оценки, типизированные ошибки, операнды из освобождённой ленты и побитное совпадение при повторном запуске.

Исполняемый пример выводит проверенные данные: маску, вероятности, выходы, результат проверки неизменности префикса, градиенты и граничные случаи:

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let evidence = ch28_causal_masking::learner_evidence()?;
    print!("{}", ch28_causal_masking::render_report(&evidence));
    Ok(())
}

Выполните cargo run --quiet --locked -p ch28-causal-masking, чтобы увидеть полный результат примера.

Проследите путь внимания по нижнему треугольнику

Треугольные таблицы показывают, какие оценки сохраняются до нормировки. Сплошная, штриховая и двойная рамки различают доступные, закрытые и диагональные ячейки, не полагаясь только на цвет.

Если изменить только

то $o_0$ и $o_1$ побитно совпадут с исходными, а

Сначала сделайте прогноз, затем откройте ответы

1. Запишите множество индексов доступных ключей для строк запросов $0$, $1$ и $2$.
2. Предскажите, может ли замена только $k_2$ и $v_2$ изменить $o_0$ или $o_1$.
3. Объясните, почему все три диагональные ячейки доступны, хотя модель не может получить доступ к токену, который нужно предсказать.
4. Предскажите градиенты запроса и ключа для каузальной головы самовнимания с одним токеном.
5. Объясните, почему после обнуления будущих вероятностей вслед за обычным softmax сумма строки перестаёт быть равной единице.
6. Определите, позволяет ли маска параллельно обрабатывать позиции токенов при обучении, когда целевые токены известны, при авторегрессионной генерации, в обоих случаях или ни в одном.
7. Назовите отдельный механизм, который появится в главе 29, не изменяя нижнетреугольную границу.

Сохраняйте границу префикса по мере роста декодера

Теперь накопительный декодер создаёт выход в позиции $i$, используя только ключи и значения до позиции $i$ включительно. Если функция потерь зависит лишь от первых двух позиций, градиент по будущему суффиксу строго равен нулю:

Именно такая информационная граница нужна авторегрессионному декодеру. В главе 29 появится относительная информация о позициях, но граница не расширится.