From transition counts to a bigram model

Turn seven arrows into one prediction row

Chapter 5 built overlapping input–target views for training a sequence model. A bigram table needs a different view of the same data: the original training documents. If we counted through overlapping windows, a transition appearing in two windows would receive two votes. Here every source transition must contribute exactly once.

Use this fixed vocabulary order throughout the chapter:

C belongs to the vocabulary but does not occur in these two training documents. The documents remain separate:

d1: BOS(0) → A(2) → A(2) → B(3) → EOS(1)
d2: BOS(0) → A(2) → B(3) → EOS(1)

Before reading on, predict four things:

1. How many arrows will be counted?
2. Which successors occur after A, and how often?
3. Which token should be most likely after A?
4. Is “C never followed A” the same case as “nothing was observed after C”?

The first document contributes four transitions and the second contributes three:

d1: BOS→A, A→A, A→B, B→EOS
d2: BOS→A, A→B, B→EOS

After aggregation, the only nonzero counts are $\operatorname{BOS}\to A=2$, $A\to A=1$, $A\to B=2$, and $B\to\operatorname{EOS}=2$. For current token A, keep the vocabulary column order [BOS, EOS, A, B, C]:

counts after A = [0, 0, 1, 2, 0]
row total       = 3

This is a bigram model: the current token selects one table row, and every vocabulary token is a possible next-token column. It deliberately forgets everything before the current token. Normalizing the observed A row gives the maximum-likelihood estimate (MLE):

So B is the unique most likely successor. The entry for A→C is exactly zero, but the row itself is well defined because three transitions leave A.

Now try to normalize row C. Its counts are [0,0,0,0,0] and its total is zero. Dividing every cell by zero does not produce a probability distribution; the entire MLE row is undefined. This is different from the defined zero in the A row.

Add-one smoothing ($\alpha=1$) adds one pseudocount to every candidate before normalization. For A:

For C, the five equal pseudocounts give denominator $0+1\times5=5$ and the uniform row $[1/5,1/5,1/5,1/5,1/5]$. That row is a fallback imposed by the smoothing rule. It is not evidence that every successor is genuinely equally likely after C, and it has no unique winner.

Finally, resist the tempting shortcut of concatenating the documents. The flat sequence [0,2,2,3,1,0,2,3,1] contains eight adjacent pairs and invents exactly one observation: EOS(1)→BOS(0). That pair belongs to neither source document.

Describe document-local counting and row normalization

The complete procedure is:

Read the first sum from the outside in. It visits each original training document $d$ separately. Positions are zero-based, just like Rust slice indices. In a document of length $|d|$, position $t$ stops at $|d|-2$, so $t+1$ is the final valid position at most. The superscript $(d)$ matters: both members of a pair come from the same document.

Current-token ID $i$ chooses a row, and candidate next-token ID $j$ chooses a column. The indicator contributes one only when the adjacent IDs at positions $t$ and $t+1$ equal $i$ and $j$. Summing row $i$ over every vocabulary column $k$ produces its outgoing observation total $N_i$.

MLE divides observed counts by $N_i$, so it exists only when $N_i>0$. Add-$\alpha$ smoothing adds $\alpha$ to all $|V|$ cells. The denominator must therefore gain $\alpha|V|$, not merely $\alpha$; this guarantees that the smoothed row still sums to one. Because $\alpha>0$, even a zero-total row has a positive denominator.

In memory, the Rust implementation stores row $i$, column $j$ at a row-major offset. The table layout is an implementation choice; the probabilistic meaning is still “next token $j$, given current token $i$.”

Account for every symbol

Two phrases that sound similar now have precise locations in the table. An unobserved successor after A is a zero cell in row A. A context with no outgoing observations is an entire row whose total is zero. Confusing a column with a row is the central misconception to avoid.

Use a transparent classical baseline without mistaking it for a strong one

Before neural language models, n-gram models estimated the next word from a fixed-size suffix of preceding words. A bigram uses the shortest nonempty such context: one current token. It is easy to inspect and fast to query, but two histories ending in the same token always receive the same distribution.

Chen and Goodman (1996) describe the bigram maximum-likelihood estimate as a pair count divided by its context count and explain why an unobserved n-gram receiving zero probability is a serious language-modeling problem. Their empirical study compares several smoothing families. Additive smoothing is a simple member of that landscape, not the best-performing practical choice: it moves equal pseudocount mass into every cell and cannot express which unseen continuations are more plausible.

Bengio, Ducharme, and Vincent (2000) describe conventional n-gram models as tables of conditional probabilities over short contexts. Their neural approach instead learns distributed word representations so that evidence can generalize through similarity—a capability the independent cells of our count table do not have.

This course keeps add-one smoothing for a narrower reason: every numerator and denominator can be checked by hand. Its weakness is visible immediately. The smoothed A row assigns $1/8$ not only to unseen A→C, but also to A→BOS, even though BOS is reserved for document beginnings. Uniform pseudocounts know nothing about token roles or syntax.

The papers do not prescribe our token IDs, document boundaries, $\alpha=1$, train-only helper, or tie order. Those are explicit teaching and implementation choices. The runnable Rust contrast below prints both MLE and add-one estimates from the same count table so their difference is observable rather than historical prose.

Make the evidence come from one checked Rust table

The executable fixture uses exactly the two documents calculated above. It passes them as separate slices, so the outer collection preserves their segmentation. The generic fitter still trusts its caller to supply one wrapped, unpadded training document per slice; a Rust slice alone cannot prove that precondition.

pub const DOCUMENT_1: &[u32] = &[BOS, A, A, B, EOS];
pub const DOCUMENT_2: &[u32] = &[BOS, A, B, EOS];
pub const TRAINING_DOCUMENTS: [&[u32]; 2] = [DOCUMENT_1, DOCUMENT_2];

pub fn fitted_model() -> Result<BigramModel, BigramError> {
    BigramModel::fit_training_documents(VOCABULARY_SIZE, ALPHA, TRAINING_DOCUMENTS)
}

fit_training_documents validates the vocabulary, smoothing amount, table size, and every supplied token ID. Its outer loop keeps documents separate; only windows(2) inside one document produces pairs. Empty and one-token slices add no transitions, although every token ID they contain is still checked. The partition-aware entry point explicitly requests Partition::Train; validation and test documents never reach the count loop.

    /// Fits one count per adjacent pair in each caller-supplied training document.
    ///
    /// Documents remain separate: the final token of one document is never paired
    /// with the first token of the next document.
    pub fn fit_training_documents<'a, I>(
        vocabulary_size: usize,
        alpha: f64,
        training_documents: I,
    ) -> Result<Self, BigramError>
    where
        I: IntoIterator<Item = &'a [u32]>,
    {
        if vocabulary_size == 0 {
            return Err(BigramError::EmptyVocabulary);
        }
        let smoothing_mass = alpha * vocabulary_size as f64;
        if !alpha.is_finite() || alpha <= 0.0 || !smoothing_mass.is_finite() {
            return Err(BigramError::InvalidAlpha);
        }

        let cell_count = vocabulary_size
            .checked_mul(vocabulary_size)
            .ok_or(BigramError::TableTooLarge)?;
        let mut counts = Vec::new();
        counts
            .try_reserve_exact(cell_count)
            .map_err(|_| BigramError::TableTooLarge)?;
        counts.resize(cell_count, 0_u64);

        let mut model = Self {
            vocabulary_size,
            alpha,
            counts,
            fitted_documents: 0,
            fitted_transitions: 0,
        };

        for document in training_documents {
            model.fitted_documents = model
                .fitted_documents
                .checked_add(1)
                .ok_or(BigramError::TooManyDocuments)?;
            for token in document {
                model.token_index(*token)?;
            }
            for pair in document.windows(2) {
                let from = model.token_index(pair[0])?;
                let to = model.token_index(pair[1])?;
                let cell = from * vocabulary_size + to;
                model.counts[cell] = model.counts[cell]
                    .checked_add(1)
                    .ok_or(BigramError::TooManyTransitions)?;
                model.fitted_transitions = model
                    .fitted_transitions
                    .checked_add(1)
                    .ok_or(BigramError::TooManyTransitions)?;
            }
        }

        Ok(model)
    }

    /// Selects the original encoded training documents and never requests held-out data.
    pub fn fit_encoded_training_partition(
        vocabulary_size: usize,
        alpha: f64,
        partitions: &EncodedCorpusPartitions,
    ) -> Result<Self, BigramError> {
        Self::fit_training_documents(
            vocabulary_size,
            alpha,
            partitions
                .documents(Partition::Train)
                .iter()
                .map(|document| document.token_ids()),
        )
    }

The query API preserves the conceptual distinction in its return type. Some(0.0) means the selected MLE row exists and this successor received no count. None means the row total is zero, so no MLE distribution exists. Smoothed queries remain defined for every valid context. When several cells share the maximum count, most_likely_tokens reports all of them in ascending ID order; that stable reporting order does not make the first ID more probable. Additive smoothing preserves this count ordering and every tie because it adds the same $\alpha$ to each cell in a row and divides them all by the same denominator.

    /// Returns `None` when no outgoing transition was observed for `from`.
    pub fn maximum_likelihood_distribution(
        &self,
        from: u32,
    ) -> Result<Option<Vec<f64>>, BigramError> {
        let total = self.row_total(from)?;
        if total == 0 {
            return Ok(None);
        }
        Ok(Some(
            self.counts_row(from)?
                .iter()
                .map(|count| *count as f64 / total as f64)
                .collect(),
        ))
    }

    pub fn maximum_likelihood_probability(
        &self,
        from: u32,
        to: u32,
    ) -> Result<Option<f64>, BigramError> {
        self.token_index(to)?;
        let total = self.row_total(from)?;
        if total == 0 {
            return Ok(None);
        }
        Ok(Some(self.count(from, to)? as f64 / total as f64))
    }

    pub fn smoothing_denominator(&self, from: u32) -> Result<f64, BigramError> {
        Ok(self.row_total(from)? as f64 + self.alpha * self.vocabulary_size as f64)
    }

    pub fn smoothed_probability(&self, from: u32, to: u32) -> Result<f64, BigramError> {
        let numerator = self.count(from, to)? as f64 + self.alpha;
        Ok(numerator / self.smoothing_denominator(from)?)
    }

    pub fn smoothed_distribution(&self, from: u32) -> Result<Vec<f64>, BigramError> {
        self.token_index(from)?;
        (0..self.vocabulary_size)
            .map(|to| {
                let to = u32::try_from(to).map_err(|_| BigramError::TokenOutOfRange)?;
                self.smoothed_probability(from, to)
            })
            .collect()
    }

    /// Returns every count-maximizing successor in ascending token-ID order.
    pub fn most_likely_tokens(&self, from: u32) -> Result<Vec<u32>, BigramError> {
        let row = self.counts_row(from)?;
        let maximum = row
            .iter()
            .copied()
            .max()
            .expect("a vocabulary is non-empty");
        row.iter()
            .enumerate()
            .filter(|(_, count)| **count == maximum)
            .map(|(token, _)| u32::try_from(token).map_err(|_| BigramError::TokenOutOfRange))
            .collect()
    }

The demo prints counts before estimates and includes both missing-evidence cases. It also names the transition that flattening would fabricate.

    println!("tokens: BOS=0 EOS=1 A=2 B=3 C=4");
    println!("alpha: {ALPHA:.1}");
    println!("training document d1: {DOCUMENT_1:?}");
    println!("training document d2: {DOCUMENT_2:?}");
    println!("counted transitions: {}", model.fitted_transitions());
    println!(
        "A counts: {} total={}",
        format_counts(model.counts_row(A).expect("A is in the vocabulary")),
        model.row_total(A).expect("A is in the vocabulary")
    );
    println!("A MLE: {}", format_probabilities(&a_mle));
    println!(
        "A add-alpha: {} denominator={:.0}",
        format_probabilities(&a_smoothed),
        model
            .smoothing_denominator(A)
            .expect("A is in the vocabulary")
    );
    println!(
        "unseen successor A->C: MLE={:.3} add-alpha={:.3}",
        model
            .maximum_likelihood_probability(A, C)
            .expect("A and C are in the vocabulary")
            .expect("A has outgoing transitions"),
        model
            .smoothed_probability(A, C)
            .expect("A and C are in the vocabulary")
    );
    println!(
        "C counts: {} total={}",
        format_counts(model.counts_row(C).expect("C is in the vocabulary")),
        model.row_total(C).expect("C is in the vocabulary")
    );
    println!("C MLE: {c_mle}");
    println!(
        "C add-alpha: {} denominator={:.0}",
        format_probabilities(&c_smoothed),
        model
            .smoothing_denominator(C)
            .expect("C is in the vocabulary")
    );
    println!("flattening would invent: EOS({EOS})->BOS(0)");

Run the example:

./course run cargo run --quiet --locked -p ch06-bigram-baseline

The implementation preserves all five count rows, the seven-transition total, the absent EOS→BOS cell, both normalized distributions, stable tie reporting, and explicit invalid-input errors. Fitting the encoded training partition produces the same table as supplying those wrapped training documents directly.

Compare every cell without losing the document evidence

Before reading the tables, predict the answers:

1. Which cells receive a pseudocount?
2. Does smoothing change the most likely successor after A?
3. Does the smoothed C row gain a unique winner?
4. Which structurally forbidden successor of A nevertheless receives probability mass?

Every document and number in the figure comes from a deterministic trace emitted by the same Rust fixture. Text labels, row totals, table structure, and the crossed boundary arrow make the calculation independently inspectable without relying on color.

Now check the predictions against the complete tables: every cell receives $\alpha$; B remains the sole maximum after A; every candidate ties after C; and even the structurally forbidden continuation A→BOS receives nonzero mass.

Predict the result, then expose the arithmetic

Work from the two separate documents and the fixed column order [BOS, EOS, A, B, C]. Do not run the demo until you have written down each prediction.

1. Enumerate all seven observed transitions, preserving their source document.
2. Derive the complete count, MLE, and add-one rows for context A.
3. Explain why $\operatorname{MLE}(A\to C)=0$ while the MLE row for context C is undefined.
4. Concatenate the documents mentally. Which exact transition is fabricated, and how many pairs would the flattened sequence contain?
5. Why would fitting from Chapter 5’s overlapping $T=3$, $S=1$ windows miscount document d1?
6. Recalculate the smoothed A row with $\alpha=1/2$.
7. Verify that both add-one rows shown in this chapter sum to one.
8. State the $\arg\max$ result after A and after C, including every tie.
9. Suppose validation data contains many A→C transitions. Should the frozen count table change before Chapter 7 scores validation?

Misconception check: “C was not observed after A, therefore row C is undefined” confuses a column with a row. The missing A→C observation is one cell in the defined A row. Row C is undefined under MLE because no training transition has C as its current token.

Freeze the first complete next-token model for scoring

The course now has a model that accepts a current token ID $i$ and returns one probability for every candidate next token $j$. It is intentionally weak—it sees only one-token context and smoothing cannot learn similarity—but its output has the same essential shape as a decoder’s output distribution.

Chapter 7 will look up the probability assigned to each actual next token and aggregate those values into negative log-likelihood and perplexity. The count table remains frozen: scoring training examples does not add a second set of counts, validation is evaluation-only, and Chapter 7’s metric interface can select only training or validation. No gradient is needed to fit this count table; optimization of learned parameters belongs to later chapters.