Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Routing, pruning and distillation for reasoning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Learns token boundaries end-to-end inside an autoregressive U-Net using a score-function (REINFORCE-style) estimator with RL variance reduction, outperforming straight-through estimator methods at ~147M parameters.

Problem

Tokenization is a hardcoded preprocessing step in LLMs; prior learned-boundary methods rely on straight-through estimators (heuristic, no guarantees) or fixed heuristics such as whitespace or entropy. The paper asks whether token boundaries can be learned end-to-end from the cross-entropy loss with

Approach

Treat token boundary placement as a stochastic policy and train it with a score-function (REINFORCE) estimator, reducing variance via an early-exit baseline, time discounting (gamma=0.99), batch-relative advantages, and a target downsample-rate penalty.

Key results

• Learned score-function policy achieves FineWeb test bits-per-byte of 1.297 vs 1.372 for Dynamic (Nawrot et al.) and 1.386 for H-Net at 147M parameters.
• Ours attains lowest bits-per-byte on PIQA (1.557), HellaSwag (1.212 vs 1.230 BPE guidance is slightly better; ours w=1 gets 1.199), and LAMBADA (1.737 vs 1.956 Nawrot, but BPE guidance 1.645 is better).
• BPE-guided boundaries (1.299) are essentially matched by the learned policy (1.297) without external priors.

Surprising findings

• The simpler w=1 variant (boundary logit depends only on the byte representation) performs comparably to the full sliding-window model.
• The learned policy places boundaries at whitespace-like characters without any linguistic inductive bias.

Limitations

• Experiments only at ~147M parameters and few billion bytes, far below frontier scale; downstream effects at scale are unknown.
• Only a single tokenization level and a fixed downsample rate target; aspect-ratio and variable rate are left to future work.
• Benchmark differences are small and mostly within noise at this scale; accuracy is near chance so bits-per-byte is used instead.

Experimental setup

Datasets: FineWeb CodeParrot

Benchmarks: PIQA HellaSwag ARC-Easy LAMBADA FineWeb Test

Metrics: bits-per-byte validation loss downsampling rate

Baselines: Uniform boundaries Dynamic (Nawrot et al.) straight-through estimator H-Net (Hwang et al.) BPE guidance

Builds on: REINFORCE autoregressive U-Net (MEGABYTE-style hierarchical byte models) Nawrot et al. dynamic pooling H-Net Gemma 2

Closest papers by meaning

• Reinforcement Mid-Training (ICLR 2026 · similarity 0.89)
• How reinforcement learning after next-token prediction facilitates learning (ICLR 2026 · similarity 0.87)
• Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.87)
• Getting the most out of your tokenizer for pre-training and domain adaptation (ICML 2024 · similarity 0.87)
• STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens (NeurIPS 2026 · similarity 0.87)
• Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs (ICLR 2026 · similarity 0.86)