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

Researchers propose learning tokenization boundaries in large language models using reinforcement learning and score function estimates instead of hardcoded compression. This approach directly optimizes discrete token boundaries, outperforming prior straight-through estimation methods at the 100 million parameter scale.

Tokenization has remained a static, hardcoded step in LLM training pipelines despite the broader industry shift toward fully differentiable, end-to-end architectures. This research addresses a fundamental inefficiency by making token boundary selection learnable through reinforcement learning techniques, specifically using score function estimates that directly optimize the discrete optimization problem rather than approximating it with continuous relaxations.

The innovation emerges from recognizing that straight-through estimators, which treat discrete token boundary selection as a continuous problem, suffer from optimization issues. Score function methods provide tighter theoretical guarantees because they directly sample from the discrete decision space. The authors crucially incorporate variance reduction techniques from reinforcement learning—particularly time discounting—to make the approach practically viable in training.

This work carries implications for LLM efficiency and adaptability. Learned tokenization could reduce model size requirements, improve compression ratios, and potentially allow models to dynamically adjust tokenization strategies for different domains or languages. The demonstrated superiority over straight-through estimates suggests the approach scales beyond the 100 million parameter experiments shown.

Future development paths include testing at larger scales (billion+ parameters) and exploring whether learned tokenization enables cross-lingual or multi-domain efficiency gains. The research validates that reinforcement learning principles can solve discrete architectural problems in deep learning, potentially opening similar approaches to other hardcoded compression or discretization steps. Industry adoption depends on whether these efficiency gains translate to production-scale models and whether the additional computational overhead of learned tokenization justifies its benefits.

• →Learned tokenization via RL outperforms traditional hardcoded and straight-through estimation approaches at demonstrated scales
• →Score function estimates provide tighter theoretical guarantees than continuous relaxations for discrete token boundary selection
• →Variance reduction techniques from reinforcement learning are essential to make score function estimation practical
• →Learned tokenization could enable dynamic, domain-specific token strategies rather than fixed compression schemes
• →The approach demonstrates broader applicability of RL principles to solving discrete architectural problems in LLMs