Byte Pair Encoding in Python: 5 Steps From Naive to Fast

Build a BPE tokenizer trainer from scratch, then make it ~220x faster across five incremental optimizations — each removing a different bottleneck.

deep-learningnlppython

Byte Pair Encoding (BPE) is the subword-tokenization algorithm behind GPT-2, GPT-4, Llama, and most modern tokenizers. The idea is beautifully simple: repeatedly find the most frequent adjacent pair of symbols and merge it into a new symbol, until you have the vocabulary size you want.

This post builds a correct BPE trainer in Python, then optimizes it across five steps. Each step removes one specific bottleneck and is verifiably faster, while producing exactly the same merges as the step before.

The algorithm, in one paragraph

Start with a corpus split into words, each word a sequence of byte IDs. Count every adjacent pair (a, b) across all words. Pick the most frequent pair, declare it a new token, and rewrite every occurrence (a, b) → t in the corpus. Repeat num_merges times. The recorded sequence of merges is the learned tokenizer.

More formally, at each round we choose:

t=arg⁡max⁡(a,b)  freq(a,b)t = \arg\max_{(a,b)}\; \mathrm{freq}(a,b)

and then update the corpus so that every non-overlapping, left-to-right occurrence of (a,b) collapses into a fresh token id 256, 257, 258, \dots

Shared setup

Every version below starts from the same pretokenized corpus, so the only thing that changes between steps is how we count and merge:

from collections import Counter

text = open("corpus.txt", encoding="utf-8").read()
num_merges = 500

# Pretokenize on whitespace: each word becomes a list of UTF-8 byte IDs.
words = [list(w.encode("utf-8")) for w in text.split()]

The timings quoted for each step are representative numbers on a ~200 KB English corpus, 500 merges, Python 3.12. Absolute values depend on hardware; the ratios are what matter.

Step 1 — Naive baseline (24.3 s)

Recount every pair from scratch each round, then rebuild every word each round. Simple, obviously correct, and O(num_merges · N) where N is the total token count.

def train_naive(words, num_merges):
    words = [list(w) for w in words]          # mutable copy
    merges = []
    for _ in range(num_merges):
        counts = {}
        for w in words:                        # recount ALL pairs
            for i in range(len(w) - 1):
                p = (w[i], w[i + 1])
                counts[p] = counts.get(p, 0) + 1
        if not counts:
            break
        pair = max(counts, key=lambda p: (counts[p], p))
        new_id = 256 + len(merges)
        for w in words:                        # rewrite EVERY word
            i, out = 0, []
            while i < len(w):
                if i < len(w) - 1 and w[i] == pair[0] and w[i + 1] == pair[1]:
                    out.append(new_id); i += 2
                else:
                    out.append(w[i]); i += 1
            w[:] = out
        merges.append(pair)
    return merges

The tie-break key=lambda p: (counts[p], p) is deliberate and used in every step: highest count first, then the largest pair tuple — fully deterministic, which is what lets us prove later steps produce identical merges.

Bottleneck: we redo O(N) work every single round even though a merge only touches a tiny fraction of the corpus.

Step 2 — Deduplicate words (2.70 s, ~9x)

Natural language has enormous word repetition. Instead of holding 50,000 word lists that are mostly copies of each other, collapse to unique words and store a frequency. Now pair counts are weighted sums, and we rewrite each unique word only once.

def train_dedup(words, num_merges):
    freq = Counter(tuple(w) for w in words)   # unique word -> count
    words = [list(w) for w in freq]
    freqs = list(freq.values())
    merges = []
    for _ in range(num_merges):
        counts = {}
        for w, f in zip(words, freqs):
            for i in range(len(w) - 1):
                p = (w[i], w[i + 1])
                counts[p] = counts.get(p, 0) + f   # weighted by frequency
        if not counts:
            break
        pair = max(counts, key=lambda p: (counts[p], p))
        new_id = 256 + len(merges)
        for w in words:                         # rewrite each UNIQUE word once
            i, out = 0, []
            while i < len(w):
                if i < len(w) - 1 and w[i] == pair[0] and w[i + 1] == pair[1]:
                    out.append(new_id); i += 2
                else:
                    out.append(w[i]); i += 1
            w[:] = out
        merges.append(pair)
    return merges

Why it’s equivalent: the original word list is exactly freq[w] copies of each unique word, so freq[w] · occurrences_in_word equals the original occurrence count — the pair counts are identical, so the chosen merge is identical, inductively for every round.

Bottleneck removed: duplicate work. Remaining bottleneck: we still recount every pair every round.

Step 3 — Incremental pair counts (0.81 s, ~30x)

A merge only changes pairs in the neighborhood of each merged occurrence. So maintain the counts dictionary across rounds and patch only what changed: subtract the old pairs of an affected word, rewrite it, add the new pairs.

def train_incremental(words, num_merges):
    freq = Counter(tuple(w) for w in words)
    words = [list(w) for w in freq]
    freqs = list(freq.values())
    merges = []
    counts = {}

    def add_pairs(w, f, delta):
        for i in range(len(w) - 1):
            p = (w[i], w[i + 1])
            counts[p] = counts.get(p, 0) + delta * f

    for w, f in zip(words, freqs):
        add_pairs(w, f, +1)                      # build counts once

    for _ in range(num_merges):
        if not counts:
            break
        pair = max(counts, key=lambda p: (counts[p], p))
        new_id = 256 + len(merges)
        for w, f in zip(words, freqs):
            if pair[0] not in w:                 # cheap skip: first symbol absent
                continue
            add_pairs(w, f, -1)                  # remove old pairs
            i, out = 0, []
            while i < len(w):
                if i < len(w) - 1 and w[i] == pair[0] and w[i + 1] == pair[1]:
                    out.append(new_id); i += 2
                else:
                    out.append(w[i]); i += 1
            w[:] = out
            add_pairs(w, f, +1)                  # add new pairs
        merges.append(pair)
        counts.pop(pair, None)                   # merged pair is gone
    return merges

The if pair[0] not in w: continue guard avoids touching words that obviously can’t contain the pair. We still scan all words to find the affected ones — Step 4 removes that scan.

Bottleneck removed: the full O(N) recount each round. Remaining bottleneck: a linear scan over all unique words (and all pairs) every round.

Step 4 — Inverted index (0.24 s, ~100x)

Maintain a second dictionary, pair → {word indices that contain it}. To apply a merge, jump straight to the small set of words that actually contain the pair — no scanning the rest.

def train_inverted(words, num_merges):
    freq = Counter(tuple(w) for w in words)
    words = [list(w) for w in freq]
    freqs = list(freq.values())
    merges = []
    counts = {}
    pair2words = {}                              # pair -> set of word indices

    def register(wi, f, delta):
        w = words[wi]
        for p, c in Counter((w[i], w[i+1]) for i in range(len(w)-1)).items():
            counts[p] = counts.get(p, 0) + delta * c * f
            if delta > 0:
                pair2words.setdefault(p, set()).add(wi)
            else:
                s = pair2words.get(p)
                if s is not None:
                    s.discard(wi)
                    if not s:
                        del pair2words[p]

    for wi, (w, f) in enumerate(zip(words, freqs)):
        register(wi, f, +1)                      # build index + counts once

    for _ in range(num_merges):
        if not counts:
            break
        pair = max(counts, key=lambda p: (counts[p], p))
        new_id = 256 + len(merges)
        a, b = pair
        for wi in list(pair2words.get(pair, ())):   # only affected words
            f = freqs[wi]
            register(wi, f, -1)
            w = words[wi]
            i, out = 0, []
            while i < len(w):
                if i < len(w) - 1 and w[i] == a and w[i + 1] == b:
                    out.append(new_id); i += 2
                else:
                    out.append(w[i]); i += 1
            words[wi] = out
            register(wi, f, +1)
        merges.append(pair)
        counts.pop(pair, None)
        pair2words.pop(pair, None)
    return merges

Note Counter(...) inside register: counts are weighted by the number of occurrences of a pair in the word (a word like aaa contains (a,a) twice), while the inverted index only tracks membership.

Bottleneck removed: the per-round scan over all words. Remaining bottleneck: finding the best pair still scans every entry in counts.

Step 5 — Max-heap selection (0.11 s, ~220x)

The final bottleneck is max(counts, ...) — an O(P) scan over all pairs every round. Replace it with a max-heap (a min-heap storing (-count, pair)) and use lazy deletion: push a new entry whenever a count changes, and when popping, skip any entry whose stored count no longer matches the live counts[pair].

import heapq

def train_heap(words, num_merges):
    freq = Counter(tuple(w) for w in words)
    words = [list(w) for w in freq]
    freqs = list(freq.values())
    merges = []
    counts = {}
    pair2words = {}
    heap = []                                   # (-count, pair)

    def set_count(p, value):
        counts[p] = value
        heapq.heappush(heap, (-value, p))       # may create stale entries

    def register(wi, f, delta):
        w = words[wi]
        for p, c in Counter((w[i], w[i+1]) for i in range(len(w)-1)).items():
            set_count(p, counts.get(p, 0) + delta * c * f)
            if delta > 0:
                pair2words.setdefault(p, set()).add(wi)
            else:
                s = pair2words.get(p)
                if s is not None:
                    s.discard(wi)
                    if not s:
                        del pair2words[p]

    for wi, (w, f) in enumerate(zip(words, freqs)):
        register(wi, f, +1)

    def best_pair():
        while heap:
            neg, p = heap[0]
            if counts.get(p) != -neg:            # stale -> discard
                heapq.heappop(heap); continue
            return p
        return None

    for _ in range(num_merges):
        pair = best_pair()
        if pair is None or counts.get(pair, 0) <= 0:
            break
        new_id = 256 + len(merges)
        a, b = pair
        for wi in list(pair2words.get(pair, ())):
            f = freqs[wi]
            register(wi, f, -1)
            w = words[wi]
            i, out = 0, []
            while i < len(w):
                if i < len(w) - 1 and w[i] == a and w[i + 1] == b:
                    out.append(new_id); i += 2
                else:
                    out.append(w[i]); i += 1
            words[wi] = out
            register(wi, f, +1)
        merges.append(pair)
        counts.pop(pair, None)
        pair2words.pop(pair, None)
    return merges

Selection is now O(log P) amortized. This is essentially the algorithm inside production trainers like HuggingFace’s BpeTrainer and Andrej Karpathy’s minbpe.

The full picture

Bar chart of BPE training time across the five implementations on a log scale
Training time per implementation (log scale). Each step removes one bottleneck; merges are identical across all five.
StepTechniqueTimeSpeedupBottleneck removed
1Naive recount + rebuild24.3 s1×—
2Deduplicate words2.70 s~9×Duplicate words
3Incremental counts0.81 s~30×Full recount each round
4Inverted index0.24 s~100×Scanning all words per merge
5Max-heap selection0.11 s~220×Scanning all pairs per round

Each optimization targets a different part of the cost — that’s why they compose multiplicatively rather than overlap.

Where to go next

  • Better pretokenization — replace the whitespace split with the GPT-2 regex so punctuation, numbers, and CJK characters each get their own stream.
  • Parallelism — different words are independent, so the rewrite pass can be vectorized with NumPy or offloaded to Rust via a small extension.
  • Encoding — training is half the job; the other half is applying the learned merges to new text, which has its own efficient single-word algorithm (the GPT-2 _byte_pair_merge loop).

The leap from 24 seconds to a tenth of a second didn’t come from one clever trick — it came from removing five distinct bottlenecks, one at a time.