Learning deterministic BPE merges

Give repeated bytes a reusable token

Chapter 1 separated text units from the IDs that name them. A table of whole words can give an observed word one ID, but a finite spelling list needs a fallback for an unseen word. Starting with every possible byte avoids that spelling limit: every UTF-8 document is already a sequence of those bytes. The cost is that a repeated byte sequence still occupies several token positions. Byte-pair encoding (BPE) learns larger reusable units from those repetitions.

Here a token ID names a vocabulary entry, and its byte expansion is the exact ordered sequence of bytes that entry represents. Initially, IDs 0 through 255 each represent the corresponding single byte. An ID is a label; it is not a Unicode code point or a measure of meaning. Learning new entries is tokenizer fitting. No neural network weights are optimized in this chapter.

Chapter 2 fixed the document roles before fitting. The trainer reads only the training documents and keeps a separate token sequence for each one. To explain one fit completely, begin with these two tiny training documents:

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

The ASCII byte for a is decimal 97, written 61 in hexadecimal; b is decimal 98, or hex 62. The first round counts every pair of adjacent token positions. In aaa, (97,97) occurs at positions 0–1 and 1–2. The middle token belongs to both windows, so these are two candidate occurrences. The separate document aba contributes one (97,98) and one (98,97):

The largest count belongs to (97,97). This first learned rule has merge rank zero, meaning position zero in the ordered rule table. It receives fresh ID 256, whose byte expansion is 61 61.

Replacement then scans each document once from left to right. In aaa, the match at positions 0–1 consumes both input tokens and emits 256. The scan resumes at position 2 and copies the remaining 97. The middle input token has already been consumed, so it cannot take part in another replacement. aba has no match for this rule and stays unchanged:

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

The candidate count is two and the replacement count is one. The first number chooses a rule using all adjacent windows; the second counts the matches actually consumed by the replacement pass. One does not determine the other when occurrences overlap.

The second round recounts the current sequences. Its candidates are now (97,98), (98,97), and (256,97), each occurring once. The old count for (97,97) is no longer relevant. This course resolves equal counts by comparing left numeric IDs first, then right numeric IDs. (97,98) wins because its left ID 97 is smaller than 98 and 256.

Rank one receives fresh ID 257 with byte expansion 61 62. Its one match is the first two tokens of aba, so the next sequences are:

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

This rule is learned even though its count is one; the trainer does not require a pair to occur at least twice. Across the two documents, the number of token positions falls from six to five to four. Expanding [256,97] gives 61 61 61, and expanding [257,97] gives 61 62 61: each document still has its original bytes. The vocabulary grows from 256 entries to 258 entries; vocabulary entries, token positions, and byte length count different things.

A document boundary ends the candidate sequence. For two one-byte documents a | a, the bar depicts the boundary and is not input data. Neither document has two positions, so there is no candidate (97,97). Within one document, spaces are ordinary bytes; this trainer has no separate word or character boundary rule.

State the rule for every round

The example repeats the same sequence of operations: count the current adjacent windows inside each document, choose one pair, give its joined bytes a fresh ID, and replace that pair in one pass per document. Recount only after those passes finish. The selection and new entry are written as:

The score has three parts and is compared lexicographically: the first unequal part decides. Maximize the candidate count $C(a,b)$ first. For equal counts, maximizing $-a$ prefers the smaller left numeric ID; if the left IDs also match, maximizing $-b$ prefers the smaller right ID. Thus the second round’s scores $(1,-97,-98)$, $(1,-98,-97)$, and $(1,-256,-97)$ select $(97,98)$. This tie policy is an explicit course choice.

The stars mark the selected pair. In the second expression, joining $a^{*}$ and $b^{*}$ means concatenating their byte expansions to define the new entry $m^{*}$. For example, token 256 expands to 61 61 and token 97 to 61; a later rule for (256,97) would store 61 61 61. The numeric IDs are neither added nor joined as the decimal strings "256" and "97".

The Rust code expresses the same priorities without negating unsigned IDs: windows(2) supplies the overlapping candidates, ordered numeric pairs make the smaller pair come first, and the selector changes its winner only for a strictly larger count. A separate replacement cursor consumes two input tokens on a match and one otherwise. Newly emitted tokens are output for this round; the pass does not rescan them.

The requested number of merges is an upper bound. An empty candidate map ends training early; otherwise even a count-one winner can add a rule. Requesting zero rounds leaves the original sequences and the 256-entry byte vocabulary unchanged.

Symbol glossary

Raw byte entries use IDs 0..=255. If $r$ is a rule’s zero-based rank, its new ID is $256+r$. If $k$ rounds succeed, the vocabulary contains $256+k$ entries, including all 256 original bytes. These expressions count and label vocabulary entries; they do not state how many token positions a document contains.

From closed word tables to repeated pair merging

A whole-word table makes an observed word a single lookup unit. The small Rust contrast fits the whitespace-delimited words in low lower and new newest. Sorting the four observed spellings gives lower ID 2. The table has no entry for either lowering or lowest, so both return the same unknown ID 0. That lookup loses the distinction between their spellings.

The whole-word example has its own ID space. Its unknown ID 0 is unrelated to the byte trainer’s ID 0, which represents the single byte 00.

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 “A New Algorithm for Data Compression” repeatedly substituted common adjacent byte pairs with unused bytes and stored a substitution table for expansion. Reusing one symbol for a repeated pair is the useful connection here. His finite unused-byte codes and compression procedure differ from this trainer’s fresh u32 IDs and full recount after each round.

Rico Sennrich, Barry Haddow, and Alexandra Birch adapted repeated pair merging for rare words in neural translation. Their 2016 subword paper, section 3.2 starts with characters, represents words separately with an end-of-word marker, and merges character sequences within a word. Words can therefore share learned pieces instead of each needing a separate whole-word entry.

The course variant instead starts with all 256 bytes and uses document boundaries as its barriers. A learned piece may include a space or end partway through a UTF-8 character. Neither historical source supplies our numeric tie policy, and pair merging does not guarantee that a piece is a linguistic morpheme. The whole-word contrast demonstrates a lookup limitation; it does not compare neural model quality or encode unseen text with BPE yet.

Read the same operations in Rust

The existing pair counter iterates over each document’s token vector. Its windows(2) call includes both windows in aaa; starting a new iterator for each document excludes pairs across documents. TokenPair orders candidates by the numeric left ID and then the right ID. The BTreeMap visits that order, and count > best_count replaces a winner only for a larger count. Equal counts therefore retain the smallest pair. These operations implement the three priorities in the selection formula:

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
}

The replacement function reads the input with an index and writes a separate output. A match emits the new ID and advances by two; a non-match copies one ID and advances by one. Since the cursor never returns to a consumed token or reads the output as new input, one pass cannot perform overlapping replacements:

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)
}

BpeTrainer::train accepts Chapter 2’s validated CorpusPartitions. It obtains training_documents(), records those documents’ stable IDs, and converts their bodies’ bytes to separate initial token vectors. Validation and test documents contribute no pair counts in this pipeline. The training view expresses this dependency; it is not a security barrier preventing other code from reading held-out data:

    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 training loop starts with the 256 single-byte expansions. Each successful round uses the current vocabulary length for its new ID, concatenates the winning entries’ bytes, replaces the pair separately in every document, and retains the rule and expansion. At the next iteration it counts the resulting sequences again. Each stored rule records both candidate count and replacement count, so the overlap distinction remains inspectable.

The executable performs two separate demonstrations. First it requests eight merges from the eight Chapter 2 training documents, holding out two validation and two test documents. Then it runs the two-round aaa/aba example. The tiny trace is not the first two ranks of the corpus fit:

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(())
}

The corpus fit’s rank-zero pair is (32,208), with candidate count 81, replacement count 81, fresh training ID 256, and hex expansion 20 d0. Here 20 is a space, and d0 begins a two-byte UTF-8 sequence whose continuation is absent from this piece. The expansion is therefore not valid standalone UTF-8. This follows from the byte-sequence rules in RFC 3629, sections 3–4.

One token need not be one character, one word, or even a complete UTF-8 string. The trainer retains bytes rather than displaying an invented character for the piece. Concatenating every token’s expansion in the complete document still preserves that document’s original bytes.

Compare the two rounds in one figure

Each stage has one row for train-aaa and one for train-aba; the document barrier prevents the end of one row from pairing with the start of the other. The first table’s count 2 refers to the overlapping windows in aaa, while its replacement count 1 refers to the one match consumed. The next stage shows the changed sequence before the second table recounts it.

The second table records three count-one candidates. Its selected row is (97,98) because numeric pair order resolves the tie. The final stage shows [256,97] and [257,97], so the two replacements and the preserved document identities can be read together. Selection is marked with words and a star as well as a border.

Optional reproduction and practice

To reproduce the existing evidence, use a checked-out repository with its pinned Rust toolchain and locked dependencies available. Run these commands from the repository root:

cargo test --locked -p llm-from-scratch --lib tokenizer::bpe_trainer::tests::
cargo test --locked -p ch03-learn-bpe-merges
cargo run --quiet --locked -p ch03-learn-bpe-merges | diff -u rust/demos/ch03-learn-bpe-merges/expected.txt -

The first command runs six trainer tests; the second runs one demo library test. Successful tests report no failures. The last command compares the demo’s standard output with the checked-in expected output; a successful comparison prints nothing and exits with status zero. That output contains both the eight corpus rules and the separate tiny trace explained above.

Use the checked answers to inspect or explain these consequences of the rule:

1. For the single document [97,97,97,97], compare the adjacent-window count with the number of matches consumed by one left-to-right pass.
2. Inspect the three score tuples after rank zero in the tiny example and explain why (97,98) is selected.
3. Explain the absent candidate (97,98) in separate documents [97] | [98]. Here | depicts a document boundary and is not input data.
4. Relate the vocabulary size to nine successful rounds out of twelve requested. Also state what a request for zero rounds retains.
5. Explain why changing repetition counts in a validation document supplies no new pair statistics to this trainer. This is a reasoning task; no corpus edit is needed.
6. Explain the byte expansion of a hypothetical rule (256,97) when token 256 stores 61 61 and token 97 stores 61.
7. Explain why the corpus rule with expansion 20 d0 cannot be printed as a standalone character, and why that does not invalidate the complete document.

Carry the fitted vocabulary toward the model

The result of fitting is an ordered rule table, a byte expansion for every vocabulary entry, and the transformed training sequences. The tiny example uses fewer token positions while adding vocabulary entries. That is the demonstrated tradeoff; this fixture does not establish optimal compression, lower next-token loss, or more meaningful word parts.

The eventual decoder will use a finite vocabulary of IDs for its inputs and outputs. BPE supplies reusable content units for that vocabulary. Learning the tokenizer’s table is separate from learning the neural weights that will model sequences of those IDs.

Chapter 4 freezes the learned rules and applies them to new input without refitting or counting validation and test frequencies. It also reserves BOS ID 0 and EOS ID 1, shifting every Chapter 3 content ID by two. Training ID 256 therefore becomes content ID 258. The stored byte expansions cross that handoff unchanged; the next chapter explains applying the frozen rules and recovering the exact input bytes.