Exact Byte-Level Probabilities from Tokenized Language Models for FIM-Tasks and Model Ensembles

Authors: Buu Phan, Brandon Amos, Itai Gat, Marton Havasi, Matthew J. Muckley, Karen Ullrich
Primary area: foundation or frontier models, including LLMs
Keywords: Language models, Tokenization, Probability, Sampling

Abstract:
Tokenization is associated with many poorly understood shortcomings in language models (LMs), yet remains an important component for long sequence scaling purposes. This work studies  how tokenization impacts  model performance by analyzing and comparing the stochastic behavior of tokenized models with their byte-level, or token-free, counterparts. We discover that, even when the two models are statistically equivalent, their predictive distributions over the next byte can be substantially different, a phenomenon we term as ``tokenization bias''. To fully characterize this phenomenon, we  introduce the Byte-Token Representation Lemma, a framework that establishes a mapping between the learned token distribution and its equivalent byte-level distribution.  From this result, we develop a next-byte sampling algorithm  that eliminates tokenization bias without requiring further training or optimization. In other words, this enables zero-shot conversion of tokenized LMs into statistically equivalent token-free ones. We demonstrate its broad applicability with two use cases: fill-in-the-middle (FIM) tasks and model ensembles. In FIM tasks where input prompts may terminate mid-token, leading to out-of-distribution tokenization, our method mitigates performance degradation and achieves 18\% improvement in FIM coding benchmarks, while consistently outperforming the standard token healing fix. For model ensembles where each model employs a distinct vocabulary, our approach enables seamless integration, resulting in improved performance up to 3.7\% over individual models across various standard baselines in reasoning, knowledge, and coding. Code is available at:https: //github.com/facebookresearch/Exact-Byte-Level-Probabilities-from-Tokenized-LMs.

Decision: accept (poster)
Reviewer scores: 8, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730140613388

thread_type:
official_review

<thread_object>
creator: reviewer_v6xm

<content>
summary:
Tokenization is a common pre-processing step, which shortens input sequences by mapping multiple bytes to discrete tokens. 
Prior work finds: Even for unlimited data and compute, tokenized models can achieve better cross-entropy loss vs. untokenized models
One hypothesis for why tokenization helps is that it allows models to handle longer contexts at less compute
Main problem: there are certain combinations that cannot be represented in token space, limiting the model’s abilities. This is harmful esp. for prompts that end mid-token.
Approach: this work presents a way to predict when tokenization is getting in the way and converts tokenized LMs into statistically equivalent byte-level models

strengths:
- The characterization of the tokenization bias phenomenon, with examples / and the synthetic example with a 3rd order Markov chain, is very interesting and useful to the community
- The paper is well-written and studies an important problem

weaknesses:
- The code infilling setting is an interesting example of where this work could be useful. However, it is not clear why the prefix would stop in the middle of a token/word or why new tokens would arise at inference time. Why are the <PRE>, <MID>, and <SUF> tokens needed? Why is the study restricted to two prompting formats/are others possible? Since tokenization schemes are lossless, it seems like any prefix should be representable with the set of available tokens at training time. Further explanation of the evaluation setting and why it is realistic would be helpful.
- Several prior methods can ensemble the predictions from models that do not share a vocabulary, or ensemble multiple predictions from a single model (https://arxiv.org/abs/2203.11171, https://arxiv.org/abs/2210.02441, https://arxiv.org/abs/2307.11031). It is not clear whether the ensembling method proposed in this paper would outperform these methods . The motivation that byte-representations help ensembling is not fully examined due to the lack of these comparisons.
- It would be useful to provide decoding efficiency metrics corresponding to each decoding approach in Figure 6 since efficiency is a motivation in Section 4.

questions:
In the caption for Figure 2, should the probability of sampling “A” on the LHS be $(1 - \alpha)$? It states the probability is $\alpha$.

I would be willing to change my opinion based on the responses to the weaknesses!
</content>
</thread_object>

</thread>
<thread>
creation date:
1730659907188

thread_type:
official_review

<thread_object>
creator: reviewer_kahr

<content>
summary:
This paper provides a thorough investigation into the phenomenon of tokenization bias in language models, proposing a byte-level prediction method to mitigate this bias. The approach has promising implications, especially for tasks like fill-in-the-middle (FIM) in coding and for model ensembling. The idea of combining language models with distinct tokenizations at the logit level recalls older ensemble methods, commonly used before tokenization became mainstream in modern LLMs. The proposed method, underpinned by the Byte-Token Representation Lemma and efficient sampling algorithms, allows for effective ensemble models by bypassing tokenization restrictions, facilitating seamless integration across diverse language models.

strengths:
What I really liked about this paper is that it very systematically investigates the phenomenon of tokenization bias. It develops a byte level prediction method and overcomes the bias. The applicability of this method and the findings in the paper have implications in areas such as coding where fill in the middle is crucial as well as in model ensembling. The idea of combining ensembles of foundation models at the logit layer is very appealing. Back in the day when there were no tokens and LMs were being built on whole words, ensembling using to be the norm in any practical LM based applications, such as speech recognition. With varying tokenization powering different LLMs today, the ability to freely combine any LLM with each other became restrictive. This paper proposes byte level efficient sampling algorithms and in conjuction with the byte-token representation lemma and the associated theoretical foundation, now allows the community to confidently combine predictions from any LLMs. Tokenization allows the LMs to overcome closed vocabulary limitation and from my perspective, this paper further allows us to overcome distinct vocabulary challenges associated with model ensembling.

weaknesses:
While the results are promising for fill in the middle (in coding) tasks, I would have liked to see experimentation on other tasks to ascertain this method's generalizability to a wider variety of tasks, espically those that require context senstitive token predictions. Computer use, a new usecase demonstrated by companies such as Anthropic, could be a good use case that the authors could consider as an example. 

An area that authors could delve further into is more analysis on the tokenization bias in multi step or long context tasks. 

In order to truly appreciate the need for this research, its important to know just how much diversity there exists in tokenization vocabulary across different LLMs today.

questions:
I am curious what do authors think about the generalizability of their methods to non-deterministic tokenization or in applications that involve mixed token vocabularies ? Could some experimental results be shown to further strengthen the paper ? 

How do authors explain the not-so-strong gains with ensebling (Table 1 and Table 2). It was not clear how the models are being ensembled together? I would have also liked to see other combination strategies as a baseline e.g. consensus voting or ranking over diverse samples from different LMs. Are these experimental results available ? 

In order to truly appreciate the need for this research, its important to know just how much diverseity there exists in tokenization vocabulary across different LLMs today. To that end, I would like to see some sort of distribution characterizing the popularity of tokenization and distribution of vocabulary sizes across various LLMs. If the world is already coverging to one tokenization and the vocabulary of tokens homogenizes over time, then what practical advantages does this paper bring ? We should debate whether or not this is true and I would like to see authors providing their rebuttal.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730690099133

thread_type:
official_review

<thread_object>
creator: reviewer_x7br

<content>
summary:
This paper investigates the impact of tokenization on the performance of language models, particularly highlighting a phenomenon called "tokenization bias." This bias comes from the discrepancies between the predictive distributions of tokenized LMs and their statistically equivalent byte-level distribution when predicting the next byte or token. In this paper, they introduce the Byte-Token Representation Lemma (BTR Lemma) to map between the token and byte domain, which allows the computation of exact byte-level probabilities without requiring extra modification of the model. This method identifies "cover encodings" for a given byte string which are all possible valid tokenizations of a string that "optimally" contain the string in the prompt. Experiments on FIM code completion task and model ensemble task show that the proposed method improves model performance.

strengths:
* This paper defines the tokenization bias phenomenon, which could cause performance issues in practical scenarios. 
* They evaluated the proposed method on FIM code completion task and mode ensemble task. 
* The proposed byte-level prediction method considers both the accuracy and efficiency in recovering the unbiased byte-level distribution.

weaknesses:
* This method is very similar to the ACL 2024 paper 'Token Alignment via Character Matching for Subword Completion' which addresses the same issue and proposes more evaluation scenarios. However, in this paper, they didn’t compare with the token alignment method. 

* Based on the results shown in Figure 6, the improvement of the proposed method on the FIM evaluation tasks in sampling settings (pass@10) over the token-level with token healing method is very minor. For example, from 84.0 to 84.1. Why is that? 

* The evaluation scenarios are limited. Only two types of tasks are covered. How’s the performance of the method on general completion tasks without the FIM setting? Does the tokenization bias also have an impact on natural language completion tasks?

questions:
* How is the proposed method compatible with other tokenization algorithms besides BPE and MPE?
* What do the right two sub-figures mean in Figure 6?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701813921

thread_type:
official_review

<thread_object>
creator: reviewer_ojbc

<content>
summary:
This paper reveals a fundamental issue in language models regarding tokenization bias - where tokenized and byte-level models, despite being statistically equivalent, produce different predictive distributions. The authors introduce the Byte-Token Representation Lemma, providing a framework to convert tokenized language models into statistically equivalent token-free ones without additional training. They demonstrate practical applications in fill-in-the-middle tasks and model ensembles, achieving significant improvements: 18% on FIM coding benchmarks and up to 3.7% on ensemble tasks. The paper provides both theoretical foundations and practical implementations through efficient algorithms for cover encoding search and next-byte sampling.

strengths:
The authors attempted to address a critical problem on tokenization for LLMs. They are the first to introduce tokenization bias and its effects on model behavior. They present a theory of it and introeduced efficient implementation of the algorithm.

The authors also studied the practical impact of the proposed solution through both code generation task (FIM, speficially) and model emsambles. Results indicate that this training-free approach is working well.

weaknesses:
While the current results are great, it'd be better if authors can also cover tasks other than FIM that suffer from this issue, e.g. complete a story in natural language or generic code completion task that doesn't require FIM.

The authors discussed computational overhead but doesn't thoroughly analyze the latency impact especially in a practical setting. It'd be better if authors can study these as part of the experiments authors conducted.

The interaction between the proposed byte-level prediction method and popular sampling techniques (like  top-p, top-k) have not been explored enough, especially in the ensemble context.

questions:
See weakness above. In addition, 

1. Do authors plan to open-source the implementation?

2. There is also a line of work similar to token healing and seems working well in both code and text. Could the authors study it and see whether the token alignment method could mitigate the issue especially in single-model completion tasks? 
Token Alignment via Character Matching for Subword Completion https://aclanthology.org/2024.findings-acl.929.pdf
</content>
</thread_object>

</thread>
<thread>
creation date:
1732070678566

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for your constructive feedback on our work, we’re excited to improve the draft further together. Multiple reviewers enjoy the insights our work provided as well as suggest many interesting directions which improves our work. Here, we provide a response to some general questions in the review.

**(1) Other tokenization schemes [Reviewer X7br, KAHr]:** Most common available tokenizers are using either BPE (Sentencepiece, tiktoken) or MPE (WordPiece). Our Byte-Token Representation Lemma applies to all of the aforementioned, and more specifically any deterministic tokenizer, as indicated in Section 4.1. The theoretical result that specifically applies to BPE/MPE is Proposition 1 in Section 3.2.2, i.e. invalid encodings do not occur in the token-level distribution. Leveraging this result, we can efficiently perform byte-level sampling, as shown in Section 4.2. Finally, we believe that our proof for tokenization bias (Proposition 1, Appendix A4) can be extended to future tokenization algorithms. 

**(2) New Results for Non-deterministic Tokenizer [Reviewer X7br, KAHr]:** Another class of tokenizer is non-deterministic, such as BPE Dropout, to combat model’s weakness against text fragmentation by randomly omitting tokens before text processing, theoretically training the model across multiple vocabularies. Although intended to enhance robustness, our analysis suggests that tokenization bias still exists. It is partially mitigated, nevertheless, since the model needs to minimize loss across varied vocabulary sets. The BTR lemma (Lemma 1) holds, requiring consideration of all valid encodings under this variability. In practice, given the substantial size of the vocabulary, BPE-dropout is expected to be robust to tokenization bias, i.e.assigning non-zero probabilities to invalid encodings. Yet, it is not commonly employed in training large-scale LLMs because the randomized tokenization introduces a computational bottleneck, which slows down the training process. Also, it may require greater model capacity to handle the increased complexity.


For experiments, we use the 1st order Markov chain example in the paper where $(\alpha, \beta)= (0.4,0.3)$. We train a LLM with BPE dropout where a sequence can be tokenized either with the vocabulary {$\texttt{<A>,<B>,<AA>\}$} or {$\texttt{<A>,<B>\}$}. Consider the following  input tokens and their next-token probabilities:

- $\texttt{``<B> <A>”}$ and {$\texttt{<A>}: 0.43, \texttt{<B>}: 0.57, \texttt{<AA>}: 10^{-5}$}. tokenization bias happens according to Def 1 but the probability does not go to 0.0 as the cross entropy loss is distributed between the two vocabularies. 

- $\texttt{``<B> <A><A>”}$ and {$\texttt{<A>}: 0.59, \texttt{<B>}: 0.41, \texttt{<AA>}: 10^{-5}$} which approximately equal to the original byte-level probabilities since the LLM detects that the only possible vocabulary that produces these tokens is {$\texttt{A,B}$}.
 
- $\texttt{``<B> <AA><A>”}$ and {$\texttt{<A>}: 10^{-4}, \texttt{<B>}: 0.99, \texttt{<AA>}: 10^{-5}$}: **tokenization bias occurs**, i.e. the LLM only outputs $\texttt{<B>}$. Since $\texttt{<AA>}$ only occurs in the vocabulary {$\texttt{<A>,<B>,<AA>\}$}, **the LLM detects the tokenization pattern** and infer that the sequence must be tokenized with this vocabulary.

Finally, we apply the BTR Lemma to remove the bias for the string $\texttt{``BA''}$ where $P(\texttt{<B><A>})=0.104$, $P(\texttt{<B><AA>})=0.0415$. Using $\texttt{BAA}$, we have $P(\texttt{<B><A><A>})=0.0443$ and $P(\texttt{<B><AA>})=0.0415$, which gives us
 $\alpha=0.405$ after factorization.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732508917651

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Encouragement to discuss + additional experimental results

comment:
The discussion period will be closing soon, we hence encourage all reviewers to respond to our rebuttal. We highly value your feedback and would appreciate your contributions to the discussion. Please take this opportunity to share your perspectives and engage with the community. Thank you for your participation and support!   


Additionally we would also like to share some additional experimental work, we implemented token alignment as an additional baseline for the FIM task. 

| Method | SPM (Greedy @1) | SPM (@10)  |
|---------|-------------------|-------------------|
| Token Alignment   | 63.0             | 82.9             |
| **Our**   | **63.9**             | **84.3**            | 

Both results lack behind our method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733285998099

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final comment before discussion period ends | Summary of rebuttal contributions

comment:
Reviewers,

Thank you for your feedback on our submission. We have made several improvements and clarifications based on your comments.

***Summary of Additional Contributions and Clarifications:***

**1. New Experimental Results**: Introduced comparisons with token alignment as a baseline for the FIM task, showing our method's superior performance.

**2. Theoretical Insights on Stochastic Tokenizers**: Expanded our framework to include non-deterministic tokenizers like BPE Dropout, demonstrating partial mitigation of tokenization bias.

**3. Clarifications on Token-Alignment**: Explained limitations of token-alignment methods and highlighted our approach's ability to enable exact likelihood computation and effective ensembling.

**4. Discussion on Tokenization Diversity**: Provided insights into the diversity of tokenization vocabularies across models, underscoring the practical importance of our research.

**5. Ensemble Method Comparisons**: Demonstrated our method's superior performance over simple voting strategies, especially in complex tasks like coding.

**6. Decoding Efficiency Metrics**: Provided detailed metrics on decoding efficiency, highlighting potential for optimization.

**7. Commitment to Open-Source**: Confirmed plans to open-source our implementation for reproducibility.


We believe these enhancements address your concerns and significantly strengthen our work. We hence kindly ask you to consider these improvements during the discussion period and *hope they encourage you to raise your scores*.

Thank you for considering our appeal.

Warm regards,
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper raises a problem of tokenization bias, which is a bias in the language model related to the choice of tokens it uses to encode text.  According to the reviews, the problem raised is important, and the analysis for it provided in the paper is clear and thorough. Another strength of the paper is that the provided techniques are both backed by solid theoretical arguments and have convincing empirical evidence showing its value on the selected tasks. 

The reviews raised a few concerns about the paper, asking for a better comparison to selected existing works, and suggested demonstrating the techniques on additional use cases. In the rebuttal, the authors provided a thorough positioning of their paper compared to the works raised in the reviews. The requirement for additional use-cases was not fully resolved in the rebuttal, but it was minor to begin with.  

Though there isn’t a consensus that the paper should be accepted, the techniques seem to be sufficiently motivated, and the edits required to fix the positioning issue seem quite mild. Given the impact of the problem and novelty of the analysis and technique, I think the strengths outweigh the weaknesses and the paper will be a good addition to ICLR.