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

Tokenization is a hardcoded compression step which remains in the training pipeline of Large Language Models (LLMs), despite a general trend towards architectures becoming increasingly end-to-end. Prior work has shown promising results at scale in bringing this compression step inside the LLMs' architecture with heuristics to draw token boundaries, and also attempts to learn these token boundaries with straight-through estimates, which treat the problem of drawing discrete token boundaries as a continuous one. We show that these token boundaries can instead be learned using score function estimates, which have tighter theoretical guarantees due to directly optimizing the problem of drawing discrete token boundaries to minimize loss. We observe that techniques from reinforcement learning, such as time discounting, are necessary to reduce the variance of this score function sufficiently to make it practicable. We demonstrate that the resultant method outperforms prior proposed straight-through estimates, both qualitatively and quantitatively at the $100$ million parameter scale.

Before a language model can read text, the text is first chopped into small pieces in a step called tokenization; grouping letters into larger chunks lets the model process language faster. Today this is done by a separate, hand-built program that follows fixed rules chosen in advance and, unlike the rest of a modern AI system, never learns or improves from data. We asked whether a model could instead figure out for itself where to cut the text, as part of its ordinary training. Our method lets the model try different ways of splitting the text and rewards the choices that make what comes next easier to predict — a trial-and-error approach borrowed from reinforcement learning. Remarkably, without being told anything about words or grammar, the model discovers sensible boundaries on its own, such as breaking at the spaces between words, and it does so better than previous attempts to learn this step. This brings us closer to AI systems that are learned from start to finish, with no hand-designed parts. The same idea could eventually help models handle data such as audio or video, where there is no obvious way to divide the stream into pieces.