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

TL;DR

NAT introduces a method to reduce token usage in reinforcement learning for language models by selectively updating only a subset of tokens, maintaining performance while lowering computational costs.

Contribution

The paper proposes NAT, a novel framework that employs unbiased partial-token policy-gradient estimation to efficiently scale RL with long sequences.

Findings

NAT matches full-token RL performance with only 50% token updates.

RPC reduces GPU memory by 18% and training time by 29%.

NAT enables more scalable RL for long chain-of-thought trajectories.

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.