23 · Content revision 4

Train a fixed-context neural language model

Assemble embeddings, a SwiGLU hidden layer, indexed next-token loss, mini-batches, and AdamW into a deterministic neural n-gram whose held-out loss improves.

Predict one token from the complete context

Take one complete two-token context $[z_{t-2},z_{t-1}]$ and its single following target $z_t$ . With $D=4$ , $H=8$ , and $V=266$ , predict the shapes after lookup, concatenation, SwiGLU, and vocabulary projection before inspecting any numeric output.

For batch size $B=1$ , the answer is:

[1,2]\to[1,2,4]\to[1,8]\to[1,8]\to[1,266].

The frozen prompt At encodes to the literal program IDs [67, 118]. Look up both rows before deciding what the model predicts. The initial logits are nearly equal because the matrices are freshly initialized; the exact seeded $\operatorname{argmax}$ is token ID 44, not a learned linguistic conclusion.

Concatenate embeddings before the hidden layer

Apply the chapter’s exact shared formula:

h=\operatorname{SwiGLU}([E_{z_{t-C}},\ldots,E_{z_{t-1}}]),\quad \ell=hW_o

$E\in\mathbb{R}^{V\times D}$ is one trainable table shared by every token position. Each of the $C$ selectors retrieves a row in $\mathbb{R}^{D}$ . Brackets concatenate those rows in chronological order, so SwiGLU receives $CD$ features. It returns $h\in\mathbb{R}^{H}$ , and $W_o\in\mathbb{R}^{H\times V}$ projects that hidden vector to $\ell\in\mathbb{R}^{V}$ .

For a batch, the final multiplication is $[B,H][H,V]=[B,V]$ . The output has one logit per vocabulary item, not one embedding feature per token.

Keep context, feature, and vocabulary axes distinct

$z_i$ is the integer token ID at sequence position $i$ .
$t$ is the position being predicted, and $C>0$ is its fixed context length.
$E$ maps a token ID to $D$ learned features.
$[\,\cdot\,]$ concatenates the $C$ embedding rows along the feature axis.
$\operatorname{SwiGLU}$ maps $CD$ context features to the hidden vector $h$ .
$W_o$ maps the $H$ hidden features to $V$ vocabulary scores.
$\ell$ is the vector of $V$ next-token logits.

One target $y_b$ belongs to each context row. The indexed mean loss is

L=-\frac{1}{B}\sum_{b=1}^{B}\log \frac{\exp(\ell_{b,y_b})}{\sum_{j=1}^{V}\exp(\ell_{b,j})}.

Chapter 21 stores $C$ shifted targets in every row. This model deliberately uses only $y_b=\operatorname{target\_row}(b)_{C-1}$ , the token immediately after the whole context. Using all $C$ entries would silently change both the objective and its denominator.

From sparse counts to learned contexts and attention

Classical count n-grams estimate each short context separately, so rare or unseen combinations receive little usable evidence as the number of possible sequences grows. In the historical_context_evidence calculation shown below, a bigram keyed only by final token $11$ mixes two followers, while the wider contexts $[10,11]$ and $[20,11]$ each retain one distinct follower.

Bengio et al., A Neural Probabilistic Language Model provide the neural step. Bengio et al. map a fixed context through learned distributed word features and a feed-forward network to a next-word probability distribution. Bengio et al. jointly learn distributed word vectors and a feed-forward function over a concatenated fixed context, so evidence about one word sequence can inform sequences made from nearby word representations.

Vaswani et al., Attention Is All You Need provide the later sequence-model step. Vaswani et al. replace recurrence and convolution with attention, and mask decoder self-attention so a position cannot use later positions during autoregressive prediction. Transformers later replace fixed-context mixing with masked self-attention while retaining learned embeddings, position-wise feed-forward transformations, and a vocabulary projection for next-token prediction.

The neural n-gram is an important integration point on the road to modern LLMs: embeddings share information across token identities, a learned nonlinear map combines the complete context, and next-token loss trains every matrix together. Attention later removes this fixed-context concatenation bottleneck.

The papers do not define this course’s byte-pair vocabulary, SwiGLU, dimensions, optimizer constants, seeds, final-target policy, or generation limit. The historical_context_evidence calculation isolates the language-model transition: exact count tables give way to shared learned features, and masked attention later removes the fixed concatenation window.

Contrast exact count contexts with shared learned features rust/demos/ch23-neural-ngram/src/lib.rs#historical-context-road

/// Count tables isolate exact contexts; learned embeddings can share features.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct HistoricalContextEvidence {
    pub bigram_followers: usize,
    pub first_fixed_context_followers: usize,
    pub second_fixed_context_followers: usize,
    pub neural_context_width: usize,
}

pub fn historical_context_evidence() -> HistoricalContextEvidence {
    let examples = [([10_u32, 11], 12_u32), ([20, 11], 13)];
    let mut bigram_counts = BTreeMap::<u32, BTreeMap<u32, usize>>::new();
    let mut fixed_context_counts = BTreeMap::<[u32; 2], BTreeMap<u32, usize>>::new();
    for (context, target) in examples {
        *bigram_counts
            .entry(context[1])
            .or_default()
            .entry(target)
            .or_default() += 1;
        *fixed_context_counts
            .entry(context)
            .or_default()
            .entry(target)
            .or_default() += 1;
    }
    HistoricalContextEvidence {
        bigram_followers: bigram_counts.get(&11).map_or(0, BTreeMap::len),
        first_fixed_context_followers: fixed_context_counts.get(&[10, 11]).map_or(0, BTreeMap::len),
        second_fixed_context_followers: fixed_context_counts
            .get(&[20, 11])
            .map_or(0, BTreeMap::len),
        neural_context_width: CONTEXT_LENGTH * EMBEDDING_WIDTH,
    }
}

Own one parameter set across every update

Configuration and errors make invalid model widths, derived-size overflow, parameter mismatches, bad batches, out-of-range targets, and non-finite greedy scores observable before they can be mistaken for training evidence:

Validate model dimensions and typed failure boundaries rust/crates/llm-from-scratch/src/models/neural_ngram.rs#neural-ngram-config-and-errors

/// The four widths that determine every fixed-context model shape.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NeuralNgramConfig {
    vocabulary_size: usize,
    context_length: usize,
    embedding_width: usize,
    hidden_width: usize,
    context_feature_width: usize,
    parameter_count: usize,
}

impl NeuralNgramConfig {
    pub fn new(
        vocabulary_size: usize,
        context_length: usize,
        embedding_width: usize,
        hidden_width: usize,
    ) -> Result<Self, NeuralNgramError> {
        if vocabulary_size == 0 {
            return Err(NeuralNgramError::EmptyVocabulary);
        }
        if context_length == 0 {
            return Err(NeuralNgramError::ZeroContextLength);
        }
        if embedding_width == 0 {
            return Err(NeuralNgramError::ZeroEmbeddingWidth);
        }
        if hidden_width == 0 {
            return Err(NeuralNgramError::ZeroHiddenWidth);
        }
        let context_feature_width = context_length.checked_mul(embedding_width).ok_or(
            NeuralNgramError::ContextFeatureWidthOverflow {
                context_length,
                embedding_width,
            },
        )?;
        let embedding_parameters = vocabulary_size.checked_mul(embedding_width);
        let branch_parameters = context_feature_width.checked_mul(hidden_width);
        let down_parameters = hidden_width.checked_mul(hidden_width);
        let output_parameters = hidden_width.checked_mul(vocabulary_size);
        let parameter_count = embedding_parameters
            .and_then(|count| branch_parameters.and_then(|branch| count.checked_add(branch)))
            .and_then(|count| branch_parameters.and_then(|branch| count.checked_add(branch)))
            .and_then(|count| down_parameters.and_then(|down| count.checked_add(down)))
            .and_then(|count| output_parameters.and_then(|output| count.checked_add(output)))
            .ok_or(NeuralNgramError::ParameterCountOverflow)?;
        Ok(Self {
            vocabulary_size,
            context_length,
            embedding_width,
            hidden_width,
            context_feature_width,
            parameter_count,
        })
    }

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

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

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

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

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

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

/// A rejected model shape, parameter set, batch, selector, or delegated operation.
#[derive(Clone, Debug, PartialEq)]
pub enum NeuralNgramError {
    EmptyVocabulary,
    ZeroContextLength,
    ZeroEmbeddingWidth,
    ZeroHiddenWidth,
    ContextFeatureWidthOverflow {
        context_length: usize,
        embedding_width: usize,
    },
    ParameterCountOverflow,
    EmptyBatch,
    ContextTokenCountOverflow {
        batch_size: usize,
        context_length: usize,
    },
    ContextTokenCountMismatch {
        expected: usize,
        actual: usize,
    },
    BatchContextLengthMismatch {
        expected: usize,
        actual: usize,
    },
    ParameterCountMismatch {
        expected: usize,
        actual: usize,
    },
    ParameterNameMismatch {
        index: usize,
        expected: &'static str,
        actual: String,
    },
    ParameterShapeMismatch {
        name: String,
        expected: Vec<usize>,
        actual: Vec<usize>,
    },
    MissingTargetRow {
        row: usize,
    },
    TargetIdOutOfBounds {
        row: usize,
        id: u32,
        vocabulary_size: usize,
    },
    TargetIdDoesNotFitUsize {
        row: usize,
        id: u32,
    },
    MaskedTokenOutOfBounds {
        id: u32,
        vocabulary_size: usize,
    },
    NoUnmaskedToken,
    NonFiniteLogit {
        token_id: usize,
        value: f64,
    },
    Initialization(InitializationError),
    Embedding(EmbeddingError),
    SwiGlu(SwiGluError),
    Linear(LinearError),
    Autodiff(TensorAutodiffError),
}

impl fmt::Display for NeuralNgramError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EmptyVocabulary => {
                formatter.write_str("neural n-gram vocabulary must not be empty")
            }
            Self::ZeroContextLength => {
                formatter.write_str("neural n-gram context length must be greater than zero")
            }
            Self::ZeroEmbeddingWidth => {
                formatter.write_str("neural n-gram embedding width must be greater than zero")
            }
            Self::ZeroHiddenWidth => {
                formatter.write_str("neural n-gram hidden width must be greater than zero")
            }
            Self::ContextFeatureWidthOverflow {
                context_length,
                embedding_width,
            } => write!(
                formatter,
                "context length {context_length} times embedding width {embedding_width} overflows usize"
            ),
            Self::ParameterCountOverflow => {
                formatter.write_str("neural n-gram parameter count overflows usize")
            }
            Self::EmptyBatch => formatter.write_str("neural n-gram batch must contain a row"),
            Self::ContextTokenCountOverflow {
                batch_size,
                context_length,
            } => write!(
                formatter,
                "batch size {batch_size} times context length {context_length} overflows usize"
            ),
            Self::ContextTokenCountMismatch { expected, actual } => write!(
                formatter,
                "neural n-gram forward needs {expected} context IDs, but received {actual}"
            ),
            Self::BatchContextLengthMismatch { expected, actual } => write!(
                formatter,
                "mini-batch context length must equal model context length {expected}, got {actual}"
            ),
            Self::ParameterCountMismatch { expected, actual } => write!(
                formatter,
                "neural n-gram needs {expected} named parameters, but received {actual}"
            ),
            Self::ParameterNameMismatch {
                index,
                expected,
                actual,
            } => write!(
                formatter,
                "parameter {index} must be named {expected}, got {actual}"
            ),
            Self::ParameterShapeMismatch {
                name,
                expected,
                actual,
            } => write!(
                formatter,
                "parameter {name} must have shape {expected:?}, got {actual:?}"
            ),
            Self::MissingTargetRow { row } => {
                write!(formatter, "mini-batch target row {row} is missing")
            }
            Self::TargetIdOutOfBounds {
                row,
                id,
                vocabulary_size,
            } => write!(
                formatter,
                "target ID {id} in row {row} is out of bounds for vocabulary size {vocabulary_size}"
            ),
            Self::TargetIdDoesNotFitUsize { row, id } => write!(
                formatter,
                "target ID {id} in row {row} does not fit this platform's selector type"
            ),
            Self::MaskedTokenOutOfBounds {
                id,
                vocabulary_size,
            } => write!(
                formatter,
                "masked token ID {id} is out of bounds for vocabulary size {vocabulary_size}"
            ),
            Self::NoUnmaskedToken => {
                formatter.write_str("greedy selection needs at least one unmasked token")
            }
            Self::NonFiniteLogit { token_id, value } => {
                write!(
                    formatter,
                    "logit for token {token_id} is not finite: {value}"
                )
            }
            Self::Initialization(error) => error.fmt(formatter),
            Self::Embedding(error) => error.fmt(formatter),
            Self::SwiGlu(error) => error.fmt(formatter),
            Self::Linear(error) => error.fmt(formatter),
            Self::Autodiff(error) => error.fmt(formatter),
        }
    }
}

impl Error for NeuralNgramError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Initialization(error) => Some(error),
            Self::Embedding(error) => Some(error),
            Self::SwiGlu(error) => Some(error),
            Self::Linear(error) => Some(error),
            Self::Autodiff(error) => Some(error),
            _ => None,
        }
    }
}

The model owns exactly five matrix nodes. Its ordered parameter registry and its persistent embedding, SwiGLU, and output-projection handles all refer to those same nodes. AdamW writes checked new values into the existing nodes, so the next forward observes the update through the already-owned layers. No layer needs to be rebuilt after an optimizer step:

Share one five-node parameter set between the registry and persistent layers rust/crates/llm-from-scratch/src/models/neural_ngram.rs#neural-ngram-parameter-owner

/// One fixed-context language model whose layers share one live parameter registry.
#[derive(Debug)]
pub struct NeuralNgram {
    config: NeuralNgramConfig,
    embedding: Embedding,
    feed_forward: SwiGlu,
    output: Linear,
    parameters: Vec<NamedParameter>,
}

impl NeuralNgram {
    /// Initializes all five matrices transactionally from one deterministic stream.
    pub fn new(config: NeuralNgramConfig, rng: &mut SplitMix64) -> Result<Self, NeuralNgramError> {
        let mut trial = rng.clone();
        let embedding = Embedding::new(
            NEURAL_NGRAM_PARAMETER_NAMES[0],
            config.vocabulary_size,
            config.embedding_width,
            &mut trial,
        )
        .map_err(NeuralNgramError::Embedding)?;
        let feed_forward = SwiGlu::new(
            "ngram.ffn",
            config.context_feature_width,
            config.hidden_width,
            config.hidden_width,
            &mut trial,
        )
        .map_err(NeuralNgramError::SwiGlu)?;
        let output = Linear::new(
            "ngram.output",
            config.hidden_width,
            config.vocabulary_size,
            false,
            &mut trial,
        )
        .map_err(NeuralNgramError::Linear)?;
        let model = Self::from_parts(config, embedding, feed_forward, output)?;
        *rng = trial;
        Ok(model)
    }

    /// Validates one ordered parameter set and builds persistent aliased layer handles.
    pub fn from_parameters(
        config: NeuralNgramConfig,
        parameters: Vec<NamedParameter>,
    ) -> Result<Self, NeuralNgramError> {
        validate_parameter_set(config, &parameters)?;
        let embedding = Embedding::from_parameter(parameters[0].clone())
            .map_err(NeuralNgramError::Embedding)?;
        let feed_forward = SwiGlu::from_parameters(
            parameters[1].clone(),
            parameters[2].clone(),
            parameters[3].clone(),
        )
        .map_err(NeuralNgramError::SwiGlu)?;
        let output = Linear::from_parameters(parameters[4].clone(), None)
            .map_err(NeuralNgramError::Linear)?;
        Ok(Self {
            config,
            embedding,
            feed_forward,
            output,
            parameters,
        })
    }

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

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

    fn from_parts(
        config: NeuralNgramConfig,
        embedding: Embedding,
        feed_forward: SwiGlu,
        output: Linear,
    ) -> Result<Self, NeuralNgramError> {
        let parameters = embedding
            .parameters()
            .iter()
            .chain(feed_forward.parameters())
            .chain(output.parameters())
            .cloned()
            .collect::<Vec<_>>();
        validate_parameter_set(config, &parameters)?;
        Ok(Self {
            config,
            embedding,
            feed_forward,
            output,
            parameters,
        })
    }
}

fn validate_parameter_set(
    config: NeuralNgramConfig,
    parameters: &[NamedParameter],
) -> Result<(), NeuralNgramError> {
    if parameters.len() != NEURAL_NGRAM_PARAMETER_NAMES.len() {
        return Err(NeuralNgramError::ParameterCountMismatch {
            expected: NEURAL_NGRAM_PARAMETER_NAMES.len(),
            actual: parameters.len(),
        });
    }
    let expected_shapes = [
        vec![config.vocabulary_size, config.embedding_width],
        vec![config.context_feature_width, config.hidden_width],
        vec![config.context_feature_width, config.hidden_width],
        vec![config.hidden_width, config.hidden_width],
        vec![config.hidden_width, config.vocabulary_size],
    ];
    for (index, ((parameter, expected_name), expected_shape)) in parameters
        .iter()
        .zip(NEURAL_NGRAM_PARAMETER_NAMES)
        .zip(expected_shapes)
        .enumerate()
    {
        if parameter.name() != expected_name {
            return Err(NeuralNgramError::ParameterNameMismatch {
                index,
                expected: expected_name,
                actual: parameter.name().to_owned(),
            });
        }
        let actual = parameter.tensor().shape();
        if actual != expected_shape {
            return Err(NeuralNgramError::ParameterShapeMismatch {
                name: expected_name.to_owned(),
                expected: expected_shape,
                actual,
            });
        }
    }
    Ok(())
}

The forward path enforces the exact shape chain, selects only the final shifted target, applies indexed mean NLL on vocabulary axis $1$ , and gives greedy generation a stable lower-ID tie break:

Run fixed-context prediction, loss, and deterministic maximum selection rust/crates/llm-from-scratch/src/models/neural_ngram.rs#neural-ngram-forward

impl NeuralNgram {
    /// Looks up and concatenates each complete context before prediction.
    pub fn forward(
        &self,
        context_ids: &[u32],
        batch_size: usize,
    ) -> Result<NeuralNgramForward, NeuralNgramError> {
        if batch_size == 0 {
            return Err(NeuralNgramError::EmptyBatch);
        }
        let expected = batch_size.checked_mul(self.config.context_length).ok_or(
            NeuralNgramError::ContextTokenCountOverflow {
                batch_size,
                context_length: self.config.context_length,
            },
        )?;
        if context_ids.len() != expected {
            return Err(NeuralNgramError::ContextTokenCountMismatch {
                expected,
                actual: context_ids.len(),
            });
        }
        let embeddings = self
            .embedding
            .forward(context_ids, &[batch_size, self.config.context_length])
            .map_err(NeuralNgramError::Embedding)?;
        let concatenated = embeddings
            .reshape(&[batch_size, self.config.context_feature_width])
            .map_err(NeuralNgramError::Autodiff)?;
        let hidden = self
            .feed_forward
            .forward(&concatenated)
            .map_err(NeuralNgramError::SwiGlu)?;
        let logits = self
            .output
            .forward(&hidden)
            .map_err(NeuralNgramError::Linear)?;
        Ok(NeuralNgramForward {
            embeddings,
            concatenated,
            hidden,
            logits,
        })
    }

    /// Scores only the token following each complete context row.
    pub fn loss(&self, batch: &MiniBatch) -> Result<TensorValue, NeuralNgramError> {
        if batch.context_length() != self.config.context_length {
            return Err(NeuralNgramError::BatchContextLengthMismatch {
                expected: self.config.context_length,
                actual: batch.context_length(),
            });
        }
        let batch_size = batch.batch_width();
        if batch_size == 0 {
            return Err(NeuralNgramError::EmptyBatch);
        }
        let mut targets = Vec::with_capacity(batch_size);
        for row in 0..batch_size {
            let target_row = batch
                .target_row(row)
                .ok_or(NeuralNgramError::MissingTargetRow { row })?;
            let id = target_row[self.config.context_length - 1];
            let target = usize::try_from(id)
                .map_err(|_| NeuralNgramError::TargetIdDoesNotFitUsize { row, id })?;
            if target >= self.config.vocabulary_size {
                return Err(NeuralNgramError::TargetIdOutOfBounds {
                    row,
                    id,
                    vocabulary_size: self.config.vocabulary_size,
                });
            }
            targets.push(target);
        }
        self.forward(batch.inputs(), batch_size)?
            .into_logits()
            .indexed_mean_nll(1, &targets)
            .map_err(NeuralNgramError::Autodiff)
    }

    /// Selects the greatest finite next-token logit, breaking exact ties by ID.
    pub fn greedy_next(
        &self,
        context_ids: &[u32],
        masked_ids: &[u32],
    ) -> Result<u32, NeuralNgramError> {
        let mut masked = vec![false; self.config.vocabulary_size];
        for &id in masked_ids {
            let index = usize::try_from(id)
                .ok()
                .filter(|index| *index < masked.len())
                .ok_or(NeuralNgramError::MaskedTokenOutOfBounds {
                    id,
                    vocabulary_size: self.config.vocabulary_size,
                })?;
            masked[index] = true;
        }
        let forward = self.forward(context_ids, 1)?;
        let logits = forward.logits().value();
        let mut best: Option<(usize, f64)> = None;
        for (token_id, &value) in logits.as_slice().iter().enumerate() {
            if !value.is_finite() {
                return Err(NeuralNgramError::NonFiniteLogit { token_id, value });
            }
            if masked[token_id] {
                continue;
            }
            if best.is_none_or(|(_, best_value)| value > best_value) {
                best = Some((token_id, value));
            }
        }
        let (token_id, _) = best.ok_or(NeuralNgramError::NoUnmaskedToken)?;
        u32::try_from(token_id).map_err(|_| NeuralNgramError::NoUnmaskedToken)
    }
}

The fixture learns eight BPE ranks only from training documents, never requests test text, materializes one batch order with seed 23, and uses its first $15$ batches. AdamW uses each raw gradient to compute the update and retains that gradient on the same live parameter node. After the update, the fixture explicitly calls zero_grad() on all five nodes before the next forward pass. Complete train and validation objectives are evaluated at steps $0$ , $8$ , and $15$ with actual row-count weighting. Two independently initialized runs must match every checked value bit for bit:

Train and replay the frozen held-out-loss fixture rust/demos/ch23-neural-ngram/src/lib.rs#chapter-neural-ngram-fixture

/// Trains two independent seeded runs and keeps one complete evidence record.
pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
    let first = run_once()?;
    let replay = run_once()?;
    let replay_bitwise = replay_equal(&first, &replay);
    require(
        replay_bitwise,
        "same-seed training did not replay bit for bit",
    )?;
    require(
        first.parameter_nodes_preserved,
        "AdamW did not preserve every model parameter node",
    )?;
    require(
        first.gradients_cleared,
        "training did not clear every post-update parameter gradient",
    )?;
    let initial = first
        .checkpoints
        .first()
        .ok_or(FixtureError::Invariant("initial checkpoint is missing"))?;
    let final_checkpoint = first
        .checkpoints
        .last()
        .ok_or(FixtureError::Invariant("final checkpoint is missing"))?;
    require(
        final_checkpoint.train_loss < initial.train_loss,
        "training loss did not improve",
    )?;
    require(
        final_checkpoint.validation_loss + 0.01 < initial.validation_loss,
        "validation loss did not improve by 0.01 nat",
    )?;
    Ok(LearnerEvidence {
        checkpoints: first.checkpoints,
        probe: first.probe,
        gradient_l1: first.gradient_l1,
        generation: first.generation,
        parameter_nodes_preserved: first.parameter_nodes_preserved,
        gradients_cleared: first.gradients_cleared,
        replay_bitwise,
        test_text_encoded_or_scored: first.test_text_encoded_or_scored,
    })
}

The executable prints the exact learner report:

Print the deterministic Chapter 23 learner report rust/demos/ch23-neural-ngram/src/main.rs#learner-neural-ngram-output

fn main() -> Result<(), Box<dyn std::error::Error>> {
    print!("{}", ch23_neural_ngram::learner_report()?);
    Ok(())
}

Run cargo run --quiet --locked -p ch23-neural-ngram. Its stdout is frozen in rust/demos/ch23-neural-ngram/expected.txt, so the shapes, losses, generated IDs, and replay evidence can be compared exactly with the values explained here.

Follow one context and the held-out loss

The exact trace records all five initial forward stages, three complete-objective checkpoints, the fixed final comparison, generated token IDs, and proof tokens. Reading these records together connects one forward pass to the measured held-out improvement without mixing in a second set of hand-computed values:

Emit exact pipeline, held-out-loss, generation, and replay evidence rust/demos/ch23-neural-ngram/src/diagram_trace.rs#neural-ngram-trace

pub fn diagram_trace() -> Result<String, FixtureError> {
    let evidence = learner_evidence()?;
    let initial = evidence.checkpoints[0];
    let final_checkpoint = *evidence
        .checkpoints
        .last()
        .expect("fixture always records a final checkpoint");
    let mut lines = vec![
        format!(
            "CONFIG|vocabulary={VOCABULARY_SIZE}|merges={REQUESTED_MERGES}|context={CONTEXT_LENGTH}|embedding={EMBEDDING_WIDTH}|concatenated={}|swiglu_inner={HIDDEN_WIDTH}|hidden={HIDDEN_WIDTH}|parameters=3384|batch={BATCH_SIZE}|evaluation_batch={EVALUATION_BATCH_SIZE}|init_seed={INIT_SEED}|shuffle_seed={SHUFFLE_SEED}|max_steps={MAX_STEPS}|lr={LEARNING_RATE:.6}|beta1={BETA1:.6}|beta2={BETA2:.6}|epsilon={EPSILON:.9}|weight_decay={WEIGHT_DECAY:.6}",
            CONTEXT_LENGTH * EMBEDDING_WIDTH,
        ),
        format!(
            "SPLIT|train_documents=8|validation_documents=2|test_text_used=no|train_contexts={TRAIN_CONTEXTS}|validation_contexts={VALIDATION_CONTEXTS}|train_batches={TRAIN_BATCHES}|train_evaluation_batches={TRAIN_EVALUATION_BATCHES}|validation_evaluation_batches={VALIDATION_EVALUATION_BATCHES}"
        ),
        format!(
            "STAGE|index=0|name=context_ids|shape=[1, 2]|ids={}",
            format_ids(&evidence.probe.context_ids)
        ),
        format!(
            "STAGE|index=1|name=embeddings|shape=[1, 2, 4]|values={}",
            format_values(&evidence.probe.embeddings)
        ),
        format!(
            "STAGE|index=2|name=concatenated|shape=[1, 8]|values={}",
            format_values(&evidence.probe.concatenated)
        ),
        format!(
            "STAGE|index=3|name=hidden|shape=[1, 8]|values={}",
            format_values(&evidence.probe.hidden)
        ),
        format!(
            "STAGE|index=4|name=logits|shape=[1, 266]|preview={}|argmax={}|argmax_logit={:.6}",
            format_values(&evidence.probe.logits_preview),
            evidence.probe.argmax,
            evidence.probe.argmax_logit,
        ),
    ];
    lines.extend(evidence.checkpoints.iter().map(|checkpoint| {
        format!(
            "LOSS|step={}|train={:.6}|validation={:.6}",
            checkpoint.step, checkpoint.train_loss, checkpoint.validation_loss
        )
    }));
    lines.extend([
        format!(
            "RESULT|step={}|initial_validation={:.6}|final_validation={:.6}|improvement={:.6}",
            final_checkpoint.step,
            initial.validation_loss,
            final_checkpoint.validation_loss,
            initial.validation_loss - final_checkpoint.validation_loss,
        ),
        format!(
            "GENERATE|prompt=At|prompt_ids={}|ids={}|stop={}",
            format_ids(&evidence.generation.prompt_ids),
            format_ids(&evidence.generation.generated_ids),
            evidence.generation.stop.label(),
        ),
        format!(
            "PROOF|replay={}|test_text={}|target=final_shifted|gradient_l1={}|parameter_nodes={}|gradients={}|generation=deterministic",
            if evidence.replay_bitwise { "bitwise" } else { "changed" },
            if evidence.test_text_encoded_or_scored {
                "encoded_or_scored"
            } else {
                "not_encoded_or_scored"
            },
            if evidence
                .gradient_l1
                .iter()
                .all(|value| value.is_finite() && *value > 0.0)
            {
                "five_positive_finite"
            } else {
                "invalid"
            },
            if evidence.parameter_nodes_preserved {
                "preserved"
            } else {
                "replaced"
            },
            if evidence.gradients_cleared {
                "cleared"
            } else {
                "retained"
            },
        ),
    ]);
    Ok(lines.join("\n") + "\n")
}

Follow one context through training

Follow exact Rust-authored token IDs through embeddings, concatenation, a hidden state, and vocabulary logits, then compare complete training and validation losses.

Fixed model configuration: $V=266,\ C=2,\ D=4,\ H=8$
Fixed data split: $N_{\mathrm{train}}=1836,\ N_{\mathrm{val}}=467$
Fixed optimizer settings: $\eta=0.010000,\ t_{\mathrm{max}}=15$

1 · Project one complete context

Project one complete context

Every tensor and displayed value comes from the initial seeded forward pass. Concatenation changes layout, not values.

0

Integer selectors

Context IDs
Shape

$[1, 2]$
Rust values
1. $67$
2. $118$
1

Learned features

Embedding rows
Shape

$[1, 2, 4]$
Rust values
1. $0.064154$
2. $0.021328$
3. $0.083333$
4. $-0.012260$
5. $0.057176$
6. $0.111494$
7. $-0.126703$
8. $-0.068284$
2

Learned features

Concatenated features
Shape

$[1, 8]$
Rust values
1. $0.064154$
2. $0.021328$
3. $0.083333$
4. $-0.012260$
5. $0.057176$
6. $0.111494$
7. $-0.126703$
8. $-0.068284$
3

Learned features

SwiGLU hidden state
Shape

$[1, 8]$
Rust values
1. $-0.002448$
2. $-0.000051$
3. $0.003220$
4. $0.003477$
5. $0.002033$
6. $0.004016$
7. $0.003727$
8. $0.003874$
4

Class scores

Vocabulary logits
Shape

$[1, 266]$
Rust values
1. $0.000075$
2. $-0.000037$
3. $0.000496$
4. $-0.001047$
5. $-0.000055$
6. $-0.001032$
Initial maximum-score token

$44$

Initial maximum logit

$0.002350$

2 · Compare complete objectives

Compare complete objectives

Each loss covers its complete partition and weights every batch by its actual row count. The final checkpoint is fixed before evaluation.

Measured checkpoint

Optimizer step $0$

Training loss

$L_{\mathrm{train}}=5.583505$

Validation loss

$L_{\mathrm{val}}=5.583482$
Measured checkpoint

Optimizer step $8$

Training loss

$L_{\mathrm{train}}=5.580106$

Validation loss

$L_{\mathrm{val}}=5.580365$
Fixed final step

Optimizer step $15$

Training loss

$L_{\mathrm{train}}=5.555850$

Validation loss

$L_{\mathrm{val}}=5.557362$

3 · Measure held-out improvement

Measure held-out improvement

Initial validation loss: $5.583482$
Final validation loss: $5.557362$
Validation improvement: $\Delta L_{\mathrm{val}}=0.026120$

4 · Continue from the trained model

Continue from the trained model

These are token IDs and preserved bytes from an intentionally tiny undertrained model, not a claim of readable text.

Prompt: At
Prompt IDs: $[67, 118]$
Stop reason: limit

$259$
$211$
$211$
$211$
$211$
$211$
$211$
$211$
$211$
$211$
$211$
$211$

5 · Verify the training boundary

Verify the training boundary

Independent replay: bitwise
Test-text boundary: not_encoded_or_scored
Target policy: final_shifted
First-step matrix-gradient check: five_positive_finite
Parameter-node identity: preserved
Gradients after explicit clearing: cleared
Generation policy: deterministic

The initial context rows concatenate without changing their numerical order. After $15$ AdamW updates, complete training loss falls from $5.583505$ to $5.555850$ , while held-out validation loss falls from $5.583482$ to $5.557362$ . The improvement is $0.026120$ nat at displayed precision. This does not claim monotonic loss or choose the best checkpoint: step $15$ was fixed before validation was inspected.

Greedy continuation emits token IDs $[259,211,211,211,211,211,211,211,211,211,211,211]$ and stops at the twelve-token limit. The preserved bytes are not valid UTF-8 after the first token, so the diagram labels them as IDs from an intentionally tiny undertrained model rather than presenting invented readable text.

Predict before training

Predict every shape for $B=64$ , $C=2$ , $D=4$ , $H=8$ , and $V=266$ .
Pick the target used from a shifted row of width $C$ .
Predict which of the five matrices should have a positive finite first-step gradient $L_1$ norm.
Explain how a $300$ -row final evaluation batch enters the complete train mean.
Decide whether validation loss must fall at every checkpoint.
Decide what a persistent layer handle and a cloned NamedParameter handle observe after AdamW updates their shared node, and whether AdamW clears that node’s gradient.
Predict what greedy generation does if BOS has the greatest logit.
Identify which operation would leak the test partition.

Check the predictions

The chain is $[64,2]\to[64,2,4]\to[64,8]\to[64,8]\to[64,266]$ .
Use only index $C-1$ of the target row: the token following the complete context.
The embedding, gate, up, down, and output matrices all have positive finite gradient $L_1$ norms.
Multiply its batch mean by $300$ , add the other weighted sums, then divide by all $1836$ contexts.
No. This fixture requires the fixed final validation loss to beat initialization by more than $0.01$ nat; it does not impose monotonicity.
Cloning a NamedParameter copies its handle, not its tensor data or node. The copied handle and the persistent layer handle still point to the same TensorValue node, so both observe the new value. AdamW leaves the accumulated gradient on that node, and the training loop explicitly calls zero_grad() after the update. Independent replay constructs a separate seeded model.
BOS is masked, so $\operatorname{argmax}$ selects the greatest remaining finite logit with the lower ID winning an exact tie.
Reading test document text for fitting, encoding, selection, early stopping, scoring, or display would invalidate the boundary.

Replace fixed context with causal information mixing next

The components built so far now train a complete fixed-context next-token model with frozen data partitions and AdamW, then produce deterministic greedy tokens. Chapters 24–32 replace fixed-context concatenation with residual, normalized, attention-based causal information mixing between sequence positions.

This model is functional but deliberately narrow: every prediction sees exactly $C$ preceding IDs, regardless of how much earlier context exists. Residual paths come next, before normalization and masked self-attention expand the sequence computation.