32 · Content revision 4

Stack a decoder and tie its vocabulary head

Learn how token embeddings, repeated causal blocks, final RMSNorm, and one tied vocabulary table produce differentiable next-token logits.

Predict the axes before tracing the decoder

Assemble token lookup, repeated pre-normalized decoder blocks, final RMSNorm, and one genuinely tied vocabulary projection into differentiable logits. The fixture sends token IDs $z=[0,1,2]$ through a two-block decoder with vocabulary size $V=5$ , model width $d_{\mathrm{model}}=4$ , two attention heads, feed-forward width $d_{\mathrm{ff}}=4$ , and context capacity $4$ . Its targets are $y=[1,2,3]$ .

Before running it, predict the five structural results:

lookup returns shape $[1,3,4]$ ;
both decoder blocks and final RMSNorm preserve $[1,3,4]$ ;
the vocabulary projection returns $[1,3,5]$ ;
token_embedding.weight appears once, even though two tape paths use it;
changing only token position $2$ cannot change logit rows $0$ or $1$ .

The deterministic fixture then supplies numbers. At token position $1$ , the logit row is

[-0.862249,\ 0.967613,\ -0.991545,\ -0.446363,\ 1.234533].

Its largest entry is vocabulary ID $4$ . Across the three target positions, the mean indexed negative log likelihood is $2.045535$ . This is evidence from an untrained tiny fixture, not evidence that it has learned language.

Reuse the embedding table at the far end of the stack

The complete forward formula is

\ell=\operatorname{RMSNorm}(B_N(\cdots B_1(E[z])\cdots))E^\top

This notation describes a positive depth $N\geq 1$ . At the valid zero-block boundary, the empty block composition is the identity, so the formula reduces to

\ell=\operatorname{RMSNorm}(E[z])E^\top.

The token selectors satisfy $z\in\{0,\ldots,V-1\}^{B\times T}$ . Lookup through $E\in\mathbb{R}^{V\times d_{\mathrm{model}}}$ produces a residual stream in $\mathbb{R}^{B\times T\times d_{\mathrm{model}}}$ . Every $B_i$ preserves that shape. Final $\operatorname{RMSNorm}$ also preserves it. Multiplication by $E^\top$ changes only the last axis, so

\ell\in\mathbb{R}^{B\times T\times V}.

Here $\ell$ is the logit tensor, not the scalar training loss. For targets $y$ , the fixture computes

\mathcal{L}=-\frac{1}{BT}\sum_{b=1}^{B}\sum_{t=1}^{T} \log\operatorname{softmax}(\ell_{b,t,:})_{y_{b,t}}.

With bias-free components, one block owns $4d_{\mathrm{model}}^2+3d_{\mathrm{model}}d_{\mathrm{ff}}+2d_{\mathrm{model}}$ scalars. The complete tied model owns

Vd_{\mathrm{model}}+N\left(4d_{\mathrm{model}}^2 +3d_{\mathrm{model}}d_{\mathrm{ff}}+2d_{\mathrm{model}}\right) +d_{\mathrm{model}}.

For this fixture, the count is $20+2(64+48+8)+4=264$ . A separate vocabulary matrix would add $Vd_{\mathrm{model}}=20$ scalars and raise the total to $284$ .

Distinguish one parameter from its two uses

$z$ is the rank-two batch of integer token IDs.
$E$ is one trainable vocabulary-by-feature table.
$E[z]$ gathers token rows from that table for the lookup role.
$B_i$ is the $i$ -th pre-normalized causal decoder block. Equal configuration does not mean shared block parameters.
$N$ is decoder depth. This implementation accepts zero, one, or more blocks.
Final $\operatorname{RMSNorm}$ owns a gain separate from every block gain.
$E^\top$ is a differentiable transpose view of $E$ , not another parameter.
$B$ is batch size, $T$ is sequence length, and $V$ is vocabulary size.
$d_{\mathrm{model}}$ is residual-stream width, while $d_{\mathrm{ff}}$ is the hidden width of each SwiGLU branch.
$y$ contains one target vocabulary ID for each pair $(b,t)$ .
$\ell$ contains vocabulary logits; $\mathcal{L}$ is the scalar mean loss.

Weight tying is stronger than equal initialization. There is no lm_head.weight, no copied matrix, and no synchronization step. Reverse mode reaches token_embedding.weight once through gathered input rows and once through the output projection, then adds both contributions on the same leaf:

\bar E=\bar E_{\mathrm{lookup}}+\bar E_{\mathrm{output}}.

From separate recurrent components to one decoder stack

Recurrent language models with separate input and output tables, followed by the transition to tied vocabulary weights and repeated causal Transformer layers. This is the bounded road to the modern LLM boundary assembled here.

Using the Output Embedding to Improve Language Models supplies the earlier comparison. Press and Wolf describe the input-embedding, intervening-computation, and output-score-matrix roles in recurrent neural language models, recommend tying the two embeddings, and analyze the tied update as contributions from both roles. Earlier recurrent neural language models commonly separated an input embedding, stepwise recurrent computation, and a distinct vocabulary classifier, so input and output word tables could own independent parameters and updates.

Attention Is All You Need changes the sequence architecture. Vaswani and colleagues define a stacked causally masked Transformer decoder and report sharing one matrix between embeddings and the pre-softmax linear transformation. Their full translation decoder also has encoder-decoder attention and post-normalization, so it is not the same model as this decoder-only pre-norm fixture.

Language Models are Unsupervised Multitask Learners provides the next model-level ordering. Radford and colleagues describe GPT-2 as a multi-layer Transformer language model with normalization at each sub-block input and an additional normalization after the final block. This chapter uses that ordering boundary without claiming GPT-2’s scale, initialization, tokenizer, context, or trained parameters.

LLaMA gives a bounded modern family. Touvron and colleagues place pre-normalization with RMSNorm, SwiGLU, RoPE, and causal attention in a modern Transformer language-model family with model depth as an explicit architecture dimension. A decoder-only LLM maps token IDs through a repeated causal residual stack into vocabulary logits; this course uses the already-taught RMSNorm, RoPE, attention, and SwiGLU pieces while making its tied-head choice explicit rather than universal.

Weight tying made one vocabulary-feature matrix serve both roles, the Transformer organized causally masked layers into a stack, and GPT-2 placed normalization at each sub-block input plus one final normalization before the vocabulary head. The model boundary therefore evolved from a token embedding feeding stepwise recurrent state and a separate classifier toward repeated causally masked blocks whose final hidden states are normalized and projected to vocabulary logits. Weight tying reduces parameters while requiring both gradient contributions to accumulate on one leaf.

The runnable contrast stays on that LLM architecture path. The tied fixture has $264$ scalars; an otherwise identical model with a separate $[5,4]$ vocabulary classifier would have $284$ . Detaching each use in turn checks that the full table gradient is the sum of the lookup-role and output-role contributions. This isolates the parameter-sharing consequence without claiming that the tiny fixture reproduces the recurrent predecessor, Transformer, GPT-2, or LLaMA.

Separate the lookup and output paths and add their contributions on one vocabulary table rust/demos/ch32-decoder-model/src/lib.rs#tied-gradient-proof

fn tied_role_gradient(include_lookup: bool, include_head: bool) -> Vec<f64> {
    let table = NamedParameter::from_tensor(
        "token_embedding.weight",
        tensor(&[VOCABULARY, MODEL_WIDTH], &PROBE_TABLE),
    )
    .expect("probe table must be valid");
    let embedding = Embedding::from_parameter(table.clone()).expect("probe table must embed");
    let lookup = embedding
        .forward(&[0, 1], &[1, 2])
        .expect("probe lookup must be valid");
    let lookup = if include_lookup {
        lookup
    } else {
        lookup.detach()
    };
    let normalized = RmsNorm::new("final_norm.gain", MODEL_WIDTH, RMS_EPSILON)
        .expect("probe norm must be valid")
        .forward(&lookup)
        .expect("probe norm forward must be valid");
    let head_source = if include_head {
        table.tensor().clone()
    } else {
        table.tensor().detach()
    };
    let logits = normalized
        .matmul(
            &head_source
                .transpose(0, 1)
                .expect("probe head transpose must be valid"),
        )
        .expect("probe tied projection must be valid");
    logits
        .indexed_mean_nll(2, &[1, 2])
        .expect("probe indexed loss must be valid")
        .backward_with_seed(&tensor(&[], &[1.0]).view(), GraphRetention::Retain)
        .expect("probe backward must succeed");
    table.tensor().gradient().map_or_else(
        || vec![0.0; PROBE_TABLE.len()],
        |gradient| gradient.as_slice().to_vec(),
    )
}

fn tied_gradient_decomposition_error() -> f64 {
    let full = tied_role_gradient(true, true);
    let lookup = tied_role_gradient(true, false);
    let head = tied_role_gradient(false, true);
    full.iter()
        .zip(lookup.iter().zip(head.iter()))
        .map(|(full, (lookup, head))| (full - lookup - head).abs())
        .fold(0.0, f64::max)
}

Make the complete model boundary explicit in Rust

DecoderModelError identifies invalid vocabulary, width, head, feed-forward, depth, context, RoPE, RMSNorm, component, parameter-name, token, target, and tape boundaries. Delegated block, embedding, normalization, projection, and indexed loss failures retain their source instead of collapsing into one generic model error.

Attribute each complete-model failure to the boundary that rejected it rust/crates/llm-from-scratch/src/models/decoder.rs#decoder-model-errors

/// A model-owned taped operation that rejected a forward or loss calculation.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DecoderModelStage {
    TiedWeightTranspose,
    TiedVocabularyProjection,
    IndexedMeanLoss,
}

impl fmt::Display for DecoderModelStage {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::TiedWeightTranspose => "transpose tied embedding weight",
            Self::TiedVocabularyProjection => "project tied vocabulary logits",
            Self::IndexedMeanLoss => "indexed mean negative log likelihood",
        })
    }
}

/// Which normalization inside a repeated block disagreed with model config.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DecoderModelNorm {
    Attention,
    FeedForward,
}

impl fmt::Display for DecoderModelNorm {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::Attention => "attention",
            Self::FeedForward => "feed-forward",
        })
    }
}

/// A rejected decoder configuration, component assembly, input, or tape stage.
#[derive(Clone, Debug, PartialEq)]
pub enum DecoderModelError {
    EmptyVocabulary,
    ZeroModelWidth,
    ZeroHeadCount,
    ModelWidthNotDivisible {
        model_width: usize,
        heads: usize,
    },
    OddHeadWidth {
        head_width: usize,
    },
    ZeroFeedForwardWidth,
    ZeroPositionCapacity,
    InvalidRopeBase {
        value: f64,
    },
    InvalidRmsEpsilon {
        value: f64,
    },
    LayerAllocationFailed {
        layers: usize,
    },
    ParameterAllocationFailed {
        tensors: usize,
    },
    ParameterCountMismatch {
        expected: usize,
        actual: usize,
    },
    TargetAllocationFailed {
        targets: usize,
    },
    Embedding(EmbeddingError),
    Block {
        layer: usize,
        source: DecoderBlockError,
    },
    FinalNorm(RmsNormError),
    Initialization(InitializationError),
    LayerCountMismatch {
        expected: usize,
        actual: usize,
    },
    EmbeddingVocabularyMismatch {
        expected: usize,
        actual: usize,
    },
    EmbeddingWidthMismatch {
        expected: usize,
        actual: usize,
    },
    BlockModelWidthMismatch {
        layer: usize,
        expected: usize,
        actual: usize,
    },
    BlockHeadCountMismatch {
        layer: usize,
        expected: usize,
        actual: usize,
    },
    BlockFeedForwardWidthMismatch {
        layer: usize,
        expected: usize,
        actual: usize,
    },
    BlockPositionCapacityMismatch {
        layer: usize,
        expected: usize,
        actual: usize,
    },
    BlockRopeBaseMismatch {
        layer: usize,
        expected: f64,
        actual: f64,
    },
    BlockRmsEpsilonMismatch {
        layer: usize,
        norm: DecoderModelNorm,
        expected: f64,
        actual: f64,
    },
    FinalNormWidthMismatch {
        expected: usize,
        actual: usize,
    },
    FinalNormEpsilonMismatch {
        expected: f64,
        actual: f64,
    },
    ParameterNameMismatch {
        index: usize,
        expected: String,
        actual: String,
    },
    TokenRank {
        rank: usize,
    },
    EmptyBatch,
    EmptyTokens,
    TokenCountOverflow {
        batch: usize,
        tokens: usize,
    },
    TokenCountMismatch {
        expected: usize,
        actual: usize,
    },
    ContextLengthExceeded {
        tokens: usize,
        max_positions: usize,
    },
    TargetCountMismatch {
        expected: usize,
        actual: usize,
    },
    TargetIdOutOfBounds {
        position: usize,
        id: u32,
        vocabulary_size: usize,
    },
    Autodiff {
        stage: DecoderModelStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for DecoderModelError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EmptyVocabulary => {
                formatter.write_str("decoder vocabulary must contain at least one token")
            }
            Self::ZeroModelWidth => formatter.write_str("decoder model width must be nonzero"),
            Self::ZeroHeadCount => formatter.write_str("decoder head count must be nonzero"),
            Self::ModelWidthNotDivisible { model_width, heads } => write!(
                formatter,
                "decoder model width {model_width} must be divisible by head count {heads}"
            ),
            Self::OddHeadWidth { head_width } => write!(
                formatter,
                "decoder per-head width must be even for RoPE, got {head_width}"
            ),
            Self::ZeroFeedForwardWidth => {
                formatter.write_str("decoder feed-forward width must be nonzero")
            }
            Self::ZeroPositionCapacity => {
                formatter.write_str("decoder context capacity must be nonzero")
            }
            Self::InvalidRopeBase { value } => {
                write!(
                    formatter,
                    "decoder RoPE base must be finite and positive, got {value:?}"
                )
            }
            Self::InvalidRmsEpsilon { value } => write!(
                formatter,
                "decoder RMSNorm epsilon must be finite and nonnegative, got {value:?}"
            ),
            Self::LayerAllocationFailed { layers } => {
                write!(formatter, "could not reserve {layers} decoder blocks")
            }
            Self::ParameterAllocationFailed { tensors } => {
                write!(
                    formatter,
                    "could not reserve {tensors} decoder parameter handles"
                )
            }
            Self::ParameterCountMismatch { expected, actual } => write!(
                formatter,
                "decoder parameter count must be {expected}, got {actual}"
            ),
            Self::TargetAllocationFailed { targets } => {
                write!(
                    formatter,
                    "could not reserve {targets} decoder target indices"
                )
            }
            Self::Embedding(source) => write!(formatter, "token embedding: {source}"),
            Self::Block { layer, source } => {
                write!(formatter, "decoder block {layer}: {source}")
            }
            Self::FinalNorm(source) => write!(formatter, "final RMSNorm: {source}"),
            Self::Initialization(source) => source.fmt(formatter),
            Self::LayerCountMismatch { expected, actual } => write!(
                formatter,
                "decoder config needs {expected} blocks, but received {actual}"
            ),
            Self::EmbeddingVocabularyMismatch { expected, actual } => write!(
                formatter,
                "token embedding vocabulary must be {expected}, got {actual}"
            ),
            Self::EmbeddingWidthMismatch { expected, actual } => write!(
                formatter,
                "token embedding width must be {expected}, got {actual}"
            ),
            Self::BlockModelWidthMismatch {
                layer,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder block {layer} model width must be {expected}, got {actual}"
            ),
            Self::BlockHeadCountMismatch {
                layer,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder block {layer} head count must be {expected}, got {actual}"
            ),
            Self::BlockFeedForwardWidthMismatch {
                layer,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder block {layer} feed-forward width must be {expected}, got {actual}"
            ),
            Self::BlockPositionCapacityMismatch {
                layer,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder block {layer} position capacity must be {expected}, got {actual}"
            ),
            Self::BlockRopeBaseMismatch {
                layer,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder block {layer} RoPE base must be {expected:?}, got {actual:?}"
            ),
            Self::BlockRmsEpsilonMismatch {
                layer,
                norm,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder block {layer} {norm} RMSNorm epsilon must be {expected:?}, got {actual:?}"
            ),
            Self::FinalNormWidthMismatch { expected, actual } => write!(
                formatter,
                "final RMSNorm width must be {expected}, got {actual}"
            ),
            Self::FinalNormEpsilonMismatch { expected, actual } => write!(
                formatter,
                "final RMSNorm epsilon must be {expected:?}, got {actual:?}"
            ),
            Self::ParameterNameMismatch {
                index,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder parameter {index} must be named {expected:?}, got {actual:?}"
            ),
            Self::TokenRank { rank } => write!(
                formatter,
                "decoder token shape must have rank two [batch, tokens], got rank {rank}"
            ),
            Self::EmptyBatch => formatter.write_str("decoder batch must be nonempty"),
            Self::EmptyTokens => formatter.write_str("decoder token sequence must be nonempty"),
            Self::TokenCountOverflow { batch, tokens } => write!(
                formatter,
                "decoder token count overflows for batch {batch} and length {tokens}"
            ),
            Self::TokenCountMismatch { expected, actual } => write!(
                formatter,
                "decoder token shape needs {expected} IDs, but received {actual}"
            ),
            Self::ContextLengthExceeded {
                tokens,
                max_positions,
            } => write!(
                formatter,
                "decoder sequence length {tokens} exceeds context capacity {max_positions}"
            ),
            Self::TargetCountMismatch { expected, actual } => write!(
                formatter,
                "decoder loss needs {expected} targets, but received {actual}"
            ),
            Self::TargetIdOutOfBounds {
                position,
                id,
                vocabulary_size,
            } => write!(
                formatter,
                "target token ID {id} at flat position {position} is out of bounds for vocabulary size {vocabulary_size}"
            ),
            Self::Autodiff { stage, source } => write!(formatter, "decoder {stage}: {source}"),
        }
    }
}

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

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

fn autodiff_error(
    stage: DecoderModelStage,
) -> impl FnOnce(TensorAutodiffError) -> DecoderModelError {
    move |source| DecoderModelError::Autodiff { stage, source }
}

DecoderModelConfig validates all architecture settings even when $N=0$ . An empty stack therefore cannot hide an invalid head layout, zero vocabulary, bad context capacity, nonfinite RoPE base, or invalid RMSNorm epsilon.

Validate vocabulary, dimensions, depth, context, RoPE, and normalization before assembly rust/crates/llm-from-scratch/src/models/decoder.rs#decoder-model-config

/// Every dimension and numerical constant owned by one decoder model.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DecoderModelConfig {
    vocabulary_size: usize,
    model_width: usize,
    heads: usize,
    feed_forward_width: usize,
    layers: usize,
    max_positions: usize,
    rope_base: f64,
    rms_epsilon: f64,
}

impl DecoderModelConfig {
    #[allow(clippy::too_many_arguments)]
    pub const fn new(
        vocabulary_size: usize,
        model_width: usize,
        heads: usize,
        feed_forward_width: usize,
        layers: usize,
        max_positions: usize,
        rope_base: f64,
        rms_epsilon: f64,
    ) -> Self {
        Self {
            vocabulary_size,
            model_width,
            heads,
            feed_forward_width,
            layers,
            max_positions,
            rope_base,
            rms_epsilon,
        }
    }

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

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

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

    pub const fn head_width(self) -> Option<usize> {
        if self.heads == 0 || !self.model_width.is_multiple_of(self.heads) {
            None
        } else {
            Some(self.model_width / self.heads)
        }
    }

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

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

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

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

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

    pub const fn block_config(self) -> DecoderBlockConfig {
        DecoderBlockConfig::new(
            self.model_width,
            self.heads,
            self.feed_forward_width,
            self.max_positions,
            self.rope_base,
            self.rms_epsilon,
        )
    }
}

fn validate_config(config: DecoderModelConfig) -> Result<(), DecoderModelError> {
    if config.vocabulary_size == 0 {
        return Err(DecoderModelError::EmptyVocabulary);
    }
    if config.model_width == 0 {
        return Err(DecoderModelError::ZeroModelWidth);
    }
    if config.heads == 0 {
        return Err(DecoderModelError::ZeroHeadCount);
    }
    if !config.model_width.is_multiple_of(config.heads) {
        return Err(DecoderModelError::ModelWidthNotDivisible {
            model_width: config.model_width,
            heads: config.heads,
        });
    }
    let head_width = config.model_width / config.heads;
    if !head_width.is_multiple_of(2) {
        return Err(DecoderModelError::OddHeadWidth { head_width });
    }
    if config.feed_forward_width == 0 {
        return Err(DecoderModelError::ZeroFeedForwardWidth);
    }
    if config.max_positions == 0 {
        return Err(DecoderModelError::ZeroPositionCapacity);
    }
    if !config.rope_base.is_finite() || config.rope_base <= 0.0 {
        return Err(DecoderModelError::InvalidRopeBase {
            value: config.rope_base,
        });
    }
    if !config.rms_epsilon.is_finite() || config.rms_epsilon < 0.0 {
        return Err(DecoderModelError::InvalidRmsEpsilon {
            value: config.rms_epsilon,
        });
    }
    Ok(())
}

fn expected_parameter_tensors(layers: usize) -> Result<usize, DecoderModelError> {
    layers
        .checked_mul(BLOCK_PARAMETER_SUFFIXES.len())
        .and_then(|count| count.checked_add(2))
        .ok_or(DecoderModelError::ParameterAllocationFailed {
            tensors: usize::MAX,
        })
}

Construction owns token_embedding.weight, then nine stable parameters for each entry under blocks.N, then final_norm.gain. The two-block fixture has $20$ tensors and $264$ scalars. There is deliberately no output-head parameter. Initialization uses a trial random stream and commits the caller’s stream only after all components and cross-component invariants pass.

Stable order also defines one reusable decoder parameter layout. Before from_parameters binds an existing parameter list to model components, validate_parameter_layout checks the configuration, the exact tensor count, the stable name at every list index, and the shape required by the corresponding component. At list index zero, token_embedding.weight is the sole embedding and output-projection table slot; there is no separate output-projection slot.

The layout check borrows each name and tensor only for the duration of its inspection. It does not create parameter nodes, component handles, or tensor copies. Passing the check therefore proves only that the list has the required layout. from_parameters then gives the embedding, blocks, and final RMSNorm shared handles to the validated parameter leaves. Only this binding step makes lookup and output projection use the same live embedding node.

forward accepts only nonempty rank-two token IDs with $T\leq\mathrm{max\_positions}$ . It looks up rows, calls each block at position offset zero, applies final RMSNorm, transposes the original table leaf, and multiplies directly to obtain logits. loss first checks one valid target per token, then applies indexed mean negative log likelihood on vocabulary axis $2$ .

Compose lookup, distinct blocks, final RMSNorm, tied logits, indexed loss, and stable parameters rust/crates/llm-from-scratch/src/models/decoder.rs#decoder-model-layer

/// Inspectable values from lookup through the tied vocabulary projection.
#[derive(Clone, Debug)]
pub struct DecoderModelForward {
    embedding: TensorValue,
    blocks: Vec<DecoderBlockForward>,
    final_norm: RmsNormForward,
    logits: TensorValue,
}

impl DecoderModelForward {
    pub fn embedding(&self) -> &TensorValue {
        &self.embedding
    }

    pub fn blocks(&self) -> &[DecoderBlockForward] {
        &self.blocks
    }

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

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

    pub fn into_logits(self) -> TensorValue {
        self.logits
    }
}

/// Token lookup, repeated decoder blocks, final RMSNorm, and one tied head.
#[derive(Clone, Debug)]
pub struct DecoderModel {
    config: DecoderModelConfig,
    embedding: Embedding,
    blocks: Vec<DecoderBlock>,
    final_norm: RmsNorm,
    parameters: NamedParameters,
}

impl DecoderModel {
    /// Initializes the full model transactionally from one deterministic stream.
    pub fn new(
        config: DecoderModelConfig,
        rng: &mut SplitMix64,
    ) -> Result<Self, DecoderModelError> {
        validate_config(config)?;
        let mut trial = rng.clone();
        let embedding = Embedding::new(
            "token_embedding.weight",
            config.vocabulary_size,
            config.model_width,
            &mut trial,
        )
        .map_err(DecoderModelError::Embedding)?;
        let mut blocks = Vec::new();
        blocks.try_reserve_exact(config.layers).map_err(|_| {
            DecoderModelError::LayerAllocationFailed {
                layers: config.layers,
            }
        })?;
        for layer in 0..config.layers {
            blocks.push(
                DecoderBlock::new(format!("blocks.{layer}"), config.block_config(), &mut trial)
                    .map_err(|source| DecoderModelError::Block { layer, source })?,
            );
        }
        let final_norm = RmsNorm::new("final_norm.gain", config.model_width, config.rms_epsilon)
            .map_err(DecoderModelError::FinalNorm)?;
        let model = Self::from_parts(config, embedding, blocks, final_norm)?;
        *rng = trial;
        Ok(model)
    }

    // region:decoder-parameter-rebuild
    /// Rebuilds every component handle from one exact stable-order parameter set.
    ///
    /// State restoration uses this construction boundary to create an isolated
    /// decoder. Ordinary optimizer steps instead update the existing leaves, so
    /// the registry, components, and tied embedding keep their live aliases.
    pub fn from_parameters(
        config: DecoderModelConfig,
        parameters: Vec<NamedParameter>,
    ) -> Result<Self, DecoderModelError> {
        validate_parameter_layout(config, parameters.as_slice())?;
        let expected = parameters.len();

        let embedding = Embedding::from_parameter(parameters[0].clone())
            .map_err(DecoderModelError::Embedding)?;
        let mut blocks = Vec::new();
        blocks.try_reserve_exact(config.layers).map_err(|_| {
            DecoderModelError::LayerAllocationFailed {
                layers: config.layers,
            }
        })?;
        for layer in 0..config.layers {
            let start = 1 + layer * BLOCK_PARAMETER_SUFFIXES.len();
            let attention_norm = RmsNorm::from_gain(parameters[start].clone(), config.rms_epsilon)
                .map_err(|source| DecoderModelError::Block {
                    layer,
                    source: DecoderBlockError::AttentionNorm(source),
                })?;
            let attention = MultiHeadAttention::from_parameters(
                parameters[start + 1].clone(),
                parameters[start + 2].clone(),
                parameters[start + 3].clone(),
                parameters[start + 4].clone(),
                config.heads,
                config.max_positions,
                config.rope_base,
            )
            .map_err(|source| DecoderModelError::Block {
                layer,
                source: DecoderBlockError::Attention(source),
            })?;
            let feed_forward_norm =
                RmsNorm::from_gain(parameters[start + 5].clone(), config.rms_epsilon).map_err(
                    |source| DecoderModelError::Block {
                        layer,
                        source: DecoderBlockError::FeedForwardNorm(source),
                    },
                )?;
            let feed_forward = SwiGlu::from_parameters(
                parameters[start + 6].clone(),
                parameters[start + 7].clone(),
                parameters[start + 8].clone(),
            )
            .map_err(|source| DecoderModelError::Block {
                layer,
                source: DecoderBlockError::FeedForward(source),
            })?;
            blocks.push(
                DecoderBlock::from_parts(
                    attention_norm,
                    attention,
                    feed_forward_norm,
                    feed_forward,
                )
                .map_err(|source| DecoderModelError::Block { layer, source })?,
            );
        }
        let final_norm = RmsNorm::from_gain(parameters[expected - 1].clone(), config.rms_epsilon)
            .map_err(DecoderModelError::FinalNorm)?;
        Self::from_parts(config, embedding, blocks, final_norm)
    }
    // endregion:decoder-parameter-rebuild

    /// Assembles exact named components and rejects any configuration drift.
    pub fn from_parts(
        config: DecoderModelConfig,
        embedding: Embedding,
        blocks: Vec<DecoderBlock>,
        final_norm: RmsNorm,
    ) -> Result<Self, DecoderModelError> {
        validate_config(config)?;
        if embedding.vocabulary_size() != config.vocabulary_size {
            return Err(DecoderModelError::EmbeddingVocabularyMismatch {
                expected: config.vocabulary_size,
                actual: embedding.vocabulary_size(),
            });
        }
        if embedding.embedding_width() != config.model_width {
            return Err(DecoderModelError::EmbeddingWidthMismatch {
                expected: config.model_width,
                actual: embedding.embedding_width(),
            });
        }
        if blocks.len() != config.layers {
            return Err(DecoderModelError::LayerCountMismatch {
                expected: config.layers,
                actual: blocks.len(),
            });
        }
        for (layer, block) in blocks.iter().enumerate() {
            if block.model_width() != config.model_width {
                return Err(DecoderModelError::BlockModelWidthMismatch {
                    layer,
                    expected: config.model_width,
                    actual: block.model_width(),
                });
            }
            if block.attention().heads() != config.heads {
                return Err(DecoderModelError::BlockHeadCountMismatch {
                    layer,
                    expected: config.heads,
                    actual: block.attention().heads(),
                });
            }
            if block.feed_forward().hidden_width() != config.feed_forward_width {
                return Err(DecoderModelError::BlockFeedForwardWidthMismatch {
                    layer,
                    expected: config.feed_forward_width,
                    actual: block.feed_forward().hidden_width(),
                });
            }
            if block.attention().rope().max_positions() != config.max_positions {
                return Err(DecoderModelError::BlockPositionCapacityMismatch {
                    layer,
                    expected: config.max_positions,
                    actual: block.attention().rope().max_positions(),
                });
            }
            if block.attention().rope().base().to_bits() != config.rope_base.to_bits() {
                return Err(DecoderModelError::BlockRopeBaseMismatch {
                    layer,
                    expected: config.rope_base,
                    actual: block.attention().rope().base(),
                });
            }
            for (norm, epsilon) in [
                (
                    DecoderModelNorm::Attention,
                    block.attention_norm().epsilon(),
                ),
                (
                    DecoderModelNorm::FeedForward,
                    block.feed_forward_norm().epsilon(),
                ),
            ] {
                if epsilon.to_bits() != config.rms_epsilon.to_bits() {
                    return Err(DecoderModelError::BlockRmsEpsilonMismatch {
                        layer,
                        norm,
                        expected: config.rms_epsilon,
                        actual: epsilon,
                    });
                }
            }
        }
        if final_norm.feature_width() != config.model_width {
            return Err(DecoderModelError::FinalNormWidthMismatch {
                expected: config.model_width,
                actual: final_norm.feature_width(),
            });
        }
        if final_norm.epsilon().to_bits() != config.rms_epsilon.to_bits() {
            return Err(DecoderModelError::FinalNormEpsilonMismatch {
                expected: config.rms_epsilon,
                actual: final_norm.epsilon(),
            });
        }

        let parameter_tensors = expected_parameter_tensors(config.layers)?;
        let mut listed = Vec::new();
        listed.try_reserve_exact(parameter_tensors).map_err(|_| {
            DecoderModelError::ParameterAllocationFailed {
                tensors: parameter_tensors,
            }
        })?;
        listed.push(embedding.table().clone());
        for block in &blocks {
            listed.extend(block.parameters().iter().cloned());
        }
        listed.push(final_norm.gain().clone());
        validate_parameter_names(listed.as_slice(), config.layers)?;
        let parameters = NamedParameters::try_new(listed)?;

        Ok(Self {
            config,
            embedding,
            blocks,
            final_norm,
            parameters,
        })
    }

    /// Runs the complete model while retaining one evidence record per layer.
    pub fn forward(
        &self,
        token_ids: &[u32],
        token_shape: &[usize],
    ) -> Result<DecoderModelForward, DecoderModelError> {
        self.validate_tokens(token_ids, token_shape)?;
        let embedding = self
            .embedding
            .forward(token_ids, token_shape)
            .map_err(DecoderModelError::Embedding)?;
        let mut current = embedding.clone();
        let mut block_forwards = Vec::new();
        block_forwards
            .try_reserve_exact(self.blocks.len())
            .map_err(|_| DecoderModelError::LayerAllocationFailed {
                layers: self.blocks.len(),
            })?;
        for (layer, block) in self.blocks.iter().enumerate() {
            let forward = block
                .forward(&current, 0)
                .map_err(|source| DecoderModelError::Block { layer, source })?;
            current = forward.output().clone();
            block_forwards.push(forward);
        }
        let final_norm = self
            .final_norm
            .forward_with_intermediates(&current)
            .map_err(DecoderModelError::FinalNorm)?;
        let tied_weight = self
            .embedding
            .table()
            .tensor()
            .transpose(0, 1)
            .map_err(autodiff_error(DecoderModelStage::TiedWeightTranspose))?;
        let logits = final_norm
            .output()
            .matmul(&tied_weight)
            .map_err(autodiff_error(DecoderModelStage::TiedVocabularyProjection))?;

        Ok(DecoderModelForward {
            embedding,
            blocks: block_forwards,
            final_norm,
            logits,
        })
    }

    /// Computes one mean next-token loss over the vocabulary axis.
    pub fn loss(
        &self,
        token_ids: &[u32],
        token_shape: &[usize],
        targets: &[u32],
    ) -> Result<TensorValue, DecoderModelError> {
        let expected = self.validate_tokens(token_ids, token_shape)?;
        if targets.len() != expected {
            return Err(DecoderModelError::TargetCountMismatch {
                expected,
                actual: targets.len(),
            });
        }
        let mut target_indices = Vec::new();
        target_indices
            .try_reserve_exact(expected)
            .map_err(|_| DecoderModelError::TargetAllocationFailed { targets: expected })?;
        for (position, &id) in targets.iter().enumerate() {
            let valid = usize::try_from(id)
                .ok()
                .is_some_and(|index| index < self.config.vocabulary_size);
            if !valid {
                return Err(DecoderModelError::TargetIdOutOfBounds {
                    position,
                    id,
                    vocabulary_size: self.config.vocabulary_size,
                });
            }
            target_indices
                .push(usize::try_from(id).expect("validated target token ID must fit usize"));
        }
        self.forward(token_ids, token_shape)?
            .into_logits()
            .indexed_mean_nll(2, &target_indices)
            .map_err(autodiff_error(DecoderModelStage::IndexedMeanLoss))
    }

    fn validate_tokens(
        &self,
        token_ids: &[u32],
        token_shape: &[usize],
    ) -> Result<usize, DecoderModelError> {
        if token_shape.len() != 2 {
            return Err(DecoderModelError::TokenRank {
                rank: token_shape.len(),
            });
        }
        let batch = token_shape[0];
        let tokens = token_shape[1];
        if batch == 0 {
            return Err(DecoderModelError::EmptyBatch);
        }
        if tokens == 0 {
            return Err(DecoderModelError::EmptyTokens);
        }
        if tokens > self.config.max_positions {
            return Err(DecoderModelError::ContextLengthExceeded {
                tokens,
                max_positions: self.config.max_positions,
            });
        }
        let expected = batch
            .checked_mul(tokens)
            .ok_or(DecoderModelError::TokenCountOverflow { batch, tokens })?;
        if token_ids.len() != expected {
            return Err(DecoderModelError::TokenCountMismatch {
                expected,
                actual: token_ids.len(),
            });
        }
        Ok(expected)
    }

    pub const fn config(&self) -> DecoderModelConfig {
        self.config
    }

    pub const fn embedding(&self) -> &Embedding {
        &self.embedding
    }

    pub fn blocks(&self) -> &[DecoderBlock] {
        &self.blocks
    }

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

    /// The sole table used for both token lookup and vocabulary projection.
    pub const fn tied_embedding(&self) -> &NamedParameter {
        self.embedding.table()
    }

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

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

The focused numerical proof checks all $20$ coordinates of the tied table and all $4$ coordinates of the final gain against central differences. Its bound is

\tau=2\times10^{-5}.

The full two-block loss also gives finite gradients to all $20$ registered parameter tensors.

Check the tied table and final gain numerically, then inspect every stacked parameter gradient rust/demos/ch32-decoder-model/src/lib.rs#gradient-checks

#[derive(Clone, Debug, PartialEq)]
pub struct GradientEvidence {
    pub tied_table_checks: usize,
    pub final_norm_checks: usize,
    pub tolerance: f64,
    pub passed: bool,
    pub stack_parameter_tensors: usize,
    pub stack_gradient_tensors: usize,
    pub decomposition_error: f64,
}

fn gradient_evidence(model: &DecoderModel) -> Result<GradientEvidence, Box<dyn Error>> {
    let probe = zero_layer_model(
        tensor(&[VOCABULARY, MODEL_WIDTH], &PROBE_TABLE),
        tensor(&[MODEL_WIDTH], &PROBE_GAIN),
    );
    probe.loss(&[0, 1], &[1, 2], &[1, 2])?.backward()?;
    let table_gradient = probe
        .tied_embedding()
        .tensor()
        .gradient()
        .expect("probe tied table must receive a gradient");
    let gain_gradient = probe
        .final_norm()
        .gain()
        .tensor()
        .gradient()
        .expect("probe final gain must receive a gradient");
    let table_report = sampled_tensor_gradient_check(
        &mut tensor(&[VOCABULARY, MODEL_WIDTH], &PROBE_TABLE),
        &table_gradient.view(),
        STEP,
        TOLERANCE,
        PROBE_TABLE.len(),
        |candidate| zero_layer_loss(candidate.clone(), tensor(&[MODEL_WIDTH], &PROBE_GAIN)),
    )?;
    let gain_report = sampled_tensor_gradient_check(
        &mut tensor(&[MODEL_WIDTH], &PROBE_GAIN),
        &gain_gradient.view(),
        STEP,
        TOLERANCE,
        PROBE_GAIN.len(),
        |candidate| {
            zero_layer_loss(
                tensor(&[VOCABULARY, MODEL_WIDTH], &PROBE_TABLE),
                candidate.clone(),
            )
        },
    )?;

    model
        .loss(&TOKEN_IDS, &[BATCH, TOKENS], &TARGET_IDS)?
        .backward()?;
    let stack_gradient_tensors = model
        .parameters()
        .iter()
        .filter(|parameter| {
            parameter
                .tensor()
                .gradient()
                .is_some_and(|gradient| gradient.as_slice().iter().all(|value| value.is_finite()))
        })
        .count();

    Ok(GradientEvidence {
        tied_table_checks: table_report.checks.len(),
        final_norm_checks: gain_report.checks.len(),
        tolerance: TOLERANCE,
        passed: table_report.passed
            && gain_report.passed
            && stack_gradient_tensors == model.parameters().len(),
        stack_parameter_tensors: model.parameters().len(),
        stack_gradient_tensors,
        decomposition_error: tied_gradient_decomposition_error(),
    })
}

The deterministic fixture collects exact stage values, logits, mean loss, parameter ownership, depth checks, named invalid inputs, causality, gradient checks, and bitwise replay.

Collect one coherent evidence record for the complete decoder boundary rust/demos/ch32-decoder-model/src/lib.rs#learner-evidence

pub fn learner_evidence() -> Result<LearnerEvidence, Box<dyn Error>> {
    let model = initialized_model(LAYERS)?;
    let forward = model.forward(&TOKEN_IDS, &[BATCH, TOKENS])?;
    let mut stages = Vec::with_capacity(LAYERS + 2);
    stages.push(stage("embedding", forward.embedding()));
    for (layer, block) in forward.blocks().iter().enumerate() {
        stages.push(stage(format!("block-{layer}"), block.output()));
    }
    stages.push(stage("final-norm", forward.final_norm().output()));
    let logits = forward.logits().value_snapshot();
    let loss = model
        .loss(&TOKEN_IDS, &[BATCH, TOKENS], &TARGET_IDS)?
        .value()
        .as_slice()[0];
    let replay = initialized_model(LAYERS)?
        .forward(&TOKEN_IDS, &[BATCH, TOKENS])?
        .logits()
        .value_snapshot();
    let expected_names = model
        .parameters()
        .iter()
        .map(|parameter| parameter.name().to_owned())
        .collect::<Vec<_>>();
    let untied_parameter_scalars = model.parameter_count() + VOCABULARY * MODEL_WIDTH;

    Ok(LearnerEvidence {
        stages,
        predictions: predictions(&logits),
        logits: logits.clone(),
        loss,
        parameter_names: expected_names.clone(),
        parameter_scalars: model.parameter_count(),
        untied_parameter_scalars,
        tied_parameter_name: model.tied_embedding().name().to_owned(),
        tied_lookup_and_head: model
            .embedding()
            .table()
            .tensor()
            .is_same_node(model.tied_embedding().tensor()),
        bias_free: expected_names.iter().all(|name| !name.contains("bias")),
        stable_order: expected_names
            .first()
            .is_some_and(|name| name == "token_embedding.weight")
            && expected_names
                .last()
                .is_some_and(|name| name == "final_norm.gain"),
        depths_valid: shape_evidence()?,
        errors: error_evidence(),
        causality: causality_evidence(&model)?,
        gradients: gradient_evidence(&model)?,
        replay_bitwise: logits == replay,
    })
}

Format the exact decoder-model report while deferring optimization to Chapter 33 rust/demos/ch32-decoder-model/src/lib.rs#learner-report

pub fn render_report(evidence: &LearnerEvidence) -> String {
    let token_one = &evidence.logits.as_slice()[VOCABULARY..2 * VOCABULARY];
    format!(
        concat!(
            "chapter=32-decoder-model\n",
            "config=batch:{batch} tokens:{tokens} vocabulary:{vocabulary} model_width:{model_width} layers:{layers} heads:{heads} head_width:{head_width} feed_forward_width:{feed_forward_width} context:{max_positions}\n",
            "shape=embedding:{embedding:?} block_0:{block_0:?} block_1:{block_1:?} final_norm:{final_norm:?} logits:{logits:?}\n",
            "token_1_logits={token_one}\n",
            "targets=[1,2,3] mean_loss:{loss:.6}\n",
            "prediction=token_0:{prediction_0} token_1:{prediction_1} token_2:{prediction_2}\n",
            "tying=name:{tied_name} lookup_and_head:{tied} gradient_roles:lookup+output decomposition_error:{decomposition_error:.12}\n",
            "parameters=tensors:{parameter_tensors} scalars:{parameter_scalars} untied_scalars:{untied_scalars} saved:{saved} bias_free:{bias_free} stable_order:{stable_order}\n",
            "depths=zero_one_two:{depths} configuration_errors:{configuration} context_limit:{context} vocabulary_errors:{vocabulary_error} target_errors:{target}\n",
            "causality=prefix_0_bitwise:{prefix_0} prefix_1_bitwise:{prefix_1} suffix_changed:{suffix}\n",
            "gradcheck=tied_table:{table_checks} final_norm:{gain_checks} total:{total_checks} tolerance:{tolerance:.6} passed:{passed} stack_gradients:{stack_gradients}/{stack_parameters}\n",
            "replay=bitwise:{replay}\n",
            "next=train this decoder and select a state with validation loss only\n",
        ),
        batch = BATCH,
        tokens = TOKENS,
        vocabulary = VOCABULARY,
        model_width = MODEL_WIDTH,
        layers = LAYERS,
        heads = HEADS,
        feed_forward_width = FEED_FORWARD_WIDTH,
        max_positions = MAX_POSITIONS,
        head_width = MODEL_WIDTH / HEADS,
        embedding = evidence.stages[0].shape,
        block_0 = evidence.stages[1].shape,
        block_1 = evidence.stages[2].shape,
        final_norm = evidence.stages[3].shape,
        logits = evidence.logits.shape(),
        token_one = values_text(token_one),
        loss = evidence.loss,
        prediction_0 = evidence.predictions[0],
        prediction_1 = evidence.predictions[1],
        prediction_2 = evidence.predictions[2],
        tied_name = evidence.tied_parameter_name,
        tied = evidence.tied_lookup_and_head,
        decomposition_error = evidence.gradients.decomposition_error,
        parameter_tensors = evidence.parameter_names.len(),
        parameter_scalars = evidence.parameter_scalars,
        untied_scalars = evidence.untied_parameter_scalars,
        saved = evidence.untied_parameter_scalars - evidence.parameter_scalars,
        bias_free = evidence.bias_free,
        stable_order = evidence.stable_order,
        depths = evidence.depths_valid,
        configuration = evidence.errors.configuration,
        context = evidence.errors.context,
        vocabulary_error = evidence.errors.vocabulary,
        target = evidence.errors.target,
        prefix_0 = evidence.causality.prefix_0_bitwise,
        prefix_1 = evidence.causality.prefix_1_bitwise,
        suffix = evidence.causality.suffix_changed,
        table_checks = evidence.gradients.tied_table_checks,
        gain_checks = evidence.gradients.final_norm_checks,
        total_checks = evidence.gradients.tied_table_checks + evidence.gradients.final_norm_checks,
        tolerance = evidence.gradients.tolerance,
        passed = evidence.gradients.passed,
        stack_gradients = evidence.gradients.stack_gradient_tensors,
        stack_parameters = evidence.gradients.stack_parameter_tensors,
        replay = evidence.replay_bitwise,
    )
}

The executable prints only the frozen report:

Print the frozen Chapter 32 decoder-model learner report rust/demos/ch32-decoder-model/src/main.rs

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

Run cargo run --quiet --locked -p ch32-decoder-model. Its standard output matches rust/demos/ch32-decoder-model/expected.txt byte for byte, including the final newline.

Inspect one table at both ends of the forward path

The deterministic run records the model configuration, token IDs, every feature row after lookup and each stacked stage, all vocabulary logits, targets, predictions, mean loss, tying evidence, parameter counts, input boundaries, causality, gradient checks, and replay status.

Emit the exact twenty-eight-line decoder-model evidence record rust/demos/ch32-decoder-model/src/diagram_trace.rs

//! Locale-neutral static trace consumed by the Chapter 32 diagram parser.

use std::fmt::Write;

use crate::{
    BATCH, FEED_FORWARD_WIDTH, HEADS, LAYERS, LearnerEvidence, MAX_POSITIONS, MODEL_WIDTH,
    TARGET_IDS, TOKEN_IDS, TOKENS, VOCABULARY,
};

fn values(values: &[f64]) -> String {
    format!(
        "[{}]",
        values
            .iter()
            .map(|value| format!("{value:.6}"))
            .collect::<Vec<_>>()
            .join(",")
    )
}

fn shape(shape: &[usize]) -> String {
    format!(
        "[{}]",
        shape
            .iter()
            .map(usize::to_string)
            .collect::<Vec<_>>()
            .join(",")
    )
}

pub fn render_trace(evidence: &LearnerEvidence) -> String {
    let mut output = String::new();
    writeln!(output, "DECODER_MODEL_TRACE_V1").unwrap();
    writeln!(
        output,
        "config batch={BATCH} tokens={TOKENS} vocabulary={VOCABULARY} model_width={MODEL_WIDTH} layers={LAYERS} heads={HEADS} feed_forward_width={FEED_FORWARD_WIDTH} context={MAX_POSITIONS}"
    )
    .unwrap();
    writeln!(
        output,
        "tokens shape=[{BATCH},{TOKENS}] values=[{},{},{}]",
        TOKEN_IDS[0], TOKEN_IDS[1], TOKEN_IDS[2]
    )
    .unwrap();
    writeln!(
        output,
        "targets values=[{},{},{}]",
        TARGET_IDS[0], TARGET_IDS[1], TARGET_IDS[2]
    )
    .unwrap();
    for stage in &evidence.stages {
        for token in 0..TOKENS {
            let start = token * MODEL_WIDTH;
            writeln!(
                output,
                "stage name={} shape={} token={} values={}",
                stage.name,
                shape(&stage.shape),
                token,
                values(&stage.values.as_slice()[start..start + MODEL_WIDTH])
            )
            .unwrap();
        }
    }
    for token in 0..TOKENS {
        let start = token * VOCABULARY;
        writeln!(
            output,
            "logits token={} values={}",
            token,
            values(&evidence.logits.as_slice()[start..start + VOCABULARY])
        )
        .unwrap();
    }
    writeln!(
        output,
        "predictions values=[{},{},{}]",
        evidence.predictions[0], evidence.predictions[1], evidence.predictions[2]
    )
    .unwrap();
    writeln!(output, "loss mean={:.6}", evidence.loss).unwrap();
    writeln!(
        output,
        "tying name={} lookup_and_head={} gradient_roles=lookup+output decomposition_error={:.12}",
        evidence.tied_parameter_name,
        evidence.tied_lookup_and_head,
        evidence.gradients.decomposition_error
    )
    .unwrap();
    writeln!(
        output,
        "parameters tensors={} scalars={} untied_scalars={} saved={} bias_free={} stable_order={}",
        evidence.parameter_names.len(),
        evidence.parameter_scalars,
        evidence.untied_parameter_scalars,
        evidence.untied_parameter_scalars - evidence.parameter_scalars,
        evidence.bias_free,
        evidence.stable_order
    )
    .unwrap();
    writeln!(
        output,
        "depths zero_one_two={} configuration_errors={} context_limit={} vocabulary_errors={} target_errors={}",
        evidence.depths_valid,
        evidence.errors.configuration,
        evidence.errors.context,
        evidence.errors.vocabulary,
        evidence.errors.target
    )
    .unwrap();
    writeln!(
        output,
        "causality prefix_0_bitwise={} prefix_1_bitwise={} suffix_changed={}",
        evidence.causality.prefix_0_bitwise,
        evidence.causality.prefix_1_bitwise,
        evidence.causality.suffix_changed
    )
    .unwrap();
    writeln!(
        output,
        "gradcheck tied_table={} final_norm={} total={} tolerance={:.6} passed={} stack_gradients={}/{}",
        evidence.gradients.tied_table_checks,
        evidence.gradients.final_norm_checks,
        evidence.gradients.tied_table_checks + evidence.gradients.final_norm_checks,
        evidence.gradients.tolerance,
        evidence.gradients.passed,
        evidence.gradients.stack_gradient_tensors,
        evidence.gradients.stack_parameter_tensors
    )
    .unwrap();
    writeln!(output, "replay bitwise={}", evidence.replay_bitwise).unwrap();
    writeln!(output, "END_DECODER_MODEL_TRACE").unwrap();
    output
}

Follow one tied vocabulary table through a complete decoder

Read exact Rust-authored token rows through lookup, two distinct causal blocks, final RMSNorm, and the transpose projection back to five vocabulary logits.

Double edge: repeated blocks with distinct weights
Dashed edge: two uses of one parameter
Double underline: verified Rust evidence

One decoder path with one table in two roles

The residual width remains four through lookup, both blocks, and final normalization; only the tied projection changes the final axis to vocabulary width five.

Token IDs

z=[0,1,2]

[1,3]

Embedding lookup

E[z]

[1,3,4]

Two distinct decoder blocks

$B_1$ $B_2$

blocks.0, blocks.1

[1,3,4]

Final RMSNorm

\operatorname{RMSNorm}

[1,3,4]

Tied projection

E^\top

Vocabulary logits

\ell

[1,3,5]

One parameter, used twice

token_embedding.weight

Lookup role: $E[z]$ Transpose projection role: $E^\top$

Inspect representative hidden rows and every logit

Every stage row and logit comes from the same deterministic run, so the tables preserve one coherent path from lookup to loss.

Scrollable representative decoder-stage evidence table
Stage	Token position	Shape	Exact feature row
Embedding lookup	$t=1$	$[1,3,4]$	$[0.104430,$ $0.699156,$ $0.051971,$ $-0.160217]$
After decoder block 1	$t=1$	$[1,3,4]$	$[-0.772509,$ $0.841246,$ $1.816030,$ $0.156071]$
After decoder block 2	$t=1$	$[1,3,4]$	$[-0.183854,$ $0.605829,$ $0.709748,$ $-0.521934]$
After final RMSNorm	$t=1$	$[1,3,4]$	$[-0.338936,$ $1.116846,$ $1.308420,$ $-0.962186]$

Scrollable vocabulary-logit evidence table
Token position	$v=0$	$v=1$	$v=2$	$v=3$	$v=4$	Prediction	Target
$t=0$	$0.386115$	$0.276091$	$0.301266$	$-0.460642$	$-0.735173$	$\operatorname*{argmax}=0$	$y=1$
$t=1$	$-0.862249$	$0.967613$	$-0.991545$	$-0.446363$	$1.234533$	$\operatorname*{argmax}=4$	$y=2$
$t=2$	$-0.220241$	$0.332921$	$-0.193218$	$-0.267554$	$-0.500487$	$\operatorname*{argmax}=1$	$y=3$

Mean indexed loss: $\mathcal{L}=2.045535$

Verify ownership, tied gradients, causality, and depth boundaries

Parameter counts, role decomposition, prefix invariance, valid depths, and named input failures all come from the same deterministic fixture.

One stable parameter list

$N_\theta=264$

Scalars saved by tying: $284-264=20$

tensors=20, bias_free=true

Both roles accumulate on one leaf

$\bar E=\bar E_{\mathrm{lookup}}+\bar E_{\mathrm{output}}$

$\Delta=0.000000000000$

checked=24, stack=20/20

Changing only the suffix

$t=0$ : Bitwise unchanged

$t=1$ : Bitwise unchanged

$t=2$ : Numerically changed

Depth and boundary checks

$N\in\{0,1,2\}$ : Double underline: verified Rust evidence

configuration=true context=true vocabulary=true targets=true

The figure contains one card for token_embedding.weight. The lookup relation uses $E[z]$ near the beginning of the path, and the projection relation uses $E^\top$ after final normalization. That single-card structure matters: two separate matrix cards would suggest copied parameters. Double and dashed cues retain the repeated-versus-tied distinction without color.

Read token position $1$ downward through the evidence table. Lookup begins at

[0.104430,\ 0.699156,\ 0.051971,\ -0.160217].

The two distinct blocks produce

[-0.772509,\ 0.841246,\ 1.816030,\ 0.156071]

and then

[-0.183854,\ 0.605829,\ 0.709748,\ -0.521934].

Final RMSNorm gives

[-0.338936,\ 1.116846,\ 1.308420,\ -0.962186],

which the tied transpose maps to the five logits shown earlier. These rows, predictions, the mean loss, and the tied-gradient proof all belong to that same forward and backward run; none is an independently estimated illustration.

Test axes, ownership, causality, and failures

For $B=2$ , $T=6$ , $d_{\mathrm{model}}=12$ , and $V=40$ , write the shape after lookup, each block, final RMSNorm, and the tied projection.
For $V=40$ , $d_{\mathrm{model}}=12$ , $d_{\mathrm{ff}}=32$ , and $N=3$ , compute the tied model’s parameter count. How many scalars would an untied vocabulary head add?
List the first, tenth, and last parameter names for a two-block model. Why is there no lm_head.weight?
Draw both tape paths that reach token_embedding.weight. Which rows can the lookup role update, and which rows can the classifier role update?
Change only the final token in a three-token input. Which logits must remain bitwise unchanged?
Classify these failures: zero heads, token length above capacity, token ID $V$ , and one missing target.

Check the structural answers

Lookup, every block, and final RMSNorm have shape $[2,6,12]$ . The tied projection has shape $[2,6,40]$ .
The tied count is $40(12)+3(4(12)^2+3(12)(32)+2(12))+12=5748$ . An untied head adds $40(12)=480$ scalars.
The first is token_embedding.weight, the tenth is blocks.0.ffn.down.weight, and the last is final_norm.gain. The transpose of the first leaf serves as the head, so no head parameter exists.
Lookup directly touches rows selected by input IDs. The output role can contribute to every vocabulary row because every row receives a logit. Reverse mode adds both contributions on the one table.
Logit rows $0$ and $1$ remain bitwise unchanged; row $2$ may change. Causal masking inside every repeated block preserves that prefix boundary.
Zero heads is a configuration failure, excess token length is a context failure, ID $V$ is a vocabulary-input failure, and a missing target is a target-shape failure.

Misconception: tying copies the embedding matrix into an output head. Correction: there is one table leaf and two differentiable uses. A separate but equal matrix would have separate identity, parameter ownership, and gradients.

Train this exact model at the next boundary

The cumulative implementation now produces differentiable next-token logits and mean indexed loss from token IDs; Chapter 33 will train this exact model with a bounded deterministic loop and choose one state using validation loss only.

This chapter stops before optimization. It owns architecture, parameter identity, logits, indexed loss, causality, and gradients, but it does not clip gradients, apply AdamW, zero gradients between steps, evaluate validation data, select a checkpoint, inspect test data, generate tokens, or manage a KV cache. Those responsibilities stay visible when Chapter 33 adds the training loop.