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

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

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

Проследите оба остаточных пути до запуска блока

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

В примере $B=1$ , $T=3$ , $d_{\mathrm{model}}=4$ , $h=2$ и $d_{\mathrm{ff}}=4$ . Входной тензор равен

X= \begin{bmatrix} 2&0&0&0\\ 0&2&0&0\\ 0&0&2&0 \end{bmatrix}.

Оба вектора масштабов RMSNorm равны $[1,1,1,1]$ , а в учебном примере $\varepsilon=0$ . Поэтому среднеквадратичное значение каждой входной строки равно $1$ . Матрицы запросов, ключей, значений и выхода внимания — единичные. Матрицы вентильной ветви, ветви расширения и проекции сжатия SwiGLU тоже единичные.

Такой выбор раскрывает порядок операций, но не превращает весь блок в тождественную функцию. Каузальное внимание смешивает видимые строки токенов. Промежуточный остаточный поток нормализуется заново, а SwiGLU применяет $\operatorname{SiLU}(z)\odot z$ к каждой строке токена.

До вычисления результата выпишите шесть преобразований и два обходных пути:

X\xrightarrow{\operatorname{RMSNorm}_a}N_a \xrightarrow{\operatorname{MHA}}A, \qquad X'=X+A,

X'\xrightarrow{\operatorname{RMSNorm}_f}N_f \xrightarrow{\operatorname{FFN}}F, \qquad Y=X'+F.

Все семь названных тензоров ширины модели имеют форму $[1,3,4]$ . Только вероятности внимания имеют форму $[1,2,3,3]$ . Для позиции токена $1$ предскажите, какое значение без изменений попадёт в первый тождественный путь, а какое достигнет внимания лишь после $\operatorname{RMSNorm}_a$ .

Складывайте каждую ветвь с потоком на её входе

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

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

Два применения $\operatorname{RMSNorm}$ владеют разными обучаемыми векторами масштабов. Первая ветвь использует $\operatorname{MHA}$ и смешивает токены, а вторая через $\operatorname{FFN}$ независимо применяет одно и то же обучаемое преобразование признаков к каждой позиции. Выход обоих преобразований должен возвращаться к ширине модели, иначе остаточное сложение не определено:

x,x',y\in\mathbb{R}^{B\times T\times d_{\mathrm{model}}}.

При проекциях без смещений два вектора масштабов, четыре матрицы внимания и три матрицы SwiGLU содержат

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

Для выбранных ширин $N_\theta=8+64+48=120$ . В примере обратное распространение неравномерной скалярной функции потерь проходит через все $12$ координат входа и все $120$ координат параметров. В итоге $132$ аналитических градиента сравниваются с центральными разностями.

Не смешивайте значения ветвей со значениями остаточного потока

$x$ — остаточный поток ширины модели на входе блока.
$x'$ — поток после сложения $x$ с результатом каузального внимания.
$y$ — выход блока после сложения результата сети прямого распространения с $x'$ .
$\operatorname{RMSNorm}_a$ и $\operatorname{RMSNorm}_f$ выполняют одну операцию, но владеют разными векторами масштабов.
$\operatorname{MHA}$ — преобразование многоголового каузального внимания с RoPE без смещений из главы 30.
$\operatorname{FFN}$ — преобразование SwiGLU без смещений из главы 20.
$B$ — размер батча, $T$ — число токенов, $d_{\mathrm{model}}$ — ширина остаточного потока, а $d_{\mathrm{ff}}$ — скрытая ширина SwiGLU.
$N_a$ и $N_f$ — нормализованные входы ветвей, а $A$ и $F$ — их обучаемые выходы. Ни одно из этих четырёх обозначений не относится к тождественному пути.

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

От рекуррентного состояния и Post-LN к декодерам с Pre-LN

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

Long Short-Term Memory описывает первый рассматриваемый этап. Хохрайтер и Шмидхубер предлагают явно рекуррентную архитектуру для обучения на длительных временных зависимостях. Здесь она задаёт предшественника с последовательным состоянием, но не описывает устройство всех более поздних языковых моделей на LSTM. Рекуррентные языковые модели на LSTM передают состояние от одного шага токена к следующему, а исходный Transformer применял LayerNorm после каждого остаточного сложения; ни одна из этих схем не совпадает с собираемым здесь каузальным блоком декодера с предварительной нормализацией.

Attention Is All You Need описывает следующий этап развития архитектуры. Васвани и соавторы определяют выход подслоя исходного Transformer как остаточное сложение с последующей LayerNorm и маскируют в самовнимании декодера будущие позиции. Первый подслой можно кратко записать как $\operatorname{LayerNorm}(x+\operatorname{MHA}(x))$ : это схема Post-LN.

On Layer Normalization in the Transformer Architecture даёт названия двум схемам и исследует переход между ними. Сюн и соавторы различают схемы Post-LN и Pre-LN и исследуют, как размещение нормализации внутри остаточного блока меняет поведение градиентов при инициализации. В схеме Pre-LN нормализация переносится на вход каждого подслоя, а LLaMA служит конкретным современным примером языковой модели, где предварительная нормализация входа сочетается с RMSNorm, каузальным вниманием, RoPE и SwiGLU.

LLaMA служит таким современным примером. Туврон и соавторы описывают каузальную языковую модель Transformer, которая нормализует вход каждого подслоя с помощью RMSNorm и использует SwiGLU и RoPE. Блок сохраняет форму остаточного потока и чередует каузальное внимание, смешивающее токены, с преобразованием признаков каждого токена; в главе 32 такие повторяемые блоки образуют стек языковой модели только с декодером.

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

Исполняемое сравнение сначала вычисляет небольшое переносимое состояние в стиле RNN, а затем пропускает одни и те же зафиксированные компоненты Transformer через первый остаточный этап главы в порядке Pre-LN и через вариант, где нормализация следует за сложением. Разные численные результаты показывают роль порядка операций, но не означают, что вспомогательная функция воспроизводит LSTM, полную историческую модель или обученные веса какой-либо статьи.

Сравните последовательное переносимое состояние с фиксированными остаточными порядками Post-LN и Pre-LN rust/demos/ch31-decoder-block/src/lib.rs#historical-block-order-contrast

#[derive(Clone, Debug, PartialEq)]
pub struct HistoryEvidence {
    pub rnn_style_states: Vec<f64>,
    pub sequential_recurrence: bool,
    pub original_post_norm: bool,
    pub modern_pre_norm: bool,
    pub numeric_order_contrast: bool,
}

/// A bounded serial recurrence; this is RNN-style evidence, not an LSTM gate implementation.
pub fn rnn_style_states(inputs: &[f64]) -> Vec<f64> {
    let mut state = 0.0_f64;
    inputs
        .iter()
        .map(|input| {
            state = (0.5 * input + 0.75 * state).tanh();
            state
        })
        .collect()
}

/// Unit-gain, zero-bias LayerNorm rows for the original post-norm ordering contrast.
pub fn layer_norm_rows(input: &Tensor, epsilon: f64) -> Tensor {
    let width = *input
        .shape()
        .last()
        .expect("history input has a feature axis");
    let mut normalized = Vec::with_capacity(input.len());
    for row in input.as_slice().chunks_exact(width) {
        let mean = row.iter().sum::<f64>() / width as f64;
        let variance = row
            .iter()
            .map(|value| {
                let centered = value - mean;
                centered * centered
            })
            .sum::<f64>()
            / width as f64;
        let inverse_standard_deviation = (variance + epsilon).sqrt().recip();
        normalized.extend(
            row.iter()
                .map(|value| (value - mean) * inverse_standard_deviation),
        );
    }
    Tensor::from_vec(input.shape().to_vec(), normalized).expect("history shape is unchanged")
}

Соедините проверенные части, не скрывая их границы

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

Свяжите каждую недопустимую границу с конфигурацией, входом, позицией, лентой операций или этапом, на котором возникла ошибка rust/crates/llm-from-scratch/src/models/decoder_block.rs#decoder-block-errors

/// A component whose feature width is incompatible with the residual stream.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DecoderBlockComponent {
    AttentionNorm,
    FeedForwardNorm,
    FeedForwardInput,
    FeedForwardOutput,
}

impl fmt::Display for DecoderBlockComponent {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::AttentionNorm => "attention RMSNorm",
            Self::FeedForwardNorm => "feed-forward RMSNorm",
            Self::FeedForwardInput => "feed-forward input",
            Self::FeedForwardOutput => "feed-forward output",
        })
    }
}

/// A rejected component assembly or stage of one decoder-block forward pass.
#[derive(Clone, Debug, PartialEq)]
pub enum DecoderBlockError {
    AttentionNorm(RmsNormError),
    Attention(MultiHeadAttentionError),
    AttentionResidual(ResidualError),
    FeedForwardNorm(RmsNormError),
    FeedForward(SwiGluError),
    FeedForwardResidual(ResidualError),
    ComponentWidthMismatch {
        component: DecoderBlockComponent,
        expected: usize,
        actual: usize,
    },
    Initialization(InitializationError),
}

impl fmt::Display for DecoderBlockError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::AttentionNorm(source) => write!(formatter, "attention RMSNorm: {source}"),
            Self::Attention(source) => write!(formatter, "causal multi-head attention: {source}"),
            Self::AttentionResidual(source) => {
                write!(formatter, "attention residual merge: {source}")
            }
            Self::FeedForwardNorm(source) => {
                write!(formatter, "feed-forward RMSNorm: {source}")
            }
            Self::FeedForward(source) => write!(formatter, "SwiGLU feed-forward: {source}"),
            Self::FeedForwardResidual(source) => {
                write!(formatter, "feed-forward residual merge: {source}")
            }
            Self::ComponentWidthMismatch {
                component,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder-block {component} width must be {expected}, got {actual}"
            ),
            Self::Initialization(source) => source.fmt(formatter),
        }
    }
}

impl Error for DecoderBlockError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::AttentionNorm(source) | Self::FeedForwardNorm(source) => Some(source),
            Self::Attention(source) => Some(source),
            Self::AttentionResidual(source) | Self::FeedForwardResidual(source) => Some(source),
            Self::FeedForward(source) => Some(source),
            Self::Initialization(source) => Some(source),
            Self::ComponentWidthMismatch { .. } => None,
        }
    }
}

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

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

Публичный метод forward последовательно вызывает нормализацию, внимание и первое остаточное сложение, а затем — вторую нормализацию, SwiGLU и второе остаточное сложение. Возвращаемые данные сохраняют все шесть значений этапов, полные данные внимания и сети прямого распространения на одной дифференцируемой ленте операций.

Выполните два независимо нормализованных преобразования и сложите каждый результат с остаточным потоком на входе его ветви rust/crates/llm-from-scratch/src/models/decoder_block.rs#decoder-block-layer

/// Every inspectable value produced by the two pre-normalized residual paths.
#[derive(Clone, Debug)]
pub struct DecoderBlockForward {
    attention_norm: RmsNormForward,
    attention: MultiHeadAttentionForward,
    after_attention: TensorValue,
    feed_forward_norm: RmsNormForward,
    feed_forward: SwiGluForward,
    output: TensorValue,
}

impl DecoderBlockForward {
    pub fn attention_norm(&self) -> &RmsNormForward {
        &self.attention_norm
    }

    pub fn attention(&self) -> &MultiHeadAttentionForward {
        &self.attention
    }

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

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

    pub fn feed_forward_norm(&self) -> &RmsNormForward {
        &self.feed_forward_norm
    }

    pub fn feed_forward(&self) -> &SwiGluForward {
        &self.feed_forward
    }

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

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

/// RMSNorm → causal MHA → residual, then RMSNorm → SwiGLU → residual.
#[derive(Clone, Debug)]
pub struct DecoderBlock {
    attention_norm: RmsNorm,
    attention: MultiHeadAttention,
    feed_forward_norm: RmsNorm,
    feed_forward: SwiGlu,
    parameters: NamedParameters,
    model_width: usize,
}

impl DecoderBlock {
    /// Initializes every matrix transactionally from one deterministic stream.
    pub fn new(
        parameter_prefix: impl Into<String>,
        config: DecoderBlockConfig,
        rng: &mut SplitMix64,
    ) -> Result<Self, DecoderBlockError> {
        let parameter_prefix = parameter_prefix.into();
        let attention_norm = RmsNorm::new(
            format!("{parameter_prefix}.attention_norm.gain"),
            config.model_width,
            config.rms_epsilon,
        )
        .map_err(DecoderBlockError::AttentionNorm)?;
        let feed_forward_norm = RmsNorm::new(
            format!("{parameter_prefix}.ffn_norm.gain"),
            config.model_width,
            config.rms_epsilon,
        )
        .map_err(DecoderBlockError::FeedForwardNorm)?;

        let mut trial = rng.clone();
        let attention = MultiHeadAttention::new(
            format!("{parameter_prefix}.attention"),
            config.model_width,
            config.heads,
            config.max_positions,
            config.rope_base,
            &mut trial,
        )
        .map_err(DecoderBlockError::Attention)?;
        let feed_forward = SwiGlu::new(
            format!("{parameter_prefix}.ffn"),
            config.model_width,
            config.feed_forward_width,
            config.model_width,
            &mut trial,
        )
        .map_err(DecoderBlockError::FeedForward)?;
        let block = Self::from_parts(attention_norm, attention, feed_forward_norm, feed_forward)?;
        *rng = trial;
        Ok(block)
    }

    /// Assembles already named deterministic components after cross-width checks.
    pub fn from_parts(
        attention_norm: RmsNorm,
        attention: MultiHeadAttention,
        feed_forward_norm: RmsNorm,
        feed_forward: SwiGlu,
    ) -> Result<Self, DecoderBlockError> {
        let model_width = attention.model_width();
        for (component, actual) in [
            (
                DecoderBlockComponent::AttentionNorm,
                attention_norm.feature_width(),
            ),
            (
                DecoderBlockComponent::FeedForwardNorm,
                feed_forward_norm.feature_width(),
            ),
            (
                DecoderBlockComponent::FeedForwardInput,
                feed_forward.input_width(),
            ),
            (
                DecoderBlockComponent::FeedForwardOutput,
                feed_forward.output_width(),
            ),
        ] {
            if actual != model_width {
                return Err(DecoderBlockError::ComponentWidthMismatch {
                    component,
                    expected: model_width,
                    actual,
                });
            }
        }

        let mut listed = Vec::with_capacity(9);
        listed.extend(attention_norm.parameters().iter().cloned());
        listed.extend(attention.parameters().iter().cloned());
        listed.extend(feed_forward_norm.parameters().iter().cloned());
        listed.extend(feed_forward.parameters().iter().cloned());
        let parameters = NamedParameters::try_new(listed)?;

        Ok(Self {
            attention_norm,
            attention,
            feed_forward_norm,
            feed_forward,
            parameters,
            model_width,
        })
    }

    /// Runs the two transformation branches in exact pre-normalized order.
    pub fn forward(
        &self,
        input: &TensorValue,
        position_offset: usize,
    ) -> Result<DecoderBlockForward, DecoderBlockError> {
        let attention_norm = self
            .attention_norm
            .forward_with_intermediates(input)
            .map_err(DecoderBlockError::AttentionNorm)?;
        let attention = self
            .attention
            .forward(attention_norm.output(), position_offset)
            .map_err(DecoderBlockError::Attention)?;
        let after_attention = residual_add(input, attention.output())
            .map_err(DecoderBlockError::AttentionResidual)?;
        let feed_forward_norm = self
            .feed_forward_norm
            .forward_with_intermediates(&after_attention)
            .map_err(DecoderBlockError::FeedForwardNorm)?;
        let feed_forward = self
            .feed_forward
            .forward_with_intermediates(feed_forward_norm.output())
            .map_err(DecoderBlockError::FeedForward)?;
        let output = residual_add(&after_attention, feed_forward.output())
            .map_err(DecoderBlockError::FeedForwardResidual)?;

        Ok(DecoderBlockForward {
            attention_norm,
            attention,
            after_attention,
            feed_forward_norm,
            feed_forward,
            output,
        })
    }

    pub fn attention_norm(&self) -> &RmsNorm {
        &self.attention_norm
    }

    pub fn attention(&self) -> &MultiHeadAttention {
        &self.attention
    }

    pub fn feed_forward_norm(&self) -> &RmsNorm {
        &self.feed_forward_norm
    }

    pub fn feed_forward(&self) -> &SwiGlu {
        &self.feed_forward
    }

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

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

    pub fn parameter_count(&self) -> usize {
        self.parameters
            .as_slice()
            .iter()
            .map(|parameter| parameter.tensor().value().len())
            .sum()
    }
}

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

\delta=10^{-6},\qquad \tau=2\times10^{-5}.

Для каждой из $132$ проверяемых координат аналитический градиент и численное сравнение конечны, а расхождение укладывается в этот допуск. Учебный отчёт фиксирует выводы, но не подменяет сами проверки:

Сформируйте точный отчёт о примере, порядке, каузальности, параметрах, градиентах, ошибках, повторе и следующей границе rust/demos/ch31-decoder-block/src/lib.rs#learner-report

pub fn render_report(evidence: &LearnerEvidence) -> String {
    let primary = &evidence.primary;
    let parameters = &evidence.parameters;
    let errors = &evidence.errors;
    let gradients = &evidence.gradients;
    let history = &evidence.history;
    [
        "chapter=31-decoder-block".to_owned(),
        format!(
            "config=batch:{BATCH} tokens:{TOKENS} model_width:{MODEL_WIDTH} heads:{HEADS} head_width:{HEAD_WIDTH} feed_forward_width:{FEED_FORWARD_WIDTH} epsilon:{RMS_EPSILON:.6}"
        ),
        format!(
            "shape=input:{} attention_norm:{} attention_weights:{} attention_branch:{} after_attention:{} feed_forward_norm:{} feed_forward_branch:{} output:{} probe_logits:{}",
            format_shape(&evidence.shapes.input),
            format_shape(&evidence.shapes.attention_norm),
            format_shape(&evidence.shapes.attention_weights),
            format_shape(&evidence.shapes.attention_branch),
            format_shape(&evidence.shapes.after_attention),
            format_shape(&evidence.shapes.feed_forward_norm),
            format_shape(&evidence.shapes.feed_forward_branch),
            format_shape(&evidence.shapes.output),
            format_shape(&evidence.shapes.probe_logits),
        ),
        format!(
            "order=attention_norm->attention->residual->feed_forward_norm->feed_forward->residual pre_norm:{} post_norm_differs:{}",
            primary.pre_norm_order, primary.post_norm_differs
        ),
        format!(
            "causality=prefix_0_bitwise:{} prefix_1_bitwise:{} suffix_changed:{} future_probabilities_zero:{}",
            primary.prefix_zero_unchanged,
            primary.prefix_one_unchanged,
            primary.suffix_changed,
            primary.future_probabilities_zero,
        ),
        format!(
            "parameters=tensors:{} scalars:{} bias_free:{} stable_order:{} distinct:{}",
            parameters.tensors,
            parameters.scalars,
            parameters.bias_free,
            parameters.stable_order,
            parameters.node_distinct,
        ),
        format!(
            "gradcheck=input:{} parameters:{} total:{} tolerance:{GRADIENT_TOLERANCE:.6} passed:{} tape_finite:{}",
            gradients.input_checks,
            gradients.parameter_checks,
            gradients.input_checks + gradients.parameter_checks,
            gradients.passed,
            primary.tape_finite,
        ),
        format!(
            "errors=configuration:{} component_width:{} input_rank:{} input_width:{} empty_tokens:{} position_range:{} released_input:{}",
            errors.configuration_rejected,
            errors.component_width_rejected,
            errors.input_rank_rejected,
            errors.input_width_rejected,
            errors.empty_tokens_rejected,
            errors.position_range_rejected,
            errors.released_input_rejected,
        ),
        format!(
            "history=sequential_recurrence:{} original_post_norm:{} modern_pre_norm:{} numeric_order_contrast:{}",
            history.sequential_recurrence,
            history.original_post_norm,
            history.modern_pre_norm,
            history.numeric_order_contrast,
        ),
        format!(
            "replay={}",
            if evidence.replay_bitwise {
                "bitwise"
            } else {
                "mismatch"
            }
        ),
        "next=stack these blocks between token embeddings and a tied vocabulary head".to_owned(),
    ]
    .join("\n")
        + "\n"
}

Точка входа исполняемого примера печатает только этот отчёт:

Напечатайте зафиксированный учебный отчёт о блоке декодера rust/demos/ch31-decoder-block/src/main.rs

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

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

Рассмотрите оба обходных пути и обе ветви преобразования

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

Выведите строгую трассировку блока декодера из тридцати трёх строк rust/demos/ch31-decoder-block/src/diagram_trace.rs#decoder-block-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 history = &evidence.history;
    let mut lines = vec![
        String::from(
            "CONFIG|batch=1|tokens=3|model_width=4|heads=2|head_width=2|feed_forward_width=4|epsilon=0.000000|stage_order=[attention-norm,attention,residual-1,feed-forward-norm,feed-forward,residual-2]",
        ),
        format!("SHAPE|stage=input|value={}", format_shape(&shapes.input)),
        format!(
            "SHAPE|stage=attention-norm|value={}",
            format_shape(&shapes.attention_norm)
        ),
        format!(
            "SHAPE|stage=attention-weights|value={}",
            format_shape(&shapes.attention_weights)
        ),
        format!(
            "SHAPE|stage=attention-branch|value={}",
            format_shape(&shapes.attention_branch)
        ),
        format!(
            "SHAPE|stage=after-attention|value={}",
            format_shape(&shapes.after_attention)
        ),
        format!(
            "SHAPE|stage=feed-forward-norm|value={}",
            format_shape(&shapes.feed_forward_norm)
        ),
        format!(
            "SHAPE|stage=feed-forward-branch|value={}",
            format_shape(&shapes.feed_forward_branch)
        ),
        format!("SHAPE|stage=output|value={}", format_shape(&shapes.output)),
        format!(
            "SHAPE|stage=probe-logits|value={}",
            format_shape(&shapes.probe_logits)
        ),
        stage_record("input", primary.input.as_slice()),
        stage_record("attention-norm", primary.attention_norm.as_slice()),
        stage_record("attention-branch", primary.attention_branch.as_slice()),
        stage_record("after-attention", primary.after_attention.as_slice()),
        stage_record("feed-forward-norm", primary.feed_forward_norm.as_slice()),
        stage_record(
            "feed-forward-branch",
            primary.feed_forward_branch.as_slice(),
        ),
        stage_record("output", primary.output.as_slice()),
    ];
    lines.extend(
        (0..HEADS)
            .flat_map(|head| (0..TOKENS).map(move |query| (head, query)))
            .map(|(head, query)| weight_record(evidence, head, query)),
    );
    lines.push(format!(
        "MERGE|name=attention|identity=input|branch=attention-branch|result=after-attention|exact={}",
        primary.first_residual_exact
    ));
    lines.push(format!(
        "MERGE|name=feed-forward|identity=after-attention|branch=feed-forward-branch|result=output|exact={}",
        primary.second_residual_exact
    ));
    lines.extend((0..TOKENS).map(|token| {
        let start = token * 3;
        format!(
            "PROBE|token={token}|values={}",
            format_vector(&primary.probe_logits.as_slice()[start..start + 3])
        )
    }));
    lines.push(format!(
        "ORDER_PROOF|pre_norm={}|post_norm_differs={}|post_norm_token_1={}|pre_norm_token_1={}",
        primary.pre_norm_order,
        primary.post_norm_differs,
        format_vector(token_row(primary.post_norm_first_stage.as_slice(), 1)),
        format_vector(token_row(primary.after_attention.as_slice(), 1)),
    ));
    lines.push(format!(
        "CAUSAL_PROOF|position_0={}|position_1={}|position_2={}|future_probabilities={}",
        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.future_probabilities_zero {
            "exact-zero"
        } else {
            "nonzero"
        },
    ));
    lines.push(format!(
        "PARAMETERS|tensors={}|scalars={}|bias={}|stable_order={}|distinct={}|names=[{}]",
        parameters.tensors,
        parameters.scalars,
        !parameters.bias_free,
        parameters.stable_order,
        parameters.node_distinct,
        parameters.names.join(",")
    ));
    lines.push(format!(
        "GRADIENTS|input={}|parameters={}|total={}|tolerance={GRADIENT_TOLERANCE:.6}|passed={}|tape_finite={}",
        gradients.input_checks,
        gradients.parameter_checks,
        gradients.input_checks + gradients.parameter_checks,
        gradients.passed,
        primary.tape_finite,
    ));
    lines.push(format!(
        "HISTORY|rnn_style_states={}|sequential={}|original_post_norm={}|modern_pre_norm={}|numeric_order_contrast={}",
        format_vector(&history.rnn_style_states),
        history.sequential_recurrence,
        history.original_post_norm,
        history.modern_pre_norm,
        history.numeric_order_contrast,
    ));
    debug_assert_eq!(lines.len(), 33);
    lines.join("\n") + "\n"
}

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

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

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

Один блок, сохраняющий форму, в порядке вычислений

$x'=x+\operatorname{MHA}(\operatorname{RMSNorm}(x))$ $y=x'+\operatorname{FFN}(\operatorname{RMSNorm}(x'))$

Входной остаточный поток: $[1,3,4]$
Строки после RMSNorm перед вниманием: $[1,3,4]$
Веса каузального внимания: $[1,2,3,3]$
Выход ветви внимания: $[1,3,4]$
После остаточного сложения с вниманием: $[1,3,4]$
Строки после RMSNorm перед сетью: $[1,3,4]$
Выход ветви SwiGLU: $[1,3,4]$
Выход блока: $[1,3,4]$
Проверочные логиты: $[1,3,3]$

Нормализуйте вход перед ветвью каузального внимания

Тождественная копия проходит в обход RMSNorm и внимания; позиции видимых токенов смешивает только пунктирная ветвь.

Входной остаточный поток

x=[0.000000,

2.000000,

0.000000,

0.000000]

[1,3,4]

Тождественный путь

x=[0.000000,

2.000000,

0.000000,

0.000000]

Вход RMSNorm перед вниманием

N_a=[0.000000,

2.000000,

0.000000,

0.000000]

Каузальное многоголовое внимание

A=[0.010881,

1.989119,

0.000000,

0.000000]

Двойная рамка: остаточное сложение

+

После остаточного сложения с вниманием

x'=[0.010881,

3.989119,

0.000000,

0.000000]

[1,3,4]

Нормализуйте промежуточный поток перед SwiGLU

Вторая тождественная копия проходит в обход отдельной RMSNorm и преобразования признаков каждого токена.

После остаточного сложения с вниманием

x'=[0.010881,

3.989119,

0.000000,

0.000000]

[1,3,4]

Тождественный путь

x'=[0.010881,

3.989119,

0.000000,

0.000000]

Вход RMSNorm перед сетью

N_f=[0.005455,

1.999993,

0.000000,

0.000000]

Ветвь SwiGLU

F=[0.000015,

3.523159,

0.000000,

0.000000]

Двойная рамка: остаточное сложение

+

Выход блока

y=[0.010896,

7.512278,

0.000000,

0.000000]

[1,3,4]

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

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

Прокручиваемая таблица значений этапов блока декодера: Нормализуйте вход перед ветвью каузального внимания
Позиция токена	Входной остаточный поток	Строки после RMSNorm перед вниманием	Выход ветви внимания	После остаточного сложения с вниманием
$t=0$	$[2.000000,0.000000,0.000000,0.000000]$	$[2.000000,0.000000,0.000000,0.000000]$	$[2.000000,0.000000,0.000000,0.000000]$	$[4.000000,0.000000,0.000000,0.000000]$
$t=1$	$[0.000000,2.000000,0.000000,0.000000]$	$[0.000000,2.000000,0.000000,0.000000]$	$[0.010881,1.989119,0.000000,0.000000]$	$[0.010881,3.989119,0.000000,0.000000]$
$t=2$	$[0.000000,0.000000,2.000000,0.000000]$	$[0.000000,0.000000,2.000000,0.000000]$	$[0.666667,0.666667,1.788570,0.000000]$	$[0.666667,0.666667,3.788570,0.000000]$

Прокручиваемая таблица значений этапов блока декодера: Нормализуйте промежуточный поток перед SwiGLU
Позиция токена	Строки после RMSNorm перед сетью	Выход ветви SwiGLU	Выход блока
$t=0$	$[2.000000,0.000000,0.000000,0.000000]$	$[3.523188,0.000000,0.000000,0.000000]$	$[7.523188,0.000000,0.000000,0.000000]$
$t=1$	$[0.005455,1.999993,0.000000,0.000000]$	$[0.000015,3.523159,0.000000,0.000000]$	$[0.010896,7.512278,0.000000,0.000000]$
$t=2$	$[0.341520,0.341520,1.940806,0.000000]$	$[0.068180,0.068180,3.293781,0.000000]$	$[0.734847,0.734847,7.082351,0.000000]$

Прокручиваемая таблица значений этапов блока декодера: Сверьте точные значения этапов, каузальные строки, порядок, параметры и градиенты
Позиция токена	Проверочные логиты
$t=0$	$[7.523188,0.000000,-7.523188]$
$t=1$	$[0.010896,7.512278,-7.523174]$
$t=2$	$[7.817198,7.817198,-1.469694]$

Прокручиваемая таблица вероятностей каузального внимания
$h$	Позиция запроса	$k=0$	$k=1$	$k=2$	Сумма строки
$0$	$q=0$	$1.000000$ Доступен	$0.000000$ Скрыт маской	$0.000000$ Скрыт маской	$1.000000$
$0$	$q=1$	$0.005440$ Доступен	$0.994560$ Доступен	$0.000000$ Скрыт маской	$1.000000$
$0$	$q=2$	$0.333333$ Доступен	$0.333333$ Доступен	$0.333333$ Доступен	$1.000000$
$1$	$q=0$	$1.000000$ Доступен	$0.000000$ Скрыт маской	$0.000000$ Скрыт маской	$1.000000$
$1$	$q=1$	$0.500000$ Доступен	$0.500000$ Доступен	$0.000000$ Скрыт маской	$1.000000$
$1$	$q=2$	$0.052857$ Доступен	$0.052857$ Доступен	$0.894285$ Доступен	$1.000000$

Результаты Pre-LN и Post-LN различаются

Численно изменилось

$\operatorname{PostNorm}_1=[-0.573144,$ $1.732042,$ $-0.579449,$ $-0.579449]$

$x'=[0.010881,$ $3.989119,$ $0.000000,$ $0.000000]$

Изменённый последний токен не влияет на предыдущие выходы

$t=0$ : Побитово не изменилось

$t=1$ : Побитово не изменилось

$t=2$ : Численно изменилось

Стабильная принадлежность параметров

$N_\theta=120$

tensors=9

bias=false

Полная проверка центральными разностями

$n=132$

$\tau=0.000020$

Подтверждено примером на Rust

Строки этапов показывают, какое значение идёт по каждому тождественному пути и какое возвращается из обучаемой ветви. Каузальная таблица выделяет единственное преобразование, смешивающее токены, а численное сравнение ставит рядом результаты Pre-LN и Post-LN. Сплошные тождественные пути, пунктирные ветви преобразования и двойные остаточные сложения различают эти роли без опоры на цвет.

Проследите позицию токена $1$ в первой ветви. Входная строка равна $[0,2,0,0]$ , а результат внимания — $[0.010881,1.989119,0,0]$ . Первое остаточное сложение даёт

x'_1=[0.010881,3.989119,0,0].

Если нормализовать уже результат сложения, в фиксированном сравнении получится $[-0.573144,1.732042,-0.579449,-0.579449]$ . Различие служит прямым численным доказательством того, что порядки Pre-LN и Post-LN не взаимозаменяемы.

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

Расположите операции по порядку: первое остаточное сложение, RMSNorm перед сетью прямого распространения, RMSNorm перед вниманием, SwiGLU, каузальное внимание, второе остаточное сложение.
При $B=2$ , $T=5$ и $d_{\mathrm{model}}=12$ запишите форму остаточного потока после каждого этапа. Какую ось может изменить остаточное сложение?
Если ветвь внимания возвращает ноль, выведите $x'$ . Становится ли весь блок тождественной функцией?
Измените только последний входной токен. Какие строки выхода обязаны остаться прежними и какое преобразование устанавливает эту границу?
При $d_{\mathrm{model}}=12$ и $d_{\mathrm{ff}}=32$ вычислите полное число параметров без смещений, включая оба вектора масштабов RMSNorm.
Перепишите только первый остаточный этап в порядке Post-LN и укажите, где его значение впервые может разойтись с вычислением Pre-LN.

Сверить структурные ответы

RMSNorm перед вниманием, каузальное внимание, первое остаточное сложение, RMSNorm перед сетью прямого распространения, SwiGLU, второе остаточное сложение.
Каждый остаточный поток имеет форму $[2,5,12]$ . Ни одно остаточное сложение не может изменить ось: формы обоих операндов должны полностью совпадать.
$x'=x$ . Ветвь сети прямого распространения всё ещё может изменить $y$ , поэтому весь блок не обязан быть тождественным.
Строки от $0$ до $3$ обязаны остаться побитово прежними. Каузальная маска внимания не пропускает суффикс к более ранним строкам, а последующее преобразование признаков каждого токена не смешивает разные токены.
$2(12)+4(12)^2+3(12)(32)=1752$ скалярных параметра.
Post-LN вычисляет $\operatorname{LayerNorm}(x+\operatorname{MHA}(x))$ , а Pre-LN — $x+\operatorname{MHA}(\operatorname{RMSNorm}(x))$ . Входы внимания могут различаться ещё до появления результата остаточного сложения.

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

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

Теперь в собираемой реализации есть один полный каузальный блок декодера глубины один. В главе 32 такие блоки будут повторяться между эмбеддингами токенов, завершающей RMSNorm и связанной проекцией на словарь.

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