Learning deterministic BPE merges

Predict two rounds before running them

Keep these as two documents, not one concatenated string:

train-aaa = aaa = [97, 97, 97]
train-aba = aba = [97, 98, 97]

Here 97 is byte 61 in hexadecimal (a), and 98 is byte 62 (b). Before reading the next paragraph, list every adjacent position. In aaa, the pair (97,97) starts at positions 0 and 1. Those occurrences overlap at the middle byte, but they are still two candidate positions. The second document adds one (97,98) and one (98,97):

Rank 0 therefore selects (97,97) and assigns fresh ID 256. Now make one left-to-right replacement pass. The match at positions 0–1 consumes both input tokens, so the scan resumes at position 2. It cannot reuse the middle 97 for a second replacement. The new stages are:

train-aaa = [256, 97]
train-aba = [97, 98, 97]

This is the chapter’s most important distinction: the candidate count is 2, but the replacement count is 1. Counting asks which rule to learn; replacement asks how that one rule transforms a sequence without consuming a token twice.

Predict rank 1. The candidates (97,98), (98,97), and (256,97) each occur once. This course compares numeric pairs, so (97,98) wins: its left ID 97 is smaller than 98 and 256. Fresh ID 257 represents bytes 61 62, producing:

train-aaa = [256, 97]
train-aba = [257, 97]

Finally consider two one-byte documents, a | a. The bar marks a document boundary; it is not a byte in either document. Each sequence has length one, so neither supplies a window of length two. The candidate map is empty—there is no synthetic (97,97) across the boundary.

Select one reproducible rule

For each round, the course fixes the winner and its new symbol with:

The three-part tuple is compared lexicographically. First maximize the frequency $C(a,b)$. If counts tie, maximizing $-a$ chooses the smaller left ID; if those also tie, maximizing $-b$ chooses the smaller right ID. This numeric tie rule is a course reproducibility policy, not a claim about every BPE implementation.

The second expression operates on byte expansions. If token 256 expands to 61 61 and token 97 expands to 61, merging (256,97) would store 61 61 61. It does not concatenate the decimal strings "256" and "97". If no document has an adjacent pair, training stops early; a winner with count one is otherwise valid.

Symbol glossary

Raw byte symbols occupy IDs 0..=255. A successful zero-based rank $r$ receives ID $256+r$, so $k$ successful rounds make a training-space vocabulary of $256+k$ symbols. Chapter 4 will map these content IDs into its final layout after reserving control tokens.

From closed word tables to repeated pair merging

A fixed whole-word vocabulary assigns an ID to every spelling seen during fitting. It is simple, but every unseen spelling needs a fallback. In the runnable contrast, lower has its own fitted ID while lowering and any other unseen word collapse to the same ID 0:

pub fn fit_whole_word_vocabulary(documents: &[&str]) -> BTreeMap<String, u32> {
    let words = documents
        .iter()
        .flat_map(|document| document.split_whitespace())
        .collect::<BTreeSet<_>>();
    words
        .into_iter()
        .enumerate()
        .map(|(index, word)| (word.to_owned(), index as u32 + 1))
        .collect()
}

/// Returns the fitted ID or one shared unknown-word bucket.
pub fn whole_word_id(vocabulary: &BTreeMap<String, u32>, word: &str) -> u32 {
    vocabulary.get(word).copied().unwrap_or(UNKNOWN_WORD_ID)
}

Philip Gage’s 1994 compression algorithm took a different route: repeatedly find a frequent adjacent byte pair, substitute an unused byte, and save enough of the substitution table to reverse the compression. Read the original byte-pair compression article. Its goal and finite-byte substitution machinery differ from this tokenizer, which uses u32 IDs and retains each symbol’s byte expansion.

Rico Sennrich, Barry Haddow, and Alexandra Birch later adapted repeated pair merging to open-vocabulary neural translation. Their 2016 subword paper starts from character sequences and prevents merges across word boundaries. This course deliberately implements another variant: it begins with raw UTF-8 bytes, permits a space to participate in a pair inside one document, and treats only document boundaries as hard barriers.

Those distinctions matter. “BPE” names a family of repeated-pair procedures, not one universal tie rule or boundary policy. Our numeric-smallest tie rule is stated and tested so the same corpus produces the same ranks on every run.

Key idea: BPE training is a deterministic sequence of count, tie-break, and non-overlapping replacement decisions—not a vague “merge frequent pairs” operation.

Implement the trainer without a tokenizer library

The pair counter receives a slice of token sequences rather than one flattened array. Calling windows(2) separately for each sequence counts both overlapping windows in aaa and makes crossing a document boundary impossible. A BTreeMap keeps candidates in numeric pair order; scanning it while replacing the winner only for a strictly larger count preserves the first—therefore smallest—pair on a tie:

pub fn count_adjacent_pairs(sequences: &[Vec<u32>]) -> BTreeMap<TokenPair, usize> {
    let mut counts = BTreeMap::new();
    for sequence in sequences {
        for window in sequence.windows(2) {
            let pair = TokenPair::new(window[0], window[1]);
            *counts.entry(pair).or_insert(0) += 1;
        }
    }
    counts
}

/// Selects the greatest count, breaking ties by the smallest `(left, right)` IDs.
pub fn choose_most_frequent_pair(
    counts: &BTreeMap<TokenPair, usize>,
) -> Option<(TokenPair, usize)> {
    let mut winner = None;
    for (&pair, &count) in counts {
        match winner {
            None => winner = Some((pair, count)),
            Some((_, best_count)) if count > best_count => winner = Some((pair, count)),
            Some(_) => {}
        }
    }
    winner
}

Selection does not decide how overlapping matches are rewritten. The replacement function scans with an index. A match emits one new ID and advances by two; a non-match copies one ID and advances by one:

pub fn replace_pair_left_to_right(
    sequence: &[u32],
    pair: TokenPair,
    replacement: u32,
) -> (Vec<u32>, usize) {
    let mut output = Vec::with_capacity(sequence.len());
    let mut replacements = 0;
    let mut index = 0;

    while index < sequence.len() {
        if index + 1 < sequence.len()
            && sequence[index] == pair.left
            && sequence[index + 1] == pair.right
        {
            output.push(replacement);
            replacements += 1;
            index += 2;
        } else {
            output.push(sequence[index]);
            index += 1;
        }
    }

    (output, replacements)
}

The public trainer accepts CorpusPartitions, not a free-form list that might quietly include held-out text. It records the exact stable IDs returned by training_documents() and converts only those document bodies from bytes to initial IDs:

    pub fn train(self, partitions: &CorpusPartitions<'_>) -> Result<BpeTraining, BpeTrainingError> {
        let available_merge_ids = u128::from(u32::MAX) - u128::from(u8::MAX);
        if self.max_merges as u128 > available_merge_ids {
            return Err(BpeTrainingError::new(
                "requested merge count exceeds the u32 token-ID space",
            ));
        }
        let training_documents = partitions.training_documents();
        let document_ids = training_documents
            .iter()
            .map(|document| document.id().to_owned())
            .collect::<Vec<_>>();
        let sequences = training_documents
            .iter()
            .map(|document| bytes_to_tokens(document.text().as_bytes()))
            .collect::<Vec<_>>();

        learn_from_token_sequences(self.max_merges, document_ids, sequences)
    }

The internal loop assigns the current vocabulary length as the next ID, stores the concatenated byte expansion, replaces the pair in every training sequence, and records both candidate and replacement counts. Tests cover overlap, ties, document barriers, zero rounds, token-ID overflow, exact provenance, vocabulary growth, rule uniqueness, and repeatability.

The executable then prints the eight ranks learned from the real corpus and a compact two-round trace whose counts, selections, replacements, and stopping condition can be inspected directly:

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let corpus = Corpus::from_json(CORPUS_JSON)?;
    let manifest = SplitManifest::from_json(SPLIT_MANIFEST)?;
    let partitions = manifest.partition(&corpus)?;
    let training = BpeTrainer::new(8).train(&partitions)?;

    println!("corpus checksum: {}", corpus.checksum());
    println!("statistics source: train only");
    println!("training documents: {:?}", training.training_document_ids());
    println!(
        "held out from trainer: validation={} test={}",
        partitions.documents(Partition::Validation).len(),
        partitions.documents(Partition::Test).len()
    );
    println!(
        "merge rounds: requested={} learned={}",
        training.requested_merges(),
        training.rules().len()
    );
    for rule in training.rules() {
        println!(
            "corpus rank {}: pair={} count={} replacements={} token={} bytes={:02x?}",
            rule.rank(),
            format_pair(rule.pair()),
            rule.candidate_count(),
            rule.replacement_count(),
            rule.token_id(),
            training
                .token_bytes(rule.token_id())
                .expect("learned token has bytes")
        );
    }

    let mut fixture = vec![bytes_to_tokens(b"aaa"), bytes_to_tokens(b"aba")];
    let mut fixture_vocabulary = (u8::MIN..=u8::MAX)
        .map(|byte| vec![byte])
        .collect::<Vec<_>>();
    println!("TRACE bpe-merges-v1 BEGIN");
    print_trace_stage(0, &fixture);
    for rank in 0..2 {
        let counts = count_adjacent_pairs(&fixture);
        let (winner, count) = choose_most_frequent_pair(&counts).expect("fixture has a pair");
        let token_id = BYTE_TOKEN_COUNT + rank;
        for (pair, candidate_count) in &counts {
            let selected = if *pair == winner { "yes" } else { "no" };
            println!(
                "CANDIDATE rank={rank} left={} right={} count={candidate_count} winner={selected}",
                pair.left(),
                pair.right()
            );
        }
        let mut merged_bytes = fixture_vocabulary[winner.left() as usize].clone();
        merged_bytes.extend_from_slice(&fixture_vocabulary[winner.right() as usize]);
        let mut replacements = 0;
        for sequence in &mut fixture {
            let (next, replaced) = replace_pair_left_to_right(sequence, winner, token_id);
            *sequence = next;
            replacements += replaced;
        }
        println!(
            "MERGE rank={rank} left={} right={} count={count} replacements={replacements} token={token_id} bytes_hex={}",
            winner.left(),
            winner.right(),
            format_hex(&merged_bytes)
        );
        fixture_vocabulary.push(merged_bytes);
        print_trace_stage(rank as usize + 1, &fixture);
    }
    println!("TRACE bpe-merges-v1 END");
    println!(
        "document barrier candidates for A=\"a\" B=\"a\": {}",
        count_adjacent_pairs(&[bytes_to_tokens(b"a"), bytes_to_tokens(b"a")]).len()
    );

    let words = fit_whole_word_vocabulary(&["low lower", "new newest"]);
    println!("historical whole-word types: {}", words.len());
    println!(
        "historical lookup lower: {}",
        whole_word_id(&words, "lower")
    );
    println!(
        "historical lookup lowering: {} (unknown)",
        whole_word_id(&words, "lowering")
    );
    println!("chapter 4 handoff: freeze ranks and encode arbitrary bytes");

    Ok(())
}

Run it from the repository root:

cargo test --workspace --locked
cargo run --quiet --locked -p ch03-learn-bpe-merges | diff -u rust/demos/ch03-learn-bpe-merges/expected.txt -

A successful diff is silent. On the real training partition, rank 0 is (32,208) with count 81, new ID 256, and bytes 20 d0. Byte 20 is a space; d0 begins many two-byte Cyrillic UTF-8 encodings but is incomplete alone. Consequently 20 d0 is not a standalone character or valid standalone UTF-8 string. Preserve and display it as bytes or hex.

Implementation takeaway: keep byte expansions attached to each new ID; a merge token is useful because the tokenizer can later expand it exactly.

Inspect the same trace as a static figure

The figure follows the same deterministic two-round trace. Each stage keeps both documents visible; each candidate table distinguishes the overlapping count from the non-overlapping replacement count. The selected row has text, a star, weight, and a border, so color is not required. Numeric and hex lanes remain left-to-right in every spoken language, and the layout stacks at a narrow viewport.

Read stage 0 again and point to the shared middle 97 in aaa. Then compare the rank-0 count 2 with replacements 1. At rank 1, ignore display position and apply the numeric tuple rule: $(97,98)$ must beat $(98,97)$ and $(256,97)$.

Predict, then check

1. For one document [97,97,97,97], compute $C(97,97)$ and the number of replacements in one left-to-right pass.
2. After rank 0 in the worked example, write the score tuple for each of the three count-one candidates and select the winner.
3. Explain why documents [97] | [98] do not create candidate (97,98).
4. Predict the vocabulary size after 12 requested rounds if training stops after only 9 successful rounds.
5. Decide whether adding 10,000 repetitions to a validation document can alter the learned rank table.
6. Expand a hypothetical merge (256,97) when token 256 stores bytes 61 61.
7. Explain why printing real rank 0’s bytes 20 d0 as a character is incorrect.

When testing a modification, predict its first changed rank before running the demo. If the output changes, inspect the corpus provenance, boundary handling, counting policy, tie rule, and replacement policy separately rather than treating all “BPE” implementations as interchangeable.

Freeze ranks before applying them

This chapter has learned an ordered rule table, but it has not yet defined how to encode arbitrary input. The distinction is deliberate: training discovers ranks; encoding later applies the already frozen ranks without recounting candidates on the new text.

Chapter 4 will reserve BOS and EOS, shift every Chapter 3 content ID by two, apply these ranks to arbitrary UTF-8 bytes, and decode each ID through its stored byte expansion. The order matters because an earlier merge can create an operand for a later one. Validation and test bytes remain absent from every merge statistic, even when Chapter 4 tests round trips on held-out examples.