Pith. sign in

REVIEW 3 major objections 4 minor 4 references

This paper shows that token boundaries in an autoregressive byte-level language model can be learned with a score-function (REINFORCE) gradient estimate, not a straight-through heuristic, and that the learned boundaries beat prior methods.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-02 23:21 UTC pith:U7CZZPD4

load-bearing objection A well-executed paper showing REINFORCE-style token-boundary learning beats straight-through estimators at 100M scale, but the trained objective is a biased surrogate of the theory it claims to realize. the 3 major comments →

arxiv 2602.13940 v2 pith:U7CZZPD4 submitted 2026-02-15 cs.LG cs.AI

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

Sam Dauncey , Roger Wattenhofer This is my paper
classification cs.LG cs.AI
keywords tokenizationreinforcement learningscore function estimatorREINFORCEautoregressive U-netbyte-level language modelsstraight-through estimatorend-to-end learning
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper tries to prove that tokenization—the hardcoded compression step that turns text into subword tokens before training—can be learned end-to-end inside an autoregressive byte-level model by treating each token boundary as a discrete stochastic decision and optimizing it with a score-function (REINFORCE) gradient estimate. Earlier attempts used straight-through estimators that relax the discrete decision into a continuous one; the paper argues these lack a convergence guarantee, while score-function estimates directly optimize the expected next-byte loss and offer a local-optimality guarantee. To make the high-variance estimate practical, it adds an early-exit baseline, time discounting, and batch-relative advantage centering. It reports that the resulting method learns boundaries that align with whitespace and semantic breaks, and beats straight-through baselines on bits-per-byte and LAMBADA at the 100M-parameter scale. If right, this suggests tokenization can be removed from the handcrafted part of LLM pipelines and folded into learning itself.

Core claim

The paper claims that token boundaries in an autoregressive U-net can be learned by directly optimizing the expected byte-level cross-entropy via a score-function estimator, rather than by relaxing the discrete boundary decisions as straight-through estimators do. The reward for each boundary is the improvement in next-byte log-probability over an early-exit byte-level prediction, discounted over future bytes and centered across the batch to reduce variance. At 147M parameters on FineWeb, this reaches 1.297 bits-per-byte versus 1.372 and 1.386 for the straight-through baselines, and lifts LAMBADA zero-shot accuracy to 0.086 versus 0.036, 0.029, and 0.020. The learned boundaries concentrate a

What carries the argument

The central object is the stochastic boundary policy πθ(ai|x≤i, a<i) = Bernoulli(σ(li)), where the logit li is computed from the byte-level representation Xi plus a sliding window of previous boundary decisions, all scaled and softcapped. The policy is trained by REINFORCE with the reward Ri = log pθ(xi|x<i, a<i) − log p_early θ(xi|x<i), discounted as Gi = Σ γ^j R_{i+j+1}, and batch-centered as Ai,b = Gi,b − mean over the batch. These variance-reduction techniques make a single-sample score-function estimate workable; the policy-gradient term carries a local-optimality guarantee that the paper argues straight-through estimates lack.

Load-bearing premise

The discounted, batch-centered reward in equations (13)–(14) is a close enough proxy for the true expected next-byte loss that optimizing it still yields a locally optimal tokenization strategy.

What would settle it

Train the same U-net architecture with the same compute, data, and downsampling rate but replace the learned boundary policy with a fixed uniform or whitespace heuristic; if the RL-learned boundaries do not reduce held-out bits-per-byte below the heuristic at equal FLOPs, the central claim that learning boundaries by reinforcement learning improves tokenization is refuted. A second probe: keep the method identical but set the discount factor to γ = 1.0; if performance does not degrade, the paper's assertion that time discounting is necessary for variance reduction would be unsupported.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • If the central claim is correct, tokenizer design can be removed from the handcrafted preprocessing stage and learned directly to minimize the model's loss.
  • The policy-gradient formulation gives a theoretical local-optimality guarantee for learned token boundaries, a property not available for straight-through estimators.
  • In the reported experiments, the learned boundaries improve held-out bits-per-byte and LAMBADA accuracy over prior straight-through baselines at the 100M-parameter scale.
  • The method remains effective across different target downsampling rates and transfers to Python code, suggesting it is not tied to a single domain or compression ratio.
  • Because the boundary policy is lightweight and architecture-agnostic, the same score-function update could be combined with other up/downsampling designs and hierarchical architectures.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Editorial extension: if the convergence guarantee survives at the much larger scales where byte-level U-nets have been shown to beat token-level baselines, learned tokenization could eventually replace BPE and hand-tuned tokenizer rules entirely.
  • Editorial extension: the early-exit baseline suggests a reusable principle—reward discrete structural decisions by how much they beat a cheap internal predictor—that could transfer to other discrete choices such as chunking long contexts, Mixture-of-Experts routing, or adaptive computation.
  • Editorial extension: an ablation that sets the discount factor γ to 1.0, or removes batch centering, would directly test whether the reported gains come from variance reduction; the paper does not isolate these components separately.
  • Editorial extension: testing on morphologically rich or multilingual byte streams would clarify whether the discovered whitespace-aligned boundaries reflect a general linguistic prior or a property of the English-dominated training data.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper proposes to learn token boundaries inside an autoregressive byte-level U-net by directly optimizing the expected next-byte cross-entropy with a score-function (REINFORCE) estimator. The authors derive the standard stochastic-computation-graph gradient identity (Eq. 10), then introduce practical variance-reduction modifications: an early-exit baseline (Eq. 12), time discounting (Eq. 13), batch-relative centering (Eq. 14), and a target-downsampling-rate regularizer (Eq. 20). They train 147M-parameter models on FineWeb and 90M-parameter models on CodeParrot, compare against uniform and two straight-through-estimator baselines (Nawrot et al. and Hwang et al.), and report improved bits-per-byte and zero-shot accuracy, alongside qualitative evidence that learned boundaries align with whitespace and other semantic structures. They also investigate sensitivity to target downsample rate on smaller models.

Significance. If the results hold, the paper makes a useful contribution to end-to-end tokenization by showing that score-function estimators, with suitable variance reduction, can be made practicable for discrete boundary learning and can outperform straight-through estimators at the 100M-parameter scale. The derivation of Eq. (10) is standard and correct, and the architectural flexibility of the approach is a strength. However, the paper's theoretical guarantee is claimed for the unbiased estimator, while the actually optimized loss uses a biased, discounted, and λ-scaled surrogate; the empirical comparisons lack error bars. These gaps materially weaken the central claims as currently stated.

major comments (3)
  1. [§2.3–§2.4, Eqs. (10), (13)–(15), (22)] The theoretical guarantee in §2.3—that gradient descent on the policy-gradient term yields a locally optimal tokenization strategy—applies to the unbiased REINFORCE gradient of Eq. (10). The loss actually trained, Eq. (22), replaces the full reward-to-go with the discounted sum G_i of Eq. (13) at γ=0.99 and scales L_π by λ_π=10^{-2}. The text acknowledges 'a small amount of bias' but gives no bound or argument that this bias is small relative to the difference between the surrogate optimum and the true objective of Eq. (9). Since boundary decisions can affect loss far downstream through the hierarchical architecture, discounting can in principle prefer myopic boundaries. The claim that the method has 'tighter theoretical guarantees' than straight-through estimators is therefore not established for the implemented algorithm. I recommend either proving a bias bound, running an ablation wit
  2. [Table 1 and §4.2] The central empirical claim—'outperforms prior straight-through estimates'—rests on Table 1, which reports single numbers with no standard errors, confidence intervals, or number of seeds. The LAMBADA differences (0.086 vs 0.036/0.029/0.020) are large in relative terms but the absolute accuracies are low and likely noisy at 147M scale. Similarly, the CodeParrot validation-loss comparison in Appendix D reports a single run per method. Without repeated runs or some measure of variance, the reader cannot assess whether the improvements are significant. Please report multiple seeds, or at least provide seed-level results and error bars.
  3. [§2.6, Eq. (20); §4.1] The target-downsample-rate loss L_target applies a single scalar pressure to the batch-mean logit and is scaled by λ_target=10^{-2}. This is another departure from the unbiased objective: the model is not optimizing the expected loss under a free boundary policy, but a penalized objective that forces a specific average rate. The paper does not discuss whether the reported semantic boundaries are a consequence of the rate pressure rather than of minimizing loss. A controlled experiment with different λ_target values, or at least a discussion of this confound, would strengthen the interpretation of the qualitative results.
minor comments (4)
  1. [Throughout] Typos and minor wording issues: 'Appendicies' (Appendix B), 'soley' (Abstract, §5), 'a mechanism to keep' (§2.6), 'onCodeParrot' (§4.3), 'an100' (caption of Table 2 context).
  2. [§2.5, Eq. (17)] The logit computation uses a sliding window of preceding boundaries; the notation 'WkXi' and 'Wj' is inconsistent (k vs j). Also, the scaling constant D=16 and softcapping are introduced without sensitivity analysis; please clarify whether results are robust to these choices.
  3. [§4 and Appendix D] The paper says the dataset size, batch size, and learning rate were selected using scaling laws from Porian et al. (2024), but does not report the specific scaling-law computation. The hyperparameter table is helpful, but a brief justification of the selected values would improve reproducibility.
  4. [§3, references] The reference 'Main Horse (pseudonym)' in the H-Net discussion is unusual for a peer-reviewed venue. If the technical details are only available in that blog post, please provide a more citable source or include the relevant specifics in the paper.

Circularity Check

0 steps flagged

No significant circularity: the derivation is a standard score-function identity applied to the model's own next-byte loss; empirical comparisons are independent external benchmarks.

full rationale

Eq. (10) is the standard stochastic-computation-graph gradient identity (Schulman et al., 2015; Williams, 1992), and the objective being optimized, Eq. (9), is the model's own expected next-byte cross-entropy. No predicted quantity is defined in terms of the fitted output, and no fitted parameter is renamed as a prediction. The variance-reduction terms in §2.4—early-exit baseline, time discounting, batch-relative centering—are control variates or reward shifts rather than targets fitted to the evaluation data. The empirical claims in Table 1 and Figure 10 are held-out external benchmarks and validation losses, independent of the theoretical derivation. There is no load-bearing self-citation: the cited prior work is external, and the autoregressive U-net is argued from architectural desiderata rather than imported from a self-authored uniqueness theorem. The acknowledged bias introduced by γ=0.99 in Eq. (13) is a potential soundness/overclaim gap between the unbiased theory of §2.3 and the practical estimator, not a circular reduction: it does not make the reported predictions equal to the inputs by construction. Consequently, no circularity is present.

Axiom & Free-Parameter Ledger

7 free parameters · 6 axioms · 0 invented entities

The method rests on standard score-function estimation plus several domain-specific assumptions about U-net architectures and the smallness of discounting bias. The main free parameters are loss weights, the discount factor, the target downsample rate, and window/scale factors; no new physical or model entities are introduced.

free parameters (7)
  • Discount factor γ = 0.99
    Hand-chosen in §2.4 to trade bias against variance; it changes the objective from exact reward-to-go to a discounted surrogate, which is central to the method.
  • Logit scaling factor D = 16
    Chosen in §2.5 to make token-boundary logits approximately uniform at initialization; affects exploration and stability.
  • Target downsample rate π̄target = 1/5 (natural language), 1/n (aspect-ratio runs)
    Set per experiment in §2.5-2.6; the downsample-rate targeting loss is designed to push the model toward this rate.
  • Policy loss weight λπ = 1e-2
    Set in §2.7; controls the strength of the policy-gradient term relative to the autoregressive loss.
  • Target-rate loss weight λtarget = 1e-2
    Set in §2.7; controls the downsample-rate pressure.
  • Early-exit loss weight λearly = 1e-1
    Set in §2.7; controls how strongly the early-exit baseline model is trained.
  • Boundary context window w = 8
    Set in §2.5; limits how many previous token boundaries influence the current boundary logit, a compute/expressivity tradeoff.
axioms (6)
  • standard math Score-function gradient identity: ∇θ E_{a∼πθ} log pθ(y|a,x) = E[∇θ log pθ + log pθ ∇θ log πθ]
    Used in §2.3/Appendix A.1 as the foundation of the method; standard stochastic computation graph result.
  • standard math E_{a∼πθ} ∇θ log πθ(a|x) = 0, so baselines independent of a_i can be added without bias
    Used in §2.4 to justify the early-exit and batch-relative baselines.
  • domain assumption Modern GPUs cannot efficiently perform sparse memory accesses, so tokenization for feed-forward models must be consistent across layers
    Used in §2.2 to argue that the autoregressive U-net structure is necessary rather than optional.
  • domain assumption The byte-level encoder X = encode(x) is independent of boundary decisions a, so the early-exit baseline p_early is tokenization-independent
    Required for the validity of the early-exit baseline in eq. (11)-(12); depends on the specific U-net architecture choice.
  • domain assumption A per-byte Bernoulli policy with a window of w=8 previous boundaries is expressive enough to represent useful tokenization heuristics
    Used in §2.5; restrictive but pragmatic. Fixed-stride strategies need access to previous boundary positions, hence the window.
  • ad hoc to paper The bias introduced by time discounting (γ=0.99) is small enough that the learned policy still approximately minimizes the true expected loss
    This is the weakest premise: eq. (13) replaces the exact reward-to-go with a discounted sum, so the 'locally optimal' guarantee from eq. (10) does not strictly apply to the trained objective.

pith-pipeline@v1.3.0-alltime-deepseek · 12165 in / 13137 out tokens · 123407 ms · 2026-08-02T23:21:00.331021+00:00 · methodology

0 comments
read the original abstract

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.

Figures

Figures reproduced from arXiv: 2602.13940 by Roger Wattenhofer, Sam Dauncey.

Figure 1
Figure 1. Figure 1: [Left] An example of the autoregressive U-net architecture (Nawrot et al., 2022), computed [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: [Left] Token boundaries learned by our 147M-parameter model on a held-out sample of [PITH_FULL_IMAGE:figures/full_fig_p007_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Token boundaries learned by our 147M-parameter model on a held-out sample of the [PITH_FULL_IMAGE:figures/full_fig_p014_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Token boundaries learned by a 147M-parameter model using the straight-through estima [PITH_FULL_IMAGE:figures/full_fig_p015_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Token boundaries learned by a 147M-parameter model using the straight-through estima [PITH_FULL_IMAGE:figures/full_fig_p016_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Token boundaries learned by our 90M-parameter model on a held-out sample of the [PITH_FULL_IMAGE:figures/full_fig_p017_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: Token boundaries learned by a 90M-parameter model using the straight-through estimator [PITH_FULL_IMAGE:figures/full_fig_p017_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Token boundaries learned by an array of models using our method with varying target [PITH_FULL_IMAGE:figures/full_fig_p018_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Token boundaries learned by an array of models using our method with varying target [PITH_FULL_IMAGE:figures/full_fig_p019_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: [left] Flops vs validation loss curves for [PITH_FULL_IMAGE:figures/full_fig_p021_10.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

4 extracted references · 3 linked inside Pith

  1. [2]

    Jessica Rumbelow and Matthew Watkins

    URLhttps://doi.org/10.48550/arXiv.2408.00118. Jessica Rumbelow and Matthew Watkins. Solidgoldmagikarp (plus, prompt generation). Less- Wrong, 2023. URLhttps://www.lesswrong.com/posts/aPeJE8bSo6rAFoLqg/ solidgoldmagikarp-plus-prompt-generation. John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computatio...

    Pith/arXiv arXiv doi:10.48550/arxiv.2408.00118 2023

  2. [3]

    Chaofan Tao, Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, and Ngai Wong

    URLhttps://proceedings.neurips.cc/paper_files/paper/2024/ file/e21955c93dede886af1d0d362c756757-Paper-Conference.pdf. Chaofan Tao, Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, and Ngai Wong. Scaling laws with vocabulary: Larger models deserve larger vocabularies. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Pa- quet, J...

    2024

  3. [2022]

    Junxiong Wang, Tushaar Gangavarapu, Jing Nathan Yan, and Alexander M Rush

    URLhttps://openreview.net/forum?id=JtBRnrlOEFN. Junxiong Wang, Tushaar Gangavarapu, Jing Nathan Yan, and Alexander M Rush. Mambabyte: Token-free selective state space model. InFirst Conference on Language Modeling, 2024. URL https://openreview.net/forum?id=X1xNsuKssb. Ronald J Williams. Simple statistical gradient-following algorithms for connectionist re...

    Pith/arXiv arXiv 2024

  4. [2024]

    URLhttps://proceedings.neurips.cc/paper_files/paper/2024/ file/b6341525cd84f3be0ef203e4d7cd8556-Paper-Conference.pdf. Morgane Rivi`ere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, L´eonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ram ´e, Johan Ferret, Peter Liu, Pouya Tafti, Abe Friesen, Michelle Casbon, Sabela Ramos...

    Pith/arXiv arXiv 2024