Extend BPE with bounded local updates

The problem: repeated scans must not change learned choices

Our byte-pair encoding trainer counts adjacent token pairs, selects one pair, replaces its occurrences and repeats. The first course’s Chapter 3 implemented that scalar trainer; Chapter 4 applied its frozen ranks and decoded the resulting tokens. Both are already in the practical course’s starting crate. Here we extend the existing BpeTrainer and BpeTokenizer.

The scalar trainer recounts every current training sequence at every merge rank. With more documents and more requested merges, those repeated scans revisit many edges that have not changed.

We can retain pair statistics and update the edges affected by a replacement. Local updates avoid recounting unaffected edges, but introduce a correctness obligation: the retained counts must describe the current sequences exactly. An old count or an out-of-order replacement can change the learned vocabulary even if the result still decodes to the original bytes.

We will preserve the course’s policy: learn from training documents only, count overlapping adjacent occurrences, break frequency ties by numeric token IDs, and replace matches left to right without sharing an input token between two replacements. A faster implementation of that policy must keep those decisions.

Work through two merge ranks

Use two manually prepared training documents, abab and abac. They are separate documents, so the last byte of one never forms a pair with the first byte of the other. The held-out text acac will be encoded after learning; it supplies no pair counts.

Keep the existing public token layout. Program IDs 0 and 1 are BOS and EOS; one-byte content IDs run from 2 through 257. The bytes a, b and c therefore have IDs 99, 100 and 101. The first learned content ID is 258. There is no PAD token. BOS and EOS are structural controls and do not enter training pair counts.

Before rank $0$, the current sequences are [99,100,99,100] and [99,100,99,101]. The pair a,b occurs $3$ times, b,a occurs $2$ times, and a,c occurs once. Select a,b, assign its merged bytes ab the new ID 258, and replace its nonoverlapping occurrences. The sequences become [258,258] and [258,99,101].

Before rank $1$, there are only three adjacent occurrences: (258,258), (258,99) and (99,101), each with count $1$. Numeric pair order compares the left IDs first. The pair (99,101) wins the tie, so the bytes ac receive ID 259. The two training sequences are now [258,258] and [258,259].

The frozen rules encode held-out acac as [259,259]. Decoding those content tokens reconstructs acac; document framing would instead give [0,259,259,1]. Encoding held-out text applies the learned rules and does not learn another rule.

Select from the current counts

$p_t=\min_{\mathrm{lex}}\operatorname*{argmax}_{(a,b)}c_t(a,b)$

First collect all pairs having the largest current count. Then choose the numerically smallest pair from that set: smallest left ID, then smallest right ID. If there are no adjacent pairs, learning stops. The prepared-training policy can require an exact merge count. With that policy enabled, training fails if the documents run out of pairs before that count is reached; it does not invent vocabulary entries. Without it, the requested count is an upper limit, as in the original scalar trainer.

What the symbols count

• $t$ is the zero-based merge rank, counting completed earlier merge rounds.
• $a,b$ are the left and right content token IDs of an adjacent pair in a current training document.
• $c_t(a,b)$ is the integer number of adjacent occurrences of that pair in the current training sequences at rank $t$; overlaps count and document boundaries never contribute.
• $p_t$ is the selected pair: maximum current count, with ties resolved by the smallest left ID and then the smallest right ID.

For aaa, two overlapping a,a occurrences contribute to the count, but left-to-right replacement consumes the first two bytes together and leaves the third byte alone. That produces one replacement. For aaaa, three occurrences contribute to the count and two nonoverlapping replacements are possible. Frequency and replacement count answer different questions.

Keep live edges, then apply learned ranks

BpeTrainer::train_prepared uses Chapter 2’s caller-supplied prepared-corpus reader. The caller supplies the training-only stream, its expected SHA-256 and the identity of the accepted corpus evidence. Training verifies the exact stream bytes and rejects any supplied split metadata that is not train. These checks do not discover overlap or establish rights, privacy or suitability; accepting that preparation evidence remains the caller’s responsibility.

The existing trainer keeps its scalar learning path. The new method adds a bounded working representation behind that same type; it does not define a second tokenizer. The existing BpeTokenizer still owns token layout, frozen rules, vocabulary expansions, document framing and byte decoding.

    /// Extend this trainer to bounded, selected prepared documents with local pair updates.
    pub fn train_prepared<R: std::io::Read>(
        self,
        reader: R,
        read_limits: crate::data::prepared_corpus::ReadLimits,
        binding: super::policy::TrainingBinding,
        policy: super::policy::TokenizerPolicy,
        limits: super::policy::TokenizerLimits,
        trace: bool,
    ) -> Result<PreparedBpeTraining, super::policy::TokenizerError> {
        if self.max_merges != policy.merge_count {
            return Err(super::policy::TokenizerError::Policy(
                "trainer and serialized merge budgets differ",
            ));
        }
        super::incremental::train_prepared(reader, read_limits, binding, policy, limits, trace)
    }

Each ingested byte becomes a token node with stable position, predecessor and successor links. The input text is then released. Consumed nodes retain their arena slots, so the node limit counts every ingested byte position. Node slots and pair indexes stay in memory under caller-selected bounds. Streaming input therefore does not mean constant-memory training or an external-memory algorithm.

A pair index records its current count and ordered live occurrence positions. A priority queue puts the greatest count first, using the smallest numeric pair when counts tie. Count updates can leave old queue entries behind. Each update changes a generation number; a popped entry is usable only if its count and generation still match the current pair state. A queue at its configured capacity is rebuilt from current pairs, preserving the same winner order.

    fn winner(&mut self) -> Option<(TokenPair, u64)> {
        while let Some(entry) = self.heap.pop() {
            let pair = entry.pair.0;
            if self.pairs.get(&pair).is_some_and(|state| {
                state.count == entry.count && state.generation == entry.generation
            }) {
                return Some((pair, entry.count));
            }
        }
        None
    }

Replacing one pair removes the edge entering the left node, the selected edge, and the edge leaving the right node, when those edges exist. The merged node then supplies at most two new edges. Updating those edges immediately keeps adjacent replacements from subtracting a shared old edge twice. Ordered live positions enforce leftmost replacement; consumed overlapping occurrences vanish.

    fn replace(&mut self, pair: TokenPair, output: u32) -> Result<u64, TokenizerError> {
        let mut replacements = 0u64;
        // Stable indices follow source order; first() gives the leftmost live match.
        while let Some(left) = self
            .pairs
            .get(&pair)
            .and_then(|s| s.occurrences.first())
            .copied()
        {
            if self.edge(left) != Some(pair) {
                return Err(TokenizerError::Invariant("stale live occurrence"));
            }
            let right = self.nodes[left]
                .next
                .ok_or(TokenizerError::Invariant("pair has no right node"))?;
            let previous = self.nodes[left].previous;
            let next = self.nodes[right].next;
            // Three distinct edge starts disappear; adjacent replacements cannot
            // remove a shared old edge twice because the lists update immediately.
            if let Some(index) = previous {
                self.remove_edge(index)?;
            }
            self.remove_edge(left)?;
            self.remove_edge(right)?;
            self.nodes[left].token = output;
            self.nodes[left].next = next;
            self.nodes[right].alive = false;
            if let Some(index) = next {
                self.nodes[index].previous = Some(left);
            }
            if let Some(index) = previous {
                self.add_edge(index)?;
            }
            self.add_edge(left)?;
            replacements = replacements
                .checked_add(self.weights[self.nodes[left].document])
                .ok_or(TokenizerError::Overflow)?;
        }
        Ok(replacements)
    }

The scalar reference reuses the original trainer’s complete learning loop: recount every current sequence, choose the winner and replace left to right. train_byte_fixture supplies the same two byte documents to that loop. The demo adapts its rules and counters into the artifact metadata used by the new path; it does not write another learning loop. Both canonical payloads and their hashes appear in the output, so their agreement can be inspected alongside the different working representations.

fn scalar_fixture(
    template: &artifact::TokenizerArtifact,
    documents: &[PreparedDocument],
) -> Result<artifact::TokenizerArtifact, String> {
    let bytes = documents
        .iter()
        .map(|document| document.text.as_bytes().to_vec())
        .collect::<Vec<_>>();
    let learned = BpeTrainer::new(template.policy.merge_count)
        .train_byte_fixture(&bytes)
        .map_err(|error| error.to_string())?;
    let tokenizer = BpeTokenizer::from_training(&learned).map_err(|error| error.to_string())?;
    let merges = learned
        .rules()
        .iter()
        .zip(tokenizer.merge_rules())
        .map(|(learned, frozen)| {
            Ok(artifact::ArtifactMerge {
                candidate_count: u64::try_from(learned.candidate_count())
                    .map_err(|error| error.to_string())?,
                left: frozen.content_pair().left(),
                rank: frozen.rank(),
                replacements: u64::try_from(learned.replacement_count())
                    .map_err(|error| error.to_string())?,
                right: frozen.content_pair().right(),
                token_id: frozen.content_token_id(),
            })
        })
        .collect::<Result<Vec<_>, String>>()?;
    let count = |lengths: Vec<usize>| -> Result<u64, String> {
        lengths.into_iter().try_fold(0u64, |total, length| {
            total
                .checked_add(u64::try_from(length).map_err(|error| error.to_string())?)
                .ok_or_else(|| "scalar fixture count overflow".to_owned())
        })
    };
    let training_documents = u64::try_from(documents.len()).map_err(|error| error.to_string())?;
    let mut result = template.clone();
    result.merges = merges;
    result.counters = artifact::EfficiencyCounters {
        content_tokens: count(learned.final_sequences().iter().map(Vec::len).collect())?,
        control_tokens: training_documents
            .checked_mul(2)
            .ok_or("scalar fixture control count overflow")?,
        input_bytes: count(bytes.iter().map(Vec::len).collect())?,
        training_documents,
    };
    Ok(result)
}

Application has a different priority: the earliest learned rank among currently adjacent eligible pairs, followed by the leftmost position. The encoder rechecks adjacency and node versions before using a queued candidate, replaces that pair, and queues its new neighbors. Rules are topological: a merge’s operands already exist before that rule’s output ID is created.

    while let Some(entry) = heap.pop() {
        let Reverse((rank, left, right, _, _)) = entry;
        if candidate(&nodes, left, &ranks) != Some(entry) {
            continue;
        }
        let output = tokenizer.merge_rules()[rank].content_token_id();
        let previous = nodes[left].previous;
        let next = nodes[right].next;
        nodes[left].token = output;
        nodes[left].next = next;
        nodes[left].version = nodes[left]
            .version
            .checked_add(1)
            .ok_or(TokenizerError::Overflow)?;
        nodes[right].alive = false;
        if let Some(index) = next {
            nodes[index].previous = Some(left);
        }
        if let Some(index) = previous {
            enqueue(&mut heap, &nodes, index, &ranks, limits)?;
        }
        enqueue(&mut heap, &nodes, left, &ranks, limits)?;
    }

Read chunks are transport boundaries, not token boundaries. encode_reader_bounded reads one whole bounded document before encoding; callers invoke it separately for each document. Literal text such as <|endoftext|> remains content bytes. The existing document-framing operation inserts BOS and EOS after content encoding. Byte decoding is exact; text decoding checks UTF-8 after concatenating all token fragments, since an individual token may contain only part of a UTF-8 character.

max_nodes bounds all ingested training byte positions in one training call, and byte positions in one encoded document. Bounded decoding uses that same numeric limit as the maximum reconstructed content-byte count for one document; it checks the total before allocating the output. A short token sequence can still expand into too many bytes. The existing vocabulary constructor separately bounds each token’s expansion with max_token_bytes and the sum of all content expansions with max_vocabulary_bytes.

Preserve the artifact and its evidence boundary

The tokenizer artifact records its schema and layout, raw-byte policy, training binding, requested merge count, ordered learned rules and efficiency counters. Validation rejects unknown fields or versions, duplicate pairs, gaps, forward operands and inconsistent reduction counts. Vocabulary expansions and output serialization have explicit capacity checks.

    /// Validate metadata, then ask the existing tokenizer to build its vocabulary.
    pub fn validate(&self, limits: TokenizerLimits) -> Result<BpeTokenizer, TokenizerError> {
        limits.validate()?;
        self.policy.validate()?;
        self.training.validate()?;
        if self.schema_version != 1
            || self.layout_version != TOKENIZER_LAYOUT_VERSION
            || self.trainer_semantics != "overlapping-count-leftmost-replace-v1"
        {
            return Err(TokenizerError::Artifact(
                "unknown tokenizer format or semantics",
            ));
        }
        if self.merges.len() > limits.max_nodes {
            return Err(TokenizerError::Resource(
                "merge table exceeds working capacity",
            ));
        }
        if self.merges.len() > self.policy.merge_count
            || (self.policy.exact_merge_count && self.merges.len() != self.policy.merge_count)
        {
            return Err(TokenizerError::Artifact(
                "artifact violates requested merge count",
            ));
        }
        if self.counters.control_tokens
            != self
                .counters
                .training_documents
                .checked_mul(2)
                .ok_or(TokenizerError::Overflow)?
            || self.counters.content_tokens > self.counters.input_bytes
        {
            return Err(TokenizerError::Artifact("inconsistent efficiency counters"));
        }
        let layout = TokenizerLayout::new(self.merges.len())?;
        let mut pairs = Vec::new();
        pairs
            .try_reserve_exact(self.merges.len())
            .map_err(|_| TokenizerError::Resource("merge operand allocation failed"))?;
        let mut reduction = 0u64;
        for (rank, rule) in self.merges.iter().enumerate() {
            if rule.rank != rank
                || layout.merge_token_id(rank) != Some(rule.token_id)
                || rule.replacements == 0
                || rule.replacements > rule.candidate_count
            {
                return Err(TokenizerError::Artifact("invalid rank or merge count"));
            }
            let left = rule
                .left
                .checked_sub(CONTENT_ID_OFFSET)
                .ok_or(TokenizerError::Artifact(
                    "control ID cannot be a merge operand",
                ))?;
            let right =
                rule.right
                    .checked_sub(CONTENT_ID_OFFSET)
                    .ok_or(TokenizerError::Artifact(
                        "control ID cannot be a merge operand",
                    ))?;
            pairs.push(TokenPair::new(left, right));
            reduction = reduction
                .checked_add(rule.replacements)
                .ok_or(TokenizerError::Overflow)?;
        }
        if self.counters.input_bytes.checked_sub(reduction) != Some(self.counters.content_tokens) {
            return Err(TokenizerError::Artifact(
                "merge reductions disagree with content count",
            ));
        }
        // This single constructor owns prior-rank operands, duplicate pairs,
        // byte expansions and their bounds; there is no second vocabulary here.
        Ok(BpeTokenizer::from_merge_pairs_bounded(
            &pairs,
            limits.max_token_bytes,
            limits.max_vocabulary_bytes,
        )?)
    }

Canonical payload bytes identify the learned tokenizer and its declared training lineage. Producer timestamps, local paths and machine details belong in a separate execution receipt. The two training strategies can emit equal tokenizer payload bytes while having different producer receipts. Equal vocabulary size alone does not identify the rules or their training source.

The fixture finishes with $4$ content tokens over $8$ input bytes: $4/8=0.5$ content tokens per byte. Framing the two training documents adds $4$ structural controls, which are excluded from that numerator. This is a sequence-length ratio, not trained-model quality or throughput. Zero input bytes have no tokens-per-byte ratio. Actual memory use and speed on a larger corpus require measurements of that workload; these bounded examples supply no full-corpus result.

With Docker running, use the repository’s Rust workspace setup from the repository root. The following command runs the supplied CPU fixtures:

COURSE_CORPUS=false ./course run cargo run --offline --locked \
  -p practical-ch03-scalable-bpe-tokenizer

The JSON output includes two merge rounds, held-out content [259,259], equal scalar and incremental payload hashes, content/byte/control counters, storage counters, refusal examples and the separate policy contrast below. It loads no large corpus or model weights.

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let output =
        practical_ch03_scalable_bpe_tokenizer::fixture_output().map_err(std::io::Error::other)?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

History: subwords still need a boundary policy

A whole-word vocabulary leaves rare and unseen words difficult to represent; learned subwords reduce that problem, but their alphabet and boundary rules still determine which merges are possible.

In $2016$, Neural Machine Translation of Rare Words with Subword Units introduced BPE subword learning for neural translation. Sennrich and colleagues learn subword units by repeatedly merging frequent symbol pairs, with word boundaries retained in their character-based representation. That representation is different from treating a complete document as one raw-byte sequence.

The published GPT-2 encoder applies a reversible byte-to-Unicode mapping, an explicit pretokenization pattern and ranked merges within each resulting piece. The pretokenization boundaries determine where a learned merge can apply. Inspect the pinned encoder implementation. The Unicode symbols represent bytes; the mapping does not normalize the input text.

Isolate one letter-digit boundary in Rust

Our chapter-local Rust example implements the published byte mapping and pretokenization pattern. It reuses the course tokenizer’s existing frozen-rank engine inside each piece, with one hypothetical merge of a followed by 1. This isolates boundary behavior without implementing BPE again. The demo’s tiny vocabulary defines its displayed IDs; they are not the pretrained GPT-2 vocabulary’s IDs.

        let mut byte_order = (b'!'..=b'~')
            .chain(0xa1..=0xac)
            .chain(0xae..=0xff)
            .collect::<Vec<_>>();
        let mut symbols = byte_order
            .iter()
            .map(|&byte| char::from(byte))
            .collect::<Vec<_>>();
        let mut next_codepoint = 256u32;
        for byte in 0..=255u8 {
            if !byte_order.contains(&byte) {
                let symbol = char::from_u32(next_codepoint).ok_or("invalid byte-map symbol")?;
                byte_order.push(byte);
                symbols.push(symbol);
                next_codepoint += 1;
            }
        }
        let mut bytes = ['\0'; 256];
        let mut reverse_bytes = BTreeMap::new();
        let mut vocabulary = BTreeMap::new();
        for (id, (&byte, &symbol)) in byte_order.iter().zip(&symbols).enumerate() {
            bytes[byte as usize] = symbol;
            reverse_bytes.insert(symbol, byte);
            vocabulary.insert(symbol.to_string(), id as u32);
        }

The raw-byte fixture learns a1 from a1a1 and encodes a1 as one content token. The published GPT-2 pattern splits a1 into the letter piece a and the digit piece 1 before BPE. The proposed merge cannot cross those pieces, so this policy example emits two content tokens. Both paths reconstruct the same input bytes. Reconstruction therefore does not imply equal segmentation.

        for matched in self.pretokenizer.find_iter(text) {
            let piece = matched.map_err(|error| error.to_string())?.as_str();
            pieces.push(piece.to_owned());
            for content_id in self.tokenizer.encode_content(piece.as_bytes()) {
                let expansion = self
                    .tokenizer
                    .token_bytes(content_id)
                    .ok_or("unknown rank-engine token")?;
                let symbols = expansion
                    .iter()
                    .map(|&byte| self.bytes[byte as usize])
                    .collect::<String>();
                token_ids.push(
                    *self
                        .vocabulary
                        .get(&symbols)
                        .ok_or("missing vocabulary symbol")?,
                );
                // The published policy reverses the byte map before decoding UTF-8.
                token_bytes.push(
                    symbols
                        .chars()
                        .map(|symbol| {
                            self.reverse_bytes
                                .get(&symbol)
                                .copied()
                                .ok_or("unknown byte-map symbol")
                        })
                        .collect::<Result<Vec<_>, _>>()?,
                );
            }
        }

The regex engine executes the fixed published pattern; Rust owns the byte map, the existing rank decisions and the illustrative vocabulary IDs. The byte map is reversible and does not create a merge across separate pieces. Here the rank engine works on the pieces’ raw bytes; the resulting byte fragments are then expressed through the map and tiny vocabulary. This demonstrates selected published operations rather than reproducing the pretrained tokenizer. Its bounded results do not establish agreement with every behavior of the original Python regex runtime.

Freeze the training selection and tokenizer policy before constructing decoder token targets. Improve the implementation without changing those learned decisions or confusing byte reconstruction with equal segmentation.

Follow the changing counts and the separate boundary example

The count tables refer to the representation immediately before each rank. They change after replacement. The scalar and incremental course paths reach the same learned payload; the letter-digit example changes the boundary policy and consequently the eligible merges.

Optional practice: inspect counts, order and boundaries

1. Run the overlap test below. Explain why aaa has two counted a,a occurrences but only one replacement.
2. Run the third-rank demo test below. Explain why (258,258) wins the next tie over (258,259) after the main two-rank example.
3. Use the demo command above to inspect gpt2_policy_contrast. Explain why reconstructing a1 in both paths does not imply that both paths emit one token.

From the repository root, with Docker running and the Rust workspace setup available, run:

COURSE_CORPUS=false ./course run cargo test --offline --locked \
  -p llm-from-scratch-practical --test practical_ch03_tokenizer \
  leftmost_overlap -- --exact

Expect running 1 test and 1 passed; 0 failed. This test compares counted overlaps and nonoverlapping replacement sequences using the existing scalar trainer.

From the same working directory and setup, inspect the checked extra rank:

COURSE_CORPUS=false ./course run cargo test --offline --locked \
  -p practical-ch03-scalable-bpe-tokenizer --lib \
  tests::optional_third_rank_is_the_checked_course_exercise -- --exact

Expect running 1 test and 1 passed; 0 failed. Zero selected tests means the requested example did not run.

Hand frozen token sequences to batching

Use the frozen course tokenizer to encode each prepared document separately. Held-out encoding does not fit new merges. Chapter 4 will place those token sequences into variable-length batches without changing tokenizer identity. The embedding and output layers must continue to interpret each content ID according to that same learned vocabulary.