BPE Merge Step

Start from brute force

The obvious version handles one candidate at a time: list the distinct adjacent pairs, and for each one rescan the whole corpus to count its occurrences, keeping the best count under the tie-break. It is correct, because every candidate gets an exact count. With P distinct pairs and n symbols that is O(P · n) time, which approaches O(n²) on a varied corpus where most pairs are distinct, plus O(P) space for the candidates. A tokenizer trainer runs this step thousands of times, so the rescan is the bottleneck. Joining each word into a string and calling replace is faster but wrong: it erases symbol boundaries, so ["a", "bc"] and ["ab", "c"] both become "abc".

The approach

One pass can count every pair at once. The reference bpe_merge_step walks each word and bumps counts[(w[i], w[i + 1])] in a dict, so after reading n symbols every adjacent pair, overlapping ones included, has its exact count. If counts is empty because no word has two symbols, it returns None with copies of the words. Otherwise min(counts, key=lambda p: (-counts[p], p)) picks the winner: negating the count makes the most frequent pair the smallest key, and the tuple itself breaks ties lexicographically.

The merge is a second pass with an index walk. For each word i starts at 0; if (w[i], w[i + 1]) == best it appends the concatenation and jumps i by 2, otherwise it appends w[i] and steps by 1. Jumping by 2 is the invariant that keeps merges non-overlapping and left to right: a symbol consumed by one merge can never start another, so ["a", "a", "a"] becomes ["aa", "a"]. The TypeScript bpeMergeStep keys its Map on the two symbols joined by a NUL character and compares those keys as strings for the tie-break.

Walk through an example

On the first example two pairs tie at a count of 2, so the tie-break decides the merge.

Complexity

One step is O(n) time for n symbols: a counting pass, an O(P) scan of the P distinct pairs to pick the winner, and a merge pass. Space is O(P + n) for counts and the new corpus. Counting each candidate by rescanning is O(P · n), up to O(n²). A full run of k merges repeats the step, so it costs O(k · n).

Edge cases to test

• [["a", "a", "a"]] → (("a", "a"), [["aa", "a"]]). The pair is counted twice because the occurrences overlap, but only one non-overlapping merge fits, from the left.
• [["a", "a", "a", "a"]] → (("a", "a"), [["aa", "aa"]]). Jumping i by 2 after each merge gives two clean merges, not three.
• [["b", "a", "n", "a", "n", "a"]] → (("a", "n"), [["b", "an", "an", "a"]]). (a, n) and (n, a) tie at 2, and tuple order picks (a, n).
• [["x", "y"], ["x", "y"], ["a", "b"]] → winner ("x", "y"). A higher count beats a lexicographically smaller pair; the tie-break applies only to equal counts.
• [["q"], ["x", "y"]] → (("x", "y"), [["q"], ["xy"]]). A one-symbol word has no pairs and passes through unchanged.
• [["a"], ["b"]] → (None, [["a"], ["b"]]). No word has two symbols, so nothing merges; the reference returns None (null in TypeScript) and copies of the words.

Common mistakes

• Not recounting pair frequencies after a merge — merging changes which pairs are now adjacent.
• Merging across word boundaries when the scheme keeps words separate.
• Breaking frequency ties inconsistently, so the learned vocabulary is non-deterministic.

Follow-ups an interviewer asks next

Run the full training loop for k merges efficiently.

Recounting the whole corpus after every merge costs O(n) per merge. Practical trainers first collapse the corpus into unique words with frequencies, so a word seen 10,000 times is processed once with weight 10,000. Then they update pair counts incrementally: a merge only changes the pairs touching the merged positions, and an index from each pair to the words containing it limits the update to those words.

How do you tokenize new text with the learned merges?

Replay the merges by the order they were learned. Split the text into symbols (characters or bytes), then repeatedly merge the adjacent pair with the lowest merge rank, until no adjacent pair appears in the merge table. Ranking by training order, not by frequency in the new text, guarantees that the same text always tokenizes the same way.

Why do real tokenizers start from bytes instead of characters?

Starting from the 256 byte values means any input, in any language or with emoji, can be represented without an unknown token, because every string is a sequence of bytes. GPT-2 popularized this byte-level BPE. The cost is that a character outside ASCII starts as several byte symbols, which the merges then have to rebuild.

Should merges cross word boundaries?

Usually not. Trainers split text into words or chunks first, often keeping a leading space attached to the word, and count pairs only inside each chunk. That stops a token from spanning the end of one word and the start of the next, which would spend vocabulary on accidental neighbors. This problem enforces the same rule by keeping each word as a separate list.

Where this shows up

This merge step is how byte-pair-encoding tokenizers, the kind GPT-style models use, are trained: repeat it tens of thousands of times and frequent words collapse into single tokens. Understanding it demystifies why token counts, costs, and context limits behave the way they do, and why code and rare words fragment into many tokens.

Reference solution

def bpe_merge_step(words):
    counts = {}
    for w in words:
        for i in range(len(w) - 1):
            pair = (w[i], w[i + 1])
            counts[pair] = counts.get(pair, 0) + 1
    if not counts:
        return (None, [list(w) for w in words])
    best = min(counts, key=lambda p: (-counts[p], p))
    merged = []
    for w in words:
        out, i = [], 0
        while i < len(w):
            if i + 1 < len(w) and (w[i], w[i + 1]) == best:
                out.append(w[i] + w[i + 1])
                i += 2
            else:
                out.append(w[i])
                i += 1
        merged.append(out)
    return (best, merged)


print(bpe_merge_step([['l','o','w'], ['l','o','w','e','r']]))

function bpeMergeStep(words: string[][]): [string[] | null, string[][]] {
  const counts = new Map<string, number>();
  for (const w of words) {
    for (let i = 0; i < w.length - 1; i++) {
      const key = w[i] + '\u0000' + w[i + 1];
      counts.set(key, (counts.get(key) ?? 0) + 1);
    }
  }
  if (counts.size === 0) return [null, words.map((w) => [...w])];
  let best: string | null = null;
  let bestCount = -1;
  for (const [key, c] of counts) {
    if (c > bestCount || (c === bestCount && best !== null && key < best)) {
      best = key;
      bestCount = c;
    }
  }
  const [left, right] = (best as string).split('\u0000');
  const merged = words.map((w) => {
    const out: string[] = [];
    let i = 0;
    while (i < w.length) {
      if (i + 1 < w.length && w[i] === left && w[i + 1] === right) {
        out.push(w[i] + w[i + 1]);
        i += 2;
      } else {
        out.push(w[i]);
        i += 1;
      }
    }
    return out;
  });
  return [[left, right], merged];
}

console.log(bpeMergeStep([['l','o','w'], ['l','o','w','e','r']]));