BytePiece: A Purer, Higher-Compression Tokenizer
Currently the most popular tokenizer for LLMs is probably Google's SentencePiece, since it satisfies a number of desirable properties for a tokenizer — language-agnostic, data-driven, and so on — and because it's written in C++, tokenization is fast, which makes it well suited to efficiency-oriented scenarios. However, it also has some notable drawbacks, such as slow training (the BPE algorithm) and heavy memory usage; and precisely because it's written in C++, it's a black box to most users, which makes it inconvenient to study or extend.
In fact, training a tokenizer is essentially equivalent to the classic problem of "new word discovery," and I've previously written a series of posts on Chinese word segmentation and minimum entropy, so I've accumulated some experience in new word discovery. I've had the idea of writing my own tokenizer for a long time. These past few days I finally found the time to get a first version working — a humble imitation of SentencePiece, which I've named "BytePiece."
Github: https://github.com/bojone/bytepiece
Desired properties
Since we're rewriting a tokenizer, we should first think about what an ideal tokenizer should look like, so that we have a way to judge whether the final result meets expectations. As I see it, a tokenizer should at least have the following basic properties:
1. Lossless reconstruction: the tokenization result should be losslessly recoverable back into the original input;
2. High compression rate: for a given vocabulary size, the number of tokens produced for the same batch of data should be as small as possible;
3. Language-agnostic: being statistically based, neither training nor tokenization should introduce language-specific assumptions;
4. Data-driven: it should be possible to train directly on raw corpora in an unsupervised manner;
5. Training-friendly: training should be completable in a reasonable amount of time and with reasonable resources.
Beyond these, there are some bonus points, such as fast tokenization speed, readable code, and ease of extension — nice to have, but I don't consider them essential.
To me, SentencePiece's biggest pain points are "lossless reconstruction" and "training-friendliness." First, by default SentencePiece performs NFKC normalization, which causes irreversible changes such as converting full-width commas into half-width commas — so by default it doesn't even satisfy "lossless reconstruction." Because of this it stayed off my shortlist for a long time, until I later discovered that setting the parameter --normalization_rule_name=identity during training disables all such transformations. So SentencePiece does support lossless reconstruction, it just requires special configuration.
As for training, it's even more maddening. SentencePiece supports two mainstream algorithms, BPE and Unigram. Unigram trains at a tolerable speed but with somewhat lower compression, while BPE achieves higher compression but trains an order of magnitude slower than Unigram! And regardless of whether you use BPE or Unigram, the training process is extremely memory-hungry. In short, training a SentencePiece model on a sizable corpus is simply not a pleasant experience.
Model design
Building a new tokenizer can be broken down into three parts: 1. the basic unit; 2. the tokenization algorithm; 3. the training algorithm. Once these three are settled, everything else is just a matter of engineering. Below I go through BytePiece's thinking on each of these in turn.
Basic unit
As we know, Python 3's default string type is Unicode. If we take Unicode as the basic unit, we call this Char-based. Char-based is intuitive and convenient — a Chinese character shows up as a single character of length 1 — but there are simply too many distinct characters across different languages, and even just covering single characters would consume an enormous vocab_size, let alone incorporating whole words. So, like mainstream tokenizers, BytePiece uses the Byte as its basic unit.
Once we go back to Bytes, a lot of things suddenly become clear. Since there are only 256 distinct single bytes, as long as the vocabulary includes all 256 of them, OOV (Out of Vocabulary) issues are eliminated entirely — an obvious benefit. Furthermore, we know that the average information entropy of Chinese characters is higher than that of English letters. If we go with Char-based, then even though every character superficially has length 1, the "intrinsic" granularity differs, which biases the statistics. By comparison, the information entropy per Byte is much more uniform (for instance, most Chinese characters' UTF-8 encoding corresponds to 3 bytes, and the average information entropy of a Chinese character is roughly 2–3 times that of an English letter, which corresponds to a single byte), so statistics computed over Bytes are more unbiased, which makes the model more "language-agnostic."
BytePiece goes further than SentencePiece in being Byte-based: SentencePiece first processes text as Char-based and only falls back to Byte-based when it encounters OOV, whereas BytePiece converts the text into Bytes via text.encode() right from the start, before any further processing — a purer approach.
Tokenization algorithm
There are really only a handful of dictionary-based tokenization algorithms — maximum matching, shortest path, maximum probability path, and so on. Interested readers can refer to Matrix67's earlier post A Chat on Chinese Automatic Word Segmentation and Semantic Recognition (Part 1): Chinese Word Segmentation Algorithms.
Like jieba and other Chinese word-segmentation tools, BytePiece chooses the maximum probability path for tokenization, also known as the "unigram model," i.e. Unigram. There are three reasons for choosing Unigram: first, maximum probability under Unigram is, in other words, maximum likelihood, and this aligns with the training objective of LLMs, which is also maximum likelihood — the two are thus more consistent; second, from a compression standpoint, maximum probability is really the shortest encoding length (also called minimum description length), which embodies maximizing the compression rate, consistent with the belief that "compression is intelligence"; third, the optimal Unigram tokenization can be found via the Viterbi algorithm in linear time, which is the theoretically optimal complexity.
Of course, just as there's a "unigram model," there are naturally more complex "bigram" and "trigram" models, but the increase in complexity they bring far outweighs the benefit, so we generally don't consider these higher-order models.
Training algorithm
The reason I discuss the tokenization algorithm before the training algorithm is that only once the tokenization algorithm is fixed can we determine the training objective, and hence study a corresponding training algorithm.
As mentioned at the outset, training a tokenizer is essentially the classic problem of "new word discovery," and I've previously proposed several new-word-discovery algorithms, such as in New Word Discovery Based on Segmentation, Unsupervised Segmentation Based on a Language Model, and A Better New Word Discovery Algorithm. Looking back now, the one that fits best and has the most potential in combination with the Unigram tokenization algorithm is Unsupervised Segmentation Based on a Language Model. BytePiece's training is implemented based on this approach, which I refer to here as Byte-based N-gram Language Model (BNLM).
Specifically, for Unigram tokenization, if a byte string of length $l$, $c_1, c_2, \dots, c_l$, has optimal segmentation $w_1, w_2, \dots, w_m$, then the product of probabilities $p(w_1)p(w_2)\dots p(w_m)$ should be the largest among all possible segmentations. Let the lengths of $w_1,w_2,\cdots,w_m$ be $l_1,l_2,\cdots,l_m$ respectively; then, by the conditional decomposition formula,
\begin{equation}\prod_{i=1}^m p(w_i) = \prod_{i=1}^m \prod_{j=L_{i-1} + 1}^{j=L_{i-1} + l_i} p(c_j|c_{L_{i-1} + 1},\cdots,c_{j-1})\end{equation}
where $L_i=l_1+l_2+\cdots+l_i$. If we consider only an $n$-gram model, approximating $j\gt L_{i-1} + n$'s $p(c_j|c_{L_{i-1} + 1},\cdots,c_{j-1})$ uniformly with $p(c_j|c_{j - n + 1},\cdots,c_{j-1})$, then Unigram tokenization reduces to a character (byte) tagging problem, and training the tokenizer reduces to training an $n$-gram language model (we recommend $n=6$), which can be done directly in an unsupervised manner. For a more detailed treatment, please refer to the original post Unsupervised Segmentation Based on a Language Model.
(Note: $n=6$ merely means that BytePiece's statistics only go up to 6-grams, not that the maximum piece length it can generate is 6 — because for $6$-grams beyond $n$, we approximate the conditional probability using the 6-gram statistics, so it can in principle be extended to arbitrary order, meaning pieces of arbitrary length can in theory be generated.)
Implementation
Once the theory is settled, what's left is the tedious work of implementation. I managed to put together a working codebase:
Github: https://github.com/bojone/bytepiece
The code is quite simple — a single file containing just two classes, Trainer and Tokenizer, corresponding to the two parts: tokenization and training. Tokenization uses pyahocorasick to build an Aho-Corasick automaton for a modest speedup — it's usable, but still noticeably slower than SentencePiece, since pure Python simply can't compete with C++ on speed. Training is split into four main steps: 1. counting $n$-grams; 2. pruning $n$-grams; 3. pre-tokenization; 4. pruning the pre-tokenization results. Steps 1, 3, and 4 are all compute-intensive and parallelizable, so I implemented corresponding multiprocess versions. With enough processes running (I used 64, and each process was essentially running at full utilization), training speed can rival SentencePiece's Unigram training speed.
I should also say a word about result pruning. The most basic criteria for pruning are of course frequency and vocab_size, but that's not enough, because sometimes $p(w_1)p(w_2) > p(w_1\circ w_2)$ ($w_1\circ w_2$ denotes the concatenation of two words) and all three of $w_1,w_2,w_1\circ w_2$ appear in the vocabulary at once. In this situation, the word $w_1\circ w_2$ will never actually be produced by segmentation, so keeping it in the vocabulary is pure waste of space — so the pruning process also excludes such cases.
Benchmarks
Now for the part everyone loves — the proof of the pudding is in the eating. First, a small-scale test: I randomly sampled 100,000 entries from a previously open-sourced dataset from WuDao (the exported file is about 330MB) as the training set, and separately sampled another 1,000 entries as the test set, training a vocabulary of size 50k. The comparison results are as follows:
$$\begin{array}{c|ccc} \hline & \text{training time}\downarrow & \text{max memory usage}\downarrow & \text{compression ratio}\uparrow \\ \hline \text{SP-BPE} & \text{55.3 min} & \text{5.2GB} & 4.80 \\ \text{SP-Unigram} & \text{1.6 min} & \text{2.5GB} & 4.73 \\ \text{BytePiece} & \text{6.5 min} & \text{4.3GB} & 5.05 \\ \hline \end{array}$$
To clarify, SP-BPE and SP-Unigram refer to SentencePiece with model_type set to BPE and Unigram respectively, trained with the following code:
spm.SentencePieceTrainer.train('--input=wudao.txt --model_prefix=wudao_m --vocab_size=50000 --model_type=bpe --train_extremely_large_corpus=true --normalization_rule_name=identity')
spm.SentencePieceTrainer.train('--input=wudao.txt --model_prefix=wudao_m2 --vocab_size=50000 --model_type=unigram --train_extremely_large_corpus=true --normalization_rule_name=identity')
The compression rate is measured in "bytes/token," i.e., the average number of bytes per token. As you can see, BytePiece achieves the highest compression rate while keeping training time and memory usage relatively moderate.
Next, a larger-scale test. From a mixed Chinese-English corpus with a roughly 3:5 ratio, I extracted 100,000 samples to train a Tokenizer with vocab_size=100k. The texts in this corpus are fairly long, so the exported file for 100,000 samples already comes to 13GB. The test set consists of two parts: one is 1,000 samples drawn from the same corpus (i.e., same-source), and the other is the same 1,000 samples from the WuDao dataset used earlier (i.e., different-source). The results are as follows:
$$\begin{array}{c|cccc} \hline & \text{training time}\downarrow & \text{max memory usage}\downarrow & \text{compression ratio (homologous)}\uparrow & \text{compression ratio (heterologous)}\uparrow \\ \hline \text{SP-BPE} & \text{19.21 h} & \text{97GB} & 4.52 & 4.46 \\ \text{SP-Unigram} & \text{2.02 h} & \text{384GB} & 4.51 & 4.48 \\ \text{BytePiece} & \text{2.24 h} & \text{51GB} & 5.39 & 4.51\\ \hline \end{array}$$
Whether in terms of training time, memory, or compression rate, it looks like the larger the training data, the greater BytePiece's advantage!
To be continued
Based on the results so far, BytePiece has a clear advantage in training, and its tokenization quality is also decent. However, being pure Python takes its toll: tokenization speed is only about 1/10 that of SentencePiece, which is one direction for future optimization. I'd welcome any C/C++ experts who'd like to get involved and help improve BytePiece's tokenization speed. (Note: starting from version 0.2.0, the tokenization function has been accelerated with Cython, and BytePiece's tokenization speed is now close to that of BPE, and can even surpass BPE when the text is long enough.)
In fact, using techniques such as random sampling and dynamic pruning, BytePiece's training speed and memory usage could be further optimized. Currently, to ensure deterministic results, BytePiece doesn't perform any pruning until all statistics have been fully collected — this guarantees consistent results whether running single-process or multi-process. If the input were shuffled randomly and pruning performed periodically, memory usage could be further controlled and the statistics collection sped up, and the impact on the final results would likely be small. This is something I plan to introduce further down the line based on user feedback.
Beyond all this, there are still quite a few details in BytePiece that need polishing, and there may well be bugs I haven't yet discovered — your understanding and feedback are much appreciated.
Summary
This post introduced BytePiece, a tokenizer I developed myself. It is a Byte-based Unigram tokenizer implemented in pure Python, making it more readable and easier to extend. Thanks to its new training algorithm, it typically achieves higher compression rates than existing tokenizers, and it supports multiprocess-accelerated training. Moreover, since it operates directly on the UTF-8 bytes of the text with almost no preprocessing, it is purer and more language-agnostic.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.