Technical Summary: Tokenization with Split Trees (ToaST)

Problem Statement

Tokenization is a foundational step in natural language processing, yet existing methods often rely on greedy heuristics that may not optimize for compression or downstream performance. While Byte Pair Encoding (BPE) minimizes token counts via a greedy merge strategy, it is a heuristic approximation. Prior work has established that finding a tokenizer that maximizes compression is NP-complete, even over bounded alphabets. Furthermore, the relationship between compression and downstream task performance remains debated; while some studies suggest a correlation, others argue compression alone is insufficient to explain tokenization quality. Existing alternatives, such as UnigramLM or PathPiece, utilize different inference procedures (e.g., ablative top-down or shortest-path inference) but still face challenges in global optimization and computational scaling.

Methodology

The authors introduce Tokenization with Split Trees (ToaST), a subword tokenization method that decouples the construction of candidate tokens from the selection of the final vocabulary. The approach consists of three main stages:

1. Split Tree Construction

Unlike BPE, which builds tokens bottom-up, or UnigramLM, which learns probabilities top-down, ToaST constructs a full binary tree for every unique pretoken in the training corpus independently of any vocabulary.

• Preprocessing: Training data is split into pretokens using a regular expression.
• Splitting Logic: For a given pretoken, the algorithm recursively splits it into two non-empty parts. The split point is chosen deterministically to maximize the minimum count of the resulting left and right byte n-grams ($\max \min(c_p, c_{p'})$) based on precomputed n-gram statistics from the corpus.
• Structure: This process continues until the leaves are single bytes. The nodes of these trees form the set of all candidate tokens. This construction is independent of the final vocabulary and requires no iterative training.

2. Recursive Inference

Given a vocabulary $V$, tokenization is performed via a recursive descent of the split tree:

• Starting at the root (the full pretoken), the algorithm checks if the current node is in $V$.
• If the node is in $V$, it is emitted as a token.
• If not, the algorithm recursively processes the two children.
• Properties: This procedure ensures that any vocabulary containing all necessary single bytes yields a valid tokenization. Removing a token from the vocabulary simply forces the inference to descend further into the tree, with no cascading effects on other tokens, unlike merge-based methods.

3. Vocabulary Selection via Integer Programming

The core innovation is formulating vocabulary selection as an Integer Program (IP) to minimize the total token count over the training corpus.

• Variables: Binary variables $x_i$ indicate if candidate token $i$ is in the vocabulary. Binary variables $z_{jk}$ indicate if node $k$ in split tree $j$ is emitted during inference.
• Objective: Minimize the total weighted count of emitted tokens: $\min \sum_{j} c_j \sum_{k} z_{jk}$.
• Constraints:
  • Vocabulary Size: The sum of selected tokens must equal the target size $m$.
  • Coverage: For every path from a leaf to the root, exactly one node must be selected (ensuring the entire pretoken is covered).
  • Consistency: A node can only be emitted ($z_{jk}=1$) if it is in the vocabulary ($x_i=1$) and none of its ancestors are.
• Solving: The IP is relaxed to a Linear Program (LP). In practice, the LP relaxation is "near-integral," meaning the solution contains very few fractional values. A simple rounding heuristic is applied to obtain an integer solution that is provably near-optimal (within a negligible fraction of the global optimum).

Key Contributions

1. New Tokenization Framework: ToaST introduces a split-tree based approach that avoids the greedy approximations of BPE and the specific inference constraints of UnigramLM.
2. Global Optimization: By formulating vocabulary selection as an IP with a tight LP relaxation, ToaST finds a vocabulary that globally minimizes token counts for a given set of split trees, rather than relying on local greedy decisions.
3. Scalability: The method scales quadratically with the number of split trees, allowing the processing of large corpora (covering 99% of pretoken counts) within feasible training times.
4. Decoupled Design: The separation of tree construction (based on corpus statistics) and vocabulary selection (based on optimization) allows for flexible extensions, such as incorporating morphological or character-level preferences without retraining the optimization model.

Experimental Results

The authors evaluated ToaST on English text using the CulturaX dataset and compared it against BPE, WordPiece, and UnigramLM.

• Compression: At vocabulary sizes of 40,960 and above, ToaST reduces token counts by more than 11% compared to all baselines. This is measured in bytes per token.
• Token Distribution: ToaST significantly reduces the frequency of "Leaf" tokens (single-byte fallbacks) by 14–19x compared to baselines at a vocabulary size of 65,536. Conversely, it produces substantially more "Root" tokens (full pretokens), indicating better retention of semantic units.
• Rényi Efficiency: ToaST achieves higher Rényi efficiency (with $\alpha=2.5$), indicating a more uniform token distribution and fewer high-frequency tokens, which is often correlated with better downstream performance.
• Downstream Performance: In experiments training 1.5B parameter language models (nanochat), ToaST achieved the highest CORE score (an average of 22 benchmarks), outperforming baselines by 2.6% to 7.6%. It significantly outperformed BPE and UnigramLM, and was the best performer on 13 of 22 individual tasks.
• Optimization Quality: The LP relaxation was found to be exceptionally tight. Across 128 problem instances, the maximum relative gap between the rounded integer solution and the LP bound was $5.76 \times 10^{-7}$, with many instances being fully integral.

Significance and Claims

The paper claims that ToaST demonstrates the value of optimizing tokenization directly for compression using a global optimization framework rather than relying on greedy heuristics. The authors argue that the superior compression leads to a larger effective context window and, consequently, improved downstream performance.

The significance of the work lies in:

• Practicality: Despite the theoretical NP-completeness of optimal tokenization, the LP relaxation makes the approach computationally feasible and practical for large-scale training.
• Performance: The empirical results suggest that the 11% reduction in token count translates to measurable gains in language model performance (CORE score), supporting the hypothesis that better compression correlates with better model utility.
• Flexibility: The architecture allows for future extensions, such as handling multi-byte characters or morphological splits, by modifying the tree construction hierarchy without altering the vocabulary optimization logic.

The authors note limitations, specifically that experiments were conducted exclusively on English text, and the effectiveness of training on the top aggregated pretokens may vary for non-space-delimited languages. They also disclose the use of LLMs for code implementation and figure generation.