🤖 AI

You Can Learn Tokenization End-to-End with Reinforcement Learning

This paper proposes learning tokenization end-to-end using reinforcement learning with score function estimates and time discounting to reduce variance, demonstrating that this approach outperforms prior straight-through methods at the 100 million parameter scale.

Original authors: Sam Dauncey, Roger Wattenhofer

Published 2026-08-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sam Dauncey, Roger Wattenhofer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern computers do not naturally understand human language; they only understand numbers. To bridge this gap, artificial intelligence systems that process text must first translate words and sentences into a stream of numbers they can digest. For years, the standard method for doing this has been a rigid, pre-programmed step called tokenization. Think of this process like a librarian who, before a book can be read by a machine, must cut the pages into specific, pre-determined chunks based on a fixed rulebook. This rulebook decides where one word ends and the next begins, often grouping common letter combinations together. While this works well, it is a manual, static step that sits outside the brain of the artificial intelligence, separate from the learning process itself. As these digital minds grow larger and more capable, scientists have begun to wonder if this rigid, pre-cut approach is holding them back, and whether the machine could learn to cut the text itself in a way that makes more sense to its own internal logic.

A team of researchers at ETH Zurich has taken a significant step toward answering that question by teaching an artificial intelligence to learn how to cut its own text from scratch, without any pre-written rules. In their new work, they demonstrated that a computer model can be trained to decide exactly where to place the boundaries between pieces of text, simply by trying to minimize its mistakes while learning. Instead of following a fixed manual or using a clever guess based on spaces and punctuation, the model uses a trial-and-error approach similar to how an animal learns to navigate a maze. It stochastically guesses where a cut should go, sees how well it predicts the next part of the text, and then adjusts its strategy. If a guess leads to a better prediction, the model reinforces that decision; if it leads to confusion, it tries something else. Over time, the model discovers its own optimal way of breaking down language, effectively teaching itself the rules of tokenization as it learns the language.

The researchers found that when they let the model learn this way, it naturally began to place its cuts at the same places humans do, such as at the spaces between words or at the ends of sentences, even though no one told it to do so. This happened without the model being given any specific instructions about grammar or the structure of language. In tests using a dataset of over a hundred million parameters, this self-taught approach proved to be more effective than previous methods that tried to learn boundaries using a different, less direct mathematical technique. The new method not only produced better results but also matched the performance of systems that rely on the traditional, hand-crafted rulebooks, despite having no prior knowledge of those rules.

One of the most striking aspects of this discovery is that the model learned to be efficient. It figured out how to group bytes of data into chunks that made sense for its specific task, whether that was reading general text or understanding computer code. When the researchers tested the model on Python code, it learned to start new chunks at the beginning of function names and to keep certain common phrases together, showing an intuitive grasp of the structure of the code it was reading. This suggests that the model is not just memorizing patterns but is actively developing a strategy that aligns with the meaning of the text it processes.

The study also addressed a major hurdle that had previously made this kind of learning difficult: the noise in the learning signal. Because the decision to cut text is a binary choice—either you cut here or you don't—it is mathematically difficult to tell the model exactly how to improve. The researchers solved this by borrowing techniques from reinforcement learning, a field where agents learn through rewards and penalties. They introduced a method that smooths out the noise, allowing the model to see which specific decisions were responsible for its successes or failures over a longer stretch of text. This allowed the model to learn effectively without needing an impossibly large amount of computing power.

While the current experiments were conducted on models of a specific size, the results offer a compelling vision for the future of artificial intelligence. By removing the need for a separate, pre-programmed step to prepare text, the researchers have shown that language processing can become a fully integrated, end-to-end learning process. The model does not just learn to speak; it learns how to listen and how to parse the world around it in the most efficient way possible. This approach could eventually lead to systems that handle languages and data types for which no human-designed rulebook exists, making artificial intelligence more adaptable and inclusive. The work suggests that the rigid, artisanal design of how we feed data to machines may soon be replaced by a more fluid, self-discovered method, where the machine learns the best way to read the world on its own terms.

Technical Summary: Learning Tokenization End-to-End with Reinforcement Learning

Problem Statement
Tokenization remains a hardcoded, pre-processing step in the training pipelines of Large Language Models (LLMs), typically relying on heuristics like Byte-Pair Encoding (BPE) or hand-crafted rules (e.g., splitting digits, preserving whitespace). While the broader trend in AI architecture is moving toward end-to-end learning, tokenization has largely resisted this shift. Prior attempts to integrate tokenization into the model architecture have faced two main limitations:

Heuristic Reliance: Methods that operate directly on bytes often use fixed rules (e.g., whitespace detection) or entropy spikes to determine token boundaries, which are not learned from the data.
Theoretical Weakness of Straight-Through Estimators (STEs): Previous attempts to learn token boundaries end-to-end have relied on STEs. These methods relax discrete boundary decisions into continuous variables to allow backpropagation. However, STEs lack strong theoretical guarantees because they do not directly optimize the discrete decision problem; they rely on heuristic gradient approximations that may not converge to an optimal strategy.

Methodology
The authors propose a method to learn token boundaries end-to-end using score function estimators derived from reinforcement learning (RL), specifically within an autoregressive U-Net architecture.

Architecture: The model processes text at the byte level, stochastically predicts token boundaries, downsamples to a token-level representation, processes it, and then upsamples back to the byte level for next-byte prediction. This allows the model to reuse byte-level representations at the token level.
Score Function Estimation: Instead of STEs, the authors treat token boundary selection as a discrete stochastic decision. They formulate the learning objective as minimizing the expected next-token cross-entropy loss marginalized over all possible tokenization strategies. The gradient is estimated using the REINFORCE algorithm (Williams, 1992), which decomposes the gradient into a conditional loss term and a policy gradient term.
Variance Reduction: A direct application of REINFORCE suffers from high variance, making training unstable. The authors introduce three specific techniques to reduce variance and make the method practical:
1. Early Exit Relative Rewards: They subtract a baseline reward (estimated by an early-exit byte-level model) from the total reward to remove tokenization-independent noise.
2. Time Discounting: They apply a discount factor ( $\gamma = 0.99$ ) to future rewards when computing advantages. This decouples the influence of distant sequence parts, reducing variance by providing more independent training stimuli per sequence.
3. Batch-Relative Advantages: They center the advantage estimates across the batch to correct for biases where the final-layer model consistently outperforms the early-exit model.
Downsampling Rate Control: To prevent the model from choosing a computationally expensive strategy of separating every byte, a target downsampling rate is enforced via an additional loss term that applies pressure to the mean logit of the batch.

Key Contributions

End-to-End Learning without Inductive Bias: The authors demonstrate that tokenization strategies aligning with semantic boundaries (e.g., whitespace) can be learned purely from optimizing cross-entropy loss using a score-function estimator, without explicit linguistic priors or inductive biases.
Superiority over Straight-Through Estimators: The paper claims that their score-function approach outperforms prior STE-based methods both qualitatively (producing more semantically meaningful boundaries) and quantitatively (lower validation loss and bits-per-byte).
Robustness to Downsampling Rates: The method demonstrates robust performance across a range of downsampling rates (tokenization aspect ratios), showing that the learned strategies adapt effectively to different compression targets.

Experimental Results
Experiments were conducted on 147-million parameter models trained on the FineWeb dataset and 90-million parameter models on the CodeParrot dataset.

Qualitative Analysis: The learned token boundaries reliably align with whitespace characters (spaces, newlines) and semantic breaks, despite the model having no explicit instruction to do so. In code, boundaries align with module names and license headers.
Quantitative Performance:
- On natural language benchmarks (PIQA, HellaSwag, LAMBADA, ARC-Easy), the learned policy achieved lower bits-per-byte than uniform baselines, STE-based methods (Nawrot et al., 2023; Hwang et al., 2025), and a random baseline.
- The performance of the learned method was comparable to a BPE-guided baseline (where boundaries are fixed by a pre-trained BPE tokenizer), achieving 1.297 bits-per-byte on the FineWeb test set versus 1.299 for BPE guidance.
- The method outperformed STE-based approaches in validation loss curves throughout the training run.
Ablation: A variant with a reduced window size ( $w=1$ ) performed comparably to the full model, suggesting the sliding window context for boundary logits is not strictly necessary at this scale.

Significance and Claims
The paper posits that score function estimators, bolstered by RL variance reduction techniques, offer a theoretically sound and practically effective alternative to straight-through estimators for learning discrete token boundaries. The authors argue that this approach allows tokenization to become a fully learned component of the LLM architecture, moving away from artisanal, hand-crafted designs.

The authors remain modest regarding the scale of their experiments (100M parameter scale), noting that downstream effects on language modeling quality are difficult to assess directly at this size compared to frontier models. However, they claim the method successfully recovers the performance of BPE-guided downsampling without external priors. They identify open problems regarding the scaling of optimal tokenization aspect ratios with model size and the generalization of these strategies to variable downsampling rates. The work suggests that as models become increasingly end-to-end, score function estimates may become the preferred method for tokenization.

Technical Summary: Learning Tokenization End-to-End with Reinforcement Learning

More like this