Not all tokens are needed(NAT): token efficient reinforcement learning

Hejian Sang; Yuanda Xu; Zhengze Zhou; Ran He; Zhipeng Wang

cs.LGFeb 20, 2026

Not all tokens are needed(NAT): token efficient reinforcement learning

Authors: Hejian Sang, Yuanda Xu, Zhengze Zhou, Ran He, Zhipeng Wang

Organizations: LinkedIn Corporation, CA, USA

Abstract

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 $\pm$ 0.031	0.710 $\pm$ 0.064	0.259 $\pm$ 0.027	0.469 $\pm$ 0.037	0.126 $\pm$ 0.020	0.256 $\pm$ 0.015
URS	0.610 $\pm$ 0.050	0.741 $\pm$ 0.080	0.190 $\pm$ 0.066	0.360 $\pm$ 0.105	0.116 $\pm$ 0.051	0.229 $\pm$ 0.062
Det. Trunc.	0.361 $\pm$ 0.265	0.423 $\pm$ 0.318	0.037 $\pm$ 0.076	0.072 $\pm$ 0.159	0.027 $\pm$ 0.071	0.071 $\pm$ 0.178
RPC	0.670 $\pm$ 0.093	0.807 $\pm$ 0.135	0.170 $\pm$ 0.079	0.405 $\pm$ 0.081	0.122 $\pm$ 0.019	0.282 $\pm$ 0.065

Table 2: Token-efficient RL training results on three math benchmarks (MATH, AIME24, AIME25). We report Acc@16 and pass@16 (mean

\pm

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).

Panel of Figure 2 — 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).

Panel of Figure 3 — 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 $\pm$ 0.042	166.579 $\pm$ 6.728	291.240 $\pm$ 11.943
URS	35.758 $\pm$ 0.136	147.399 $\pm$ 34.520	259.376 $\pm$ 55.345
Det. Trunc.	24.603 $\pm$ 0.487	81.552 $\pm$ 4.322	137.211 $\pm$ 9.388
RPC	29.235 $\pm$ 0.362	121.918 $\pm$ 9.601	230.428 $\pm$ 17.573
Qwen3-8B Model

Table 3: System efficiency metrics for token-efficient RL training (mean

\pm

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

Panel of Figure 4 — 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).

Panel of Figure 5 — Figure 5: Time per step (excluding inference time) with 95% confidence interval across 5 runs for RPC (GRPO with uniform random prefix cutting).

Panel of Figure 6 — 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).

RL for Language Models

Reinforcement learning for language models uses feedback-driven policy updates to improve reasoning, alignment, and interactive behavior.

See all 433 papers

IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning
Feb 22, 2026Yinhan He, Yaochen Zhu, Mingjia Shi +5Efficient Language Model Reasoning Language Model Post-Training
University of Virginia, Charlottesville, VA, USA · LinkedIn Inc., Sunnyvale, CA, USA
You Can Learn Tokenization End-to-End with Reinforcement Learning
Feb 15, 2026Sam Dauncey, Roger WattenhoferSubword Tokenization LLM Training
Department of Electrical Engineering, ETH Zurich, Zurich, Switzerland.
Unlocking Token Rewards via Training-Free Reward Attribution
Oct 13, 2026Sitong Wu, Haoru Tan, Bin Xia +6Token-Level Credit Assignment RL for Language Model Reasoning
Chinese University of Hong Kong · University of Hong Kong · Hong Kong University of Science and Technology

Feb 22, 2026J/K move · Enter open · S save

IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning

Yinhan He, Yaochen Zhu, Mingjia Shi, Wendy Zheng+4

University of Virginia, Charlottesville, VA, USA · LinkedIn Inc., Sunnyvale, CA, USA

Efficient Language Model Reasoning Language Model Post-Training RL for Language Model Reasoning

Open paper AskPDF arXiv

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.

Abstract

RL for Language Models

Explore similar work

IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning

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

Unlocking Token Rewards via Training-Free Reward Attribution