Reinforcement learning (RL) has become a key driver of progress in large language models, but scaling RL to long chain-of-thought (CoT) trajectories is increasingly constrained by backpropagation over every generated token. Even with optimized rollout engines, full-token updates can consume a large fraction of total training cost, turning token length into a hidden tax on RL. We introduce Not All Tokens Are Needed (NAT), a unified framework that makes the token budget a first-class optimization primitive. NAT updates the policy using only a selected subset of generated tokens while preserving the learning signal of full-sequence RL. The core idea is an unbiased partial-token policy-gradient estimator via Horvitz-Thompson reweighting, which ensures statistically correct gradients despite subsampling. We instantiate NAT with two simple, plug-and-play token selection schemes: Uniform Random Sampling (URS) and Random Prefix Cutting (RPC), both of which reduce forward and backward compute and memory without modifying the reward computation or rollout pipeline. Across mathematical reasoning benchmarks, NAT matches full-token GRPO performance while using as few as 50% of tokens, providing an efficient and orthogonal pathway to scaling RL beyond the limits imposed by long trajectories. In our experiments, RPC saves 18% peak GPU memory and 29% forward and backward RL training time for Qwen3-8B.
Figures & tables
Figure 1: Barplots of Qwen3-8B RL training metrics with 95% CIs across 5 runs for GRPO (vanilla GRPO), URS (GRPO with random sampling p=0.5 ), Det. Trunc. (GRPO with deterministic prefix truncation of 50% tokens) and RPC (GRPO with uniform random prefix cutting).
Method
Unbiased?
Forward Savings
Backward Savings
Key Property
URS
Yes
No
Yes
Simple, constant p sampling
Det. Trunc.
No
Yes
Yes
Systematic bias, ignores late tokens
RPC
Yes
Yes
Yes
Structured, preserves causal context
Table 1: Comparison of token-efficient methods. RPC is the only method that provides computational savings in both passes while remaining statistically unbiased.
Method
MATH
AIME24
AIME25
Acc@16
pass@16
Acc@16
pass@16
Acc@16
pass@16
Qwen2.5-Math-7B Model
GRPO
0.610 ± 0.031
0.710 ± 0.064
0.259 ± 0.027
0.469 ± 0.037
0.126 ± 0.020
0.256 ± 0.015
URS
0.610 ± 0.050
0.741 ± 0.080
0.190 ± 0.066
0.360 ± 0.105
0.116 ± 0.051
0.229 ± 0.062
Det. Trunc.
0.361 ± 0.265
0.423 ± 0.318
0.037 ± 0.076
0.072 ± 0.159
0.027 ± 0.071
0.071 ± 0.178
RPC
0.670 ± 0.093
0.807 ± 0.135
0.170 ± 0.079
0.405 ± 0.081
0.122 ± 0.019
0.282 ± 0.065
Table 2: Token-efficient RL training results on three math benchmarks (MATH, AIME24, AIME25). We report Acc@16 and pass@16 (mean ± 95% confidence interval across runs) for two base models. Within each model, we color each cell by whether its 95% CI overlaps the GRPO baseline for the same metric ( green : overlap; red : non-overlap, grey : overlap with much high variance and lower scores).
Figure 2: Entropy curves with 95% confidence interval across 5 runs for GRPO (vanilla GRPO), URS (GRPO with random sampling p=0.5 ), Det. Trunc. (GRPO with deterministic prefix truncation of 50% of trajectory tokens) and RPC (GRPO with uniform random prefix cutting).
Figure 3: Percentage of selected tokens with 95% confidence interval across 5 runs for RPC (GRPO with uniform random prefix cutting).
Method
Avg. Peak GPU Memory (GB)
Avg. Train Time / Step (s) w/o inference
Avg. Total Time / Step (s)
Qwen2.5-Math-7B Model
GRPO
35.810 ± 0.042
166.579 ± 6.728
291.240 ± 11.943
URS
35.758 ± 0.136
147.399 ± 34.520
259.376 ± 55.345
Det. Trunc.
24.603 ± 0.487
81.552 ± 4.322
137.211 ± 9.388
RPC
29.235 ± 0.362
121.918 ± 9.601
230.428 ± 17.573
Qwen3-8B Model
Table 3: System efficiency metrics for token-efficient RL training (mean ± 95% CI across 5 runs). We report average peak GPU memory, average training time per step excluding inference, and average total time per step including inference, each averaged over training steps. Within each model block, URS overlaps with GRPO on all three metrics under 95% CIs. Green indicates better efficiency (lower mean) than GRPO; gray indicates overlap/parity with GRPO.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Gradient norm curves with 95% confidence interval across 5 runs for GRPO (vanilla GRPO), URS (GRPO with random sampling p=0.5 ), Det. Trunc. (GRPO with deterministic prefix truncation of 50% of trajectory tokens) and RPC (GRPO with uniform random prefix cutting).
Figure 5: Time per step (excluding inference time) with 95% confidence interval across 5 runs for RPC (GRPO with uniform random prefix cutting).
Figure 6: Max allocated GPU memory curves in GB with 95% confidence interval across 5 runs for GRPO (vanilla GRPO), URS (GRPO with random sampling p=0.5 ), Det. Trunc. (GRPO with deterministic prefix truncation of 50% of trajectory tokens) and RPC (GRPO with uniform random prefix cutting).
Large language models increasingly rely on long chains of thought to improve accuracy, yet such gains come with substantial inference-time costs. We revisit token-efficient post-training and argue that existing sequence-level reward-shaping methods offer limited control over how reasoning effort is allocated across tokens. To bridge the gap, we propose IAPO, an information-theoretic post-training framework that assigns token-wise advantages based on each token's conditional mutual information (MI) with the final answer. This yields an explicit, principled mechanism for identifying informative reasoning steps and suppressing low-utility exploration. We provide a theoretical analysis showing that our IAPO can induce monotonic reductions in reasoning verbosity without harming correctness. Empirically, IAPO consistently improves reasoning accuracy while reducing reasoning length by up to 36%, outperforming existing token-efficient RL methods across various reasoning datasets. Extensive empirical evaluations demonstrate that information-aware advantage shaping is a powerful and general direction for token-efficient post-training. The code is available at https://github.com/YinhanHe123/IAPO.
Yinhan He, Yaochen Zhu, Mingjia Shi +5
University of Virginia, Charlottesville, VA, USA · LinkedIn Inc., Sunnyvale, CA, USA
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.
Sam Dauncey, Roger Wattenhofer
Department of Electrical Engineering, ETH Zurich, Zurich, Switzerland.
In this paper, we propose an extremely efficient, training-free method to extract token-level reward signals directly from an existing deep reward model. Our core idea is to attribute the overall process reward to individual tokens by estimating each token's influence. This influence is defined as the change in the final macroscopic reward (e.g., the process reward) when a token is replaced with a semantically null token. Naively calculating this influence is computationally infeasible, requiring N forward passes through the PRM for an N-token sequence. We overcome this bottleneck by proposing a highly efficient gradient-based estimator. Specifically, we use a first-order Taylor approximation, which simplifies the influence calculation to the inner product of the difference between the token embedding and the null token embedding, and the gradient of the reward with respect to the token embedding. This requires only a single forward and backward pass. The resulting token-level rewards enable standard RL algorithms to perform precise credit assignment without requiring additional reward model training. Experiments on challenging reasoning benchmarks demonstrate that our method substantially improves policy optimization efficiency and enhances the generalization of LLM reasoning capabilities. Our P2T outperforms the outcome reward by +4.9% on MathVista for Qwen2.5-VL-7B-Instruct, and +11.5% on AIME24 for Qwen2.5-Math-7B. The code will be released via https://github.com/stonewst/P2T.
Sitong Wu, Haoru Tan, Bin Xia +6
Chinese University of Hong Kong · University of Hong Kong · Hong Kong University of Science and Technology