26 · Версия материала 2

Создайте представления запросов, ключей и значений

Разберите, как самовнимание Transformer создаёт из одной последовательности скрытых состояний тензоры запросов, ключей и значений с помощью трёх независимых проекций без смещения.

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

Начните с одного пакета, содержащего два состояния токенов:

X=\begin{bmatrix}1&2&-1\\0&1&2\end{bmatrix},\qquad X\in\mathbb{R}^{1\times2\times3}.

Используйте три разные матрицы весов:

W_Q=\begin{bmatrix}1&0\\0&1\\1&-1\end{bmatrix}.

W_K=\begin{bmatrix}0&1\\1&0\\-1&1\end{bmatrix}.

W_V=\begin{bmatrix}1&1\\1&-1\\0&2\end{bmatrix}.

До запуска примера умножьте первый токен на каждую матрицу весов. Ожидаемые векторы — $q_0=[0,3]$ , $k_0=[3,0]$ и $v_0=[3,-3]$ . Для второго токена должны получиться $q_1=[2,-1]$ , $k_1=[-1,2]$ и $v_1=[1,3]$ .

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

Спроецируйте последнюю ось признаков тремя способами

Полное правило прямого прохода:

Q=XW_Q,\quad K=XW_K,\quad V=XW_V

Формы тензоров должны удовлетворять условиям

X\in\mathbb{R}^{B\times T\times d_{model}}.

W_Q,W_K,W_V\in\mathbb{R}^{d_{model}\times d_{head}}.

Q,K,V\in\mathbb{R}^{B\times T\times d_{head}}.

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

Для зафиксированного примера полный результат равен

Q=\begin{bmatrix}0&3\\2&-1\end{bmatrix},\quad K=\begin{bmatrix}3&0\\-1&2\end{bmatrix},\quad V=\begin{bmatrix}3&-3\\1&3\end{bmatrix}.

Различайте роли, оси и размерности

$X$ — тензор скрытых состояний, поступающий в самовнимание.
$B$ — размер пакета, а $T$ — число позиций токенов.
$d_{model}$ — ширина входных признаков.
$d_{head}$ — ширина выхода единственной головы в этой главе.
$W_Q$ , $W_K$ и $W_V$ — независимые обучаемые матрицы весов.
$Q$ — представление запросов: с его помощью каждая позиция задаст, какую информацию следует учесть.
$K$ — представление ключей: оно опишет, по каким признакам позиции сопоставляются с запросами.
$V$ — представление значений: оно содержит информацию для последующего смешивания.

Сами обозначения не вычисляют внимание. Эта глава также не вводит ограничений делимости, связанных с несколькими головами. В главе 30 появится число голов $h$ и обычное соотношение $d_{model}=h\,d_{head}$ , используемое при разделении и объединении нескольких голов.

От обучаемого выравнивания к проекциям самовнимания

Здесь рассматривается развитие нейронного внимания на пути к современным LLM, а не история языков программирования.

Базовые рекуррентные модели энкодер–декодер сжимали исходную последовательность в один фиксированный вектор. Бахданау, Чо и Бенжио, Neural Machine Translation by Jointly Learning to Align and Translate заменяют это узкое место дифференцируемым выравниванием. На каждом шаге декодирования Бахданау, Чо и Бенжио сопоставляют каждую аннотацию энкодера с предыдущим состоянием декодера и по полученным весам формируют вектор контекста. В их обозначениях $s_{i-1}$ — предыдущее состояние декодера, а $h_j$ — одна из аннотаций энкодера, поступающих в обучаемую функцию совместимости. Состояние, играющее роль запроса, и содержимое аннотаций по-прежнему поступают из двух разных частей модели энкодера–декодера.

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

Васвани и соавторы, Attention Is All You Need описывают внимание как отображение запроса и пар ключ–значение в выход. Васвани и соавторы используют отдельные обучаемые линейные проекции для запросов, ключей и значений и определяют самовнимание внутри одной последовательности. В самовнимании вместо выравнивания двух источников одна последовательность предыдущего слоя служит входом для всех трёх ролей до вычисления оценок.

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

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

Сопоставьте две части модели с одной последовательностью, спроецированной тремя способами rust/demos/ch26-qkv-projections/src/lib.rs#historical-attention-source-contrast

fn historical_source_contrast() -> HistoryEvidence {
    HistoryEvidence {
        earlier_left: "decoder-state",
        earlier_right: "encoder-annotations",
        transformer_source: "one-sequence",
        mapping: "retrospective",
    }
}

Составьте слой из трёх существующих дифференцируемых линейных слоёв

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

Спроецируйте тензор скрытых состояний третьего ранга в независимые тензоры запросов, ключей и значений rust/crates/llm-from-scratch/src/attention/qkv.rs#qkv-layer

/// The three projected views of the same batch and token positions.
#[derive(Clone, Debug)]
pub struct QkvForward {
    query: TensorValue,
    key: TensorValue,
    value: TensorValue,
}

impl QkvForward {
    pub fn query(&self) -> &TensorValue {
        &self.query
    }

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

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

    pub fn into_parts(self) -> (TensorValue, TensorValue, TensorValue) {
        (self.query, self.key, self.value)
    }
}

/// Three independent `[model_width, head_width]` linear maps with no biases.
#[derive(Clone, Debug)]
pub struct QkvProjections {
    query: Linear,
    key: Linear,
    value: Linear,
    parameters: NamedParameters,
    model_width: usize,
    head_width: usize,
}

impl QkvProjections {
    /// Initializes Q, K, and V in that order without partially advancing `rng`.
    pub fn new(
        parameter_prefix: impl Into<String>,
        model_width: usize,
        head_width: usize,
        rng: &mut SplitMix64,
    ) -> Result<Self, QkvError> {
        let parameter_prefix = parameter_prefix.into();
        let mut trial = rng.clone();
        let query = Linear::new(
            format!("{parameter_prefix}.query"),
            model_width,
            head_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(QkvProjection::Query))?;
        let key = Linear::new(
            format!("{parameter_prefix}.key"),
            model_width,
            head_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(QkvProjection::Key))?;
        let value = Linear::new(
            format!("{parameter_prefix}.value"),
            model_width,
            head_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(QkvProjection::Value))?;
        let projections = Self::from_layers(query, key, value)?;
        *rng = trial;
        Ok(projections)
    }

    /// Gives Q/K/V semantics to three existing matrix parameters.
    pub fn from_weights(
        query_weight: NamedParameter,
        key_weight: NamedParameter,
        value_weight: NamedParameter,
    ) -> Result<Self, QkvError> {
        let query = Linear::from_parameters(query_weight, None)
            .map_err(projection_error(QkvProjection::Query))?;
        let key = Linear::from_parameters(key_weight, None)
            .map_err(projection_error(QkvProjection::Key))?;
        let value = Linear::from_parameters(value_weight, None)
            .map_err(projection_error(QkvProjection::Value))?;
        Self::from_layers(query, key, value)
    }

    fn from_layers(query: Linear, key: Linear, value: Linear) -> Result<Self, QkvError> {
        let input_widths = (query.input_width(), key.input_width(), value.input_width());
        if input_widths.0 != input_widths.1 || input_widths.0 != input_widths.2 {
            return Err(QkvError::BranchInputWidthMismatch {
                query: input_widths.0,
                key: input_widths.1,
                value: input_widths.2,
            });
        }

        let output_widths = (
            query.output_width(),
            key.output_width(),
            value.output_width(),
        );
        if output_widths.0 != output_widths.1 || output_widths.0 != output_widths.2 {
            return Err(QkvError::BranchOutputWidthMismatch {
                query: output_widths.0,
                key: output_widths.1,
                value: output_widths.2,
            });
        }

        let parameters = NamedParameters::try_new(vec![
            query.weight().clone(),
            key.weight().clone(),
            value.weight().clone(),
        ])?;
        Ok(Self {
            query,
            key,
            value,
            parameters,
            model_width: input_widths.0,
            head_width: output_widths.0,
        })
    }

    /// Projects exactly `[batch, tokens, model_width]` into three head-width views.
    pub fn forward(&self, input: &TensorValue) -> Result<QkvForward, QkvError> {
        let shape = input.shape();
        if shape.len() != 3 {
            return Err(QkvError::InputRank { rank: shape.len() });
        }
        if shape[2] != self.model_width {
            return Err(QkvError::InputWidthMismatch {
                expected: self.model_width,
                actual: shape[2],
            });
        }

        let query = self
            .query
            .forward(input)
            .map_err(projection_error(QkvProjection::Query))?;
        let key = self
            .key
            .forward(input)
            .map_err(projection_error(QkvProjection::Key))?;
        let value = self
            .value
            .forward(input)
            .map_err(projection_error(QkvProjection::Value))?;
        Ok(QkvForward { query, key, value })
    }

    pub fn query(&self) -> &Linear {
        &self.query
    }

    pub fn key(&self) -> &Linear {
        &self.key
    }

    pub fn value(&self) -> &Linear {
        &self.value
    }

    pub fn parameters(&self) -> &[NamedParameter] {
        self.parameters.as_slice()
    }

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

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

    pub const fn parameter_count(&self) -> usize {
        3 * self.model_width * self.head_width
    }
}

Обёртка принимает вход строго третьего ранга, чтобы оси пакета и токенов оставались явными. Проверки выполняются в определённом порядке: сначала ранг, затем ширина, а ветви проверяются по очереди: сначала запросы, затем ключи и значения:

Явно различайте ошибки ветвей и форм rust/crates/llm-from-scratch/src/attention/qkv.rs#qkv-errors

/// The projection branch that rejected construction or a delegated operation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum QkvProjection {
    Query,
    Key,
    Value,
}

impl fmt::Display for QkvProjection {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::Query => "query",
            Self::Key => "key",
            Self::Value => "value",
        })
    }
}

/// A rejected Q/K/V parameter set or hidden-state input.
#[derive(Clone, Debug, PartialEq)]
pub enum QkvError {
    Projection {
        projection: QkvProjection,
        source: LinearError,
    },
    InputRank {
        rank: usize,
    },
    InputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    BranchInputWidthMismatch {
        query: usize,
        key: usize,
        value: usize,
    },
    BranchOutputWidthMismatch {
        query: usize,
        key: usize,
        value: usize,
    },
    Initialization(InitializationError),
}

impl fmt::Display for QkvError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Projection { projection, source } => {
                write!(formatter, "{projection} projection: {source}")
            }
            Self::InputRank { rank } => write!(
                formatter,
                "Q/K/V input must have rank three [batch, tokens, model_width], got rank {rank}"
            ),
            Self::InputWidthMismatch { expected, actual } => write!(
                formatter,
                "Q/K/V input final width must equal model width {expected}, got {actual}"
            ),
            Self::BranchInputWidthMismatch { query, key, value } => write!(
                formatter,
                "Q/K/V model widths must match, got query {query}, key {key}, value {value}"
            ),
            Self::BranchOutputWidthMismatch { query, key, value } => write!(
                formatter,
                "Q/K/V head widths must match, got query {query}, key {key}, value {value}"
            ),
            Self::Initialization(source) => source.fmt(formatter),
        }
    }
}

impl Error for QkvError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Projection { source, .. } => Some(source),
            Self::Initialization(source) => Some(source),
            _ => None,
        }
    }
}

impl From<InitializationError> for QkvError {
    fn from(error: InitializationError) -> Self {
        Self::Initialization(error)
    }
}

fn projection_error(projection: QkvProjection) -> impl FnOnce(LinearError) -> QkvError {
    move |source| QkvError::Projection { projection, source }
}

В примере обратного прохода строится одна скалярная целевая функция:

L=\langle Q,U_Q\rangle+\langle K,U_K\rangle+\langle V,U_V\rangle.

Если черта сверху обозначает градиент обратного режима, вклады всех трёх путей суммируются в градиент общего входа:

\bar X=\bar QW_Q^{\mathsf T}+\bar KW_K^{\mathsf T}+\bar VW_V^{\mathsf T}.

Для вычисления градиентов весов объединим оси пакета и токенов в одну ось строк, обозначенную индексом $(BT)$ . После этого каждая ветвь сохраняет собственный градиент:

\bar W_Q=X_{(BT)}^{\mathsf T}\bar Q_{(BT)},\quad \bar W_K=X_{(BT)}^{\mathsf T}\bar K_{(BT)},\quad \bar W_V=X_{(BT)}^{\mathsf T}\bar V_{(BT)}.

Для входящих градиентов из примера градиент общего входа равен

\bar X=\begin{bmatrix}3&1.5&1.5\\-1.5&3.5&-5\end{bmatrix},

а у каждой ветви остаётся собственный градиент весов. Центральные разности проверяют все шесть координат каждого из тензоров $X$ , $W_Q$ , $W_K$ и $W_V$ с шагом $10^{-6}$ и допуском $2\times10^{-6}$ . Для пустых осей пакета и токенов сохраняется путь обратного распространения, а два запуска дают одинаковые битовые представления всех чисел:

Соберите точные данные прямого и обратного проходов, форм, ошибок, инициализации и градиентов rust/demos/ch26-qkv-projections/src/lib.rs#qkv-fixture

pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
    let primary = primary_once()?;
    let replay = primary_once()?;
    let (query_changed, key_unchanged, value_unchanged) = independence_evidence()?;
    let (input_checks, query_weight_checks, key_weight_checks, value_weight_checks, passed) =
        gradient_evidence(&primary)?;
    Ok(LearnerEvidence {
        replay_bitwise: same_primary_bits(&primary, &replay),
        primary,
        shapes: shape_evidence()?,
        errors: error_evidence()?,
        initialization: initialization_evidence()?,
        history: historical_source_contrast(),
        query_changed,
        key_unchanged,
        value_unchanged,
        input_checks,
        query_weight_checks,
        key_weight_checks,
        value_weight_checks,
        gradcheck_passed: passed,
    })
}

Запустите полный пример проекций запросов, ключей и значений rust/demos/ch26-qkv-projections/src/main.rs

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

Выполните cargo run --quiet --locked -p ch26-qkv-projections, чтобы увидеть те же значения прямого прохода, градиенты, проверки форм и отклонённые входы, полученные исполняемым примером.

Проследите, как одна последовательность превращается в три обучаемых представления

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

Получите из одной последовательности скрытых состояний три обучаемых представления

Проследите, как одна последовательность скрытых состояний проходит через проекции запросов, ключей и значений, а затем сравните формы, градиенты, независимость, источники данных в исторических механизмах внимания и отклонённые входы.

Рамка общего входа
Сплошная рамка — запрос
Пунктирная рамка — ключ
Двойная рамка — значение

Сохраните позиции, изменив пространство признаков

Общий вход: скрытые состояния

X=[1.000000,2.000000,-1.000000,0.000000,1.000000,2.000000]

Форма: $\operatorname{shape}(X)=[1,2,3]$
Наличие смещения: bias=false

Сплошная рамка — запрос

Проекция запросов

Запрос: какую информацию следует учесть для этой позиции?

$Q=XW_Q$

Параметр с фиксированным именем: decoder.block.0.attention.query.weight
Матрица весов проекции: $\operatorname{shape}(W_Q)=[3,2]$
Результат проекции: $\operatorname{shape}(Q)=[1,2,2]$

$W_Q$
	$0$	$1$
$0$	$1.000000$	$0.000000$
$1$	$0.000000$	$1.000000$
$2$	$1.000000$	$-1.000000$

Результат проекции
	$0$	$1$
$t_0$	$0.000000$	$3.000000$
$t_1$	$2.000000$	$-1.000000$

Пунктирная рамка — ключ

Проекция ключей

Ключ: по каким признакам эта позиция сопоставляется с запросом?

$K=XW_K$

Параметр с фиксированным именем: decoder.block.0.attention.key.weight
Матрица весов проекции: $\operatorname{shape}(W_K)=[3,2]$
Результат проекции: $\operatorname{shape}(K)=[1,2,2]$

$W_K$
	$0$	$1$
$0$	$0.000000$	$1.000000$
$1$	$1.000000$	$0.000000$
$2$	$-1.000000$	$1.000000$

Результат проекции
	$0$	$1$
$t_0$	$3.000000$	$0.000000$
$t_1$	$-1.000000$	$2.000000$

Двойная рамка — значение

Проекция значений

Значение: какое содержимое эта позиция передаёт дальше?

$V=XW_V$

Параметр с фиксированным именем: decoder.block.0.attention.value.weight
Матрица весов проекции: $\operatorname{shape}(W_V)=[3,2]$
Результат проекции: $\operatorname{shape}(V)=[1,2,2]$

$W_V$
	$0$	$1$
$0$	$1.000000$	$1.000000$
$1$	$1.000000$	$-1.000000$
$2$	$0.000000$	$2.000000$

Результат проекции
	$0$	$1$
$t_0$	$3.000000$	$-3.000000$
$t_1$	$1.000000$	$3.000000$

У трёх ветвей общие координаты, но разные веса и значения признаков.

От двух источников данных для внимания к одному

Сходство ролей запросов, ключей и значений с ранним механизмом из двух частей модели — лишь ограниченная ретроспективная аналогия.
	Предыдущее состояние декодера	Аннотации энкодера	Одна последовательность предыдущего слоя
Аддитивное внимание энкодера–декодера	`decoder-state`	`encoder-annotations`	—
Самовнимание Transformer	—	—	`one-sequence`

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

Проверьте независимость, градиенты и граничные случаи

Суммарный градиент по входу

$\bar X=\begin{bmatrix}3.000000&1.500000&1.500000\\-1.500000&3.500000&-5.000000\end{bmatrix}$

$\operatorname{shape}(\bar X)=[1,2,3]$

Градиенты весов отдельных ветвей

Запрос: какую информацию следует учесть для этой позиции?	$\bar W_Q=[1.000000,0.000000,1.000000,2.000000,-3.000000,4.000000]$
Ключ: по каким признакам эта позиция сопоставляется с запросом?	$\bar W_K=[0.500000,-1.000000,2.000000,-2.000000,1.500000,1.000000]$
Значение: какое содержимое эта позиция передаёт дальше?	$\bar W_V=[2.000000,1.000000,4.000000,1.500000,-2.000000,-2.000000]$

Измените только веса запросов

Запрос: какую информацию следует учесть для этой позиции?: Изменено
Ключ: по каким признакам эта позиция сопоставляется с запросом?: Без изменений
Значение: какое содержимое эта позиция передаёт дальше?: Без изменений

Пустая ось пакета

$[0,2,3]\to [0,2,2],[0,2,2],[0,2,2]$

Пустая ось токенов

$[2,0,3]\to [2,0,2],[2,0,2],[2,0,2]$

Отклонённые граничные случаи

Отклонено rank-two Во входе должны быть явно представлены оси пакета, токенов и признаков.
Отклонено input-width Последняя ось входа должна совпадать с шириной модели.
Отклонено branch-mismatch У всех трёх матриц проекций должна быть одинаковая входная ширина модели.

Численные проверки

$n_X=6, n_Q=6, n_K=6, n_V=6$

$\tau=0.000002$

gradcheck=true replay=bitwise names=unique initialization=transactional

Пример проверяет каждую координату входа и весов и подтверждает воспроизводимость и независимость проекций.

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

Сначала сделайте предсказания

Предскажите формы всех трёх выходов для входа $[4,7,3]$ при $d_{head}=2$ .
Вычислите $q_0$ , $k_0$ и $v_0$ для первого токена из зафиксированного примера.
Определите, может ли изменение только $W_Q$ изменить $K$ или $V$ .
Подсчитайте число обучаемых скаляров в трёх матрицах весов формы $[3,2]$ без смещения.
Предскажите формы трёх выходов для входов $[0,2,3]$ и $[2,0,3]$ .
Объясните, почему эта обёртка отклоняет вход второго ранга $[T,d_{model}]$ .
Определите, должен ли здесь $d_{model}$ делиться на $d_{head}$ .
Укажите, какой исторический механизм получает данные из двух частей модели, а какой — из одной последовательности.
Назовите вычисление, которого ещё не хватает, чтобы эти тензоры образовали выход внимания.

Сверить ответы

Каждый выход имеет форму $[4,7,2]$ .
Получаются векторы $[0,3]$ , $[3,0]$ и $[3,-3]$ .
Нет. Ветви имеют независимые веса, поэтому при изменении только матрицы запросов битовые представления ключей и значений остаются прежними.
Число параметров равно $3\times3\times2=18$ .
Для каждой ветви выходы имеют формы $[0,2,2]$ и $[2,0,2]$ .
API сохраняет оси пакета и токенов явными, поэтому требует вход третьего ранга.
Нет. Ограничение делимости появится позднее при разделении на несколько голов.
Аддитивное внимание энкодера–декодера получает данные со стороны декодера и энкодера; самовнимание проецирует одну последовательность тремя способами.
В главе 27 предстоит вычислить оценки запрос–ключ, вероятности и взвешенную смесь значений.

Далее сопоставьте запросы с ключами и смешайте значения

Накопительный декодер теперь принимает нормализованную последовательность скрытых состояний формы $[B,T,d_{model}]$ и выдаёт отдельные тензоры $Q$ , $K$ и $V$ формы $[B,T,d_{head}]$ . В главе 27 эти подготовленные представления превратятся в одну голову внимания без маски: для этого будут вычислены оценки, вероятности и смесь значений.