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

Вычислите каузальное внимание в нескольких головах, затем смешайте их выходы

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

Сначала определите границы этапов, затем вычисляйте вероятности

Возьмём один батч с тремя строками токенов, ширину модели $d_{\mathrm{model}}=4$ и число голов $h=2$ . Ширина одной головы равна

d_h=\frac{d_{\mathrm{model}}}{h}=2.

В этой главе RoPE применяется ко всей голове, поэтому ширина $d_h$ должна быть ещё и чётной. Именно поэтому $h=4$ здесь недопустимо, хотя $h=4$ делит $d_{\mathrm{model}}=4$ без остатка: в каждой голове осталась бы одна координата без пары.

Зафиксируем вход

X= \begin{bmatrix} 1&0&1&0\\ \cos(1)&-\sin(1)&0&1\\ \cos(2)&-\sin(2)&1&1 \end{bmatrix}.

Только в этом примере $W_Q=W_K=W_V=I$ . Благодаря этому группы признаков после проекции легко проследить. Для многоголового внимания это не ограничение: в обученном слое столбцы, относящиеся к любой голове, могут использовать все входные признаки.

Прежде чем смотреть на вывод программы на Rust, предскажите последовательность форм:

\begin{aligned} X:[1,3,4] &\to(Q,K,V):[1,3,4] \to\operatorname{split}(Q,K,V):[1,2,3,2],\\ \operatorname{split}(Q,K,V):[1,2,3,2] &\to A:[1,2,3,3] \to H:[1,2,3,2],\\ H:[1,2,3,2] &\to\operatorname{Concat}(H):[1,3,4] \xrightarrow{W_O}\operatorname{MHA}(X):[1,3,4]. \end{aligned}

Первые две входные координаты намеренно компенсируют поворот, задаваемый позицией. В позиции токена $t$ они равны $R(-t)[1,0]$ , поэтому поворот из главы 29 возвращает

R(t)R(-t)[1,0]=[1,0].

Таким образом, голова $0$ сравнивает три одинаковые строки запросов и ключей. До запуска программы предскажите её каузальные вероятности:

A^{(0)}= \begin{bmatrix} 1&0&0\\ \tfrac12&\tfrac12&0\\ \tfrac13&\tfrac13&\tfrac13 \end{bmatrix}.

До поворота голова $1$ получает строки $[1,0]$ , $[0,1]$ и $[1,1]$ , поэтому её разрешённые оценки не равны друг другу. При этом в обеих головах вероятность выбора ключа выше диагонали точно равна нулю.

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

H^{(0)}_1 =\tfrac12[1,0]+\tfrac12[0.540302,-0.841471] =[0.770151,-0.420735].

Для головы $1$ две видимые масштабированные оценки и их softmax равны

[s^{(1)}_{1,0},s^{(1)}_{1,1}] =[-0.595010,0.707107] \xrightarrow{\operatorname{softmax}} [0.213809,0.786191].

Её строки значений до поворота равны $[1,0]$ и $[0,1]$ , поэтому $H^{(1)}_1=[0.213809,0.786191]$ . После конкатенации двух выходов получаем $[0.770151,-0.420735,0.213809,0.786191]$ .

Наконец, предскажите, на каком этапе признаки двух голов впервые могут повлиять на одну и ту же обучаемую выходную координату. Разделение и оба вычисления внимания сохраняют головы раздельными, а конкатенация только размещает их координаты рядом. Первая возможность обучаемого смешивания голов после их вычислений внимания появляется при умножении на $W_O$ . Зафиксированная $W_O$ меняет две группы местами и тем самым показывает обучаемую маршрутизацию; в упражнении 5 внеблочный элемент создаёт настоящую сумму признаков разных голов. При этом предшествующие плотные проекции запросов, ключей и значений уже могут использовать каждую координату общей входной строки.

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

Реализация хранит три упакованные матрицы без смещений:

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

где $Q$ , $K$ и $V$ имеют форму $[B,T,d_{\mathrm{model}}]$ . Изменение формы с последующей перестановкой осей даёт $[B,h,T,d_h]$ .

Упакованная запись — лишь удобная организация столбцов отдельных голов из исходной статьи. Для запросов

W_Q=[W_1^Q\;W_2^Q\;\cdots\;W_h^Q], \qquad Q_i=XW_i^Q,

и то же соотношение выполняется для ключей и значений. Порядок операций важен: если сначала разделить исходный $X$ , каждая голова будет ограничена фиксированной частью входа. Если сначала выполнить проекцию, то все $W_i^Q$ , $W_i^K$ и $W_i^V$ могут использовать всю входную строку, а уже их выходы попадают в отдельные ветви голов.

Для головы $i$ примените RoPE только к строкам запросов и ключей:

\widetilde Q_i=\operatorname{RoPE}(Q_i),\qquad \widetilde K_i=\operatorname{RoPE}(K_i).

Затем используйте операцию каузального внимания из главы 28:

A_i=\operatorname{softmax}_{j}\!\left( \frac{\widetilde Q_i\widetilde K_i^\top}{\sqrt{d_h}}+M \right), \qquad H_i=A_iV_i.

$M$ закрывает будущие ключи, а softmax выполняется отдельно по позициям ключей в каждой голове. Одно полноразмерное вычисление внимания дало бы одну нормализованную таблицу $[T,T]$ , общую для всех признаков значений. Многоголовое внимание вместо этого строит $h$ отдельно нормализованных таблиц $[T,T]$ . Они могут различаться из-за разных проекций запросов и ключей, но при некоторых входах и параметрах их численные значения могут и совпасть.

Каждый $A_i$ имеет форму $[B,T,T]$ , а каждый $H_i$ — форму $[B,T,d_h]$ . Восстановите ширину модели конкатенацией по последней оси признаков, затем примените одну обучаемую выходную матрицу:

\operatorname{MHA}(X)=\operatorname{Concat}(H_1,\ldots,H_h)W_O

где

W_O\in\mathbb{R}^{d_{\mathrm{model}}\times d_{\mathrm{model}}}.

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

N_\theta=4d_{\mathrm{model}}^2.

При $d_{\mathrm{model}}=4$ это $64$ скалярных параметра.

Не путайте оси модели, голов и токенов

$\operatorname{MHA}$ — полный слой многоголового каузального самовнимания.
$X$ — батч входных строк скрытых состояний.
$\operatorname{Concat}$ объединяет выходы голов вдоль оси признаков в каждой позиции токена; она не нормализует и не смешивает их.
$B$ — размер батча.
$T$ — число позиций токенов в каждой последовательности.
$d_{\mathrm{model}}$ — ширина входа и выхода полного слоя.
$h$ — число отдельно нормализуемых голов внимания.
$d_h=d_{\mathrm{model}}/h$ — ширина одной головы; здесь она чётная, потому что RoPE поворачивает соседние пары координат.
$t$ — позиция токена в примере с индексацией от нуля, а $R(t)$ — двумерный поворот на $t$ радиан, используемый первой головой примера.
$W_Q$ , $W_K$ и $W_V$ — упакованные матрицы проекций ширины модели. Столбцы, относящиеся к голове $i$ , образуют $W_i^Q$ , $W_i^K$ и $W_i^V$ .
$Q_i$ , $K_i$ и $V_i$ — спроецированные строки запросов, ключей и значений головы $i$ .
$\widetilde Q_i$ и $\widetilde K_i$ — эти же запросы и ключи после RoPE; значения не поворачиваются.
$M$ — каузальная маска, закрывающая ключи справа от позиции запроса.
$j$ — индекс позиции ключа и ось, по которой нормализуется каждая строка softmax.
$A_i$ — построчно нормализованная каузальная таблица внимания головы $i$ .
$H_i=A_iV_i$ — взвешенная смесь значений головы $i$ .
$W_O$ — обучаемая выходная матрица, применяемая после конкатенации всех строк $H_i$ .
$N_\theta$ — число обучаемых скалярных параметров. $\delta$ — приращение центральной разности, $\varepsilon_g$ — допустимая погрешность этой проверки, а $\varepsilon$ — допуск для проверок сборки и инвариантов.

В формулах используется принятое в статье соглашение $i\in\{1,\ldots,h\}$ . В программе на Rust и зафиксированных данных главы индексы голов начинаются с нуля: $0,\ldots,h-1$ . Поэтому головы $0$ и $1$ в примере соответствуют $H_1$ и $H_2$ в формулах.

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

От одного рекуррентного выравнивания к параллельным проекциям внимания

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

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

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

Статья о LLaMA приводит конкретный современный пример каузальной языковой модели. Туврон и соавторы описывают модель на основе Transformer с несколькими головами внимания, RoPE в каждом слое и оптимизированной реализацией каузального многоголового внимания. Эта реализация не хранит веса внимания. В приведённом ниже плотном вычислении веса остаются видимыми для изучения; оптимизированное ядро здесь не воспроизводится.

Исполняемое сравнение не пытается воспроизвести весь рекуррентный декодер. Оно для одного зафиксированного распределения выравнивания вычисляет взвешенную сумму аннотаций энкодера — контекст исходной последовательности, запускает существующую одноголовую каузальную операцию на том же входе из трёх токенов и сравнивает её единственный тензор вероятностей $[1,3,3]$ с тензором многоголового примера $[1,2,3,3]$ . В обоих случаях строки нормализованы, но второй тензор содержит две отдельно нормализованные таблицы. Сравнение не приписывает статьям выбранную в примере упаковку матриц, правило чётной ширины, отсутствие смещений, имена и значения:

Вычислить один взвешенный контекст исходной последовательности, одну полноразмерную каузальную таблицу и две таблицы отдельных голов rust/demos/ch30-multi-head-attention/src/lib.rs#historical-multi-head-contrast

#[derive(Clone, Debug, PartialEq)]
pub struct HistoryEvidence {
    pub earlier_weighted_context: [f64; 2],
    pub earlier_distributions_per_target: usize,
    pub single_head_weight_shape: Vec<usize>,
    pub multi_head_weight_shape: Vec<usize>,
    pub single_head_tables: usize,
    pub multi_head_tables: usize,
    pub all_rows_normalized: bool,
    pub mixing_stage: &'static str,
    pub modern_example: &'static str,
    pub weight_api: &'static str,
}

fn weighted_source_context(weights: [f64; 2], annotations: [[f64; 2]; 2]) -> [f64; 2] {
    [
        weights[0] * annotations[0][0] + weights[1] * annotations[1][0],
        weights[0] * annotations[0][1] + weights[1] * annotations[1][1],
    ]
}

fn rows_are_normalized(weights: &Tensor, tokens: usize) -> bool {
    weights
        .as_slice()
        .chunks_exact(tokens)
        .all(|row| (row.iter().sum::<f64>() - 1.0).abs() <= INVARIANT_TOLERANCE)
}

pub fn historical_attention_contrast(
    primary: &PrimaryEvidence,
) -> Result<HistoryEvidence, FixtureError> {
    let earlier_weighted_context = weighted_source_context([0.25, 0.75], [[1.0, 0.0], [0.0, 1.0]]);

    let full_width = constant(&[BATCH, TOKENS, MODEL_WIDTH], primary.input.as_slice())?;
    let single_head =
        causal_scaled_dot_product_self_attention(&full_width, &full_width, &full_width)?;
    let single_head_weights = single_head.weights().value();
    let single_head_weight_shape = single_head_weights.shape().to_vec();
    let multi_head_weight_shape = primary.attention_weights.shape().to_vec();
    let single_head_tables = single_head_weights.len() / (TOKENS * TOKENS);
    let multi_head_tables = primary.attention_weights.len() / (TOKENS * TOKENS);

    Ok(HistoryEvidence {
        earlier_weighted_context,
        earlier_distributions_per_target: 1,
        single_head_weight_shape,
        multi_head_weight_shape,
        single_head_tables,
        multi_head_tables,
        all_rows_normalized: rows_are_normalized(&single_head_weights, TOKENS)
            && rows_are_normalized(&primary.attention_weights, TOKENS),
        mixing_stage: "after-concatenation",
        modern_example: "llama-causal-heads-plus-rope",
        weight_api: "dense-teaching-evidence",
    })
}

Соберите на Rust полностью дифференцируемый и проверяемый путь вычислений

Открытая функция split_heads сначала меняет форму тензора, а затем переставляет оси:

[B,T,d_{\mathrm{model}}] \to[B,T,h,d_h] \to[B,h,T,d_h].

merge_heads выполняет обратную перестановку осей и возвращает исходную форму. Обе операции остаются на ленте автоматического дифференцирования; отдельного алгоритма ручного копирования нет.

Разделить тензор на головы и объединить их операциями изменения формы и перестановки осей на ленте rust/crates/llm-from-scratch/src/attention/multi_head.rs#head-layout

/// A taped reshape/transpose stage in the public split and merge helpers.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HeadLayoutStage {
    SplitReshape,
    SplitTranspose,
    MergeTranspose,
    MergeReshape,
}

impl fmt::Display for HeadLayoutStage {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::SplitReshape => "split reshape",
            Self::SplitTranspose => "split transpose",
            Self::MergeTranspose => "merge transpose",
            Self::MergeReshape => "merge reshape",
        })
    }
}

/// A rejected rank or feature partition in `split_heads` or `merge_heads`.
#[derive(Clone, Debug, PartialEq)]
pub enum HeadLayoutError {
    SplitRank {
        rank: usize,
    },
    MergeRank {
        rank: usize,
    },
    ZeroHeadCount,
    ZeroHeadWidth,
    WidthNotDivisible {
        width: usize,
        heads: usize,
    },
    ModelWidthOverflow {
        heads: usize,
        head_width: usize,
    },
    Autodiff {
        stage: HeadLayoutStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for HeadLayoutError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::SplitRank { rank } => write!(
                formatter,
                "head split input must have rank three [batch, tokens, model_width], got rank {rank}"
            ),
            Self::MergeRank { rank } => write!(
                formatter,
                "head merge input must have rank four [batch, heads, tokens, head_width], got rank {rank}"
            ),
            Self::ZeroHeadCount => formatter.write_str("head count must be nonzero"),
            Self::ZeroHeadWidth => formatter.write_str("head width must be nonzero"),
            Self::WidthNotDivisible { width, heads } => write!(
                formatter,
                "model width {width} must be divisible by head count {heads}"
            ),
            Self::ModelWidthOverflow { heads, head_width } => write!(
                formatter,
                "merged model width overflows for {heads} heads of width {head_width}"
            ),
            Self::Autodiff { stage, source } => {
                write!(formatter, "multi-head {stage}: {source}")
            }
        }
    }
}

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

fn layout_autodiff(stage: HeadLayoutStage) -> impl FnOnce(TensorAutodiffError) -> HeadLayoutError {
    move |source| HeadLayoutError::Autodiff { stage, source }
}

/// Converts `[batch, tokens, model_width]` to `[batch, heads, tokens, head_width]`.
pub fn split_heads(input: &TensorValue, heads: usize) -> Result<TensorValue, HeadLayoutError> {
    let shape = input.shape();
    if shape.len() != 3 {
        return Err(HeadLayoutError::SplitRank { rank: shape.len() });
    }
    if heads == 0 {
        return Err(HeadLayoutError::ZeroHeadCount);
    }
    let width = shape[2];
    if width == 0 {
        return Err(HeadLayoutError::ZeroHeadWidth);
    }
    if !width.is_multiple_of(heads) {
        return Err(HeadLayoutError::WidthNotDivisible { width, heads });
    }
    let head_width = width / heads;
    let reshaped = input
        .reshape(&[shape[0], shape[1], heads, head_width])
        .map_err(layout_autodiff(HeadLayoutStage::SplitReshape))?;
    reshaped
        .transpose(1, 2)
        .map_err(layout_autodiff(HeadLayoutStage::SplitTranspose))
}

/// Converts `[batch, heads, tokens, head_width]` back to model-width rows.
pub fn merge_heads(input: &TensorValue) -> Result<TensorValue, HeadLayoutError> {
    let shape = input.shape();
    if shape.len() != 4 {
        return Err(HeadLayoutError::MergeRank { rank: shape.len() });
    }
    if shape[1] == 0 {
        return Err(HeadLayoutError::ZeroHeadCount);
    }
    if shape[3] == 0 {
        return Err(HeadLayoutError::ZeroHeadWidth);
    }
    let model_width =
        shape[1]
            .checked_mul(shape[3])
            .ok_or(HeadLayoutError::ModelWidthOverflow {
                heads: shape[1],
                head_width: shape[3],
            })?;
    let transposed = input
        .transpose(1, 2)
        .map_err(layout_autodiff(HeadLayoutStage::MergeTranspose))?;
    transposed
        .reshape(&[shape[0], shape[2], model_width])
        .map_err(layout_autodiff(HeadLayoutStage::MergeReshape))
}

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

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

Сохранить типизированные ошибки конфигурации, входа, позиций и вложенных этапов rust/crates/llm-from-scratch/src/attention/multi_head.rs#multi-head-errors

/// A Q/K/V branch whose head-layout or rotary operation failed.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MultiHeadInput {
    Query,
    Key,
    Value,
}

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

/// A cumulative tensor stage owned by the multi-head assembly.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MultiHeadStage {
    QueryLanes,
    KeyLanes,
    ValueLanes,
    RestoreWeights,
    RestoreHeadOutputs,
}

impl fmt::Display for MultiHeadStage {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::QueryLanes => "query lane flatten",
            Self::KeyLanes => "key lane flatten",
            Self::ValueLanes => "value lane flatten",
            Self::RestoreWeights => "attention-weight head restore",
            Self::RestoreHeadOutputs => "head-output restore",
        })
    }
}

/// A rejected configuration, parameter set, input, or cumulative forward stage.
#[derive(Clone, Debug, PartialEq)]
pub enum MultiHeadAttentionError {
    ZeroModelWidth,
    ZeroHeadCount,
    ModelWidthNotDivisible {
        model_width: usize,
        heads: usize,
    },
    OddHeadWidth {
        head_width: usize,
    },
    QkvProjection(QkvError),
    QkvOutputWidthMismatch {
        model_width: usize,
        projected_width: usize,
    },
    OutputProjection(LinearError),
    OutputInputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    OutputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    Initialization(InitializationError),
    InputRank {
        rank: usize,
    },
    InputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    EmptyTokens,
    PositionOffsetOverflow {
        offset: usize,
        tokens: usize,
    },
    PositionRangeExceeded {
        offset: usize,
        tokens: usize,
        max_positions: usize,
    },
    BatchHeadOverflow {
        batch: usize,
        heads: usize,
    },
    HeadLayout {
        input: MultiHeadInput,
        source: HeadLayoutError,
    },
    MergeLayout(HeadLayoutError),
    RotaryConfiguration(RopeError),
    Rotary {
        input: MultiHeadInput,
        source: RopeError,
    },
    CausalAttention(CausalMaskingError),
    Autodiff {
        stage: MultiHeadStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for MultiHeadAttentionError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::ZeroModelWidth => formatter.write_str("multi-head model width must be nonzero"),
            Self::ZeroHeadCount => formatter.write_str("multi-head head count must be nonzero"),
            Self::ModelWidthNotDivisible { model_width, heads } => write!(
                formatter,
                "multi-head model width {model_width} must be divisible by head count {heads}"
            ),
            Self::OddHeadWidth { head_width } => write!(
                formatter,
                "multi-head per-head width must be even for RoPE, got {head_width}"
            ),
            Self::QkvProjection(source) => source.fmt(formatter),
            Self::QkvOutputWidthMismatch {
                model_width,
                projected_width,
            } => write!(
                formatter,
                "Q/K/V projected width must equal model width {model_width}, got {projected_width}"
            ),
            Self::OutputProjection(source) => write!(formatter, "output projection: {source}"),
            Self::OutputInputWidthMismatch { expected, actual } => write!(
                formatter,
                "output projection input width must equal model width {expected}, got {actual}"
            ),
            Self::OutputWidthMismatch { expected, actual } => write!(
                formatter,
                "output projection width must equal model width {expected}, got {actual}"
            ),
            Self::Initialization(source) => source.fmt(formatter),
            Self::InputRank { rank } => write!(
                formatter,
                "multi-head input must have rank three [batch, tokens, model_width], got rank {rank}"
            ),
            Self::InputWidthMismatch { expected, actual } => write!(
                formatter,
                "multi-head input final width must equal model width {expected}, got {actual}"
            ),
            Self::EmptyTokens => formatter.write_str(
                "multi-head attention needs at least one token so every causal row has a key",
            ),
            Self::PositionOffsetOverflow { offset, tokens } => write!(
                formatter,
                "multi-head position interval overflows: offset {offset} plus {tokens} tokens"
            ),
            Self::PositionRangeExceeded {
                offset,
                tokens,
                max_positions,
            } => write!(
                formatter,
                "multi-head position interval [{offset}, {}) exceeds capacity {max_positions}",
                offset.saturating_add(*tokens)
            ),
            Self::BatchHeadOverflow { batch, heads } => write!(
                formatter,
                "multi-head lane count overflows for batch {batch} and {heads} heads"
            ),
            Self::HeadLayout { input, source } => {
                write!(formatter, "{input} head layout: {source}")
            }
            Self::MergeLayout(source) => write!(formatter, "head output merge: {source}"),
            Self::RotaryConfiguration(source) => {
                write!(formatter, "multi-head RoPE configuration: {source}")
            }
            Self::Rotary { input, source } => write!(formatter, "{input} RoPE: {source}"),
            Self::CausalAttention(source) => source.fmt(formatter),
            Self::Autodiff { stage, source } => {
                write!(formatter, "multi-head {stage}: {source}")
            }
        }
    }
}

impl Error for MultiHeadAttentionError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::QkvProjection(source) => Some(source),
            Self::OutputProjection(source) => Some(source),
            Self::Initialization(source) => Some(source),
            Self::HeadLayout { source, .. } => Some(source),
            Self::MergeLayout(source) => Some(source),
            Self::RotaryConfiguration(source) => Some(source),
            Self::Rotary { source, .. } => Some(source),
            Self::CausalAttention(source) => Some(source),
            Self::Autodiff { source, .. } => Some(source),
            _ => None,
        }
    }
}

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

fn multi_autodiff(
    stage: MultiHeadStage,
) -> impl FnOnce(TensorAutodiffError) -> MultiHeadAttentionError {
    move |source| MultiHeadAttentionError::Autodiff { stage, source }
}

fn head_layout(input: MultiHeadInput) -> impl FnOnce(HeadLayoutError) -> MultiHeadAttentionError {
    move |source| MultiHeadAttentionError::HeadLayout { input, source }
}

fn rotary_error(input: MultiHeadInput) -> impl FnOnce(RopeError) -> MultiHeadAttentionError {
    move |source| MultiHeadAttentionError::Rotary { input, source }
}

fn validate_configuration(
    model_width: usize,
    heads: usize,
) -> Result<usize, MultiHeadAttentionError> {
    if model_width == 0 {
        return Err(MultiHeadAttentionError::ZeroModelWidth);
    }
    if heads == 0 {
        return Err(MultiHeadAttentionError::ZeroHeadCount);
    }
    if !model_width.is_multiple_of(heads) {
        return Err(MultiHeadAttentionError::ModelWidthNotDivisible { model_width, heads });
    }
    let head_width = model_width / heads;
    if !head_width.is_multiple_of(2) {
        return Err(MultiHeadAttentionError::OddHeadWidth { head_width });
    }
    Ok(head_width)
}

Прямой проход выполняет проекции до полной ширины модели, разделяет все три тензора на головы, поворачивает только $Q$ и $K$ , объединяет оси батча и голов в независимые ветви, повторно использует проверенную одноголовую каузальную операцию, восстанавливает ось голов, объединяет признаки и применяет выходную проекцию. Возвращаемые промежуточные значения показывают границы всех этапов, а итоговый выход сохраняет форму $[B,T,d_{\mathrm{model}}]$ .

Собрать полноразмерные проекции, RoPE и каузальное внимание по головам, объединение и выходную проекцию rust/crates/llm-from-scratch/src/attention/multi_head.rs#multi-head-layer

/// Inspectable intermediate tensors from one complete multi-head forward pass.
#[derive(Clone, Debug)]
pub struct MultiHeadAttentionForward {
    projected_query_heads: TensorValue,
    projected_key_heads: TensorValue,
    projected_value_heads: TensorValue,
    rotated_query_heads: TensorValue,
    rotated_key_heads: TensorValue,
    attention_weights: TensorValue,
    head_outputs: TensorValue,
    merged: TensorValue,
    output: TensorValue,
}

impl MultiHeadAttentionForward {
    pub fn projected_query_heads(&self) -> &TensorValue {
        &self.projected_query_heads
    }

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

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

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

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

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

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

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

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

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

/// Full-width Q/K/V projections, per-head RoPE and causal attention, and W_O.
#[derive(Clone, Debug)]
pub struct MultiHeadAttention {
    qkv: QkvProjections,
    output: Linear,
    rope: RotaryEmbedding,
    parameters: NamedParameters,
    model_width: usize,
    heads: usize,
    head_width: usize,
}

impl MultiHeadAttention {
    /// Initializes Q, K, V, and O in order without partially advancing `rng`.
    pub fn new(
        parameter_prefix: impl Into<String>,
        model_width: usize,
        heads: usize,
        max_positions: usize,
        rope_base: f64,
        rng: &mut SplitMix64,
    ) -> Result<Self, MultiHeadAttentionError> {
        let head_width = validate_configuration(model_width, heads)?;
        let rope = RotaryEmbedding::new(head_width, max_positions, rope_base)
            .map_err(MultiHeadAttentionError::RotaryConfiguration)?;
        let parameter_prefix = parameter_prefix.into();
        let mut trial = rng.clone();
        let qkv = QkvProjections::new(&parameter_prefix, model_width, model_width, &mut trial)
            .map_err(MultiHeadAttentionError::QkvProjection)?;
        let output = Linear::new(
            format!("{parameter_prefix}.output"),
            model_width,
            model_width,
            false,
            &mut trial,
        )
        .map_err(MultiHeadAttentionError::OutputProjection)?;
        let layer = Self::from_validated_parts(qkv, output, rope, heads, head_width)?;
        *rng = trial;
        Ok(layer)
    }

    /// Builds an exact deterministic layer from four named full-width matrices.
    pub fn from_parameters(
        query_weight: NamedParameter,
        key_weight: NamedParameter,
        value_weight: NamedParameter,
        output_weight: NamedParameter,
        heads: usize,
        max_positions: usize,
        rope_base: f64,
    ) -> Result<Self, MultiHeadAttentionError> {
        let qkv = QkvProjections::from_weights(query_weight, key_weight, value_weight)
            .map_err(MultiHeadAttentionError::QkvProjection)?;
        let output = Linear::from_parameters(output_weight, None)
            .map_err(MultiHeadAttentionError::OutputProjection)?;
        let model_width = qkv.model_width();
        let head_width = validate_configuration(model_width, heads)?;
        if qkv.head_width() != model_width {
            return Err(MultiHeadAttentionError::QkvOutputWidthMismatch {
                model_width,
                projected_width: qkv.head_width(),
            });
        }
        if output.input_width() != model_width {
            return Err(MultiHeadAttentionError::OutputInputWidthMismatch {
                expected: model_width,
                actual: output.input_width(),
            });
        }
        if output.output_width() != model_width {
            return Err(MultiHeadAttentionError::OutputWidthMismatch {
                expected: model_width,
                actual: output.output_width(),
            });
        }
        let rope = RotaryEmbedding::new(head_width, max_positions, rope_base)
            .map_err(MultiHeadAttentionError::RotaryConfiguration)?;
        Self::from_validated_parts(qkv, output, rope, heads, head_width)
    }

    fn from_validated_parts(
        qkv: QkvProjections,
        output: Linear,
        rope: RotaryEmbedding,
        heads: usize,
        head_width: usize,
    ) -> Result<Self, MultiHeadAttentionError> {
        let model_width = qkv.model_width();
        let mut listed = Vec::with_capacity(4);
        listed.extend(qkv.parameters().iter().cloned());
        listed.push(output.weight().clone());
        let parameters = NamedParameters::try_new(listed)?;
        Ok(Self {
            qkv,
            output,
            rope,
            parameters,
            model_width,
            heads,
            head_width,
        })
    }

    /// Runs position-aware causal attention independently in every head.
    pub fn forward(
        &self,
        input: &TensorValue,
        position_offset: usize,
    ) -> Result<MultiHeadAttentionForward, MultiHeadAttentionError> {
        let shape = input.shape();
        if shape.len() != 3 {
            return Err(MultiHeadAttentionError::InputRank { rank: shape.len() });
        }
        if shape[2] != self.model_width {
            return Err(MultiHeadAttentionError::InputWidthMismatch {
                expected: self.model_width,
                actual: shape[2],
            });
        }
        if shape[1] == 0 {
            return Err(MultiHeadAttentionError::EmptyTokens);
        }
        let position_end = position_offset.checked_add(shape[1]).ok_or(
            MultiHeadAttentionError::PositionOffsetOverflow {
                offset: position_offset,
                tokens: shape[1],
            },
        )?;
        if position_end > self.rope.max_positions() {
            return Err(MultiHeadAttentionError::PositionRangeExceeded {
                offset: position_offset,
                tokens: shape[1],
                max_positions: self.rope.max_positions(),
            });
        }

        let projected = self
            .qkv
            .forward(input)
            .map_err(MultiHeadAttentionError::QkvProjection)?;
        let projected_query_heads = split_heads(projected.query(), self.heads)
            .map_err(head_layout(MultiHeadInput::Query))?;
        let projected_key_heads =
            split_heads(projected.key(), self.heads).map_err(head_layout(MultiHeadInput::Key))?;
        let projected_value_heads = split_heads(projected.value(), self.heads)
            .map_err(head_layout(MultiHeadInput::Value))?;
        let rotated_query_heads = self
            .rope
            .rotate(&projected_query_heads, position_offset)
            .map_err(rotary_error(MultiHeadInput::Query))?;
        let rotated_key_heads = self
            .rope
            .rotate(&projected_key_heads, position_offset)
            .map_err(rotary_error(MultiHeadInput::Key))?;

        let lanes =
            shape[0]
                .checked_mul(self.heads)
                .ok_or(MultiHeadAttentionError::BatchHeadOverflow {
                    batch: shape[0],
                    heads: self.heads,
                })?;
        let lane_shape = [lanes, shape[1], self.head_width];
        let query_lanes = rotated_query_heads
            .reshape(&lane_shape)
            .map_err(multi_autodiff(MultiHeadStage::QueryLanes))?;
        let key_lanes = rotated_key_heads
            .reshape(&lane_shape)
            .map_err(multi_autodiff(MultiHeadStage::KeyLanes))?;
        let value_lanes = projected_value_heads
            .reshape(&lane_shape)
            .map_err(multi_autodiff(MultiHeadStage::ValueLanes))?;
        let attended =
            causal_scaled_dot_product_self_attention(&query_lanes, &key_lanes, &value_lanes)
                .map_err(MultiHeadAttentionError::CausalAttention)?;
        let attention_weights = attended
            .weights()
            .reshape(&[shape[0], self.heads, shape[1], shape[1]])
            .map_err(multi_autodiff(MultiHeadStage::RestoreWeights))?;
        let head_outputs = attended
            .output()
            .reshape(&[shape[0], self.heads, shape[1], self.head_width])
            .map_err(multi_autodiff(MultiHeadStage::RestoreHeadOutputs))?;
        let merged = merge_heads(&head_outputs).map_err(MultiHeadAttentionError::MergeLayout)?;
        let output = self
            .output
            .forward(&merged)
            .map_err(MultiHeadAttentionError::OutputProjection)?;

        Ok(MultiHeadAttentionForward {
            projected_query_heads,
            projected_key_heads,
            projected_value_heads,
            rotated_query_heads,
            rotated_key_heads,
            attention_weights,
            head_outputs,
            merged,
            output,
        })
    }

    pub fn qkv(&self) -> &QkvProjections {
        &self.qkv
    }

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

    pub fn rope(&self) -> &RotaryEmbedding {
        &self.rope
    }

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

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

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

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

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

Неоднородная начальная сопряжённая величина на выходе даёт конечные тензоры градиентов для входа и всех четырёх матриц. Центральные разности проверяют все $12+4\cdot16=76$ координат при

\delta=10^{-6},\qquad \varepsilon_g=8\times10^{-6}.

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

Пример собирает эти значения и проверки в один воспроизводимый отчёт:

Собрать точные значения примера, инварианты, градиенты, граничные случаи и переход к следующей главе rust/demos/ch30-multi-head-attention/src/lib.rs#learner-report

pub fn render_report(evidence: &LearnerEvidence) -> String {
    let primary = &evidence.primary;
    let parameters = &evidence.parameters;
    let shapes = &evidence.shapes;
    let errors = &evidence.errors;
    let gradients = &evidence.gradients;
    let history = &evidence.history;
    [
        "chapter=30-multi-head-attention".to_owned(),
        "prediction=projection creates two learned feature lanes; each lane normalizes its own causal rows; W_O first learns how to mix the concatenated results".to_owned(),
        format!(
            "config=batch:{BATCH} tokens:{TOKENS} d_model:{MODEL_WIDTH} heads:{HEADS} d_h:{HEAD_WIDTH} offset:0 capacity:{MAX_POSITIONS} rope_base:{ROPE_BASE:.6} bias:false"
        ),
        format!(
            "input=shape:{} values:{}",
            format_shape(primary.input.shape()),
            format_vector(primary.input.as_slice())
        ),
        format!(
            "projected_query_heads=shape:{} values:{}",
            format_shape(primary.projected_query_heads.shape()),
            format_vector(primary.projected_query_heads.as_slice())
        ),
        format!(
            "projected_key_heads=shape:{} values:{}",
            format_shape(primary.projected_key_heads.shape()),
            format_vector(primary.projected_key_heads.as_slice())
        ),
        format!(
            "projected_value_heads=shape:{} values:{}",
            format_shape(primary.projected_value_heads.shape()),
            format_vector(primary.projected_value_heads.as_slice())
        ),
        format!(
            "rotated_query_heads=shape:{} values:{}",
            format_shape(primary.rotated_query_heads.shape()),
            format_vector(primary.rotated_query_heads.as_slice())
        ),
        format!(
            "rotated_key_heads=shape:{} values:{}",
            format_shape(primary.rotated_key_heads.shape()),
            format_vector(primary.rotated_key_heads.as_slice())
        ),
        format!(
            "attention_weights=shape:{} values:{}",
            format_shape(primary.attention_weights.shape()),
            format_vector(primary.attention_weights.as_slice())
        ),
        format!(
            "head_outputs=shape:{} values:{}",
            format_shape(primary.head_outputs.shape()),
            format_vector(primary.head_outputs.as_slice())
        ),
        format!(
            "merged=shape:{} values:{}",
            format_shape(primary.merged.shape()),
            format_vector(primary.merged.as_slice())
        ),
        format!(
            "output_weight=shape:{} values:{}",
            format_shape(primary.output_weight.shape()),
            format_vector(primary.output_weight.as_slice())
        ),
        format!(
            "output=shape:{} values:{}",
            format_shape(primary.output.shape()),
            format_vector(primary.output.as_slice())
        ),
        format!(
            "heads=head_0_uniform:{} head_1_distinct:{} future_probabilities_zero:{}",
            primary.uniform_head_zero,
            primary.distinct_head_weights,
            primary.future_probabilities_zero
        ),
        format!(
            "prefix_perturbed_output={} position_0_unchanged:{} position_1_unchanged:{} position_2_changed:{}",
            format_vector(primary.prefix_perturbed_output.as_slice()),
            primary.prefix_zero_unchanged,
            primary.prefix_one_unchanged,
            primary.suffix_changed
        ),
        format!(
            "layout=split_merge_bitwise:{} head_isolation_before_output:{} common_offset_weights_preserved:{} tolerance:{INVARIANT_TOLERANCE:.12}",
            primary.split_merge_bitwise,
            primary.head_isolation_before_output,
            primary.common_offset_weights_preserved
        ),
        format!(
            "parameters=names:{} shapes:{} count:{} bias_free:{} node_distinct:{}",
            format_names(&parameters.names),
            format_shapes(&parameters.shapes),
            parameters.count,
            parameters.bias_free,
            parameters.node_distinct
        ),
        format!(
            "upstream={} loss:{:.6}",
            format_vector(primary.upstream.as_slice()),
            canonical(primary.loss)
        ),
        format!("input_gradient={}", format_vector(primary.input_gradient.as_slice())),
        format!(
            "query_weight_gradient={}",
            format_vector(primary.parameter_gradients[0].as_slice())
        ),
        format!(
            "key_weight_gradient={}",
            format_vector(primary.parameter_gradients[1].as_slice())
        ),
        format!(
            "value_weight_gradient={}",
            format_vector(primary.parameter_gradients[2].as_slice())
        ),
        format!(
            "output_weight_gradient={}",
            format_vector(primary.parameter_gradients[3].as_slice())
        ),
        format!(
            "gradcheck=input:{} query:{} key:{} value:{} output:{} total:{} tolerance:{GRADIENT_TOLERANCE:.6} passed:{}",
            gradients.input_checks,
            gradients.query_checks,
            gradients.key_checks,
            gradients.value_checks,
            gradients.output_checks,
            gradients.input_checks
                + gradients.query_checks
                + gradients.key_checks
                + gradients.value_checks
                + gradients.output_checks,
            gradients.passed
        ),
        format!(
            "shapes=input:{} split:{} rotated:{} weights:{} head_output:{} merged:{} output_weight:{} output:{} empty_batch_weights:{} empty_batch_output:{} single_token_weights:{}",
            format_shape(&shapes.input),
            format_shape(&shapes.split),
            format_shape(&shapes.rotated),
            format_shape(&shapes.weights),
            format_shape(&shapes.head_output),
            format_shape(&shapes.merged),
            format_shape(&shapes.output_weight),
            format_shape(&shapes.output),
            format_shape(&shapes.empty_batch_weights),
            format_shape(&shapes.empty_batch_output),
            format_shape(&shapes.single_token_weights)
        ),
        format!(
            "errors=zero_model_width:{} zero_heads:{} nondivisible:{} odd_head_width:{} input_rank:{} input_width:{} empty_tokens:{} offset_overflow:{} position_range:{} released_input:{}",
            errors.zero_model_width_rejected,
            errors.zero_heads_rejected,
            errors.nondivisible_rejected,
            errors.odd_head_width_rejected,
            errors.input_rank_rejected,
            errors.input_width_rejected,
            errors.empty_tokens_rejected,
            errors.offset_overflow_rejected,
            errors.position_range_rejected,
            errors.released_input_rejected
        ),
        format!(
            "history=earlier_weighted_context:{} earlier_distributions_per_target:{} single_head_shape:{} multi_head_shape:{} single_head_tables:{} multi_head_tables:{} rows_normalized:{} mixing:{} modern_example:{} weight_api:{}",
            format_vector(&history.earlier_weighted_context),
            history.earlier_distributions_per_target,
            format_shape(&history.single_head_weight_shape),
            format_shape(&history.multi_head_weight_shape),
            history.single_head_tables,
            history.multi_head_tables,
            history.all_rows_normalized,
            history.mixing_stage,
            history.modern_example,
            history.weight_api
        ),
        format!(
            "proof=tape_finite:{} replay:{} heads_distinct:{} causal:{} split_merge:{} gradients:{}",
            primary.tape_finite,
            if evidence.replay_bitwise { "bitwise" } else { "mismatch" },
            primary.distinct_head_weights,
            primary.future_probabilities_zero,
            primary.split_merge_bitwise,
            gradients.passed
        ),
        "next=wrap this attention transformation in the first pre-normalized residual path".to_owned(),
    ]
    .join("\n")
        + "\n"
}

Исполняемый файл печатает этот отчёт:

Запустить полный пример многоголового внимания rust/demos/ch30-multi-head-attention/src/main.rs

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

Выполните cargo run --quiet --locked -p ch30-multi-head-attention. Стандартный вывод побайтно совпадёт с rust/demos/ch30-multi-head-attention/expected.txt, включая завершающий перевод строки.

Проследите, где отдельные головы объединяются в строку ширины модели

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

Получить значения многоголового внимания, используемые в схеме rust/demos/ch30-multi-head-attention/src/diagram_trace.rs#multi-head-trace

pub fn render_trace(evidence: &LearnerEvidence) -> String {
    let primary = &evidence.primary;
    let shapes = &evidence.shapes;
    let parameters = &evidence.parameters;
    let gradients = &evidence.gradients;
    let gradient_checks = gradients.input_checks
        + gradients.query_checks
        + gradients.key_checks
        + gradients.value_checks
        + gradients.output_checks;
    debug_assert_eq!(gradient_checks, 76);
    let mut lines = vec![
        String::from(
            "CONFIG|batch=1|tokens=3|model_width=4|heads=2|head_width=2|offset=0|max_positions=6|rope_base=100.000000|bias=false|parameter_order=[query.weight,key.weight,value.weight,output.weight]|layout=reshape-transpose",
        ),
        format!("SHAPE|stage=input|value={}", format_shape(&shapes.input)),
        format!("SHAPE|stage=split|value={}", format_shape(&shapes.split)),
        format!(
            "SHAPE|stage=rotated|value={}",
            format_shape(&shapes.rotated)
        ),
        format!(
            "SHAPE|stage=weights|value={}",
            format_shape(&shapes.weights)
        ),
        format!(
            "SHAPE|stage=head-output|value={}",
            format_shape(&shapes.head_output)
        ),
        format!("SHAPE|stage=merged|value={}", format_shape(&shapes.merged)),
        format!(
            "SHAPE|stage=output-weight|value={}",
            format_shape(&shapes.output_weight)
        ),
        format!("SHAPE|stage=output|value={}", format_shape(&shapes.output)),
    ];
    lines.extend((0..HEADS).map(|head| partition_record(evidence, head)));
    lines.extend(
        (0..HEADS)
            .flat_map(|head| (0..TOKENS).map(move |query| (head, query)))
            .map(|(head, query)| weight_record(evidence, head, query)),
    );
    lines.extend(
        (0..HEADS)
            .flat_map(|head| (0..TOKENS).map(move |token| (head, token)))
            .map(|(head, token)| head_output_record(evidence, head, token)),
    );
    lines.extend((0..TOKENS).map(|token| merged_record(evidence, token)));
    lines.extend((0..MODEL_WIDTH).map(|row| output_map_record(evidence, row)));
    lines.extend((0..TOKENS).map(|token| output_record(evidence, token)));
    lines.push(format!(
        "PREFIX_PROOF|position_0={}|position_1={}|position_2={}|split_merge={}|head_isolation={}|future_probabilities={}|common_offset={}|tolerance={INVARIANT_TOLERANCE:.12}|parameters={}|gradchecks={}",
        if primary.prefix_zero_unchanged {
            "bitwise-unchanged"
        } else {
            "changed"
        },
        if primary.prefix_one_unchanged {
            "bitwise-unchanged"
        } else {
            "changed"
        },
        if primary.suffix_changed {
            "changed"
        } else {
            "unchanged"
        },
        if primary.split_merge_bitwise {
            "bitwise"
        } else {
            "mismatch"
        },
        if primary.head_isolation_before_output {
            "before-output"
        } else {
            "failed"
        },
        if primary.future_probabilities_zero {
            "exact-zero"
        } else {
            "nonzero"
        },
        if primary.common_offset_weights_preserved {
            "preserved"
        } else {
            "changed"
        },
        parameters.count,
        gradient_checks,
    ));
    debug_assert_eq!(lines.len(), 34);
    lines.join("\n") + "\n"
}

Две головы внимания остаются раздельными до выходной проекции

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

Сплошная рамка обозначает голову 0
Штриховая рамка обозначает голову 1
Доступный ключ: сплошное подчёркивание
Закрытый ключ: штриховое подчёркивание
Разрешённая диагональ: двойная рамка

Спроецируйте признаки и разделите их по головам

Общие проекции могут использовать весь вход; единичные веса лишь показывают две группы выхода.

$[B,T,d_{\mathrm{model}}]\to[B,h,T,d_h]$

Вход: $[1,3,4]$
После разделения на головы: $[1,2,3,2]$
После RoPE: $[1,2,3,2]$
Вероятности внимания: $[1,2,3,3]$
Выходы голов: $[1,2,3,2]$
После конкатенации: $[1,3,4]$
Матрица выходной проекции: $[4,4]$
Выход слоя: $[1,3,4]$

Общие проекции могут использовать весь вход; единичные веса лишь показывают две группы выхода.
Голова и координаты спроецированных признаков	Голова 0 $0,1$	Голова 1 $2,3$
После проекции $Q$	$[1.000000,0.000000],$ $[0.540302,-0.841471],$ $[-0.416147,-0.909297]$	$[1.000000,0.000000],$ $[0.000000,1.000000],$ $[1.000000,1.000000]$
После проекции $K$	$[1.000000,0.000000],$ $[0.540302,-0.841471],$ $[-0.416147,-0.909297]$	$[1.000000,0.000000],$ $[0.000000,1.000000],$ $[1.000000,1.000000]$
После проекции $V$	$[1.000000,0.000000],$ $[0.540302,-0.841471],$ $[-0.416147,-0.909297]$	$[1.000000,0.000000],$ $[0.000000,1.000000],$ $[1.000000,1.000000]$
После RoPE $Q$	$[1.000000,0.000000],$ $[1.000000,0.000000],$ $[1.000000,0.000000]$	$[1.000000,0.000000],$ $[-0.841471,0.540302],$ $[-1.325444,0.493151]$
После RoPE $K$	$[1.000000,0.000000],$ $[1.000000,0.000000],$ $[1.000000,0.000000]$	$[1.000000,0.000000],$ $[-0.841471,0.540302],$ $[-1.325444,0.493151]$

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

У каждой головы своя матрица оценок, свой softmax по строкам и своя смесь значений, но граница каузальной видимости у них одинакова.

$A_i=\operatorname{softmax}_{j}\!\left(\widetilde Q_i\widetilde K_i^\top/\sqrt{d_h}+M\right)$ $H_i=A_iV_i$

Голова 0

Каузальные вероятности внимания: Голова 0
Позиция запроса	$k=0$	$k=1$	$k=2$	Сумма строки	Смеси значений по головам: Выход головы
$q=0$	$1.000000$ Доступный ключ · Разрешённая диагональ	$0.000000$ Закрытый ключ	$0.000000$ Закрытый ключ	$1.000000$	$[1.000000,0.000000]$
$q=1$	$0.500000$ Доступный ключ	$0.500000$ Доступный ключ · Разрешённая диагональ	$0.000000$ Закрытый ключ	$1.000000$	$[0.770151,-0.420735]$
$q=2$	$0.333333$ Доступный ключ	$0.333333$ Доступный ключ	$0.333333$ Доступный ключ · Разрешённая диагональ	$1.000000$	$[0.374718,-0.583589]$

Голова 1

Каузальные вероятности внимания: Голова 1
Позиция запроса	$k=0$	$k=1$	$k=2$	Сумма строки	Смеси значений по головам: Выход головы
$q=0$	$1.000000$ Доступный ключ · Разрешённая диагональ	$0.000000$ Закрытый ключ	$0.000000$ Закрытый ключ	$1.000000$	$[1.000000,0.000000]$
$q=1$	$0.213809$ Доступный ключ	$0.786191$ Доступный ключ · Разрешённая диагональ	$0.000000$ Закрытый ключ	$1.000000$	$[0.213809,0.786191]$
$q=2$	$0.054696$ Доступный ключ	$0.370956$ Доступный ключ	$0.574348$ Доступный ключ · Разрешённая диагональ	$1.000000$	$[0.629044,0.945304]$

Конкатенируйте выходы голов и примените обучаемую выходную проекцию

Конкатенация сохраняет обе группы признаков; выбранная выходная матрица затем меняет их местами, поэтому эти два шага нельзя спутать.

$\operatorname{MHA}(X)=\operatorname{Concat}(H_1,\ldots,H_h)W_O$

Конкатенированные выходы голов
Позиция токена	Голова 0	Голова 1	Конкатенированная строка
$t=0$	$[1.000000,0.000000]$	$[1.000000,0.000000]$	$[1.000000,0.000000,1.000000,0.000000]$
$t=1$	$[0.770151,-0.420735]$	$[0.213809,0.786191]$	$[0.770151,-0.420735,0.213809,0.786191]$
$t=2$	$[0.374718,-0.583589]$	$[0.629044,0.945304]$	$[0.374718,-0.583589,0.629044,0.945304]$

Строка матрицы выходной проекции
Строка матрицы выходной проекции	$W_O$
$r=0$	$[0.000000,0.000000,1.000000,0.000000]$
$r=1$	$[0.000000,0.000000,0.000000,1.000000]$
$r=2$	$[1.000000,0.000000,0.000000,0.000000]$
$r=3$	$[0.000000,1.000000,0.000000,0.000000]$

Строки до и после выходной проекции
Позиция токена	До выходной проекции	После выходной проекции
$t=0$	$[1.000000,0.000000,1.000000,0.000000]$	$[1.000000,0.000000,1.000000,0.000000]$
$t=1$	$[0.770151,-0.420735,0.213809,0.786191]$	$[0.213809,0.786191,0.770151,-0.420735]$
$t=2$	$[0.374718,-0.583589,0.629044,0.945304]$	$[0.629044,0.945304,0.374718,-0.583589]$

Проверьте неизменность выходов префикса и инварианты сборки

Если изменить только последнюю входную строку, строки выхода 0 и 1 останутся побитово прежними в обеих головах и после выходной проекции.

$t=0$: Без изменений · Подтверждено расчётом
$t=1$: Без изменений · Подтверждено расчётом
$t=2$: Изменилось · Подтверждено расчётом

Разделение с последующим объединением Точно восстанавливает тензор
Влияние выходов голов друг на друга Отсутствует до выходной проекции
Вероятности будущих ключей Точно равны нулю
Общий сдвиг позиций при неизменном содержимом Сохраняет вероятности внимания
Обучаемые скалярные параметры $64$
Координаты, проверенные центральными разностями $76$

Результат после изменения последнего токена $\varepsilon=0.000000000001$

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

Внимательно рассмотрите позицию токена $1$ . До $W_O$ конкатенированная строка равна

[0.770151,-0.420735,0.213809,0.786191].

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

[0.213809,0.786191,0.770151,-0.420735].

Это видимое изменение объясняет, почему конкатенацию и выходную проекцию нужно считать двумя разными этапами.

Проверяйте формы, видимость и смешивание, а не приписывайте головам роли

Для $d_{\mathrm{model}}=6$ классифицируйте $h=1$ , $h=2$ , $h=3$ , $h=4$ и $h=6$ по правилу этой главы: делимость плюс чётная ширина головы для RoPE.
Проследите все формы при $B=2$ , $T=5$ , $d_{\mathrm{model}}=12$ и $h=3$ .
Заполните треугольник каузальной видимости из трёх токенов для обеих голов.
Объясните, почему одно полноразмерное вычисление внимания не равносильно двум вычислениям внутри отдельных голов.
Примените меняющую группы местами матрицу $W_O$ из примера к символическим строкам $H^{(0)}=[a,b]$ и $H^{(1)}=[c,d]$ . Затем возьмите единичную $W_O$ , дополнительно задайте $(W_O)_{0,2}=1$ и выведите новую координату выхода $y_2$ .
Измените только токен в позиции $2$ . Предскажите, какие строки выхода могут измениться, а какие обязаны остаться побитово прежними.
Объясните, где заканчивается разделение групп в примере и почему общая плотная проекция не изолирует фиксированные части исходного входа.
Вычислите число параметров целевого слоя без смещений при $d_{\mathrm{model}}=128$ .

Проверьте структурные ответы

При $h=1$ ширина $d_h=6$ чётная; при $h=3$ ширина $d_h=2$ тоже чётная. Оба варианта допустимы. При $h=2$ получаем нечётную $d_h=3$ , а при $h=6$ — нечётную $d_h=1$ ; оба нарушают правило RoPE для полной головы. $h=4$ не делит $6$ .
Формы $Q$ , $K$ и $V$ равны $[2,5,12]$ ; после разделения — $[2,3,5,4]$ ; вероятности каждой головы имеют форму $[2,5,5]$ , а вместе — $[2,3,5,5]$ ; выходы голов равны $[2,3,5,4]$ ; после объединения и на выходе — $[2,5,12]$ .
Каждая голова сохраняет один, затем два, затем три ключа. Треугольники видимости совпадают, хотя вероятности в разрешённых ячейках могут различаться.
Одно полноразмерное вычисление создаёт одну нормализованную таблицу оценок $[5,5]$ , общую для всех признаков значений. Две головы создают две отдельно нормализованные таблицы $[5,5]$ , каждая со своим знаменателем для каждой строки запроса; при некоторых параметрах их численные значения могут совпасть.
Конкатенация даёт $[a,b,c,d]$ , а зафиксированная выходная матрица — $[c,d,a,b]$ . Если начать с единичной матрицы и добавить $(W_O)_{0,2}=1$ , третья выходная координата станет $y_2=a+c$ и получит вклады обеих голов.
Строки $0$ и $1$ не видят токен $2$ , поэтому останутся побитово прежними. Строка $2$ может измениться в обеих головах и после $W_O$ .
Единичные матрицы запросов, ключей и значений делают группы примера видимыми. Столбцы общих упакованных проекций могут использовать все входные признаки; отдельные ветви получают лишь их спроецированные выходы. $W_O$ снова может смешать эти ветви.
Число параметров равно $4\cdot128^2=65{,}536$ .

Заблуждение: многоголовое внимание разбивает входной тензор — или уже готовый результат одной головы — на части.

Исправление: обучаемые проекции выполняются до разделения. Затем каждая голова по своей группе спроецированных признаков отдельно вычисляет оценки запросов и ключей после RoPE, каузальный softmax и смесь значений. Конкатенируются только готовые выходы голов, после чего $W_O$ может смешивать их. Различие двух таблиц в примере показывает один допустимый набор параметров, но не гарантирует устойчивой семантической роли каждой головы в любой обученной модели.

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

Теперь в собираемом декодере есть полное каузальное преобразование внимания с учётом позиций. Оно принимает и возвращает $[B,T,d_{\mathrm{model}}]$ , но пока не нормализует вход и не прибавляет результат к остаточному потоку.

В главе 31 оно войдёт в первую остаточную ветвь с предварительной нормализацией,

x'=x+\operatorname{MHA}(\operatorname{RMSNorm}(x)),

после которой появится остаточная ветвь сети прямого распространения. Владение кэшем ключей и значений и инкрементальное декодирование остаются отложенными до главы 37.