BpeTrainer emits a merge for a pair that occurs zero times (negative pair_counts sign-extended by `as u64`)

[Abraham-y]

Summary

BpeTrainer can emit a merge for a pair that occurs zero times in the corpus, when continuing_subword_prefix (or end_of_word_suffix) is combined with max_token_length.

Under that combination a pair count in pair_counts: AHashMap<Pair, i32> can go negative. The value is then read through as u64, which sign-extends -1 to 18446744073709551615, and that entry wins the max-heap immediately — so the trainer merges a pair that nothing in the corpus could have produced.

Reachable with 8 words / 38 bytes. Confirmed on tokenizers==0.23.1.

This is a sibling of #2058 (positive i32 overflow); same field, opposite sign, and it needs far less data to hit.

Reproducer

Self-contained — no dependency beyond tokenizers. It trains, then replays the emitted merge list against the corpus and reports the true occurrence count of each pair at the moment it is merged. A count of 0 means the merge was unreachable.

import collections, json
from tokenizers import Tokenizer, models, pre_tokenizers, trainers

CORPUS = "##c ##cac# #ab c# #### accc#a#a#b a cb"
PREFIX, VOCAB_SIZE, MAX_TOKEN_LENGTH, SPECIALS = "##", 244, 4, ["a", "<unk>"]

def train(path):
    tok = Tokenizer(models.BPE())
    tok.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
    tok.train([path], trainers.BpeTrainer(
        vocab_size=VOCAB_SIZE, show_progress=False,
        continuing_subword_prefix=PREFIX, max_token_length=MAX_TOKEN_LENGTH,
        special_tokens=SPECIALS))
    m = json.loads(tok.to_str())["model"]["merges"]
    return [tuple(x) if isinstance(x, list) else tuple(x.split(" ", 1)) for x in m]

def replay(merges):
    counts = collections.Counter(CORPUS.split())
    words = [[[c if i == 0 else PREFIX + c for i, c in enumerate(w)], n]
             for w, n in counts.items()]
    out = []
    for rank, (a, b) in enumerate(merges):
        n = sum(c for syms, c in words
                for i in range(len(syms) - 1)
                if syms[i] == a and syms[i + 1] == b)
        out.append((rank, a, b, n))
        new = a + (b[len(PREFIX):] if b.startswith(PREFIX) else b)
        for e in words:
            syms, res, i = e[0], [], 0
            while i < len(syms):
                if i + 1 < len(syms) and syms[i] == a and syms[i + 1] == b:
                    res.append(new); i += 2
                else:
                    res.append(syms[i]); i += 1
            e[0] = res
    return out

path = "/tmp/phantom_corpus.txt"
open(path, "w").write(CORPUS + "\n")
for run in range(40):
    rep = replay(train(path))
    if any(r[3] == 0 for r in rep):
        print(f"run {run}: unreachable merge found")
        for rank, a, b, n in rep:
            print(f"  {rank:>3}  ({a!r}, {b!r})".ljust(40)
                  + f"count={n}" + ("   <-- UNREACHABLE" if n == 0 else ""))
        break

Output (tokenizers==0.23.1)

  rank  pair                        true count when merged
     0  ('#', '###')                               3
     1  ('##a', '###')                             2
     2  ('##c', '###')                             2
     3  ('##', '##c')                              2
     4  ('a', '##c')                               1
     5  ('#', '##a')                               1
     6  ('c', '##b')                               1
     7  ('c', '###')                               1
     8  ('##a', '##c#')                            1
     9  ('###', '###')                             1
    10  ('##c', '##a')                             0   <-- UNREACHABLE
    11  ('##c', '##c#')                            1
    12  ('##a#', '##b')                            1
    13  ('#a', '##b')                              1

Training with continuing_subword_prefix is itself nondeterministic (#1794, PR #2066), so this appears in a subset of runs: 4/40 runs over 4 distinct merge lists.

max_token_length is required. Sweeping it on this corpus, 30 runs each:

max_token_length	2	3	4	5	6	8	100	unset
runs with an unreachable merge	0/30	0/30	6/30	0/30	0/30	0/30	0/30	0/30

Mechanism

In tokenizers/src/models/bpe/trainer.rs:

1. Counts are signed and the delta update can subtract:

   let (mut pair_counts, mut where_to_update): (AHashMap<Pair, i32>, _) = ...
   let count = change * counts[iw] as i32;
   *pair_counts.entry(pair).or_default() += count;

   With max_token_length set, a merge that reuses an existing token id can apply the -1 neighbour delta while the length guard rejects the matching +1, so a pair already at 0 reaches -1.

2. Both queue.push sites correctly refuse non-positive counts:

   let count = pair_counts[&pair];
   if count > 0 { queue.push(Merge { pair, count: count as u64, pos }); }

3. But the staleness re-check has no such guard, and this is the path that lets a negative through:

   if top.count != pair_counts[&top.pair] as u64 {
       top.count = pair_counts[&top.pair] as u64;   // -1  ->  u64::MAX
       queue.push(top);
       continue;
   }
   
   if top.count < 1 || self.min_frequency > top.count {
       break;
   }

   An entry already on the heap whose count later goes negative is re-pushed with the sign-extended value. On the next pop the staleness check passes (both sides are the same huge number), top.count < 1 is false, and it merges — ahead of every legitimate pair, since this is a max-heap.

Suggested fix

Guard the staleness re-check the same way the two push sites are guarded, e.g. drop the entry when the live count is not positive:

let live = pair_counts[&top.pair];
if live <= 0 {
    continue;                     // spent; do not re-push
}
if top.count != live as u64 {
    top.count = live as u64;
    queue.push(top);
    continue;
}

Comparing as i32 before the cast would also work. Widening the counter to i64 (as #2058 proposes for the positive direction) does not fix this on its own — any negative value sign-extends the same way.

Impact

The emitted merge is unreachable, so it silently consumes a vocabulary slot and shifts every subsequent merge, changing the tokenizer. It affects WordPieceTrainer, which sets continuing_subword_prefix("##") in its default builder, whenever max_token_length is also set.

Environment

• tokenizers 0.23.1 (also reproduces on 0.22.2)
• Python 3.12, macOS arm64

[RohanMali2003]

Heya @Abraham-y, thanks for the incredibly detailed write-up and the reproducer!
I'd love to pick this up and submit a PR with your suggested fix, if no one else is working on it rn. Lemme know if I can go ahead

[yentur]

@RohanMali2003 sorry, I did not see your comment before I started. I have opened #2327 with the guard @Abraham-y suggested plus a deterministic Rust regression test. Happy to close mine if you would rather open yours.

[RohanMali2003]

@yentur No worries at all! I finished up my implementation and regression test as well, but I didn't quite have the chance to open the PR. Since your PR is already up, let's just go with yours so we don't duplicate effort any more- I went through it and I think it looks great, there's a couple things I could pick up to make my approach better too. Feel free to keep it open so it can get merged. Nice work! :))

[ArthurZucker]

Ty both, will have a look once we are done with #2119 👀