Paper deep dive

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

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

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 60

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%

Last extracted: 7/20/2026, 9:41:52 PM

Summary

The paper introduces Not All Tokens Are Needed (NAT), a framework for token-efficient reinforcement learning in large language models. NAT reduces computational costs by updating policies using only a subset of generated tokens, employing Horvitz-Thompson reweighting to maintain unbiased gradient estimates. It implements two selection schemes: Uniform Random Sampling (URS) and Random Prefix Cutting (RPC). Experiments on Qwen3-8B demonstrate that RPC matches full-token GRPO performance while saving 18% peak GPU memory and 29% training time.

Entities (7)

NAT · framework · 98%GRPO · algorithm · 95%Qwen3-8b · model · 95%RPC · method · 92%URS · method · 92%Horvitz-Thompson · technique · 90%Chain-of-Thought · concept · 85%

Relation Signals (7)

NAT → uses → Horvitz-Thompson

confidence 95% · The core idea is an unbiased partial-token policy-gradient estimator via Horvitz-Thompson reweighting

RPC → reduces → GPU Memory

confidence 92% · RPC saves 18% peak GPU memory

RPC → reduces → training time

confidence 92% · RPC saves ... 29% forward and backward RL training time

NAT → improves → RL Efficiency

confidence 90% · NAT ... provides an efficient and orthogonal pathway to scaling RL

NAT → instantiates → URS

confidence 90% · We instantiate NAT with two simple, plug-and-play token selection schemes: Uniform Random Sampling (URS)

NAT → instantiates → RPC

confidence 90% · We instantiate NAT with two simple, plug-and-play token selection schemes: ... Random Prefix Cutting (RPC)

GRPO → isbasefor → NAT

confidence 85% · We instantiate NAT on top of the GRPO objective

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

58,749 characters extracted from source content.

Expand or collapse full text

Not All Tokens are Needed: Token-Efficient Reinforcement Learning Hejian Sang1 * † Yuanda Xu1 * Zhengze Zhou1 * Ran He1 * Zhipeng Wang1 1 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 remains 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/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. From our experiments, RPC can save 18% peak GPU memory and 29% forward and backward RL training time for Qwen3-8B RL training. 11footnotetext: Equal contribution.22footnotetext: Corresponding author: Hejian Sang<hsang@linkedin.com> 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.5p=0.5), Det. Trunc. (GRPO with deterministic prefix truncation of 50% tokens) and RPC (GRPO with uniform random prefix cutting). 1 Introduction Recent generations of Large Language Models (LLMs) Ziegler et al. (2019); Ouyang et al. (2022); Stiennon et al. (2020) have demonstrated striking advances in mathematical reasoning Guo et al. (2025), code synthesis Pan et al. (2024), and scientific problem solving. A major driver behind these gains is post-training with reinforcement learning objectives that use verifiable signals (like unit tests for code, symbolic execution, or exact-answer checks for math) commonly referred to as Reinforcement Learning from Verifiable Rewards (RLVR) Jaech et al. (2024); Guo et al. (2025); Team et al. (2025); Yang et al. (2025). By directly optimizing for task success, RLVR can push models beyond what supervised fine-tuning alone achieves, especially on benchmarks where solutions require multi-step chains-of-thought (CoT) and careful intermediate decisions. As a result, RLVR has become a core ingredient in modern reasoning systems, including OpenAI O1 Jaech et al. (2024), DeepSeek-R1 Guo et al. (2025), Kimi 1.5 Team et al. (2025), and Qwen3 Yang et al. (2025). Despite its effectiveness, RLVR remains expensive and increasingly difficult to scale. A central reason is that standard RLVR pipelines treat all generated tokens as equally important for learning: they compute policy-gradient losses and backpropagate through every token in each rollout trajectory. However, as RLVR succeeds at eliciting deeper reasoning, trajectories often become longer and more variable, which amplifies activation memory, increases per-update FLOPs, and can trigger out-of-memory (OOM) failures or unstable optimization dynamics Guo et al. (2025); Yeo et al. (2025). In practice, the cost of RLVR grows not only with the number of rollouts, but also with the length of each rollout and long-CoT behavior can become the training bottleneck. A typical GRPO Shao et al. (2024) training cycle further highlights an efficiency mismatch across stages. Rollout generation is often the most visible bottleneck and can account for a large fraction of wall-clock time (e.g., around 70% in representative profiles Gao et al. (2025)). Consequently, recent work has aggressively optimized generation through high-throughput inference engines (Kwon et al., 2023; Zheng et al., 2024), speculative decoding (Leviathan et al., 2023) and low precision inference engine Xi et al. (2026). Yet, even when generation is accelerated, the subsequent learning phase—forward and backward passes over full trajectories—remains memory-bound and compute-intensive. This creates a practical ceiling: we can generate rollouts faster, but we cannot consume them efficiently enough to translate that throughput into scalable policy improvement. This paper takes an orthogonal approach to RLVR efficiency. Instead of optimizing how rollouts are produced, we optimize how they are used by the learner. We ask a simple but consequential question: Do we really need all tokens to train strong RL reasoners? Intuitively, not all tokens in a long CoT contribute equally to the learning signal. Many tokens are “mechanical” continuation (for example, filling in algebraic steps) reiterating previously established context, or emitting low-entropy boilerplate, while a smaller subset corresponds to high-impact decision points that steer the trajectory. If we can update the policy using only a carefully chosen subset of tokens—without distorting the training objective, we can reduce memory and compute per optimizer update, increase training throughput, and unlock longer-context RLVR without sacrificing performance. We introduce Not All Tokens are Needed (NAT), a principled framework for token-efficient RLVR. NAT performs policy optimization using only a selected subset of tokens from each rollout while keeping reward signals computed on the full response. Concretely, NAT masks tokens during the policy-gradient computation, so that only included tokens contribute to backpropagation. This truncates the activation graph, reduces memory pressure, and decreases per-step compute. Importantly, NAT is not restricted to one masking rule: it can instantiate uniform token sampling, structured prefix-based schemes, or other selection mechanisms, as long as we can correct for selection bias appropriately. The core technical ingredient of NAT is Horvitz–Thompson (HT) estimation applied to token-masked policy gradients. Each token t is assigned a inclusion probability, and NAT reweights gradients by the inverse of that probability. We prove that this HT-corrected token-masked objective yields an unbiased estimator of the original full-sequence RLVR gradient. In other words, NAT preserves the target learning signal in expectation, even though each update only backpropagates through a subset of tokens. This provides a rigorous foundation for token-efficient RLVR: partial-token updates can remain statistically aligned with optimizing the full-sequence reward objective, rather than introducing systematic drift. We instantiate NAT on top of the GRPO objective and evaluate across multiple reasoning benchmarks and model scales. Empirically, NAT achieves performance comparable to full-token GRPO while using substantially fewer tokens in the backward pass (e.g., as low as 50% token coverage), translating into lower activation memory and higher training throughput. Moreover, NAT composes naturally with system-level optimizations (faster rollout engines, batching, scheduling), because it targets a different bottleneck: the cost of consuming long trajectories during learning. Contributions. Our main contributions are: • A unified framework for token-efficient RLVR. We propose Not All Tokens are Needed (NAT), which performs RLVR updates using masked token subsets while preserving full-response reward evaluation. NAT supports a broad family of token selection strategies, including uniform random token sampling and structured prefix-based masking. • Unbiasedness via Horvitz–Thompson correction. We show that HT reweighting yields an unbiased estimator of the full-token GRPO gradient for any positive inclusion probabilities. This establishes a principled connection between token masking and faithful optimization of the original RLVR objective. • Practical speedups without sacrificing reasoning quality. Through experiments on math reasoning benchmarks, we demonstrate that NAT—especially Random Prefix Cutting (RPC)—can match full-token GRPO while reducing backpropagation volume by up to 50%, yielding direct savings: 18% less peak GPU memory and 29% less forward and backward RL training time for Qwen3-8B RL training using DAPO dataset Yu et al. (2025). 2 Preliminaries 2.1 Notation and RL Setup We consider an autoregressive policy πθ _θ that generates a response o=(o1,…,oT)o=(o_1,…,o_T) conditioned on a prompt q∼q . The policy factorizes as πθ(o∣q)=∏t=1Tπθ(ot∣q,o<t). _θ(o q)= _t=1^T _θ(o_t q,o_<t). In RL with verifiable rewards (RLVR), each completed response o receives a scalar reward R(q,o)R(q,o) computed by an automatic verifier (e.g., exact-match, unit tests, or rubric-based scoring). We distinguish the current policy πθ _θ (being optimized) from the behavior policy πθold _ _old used to collect rollouts. For each prompt q, we sample a group of G responses oii=1G∼πθold(⋅∣q)\o_i\_i=1^G _ _old(· q) and compute rewards Rii=1G\R_i\_i=1^G, where Ri:=R(q,oi)R_i:=R(q,o_i). For token t in response i, we define the standard importance ratio ri,t(θ)=πθ(oi,t∣q,oi,<t)πθold(oi,t∣q,oi,<t).r_i,t(θ)= _θ(o_i,t q,o_i,<t) _ _old(o_i,t q,o_i,<t). (1) This ratio corrects for the fact that data are collected under πθold _ _old while we optimize πθ _θ. 2.2 Group Relative Policy Optimization (GRPO) Group Relative Policy Optimization (GRPO) is a PPO-style policy-gradient method that removes the learned value function (critic) and instead constructs a baseline directly from the group of sampled responses for each prompt. Intuitively, responses in the same group compete against each other: a response is reinforced only insofar as it is better than its peers. This provides a low-variance advantage signal without training a separate critic model. Group-relative advantage. Given rewards Rii=1G\R_i\_i=1^G for prompt q, GRPO defines a normalized advantage per response: A^i=Ri−μRσR+ϵ,μR=1G∑j=1GRj,σR=1G∑j=1G(Rj−μR)2, A_i= R_i- _R _R+ε, _R= 1G _j=1^GR_j, _R= 1G _j=1^G(R_j- _R)^2, (2) where ϵ>0ε>0 is a small constant for numerical stability. GRPO then shares this response-level advantage across tokens: A^i,t:=A^i A_i,t:= A_i for all t. This design matches the RLVR setting where supervision is sequence-level (correct/incorrect), so all tokens in a completion receive the same global credit. Clipped surrogate objective. Like PPO Schulman et al. (2017), GRPO uses a clipped importance-weighted surrogate to prevent overly large policy updates. For token t in response oio_i, define the PPO-style clipped surrogate Si,t(θ)=min⁡(ri,t(θ)A^i,t,clip(ri,t(θ),1−ε,1+ε)A^i,t),S_i,t(θ)= (r_i,t(θ)\, A_i,t,\;clip (r_i,t(θ),1- ,1+ )\, A_i,t ), (3) where ε is the clip threshold. The min⁡(⋅) (·) yields a conservative (pessimistic) objective: once ri,t(θ)r_i,t(θ) leaves the trust region [1−ε, 1+ε][1- ,\,1+ ], further probability changes no longer increase the surrogate, which stabilizes optimization. Token-level GRPO loss with KL regularization. To prevent the policy from drifting too far from a fixed reference policy πref _ref (e.g., the SFT model), we incorporate a KL penalty directly into the token-level objective. We define the per-token GRPO loss as Li,tGRPO(θ)=Si,t(θ)−βDKL(πθ∥πref),L_i,t^GRPO(θ)=S_i,t(θ)-β\,D_KL\! ( _θ\,\|\, _ref ), (4) where DKL(πθ∥πref)D_KL( _θ\,\|\, _ref) denotes a per-token KL term evaluated at the same conditioning context (q,oi,<t)(q,o_i,<t), and β≥0β≥ 0 controls the strength of regularization.***Many implementations approximate this KL using the sampled-action log-prob difference log⁡πθ(oi,t∣q,oi,<t)−log⁡πref(oi,t∣q,oi,<t) _θ(o_i,t q,o_i,<t)- _ref(o_i,t q,o_i,<t); we keep the notation abstract for clarity. Although we present the general GRPO formulation with KL regularization for completeness, recent work advocates for removing KL divergence Yu et al. (2025). Sequence-level aggregation. GRPO aggregates the token-level loss by averaging over tokens within each response and then over responses in the group: JGRPO(θ)=q∼,oi∼πθold[1G∑i=1G1|oi|∑t=1|oi|Li,tGRPO(θ)].J_GRPO(θ)=E_q ,\,\o_i\ _ _old [ 1G _i=1^G 1|o_i| _t=1^|o_i|L_i,t^GRPO(θ) ]. (5) 2.3 GRPO Training Pipeline A GRPO iteration can be viewed as a three-stage pipeline with distinct system bottlenecks: (1) Rollout / Inference. For each prompt q, we sample a group of G completions oi\o_i\ from the behavior policy πθold _ _old (a lagged snapshot of the current policy). This stage is throughput-critical and is often served by high-performance inference engines (e.g., vLLM Kwon et al. (2023) or SGLang Zheng et al. (2024) ) that support continuous batching and memory-efficient KV-cache management. (2) Forward pass (scoring + log-probs). Given the collected trajectories, we evaluate each completion with the verifier to obtain rewards Ri\R_i\ and compute group-relative advantages A^i\ A_i\ via Eq. equation 2. In parallel, we run a forward pass of the current policy πθ _θ to compute token log-probabilities log⁡πθ(oi,t∣q,oi,<t) _θ(o_i,t q,o_i,<t), which are needed to form the importance ratios ri,t(θ)r_i,t(θ). Since reward evaluation is typically much cheaper than the model forward pass and can be overlapped with it, we refer to this combined stage as the forward pass for simplicity. (3) Backward pass (optimization). We backpropagate through the GRPO objective (Eq. equation 5) and update parameters using an optimizer such as AdamW Loshchilov & Hutter (2017). This stage is often memory-dominated because it must store activations for backpropagation over long chain-of-thought trajectories. This separation is useful in practice: rollout favors inference-optimized kernels and scheduling, while the training stages are dominated by activation memory and stable optimization. 3 Token-Efficient Reinforcement Learning RLVR pipelines decouple reward evaluation from optimization: rewards must be computed on the full generated response to faithfully reflect task success, but the policy-gradient update is typically formed by summing token-level log-prob terms over the entire trajectory. This creates a major efficiency opportunity. In particular, long chain-of-thought (CoT) trajectories inflate the activation footprint and FLOPs of the training stage, since standard GRPO updates backpropagate through every token. We propose a general framework that reduces the optimization cost by performing policy updates on a subset of tokens, while still using the full response to compute rewards and advantages. The key technical tool is Horvitz–Thompson (HT) estimation, which provides unbiased estimators of full-sequence objectives under randomized token selection. Conceptually, we keep the RL signal (RiR_i, A^i A_i) unchanged, and only sparsify which tokens contribute to the gradient. 3.1 General Framework: Token Masking with Horvitz–Thompson Reweighting Consider a prompt q and a sampled response oi=(oi,1,…,oi,Ti)o_i=(o_i,1,…,o_i,T_i) of length TiT_i. Recall from Section 2 that GRPO defines a token-level loss Li,tGRPO(θ)L^GRPO_i,t(θ) and optimizes the sequence-average objective by averaging over t=1,…,Tit=1,…,T_i and over samples i=1,…,Gi=1,…,G. Random masking. We introduce a binary mask mi,t∈0,1m_i,t∈\0,1\ indicating whether token t of trajectory i participates in the policy update: mi,t∼Bernoulli(pi,t),pi,t∈(0,1].m_i,t (p_i,t), p_i,t∈(0,1]. Here pi,tp_i,t is the inclusion probability. HT estimator for the per-sequence mean loss. Define the full-token per-sequence mean loss μi(θ)=1Ti∑t=1TiLi,tGRPO(θ) _i(θ)= 1T_i _t=1^T_iL^GRPO_i,t(θ). Using only the selected tokens, the Horvitz–Thompson estimator is μ^iHT(θ)=1Ti∑t=1Timi,tpi,tLi,tGRPO(θ). μ_i^HT(θ)\;=\; 1T_i _t=1^T_i m_i,tp_i,t\,L^GRPO_i,t(θ). (6) Proposition 1 (Unbiasedness of HT token masking). For any inclusion probabilities pi,tt=1Ti\p_i,t\_t=1^T_i with pi,t>0p_i,t>0, m[μ^iHT(θ)]=μi(θ)E_m[ μ_i^HT(θ)]= _i(θ). Moreover, under standard regularity conditions allowing interchange of gradient and expectation, m[∇θμ^iHT(θ)]=∇θμi(θ)E_m[ _θ μ_i^HT(θ)]= _θ _i(θ). The proof is shown in Appendix A. Consequently, replacing the full-token GRPO objective with the masked objective formed by μ^iHT μ_i^HT yields an unbiased estimator of the full-token GRPO gradient. The inclusion probabilities pi,tp_i,t define a compute–variance trade-off: smaller pi,tp_i,t reduces the expected number of backpropagated tokens, but increases estimator variance through the HT weights 1/pi,t1/p_i,t. Next we instantiate this framework with two practical masking schemes. Uniform Random Token Sampling (URS) URS sets a constant inclusion probability pi,t=p_i,t=p for all (i,t)(i,t). Intuitively, URS keeps a fraction p of tokens on average, providing approximately linear savings in the backward pass. Under HT correction, the per-token score-function contribution scales as 1/p1/p. For example, for a generic score term gi,t(θ)=∇θlog⁡πθ(oi,t∣q,oi,<t)g_i,t(θ)= _θ _θ(o_i,t q,o_i,<t), the HT-corrected contribution becomes g^i,turs(θ)=mi,tpgi,t(θ),mi,t∼Bernoulli(p). g urs_i,t(θ)\;=\; m_i,tp\,g_i,t(θ), m_i,t (p). This is unbiased: m[g^i,turs(θ)]=gi,t(θ)E_m[ g urs_i,t(θ)]=g_i,t(θ). Gradient-norm inflation under URS While URS is unbiased, HT reweighting inflates the second moment by 1/p1/p: m[‖g^i,turs(θ)‖2] _m\! [ \| g urs_i,t(θ) \|^2 ] =m[‖mi,tpgi,t(θ)‖2]=1p2m[mi,t]‖gi,t(θ)‖2 =E_m\! [ \| m_i,tpg_i,t(θ) \|^2 ]= 1p^2\,E_m[m_i,t]\,\|g_i,t(θ)\|^2 =1p‖gi,t(θ)‖2. = 1p\,\|g_i,t(θ)\|^2. Thus, decreasing p increases gradient norms (and variance) roughly as 1/p1/ p in typical regimes, which can trigger more frequent gradient clipping and slow optimization. Empirically, we find p can be reduced until clipping becomes dominant; beyond that point training becomes unstable. Limitation: URS saves backward FLOPs but not forward compute. Crucially, independent token masking does not necessarily reduce forward-pass cost in standard causal Transformers. Even if token t is masked out of the loss, computing log⁡πθ(oi,t∣q,oi,<t) _θ(o_i,t q,o_i,<t) still requires processing all preceding tokens 1,…,t−11,…,t-1 due to causal self-attention. As a result, URS primarily reduces backward computation (and some gradient-related memory), but retains near full forward cost. This motivates a structured masking scheme that is compatible with prefix-truncated computation. 4 Random Prefix Cutting for Efficient Long-Context Training Motivation. Long-CoT training is expensive not only because of backpropagation, but also due to the forward pass through long contexts. In causal Transformers, forward compute and activation memory scale with sequence length (and attention scales quadratically with length under dense attention). To reduce both forward and backward costs, we propose a structured sampler that preserves the prefix dependency structure. Random Prefix Cutting (RPC) Instead of selecting tokens independently, RPC samples a contiguous prefix of each response. For trajectory i of length TiT_i, we sample a cutoff Li∈1,…,TiL_i∈\1,…,T_i\ from a distribution qi(ℓ)=Pr⁡(Li=ℓ)q_i( )= (L_i= ), and define the mask mi,t=t≤Li.m_i,t\;=\;I\t≤ L_i\. (7) The induced inclusion probability (survival function) is pi,t=Pr⁡(Li≥t)=∑ℓ=tTiqi(ℓ),t=1,…,Ti,p_i,t\;=\; (L_i≥ t)\;=\; _ =t^T_iq_i( ), t=1,…,T_i, (8) with pi,1=1p_i,1=1, pi,t+1≤pi,tp_i,t+1≤ p_i,t, and pi,Ti=qi(Ti)>0p_i,T_i=q_i(T_i)>0. HT-corrected unbiased objective. RPC uses the same HT form as Eq. equation 6, but with the structured mask in Eq. equation 7: μ^irpc(θ)=1Ti∑t=1Timi,tpi,tLi,tGRPO(θ)=1Ti∑t=1Li1pi,tLi,tGRPO(θ). μ_i rpc(θ)\;=\; 1T_i _t=1^T_i m_i,tp_i,t\,L^GRPO_i,t(θ)\;=\; 1T_i _t=1^L_i 1p_i,t\,L^GRPO_i,t(θ). (9) By Proposition 1 and [mi,t]=pi,tE[m_i,t]=p_i,t, μ^irpc(θ) μ_i rpc(θ) is an unbiased estimator of the full-sequence mean loss, and yields an unbiased estimator of the full-token GRPO gradient. Why RPC is fundamentally different from deterministic truncation. A deterministic prefix truncation that always keeps the first Ki=⌊βTi⌋K_i= β T_i tokens corresponds to pi,t=1p_i,t=1 for t≤Kit≤ K_i and pi,t=0p_i,t=0 for t>Kit>K_i, which violates the HT requirement pi,t>0p_i,t>0 and induces systematic bias by permanently discarding late-token contributions. In reasoning tasks, late tokens often contain verification steps or final-answer formation; ignoring them can distort the learned policy and degrade convergence. RPC avoids this failure mode by ensuring every position has nonzero inclusion probability while still enabling prefix-truncated computation. Compute savings: prefix structure enables true forward truncation. Unlike URS, RPC can reduce forward compute because the model only needs to process the prefix of length LiL_i. Under dense attention, the per-sequence attention cost reduces from O(Ti2)O(T_i^2) to O(Li2)O(L_i^2), and activation memory scales with LiL_i rather than TiT_i. Thus RPC provides savings in both the forward and backward passes, while remaining statistically unbiased via HT correction. Mask dependence and variance. RPC masks are positively correlated: if mi,t=1m_i,t=1 then all earlier positions are also included. For s≤ts≤ t, Cov(mi,s,mi,t) (m_i,s,m_i,t) =Pr⁡(Li≥t)−Pr⁡(Li≥s)Pr⁡(Li≥t) = (L_i≥ t)- (L_i≥ s) (L_i≥ t) =pi,t(1−pi,s)≥ 0. =p_i,t\,(1-p_i,s)\;≥ 0. This dependence can increase variance relative to idealized independent sampling, but it is precisely what makes RPC compatible with prefix-truncated computation. Practically, stability is controlled by designing qi(ℓ)q_i( ) so that pi,tp_i,t does not decay too sharply, avoiding extreme importance weights 1/pi,t1/p_i,t near the end of the sequence. This is also to ensure that the model will get a fair amount of exposure to later parts of the trajectory. Uniform cutoff as a simple default. A robust choice is the uniform cutoff Li∼Uniform(1,…,Ti)L_i (\1,…,T_i\), i.e., qi(ℓ)=1/Tiq_i( )=1/T_i. Then pi,t=Ti−t+1Ti,[Li]=∑t=1Tipi,t=Ti+12≈Ti2.p_i,t= T_i-t+1T_i, [L_i]= _t=1^T_ip_i,t= T_i+12≈ T_i2. On average, RPC processes about half the tokens per trajectory, leading to substantial reductions in forward/backward cost while preserving unbiased access to gradients from all positions through randomization. The variance analysis of Uniform cutoff is further discussed in Appendix B. Minimum-cutoff RPC In some implementations we enforce a minimum retained prefix length C, to avoid extremely short prefixes and to stabilize the HT weights near the beginning of training. Concretely, we sample Li∼Uniform(C,…,Ti)L_i (\C,…,T_i\) (or more generally any distribution supported on [C,Ti][C,T_i]), and keep mi,t=t≤Lim_i,t=I\t≤ L_i\. Under the uniform choice, the survival probabilities become pi,t=Pr⁡(Li≥t)=1,t≤C,Ti−t+1Ti−C+1,t>C,p_i,t= (L_i≥ t)= cases1,&t≤ C,\\[2.0pt] T_i-t+1T_i-C+1,&t>C, cases which guarantees bounded importance weights 1/pi,t≤Ti−C+1Ti−t+11/p_i,t≤ T_i-C+1T_i-t+1 for all included positions and strictly prevents pathological “one-token” prefixes. The expected retained length is [Li]=C+Ti2=Ti2+C2,E[L_i]= C+T_i2\;=\; T_i2+ C2, so compared to the unconstrained uniform cutoff (which has [Li]≈Ti/2E[L_i]≈ T_i/2), enforcing a minimum cutoff increases the expected compute by an additive C/2C/2 tokens per sequence while often improving stability. Table 1: Comparison of token-efficient methods. RPC is the only method that provides computational savings in both passes while remaining statistically unbiased. 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 Summary. Token masking with HT reweighting yields unbiased gradient estimates while trading compute for variance. URS provides primarily backward savings but can inflate gradient norms as p decreases. RPC preserves prefixes, enabling true forward truncation and thus savings in both passes, while remaining unbiased and avoiding the systematic bias of deterministic truncation. 5 Experiments and Results 5.1 Experiment Setup In this section, we adopt Verl Sheng et al. (2025) as our training framework. The training dataset is DAPO-Math-17K from DAPO paper Yu et al. (2025). The models we consider are Qwen2.5-Math-7B model Yang et al. (2024) and Qwen3-8B model Yang et al. (2025). The RL algorithm we adopt is GRPO Shao et al. (2024) without KL Divergence as suggested by DAPO Yu et al. (2025). The corresponding training recipes are provided in Appendix C. All methods are trained with the same number of optimizer updates and identical optimization hyperparameters. Algorithms The algorithms we used in this experiment are: • GRPO: vanilla GRPO algorithm Shao et al. (2024) using full response tokens. • URS: GRPO with random token masking using uniform random sampling with p=0.5p=0.5. • Det. Trunc.: GRPO with deterministic prefix truncation of 50% of trajectory tokens. • RPC: GRPO with uniform random prefix cutting with a minimum of 100 tokens. Each algorithm is repeated 5 times to capture the variance of validation metrics. Given the limited GPU resources, we randomly sample 80% of DAPO-Math-17K dataset once and reuse it for all runs to speed up training. The training epoch is selected when the reward curve of GRPO is plateauing. In our case, we choose 10 epoch for Qwen2.5-Math-7B and 2 epoch for Qwen3-8B. All hyperparameters are shared across algorithms to ensure a fair comparison. Evaluation We evaluate our algorithms on standard mathematical reasoning benchmarks: AIME24, AIME25, MATH500 Hendrycks et al. (2021). For each question, we generate 16 independent responses under a decoding temperature T = 1.0, and report the average accuracy and the average number of tokens per response. 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). 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 Qwen3-8B Model GRPO 0.768± 0.003 0.911± 0.005 0.257± 0.017 0.475± 0.024 0.202± 0.011 0.386± 0.040 URS 0.768± 0.006 0.911± 0.009 0.245± 0.009 0.474± 0.056 0.207± 0.010 0.384± 0.044 Det. Trunc. 0.633± 0.004 0.902± 0.004 0.190± 0.013 0.488± 0.033 0.159± 0.014 0.374± 0.029 RPC 0.764± 0.010 0.910± 0.009 0.236± 0.014 0.442± 0.040 0.201± 0.018 0.370± 0.030 (a) Qwen2.5-Math-7B (b) Qwen3-8B Figure 2: Entropy curves with 95% confidence interval across 5 runs for GRPO (vanilla GRPO), URS (GRPO with random sampling p=0.5p=0.5), Det. Trunc. (GRPO with deterministic prefix truncation of 50% of trajectory tokens) and RPC (GRPO with uniform random prefix cutting). (a) Qwen2.5-Math-7B (b) Qwen3-8B Figure 3: Percentage of selected tokens with 95% confidence interval across 5 runs for RPC (GRPO with uniform random prefix cutting). 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. 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 GRPO 47.718± 0.215 311.502± 2.555 628.027± 7.008 URS 47.802± 0.144 312.439± 1.525 631.863± 2.692 Det. Trunc. 33.786± 0.009 183.114± 0.304 364.594± 0.934 RPC 39.234± 0.852 220.322± 2.120 400.547± 4.201 Token-Efficient RL on Accuracy. Table 5.1 shows that Det. Truncation is directionally worse than full-token GRPO on most metrics, often by a large margin in point estimates. Although several gaps are not statistically significant under the 95% CI overlap heuristic, this is primarily because Det. Truncation exhibits much higher run-to-run variance for Qwen2.5-Math-7B Model. For Qwen3-8B model, all Acc@16 metrics are significantly lower than GRPO while Pass@16 metrics are not. This is because the base model is strong and reinforcement learning algorithm does not lift pass@k at large k as discussed in (Yue et al., 2025). This behavior is consistent with the expected bias of deterministic truncation: always removing the same suffix can suppress late-step learning signals and destabilize optimization. By contrast, the unbiased sampling-based methods, URS and RPC, remain consistently on par with GRPO across MATH, AIME24, and AIME25 for both models, with overlapping 95% CIs throughout. Token-Efficient RL on Entropy. From entropy Figure 2, Det. Truncation exhibits consistently higher entropy, likely due to biased gradients and reduced effective training signal. For URS and RPC, entropy curves converge to same values as GRPO. Token-Efficient RL on GPU memory. Although absolute memory usage can vary with hardware and training configuration (e.g., batch size and maximum response length), all methods are evaluated under the same setup (Appendix C), enabling a fair comparison. To quantify memory efficiency during RL training, we use VERL’s runtime metric allocated_memory_gb Sheng et al. (2025), which records the peak GPU memory allocated at each training step. Table 5.1 show a clear pattern. Det. Truncation is the most memory-efficient method because it deterministically keeps only the first 50% of tokens, thereby avoiding activation materialization for the truncated suffix. However, this memory advantage comes with a clear quality cost: as shown in Table 5.1, Det. Truncation consistently underperforms GRPO on accuracy metrics, with high variance in several settings. By contrast, URS (uniform random token sampling) yields little to no peak-memory reduction relative to GRPO. This is expected: URS reduces the number of tokens contributing to the loss/backward signal, but does not shorten the effective forward-prefix computation induced by causal attention, so peak activation memory remains close to full-token training. RPC achieves the best trade-off. By randomly cutting contiguous prefixes and applying Horvitz–Thompson correction, it preserves statistical alignment with full-token optimization while reducing the effective sequence length used in training computation. Empirically, RPC remains on par with GRPO on accuracy across benchmarks, while reducing peak GPU memory by approximately 18–20% in our runs (e.g., 35.81→ 29.24 GB on Qwen2.5-Math-7B and 47.72→ 39.23 GB on Qwen3-8B). Overall, these results indicate that RPC provides the most practical memory–quality trade-off among the compared methods. Token-Efficient RL on Training Time. From Table 5.1, for Qwen2.5-Math-7B, both sampling-based methods are faster than full-token GRPO. Relative to GRPO, URS reduces update time (w/o inference) by about 11.5% and total step time by about 10.9%. RPC delivers larger gains, reducing update time by about 26.8% and total step time by about 20.9%. For Qwen3-8B, URS is effectively on par with GRPO in timing (differences are within roughly 1% and 95% CIs overlap), whereas RPC remains substantially faster: update time decreases by about 29.3% and total step time by about 36.2%. Overall, these results align with NAT’s design objective: token-efficient optimization mainly accelerates the learner path while leaving rollout inference unchanged. As a result, improvements are most pronounced in the no-inference metric, and convert into clear end-to-end gains when learner-side cost constitutes a meaningful portion of each step. Across both model scales, RPC provides the largest and most consistent wall-clock improvements while maintaining competitive accuracy (Table 5.1). Selected-token ratio under RPC Figure 3 shows that RPC retains approximately half of each trajectory on average, with the selected-token ratio stabilizing around ∼0.54−0.56 0.54\!-\!0.56 across training. This is slightly above 50%50\% by design: our implementation uses uniform random prefix cutting with a minimum retained prefix length C=100C=100, so the expected selection ratio is strictly larger than 0.50.5, with the additive term C/(2Ti)C/(2T_i) most visible for shorter responses. Empirically, the observed ∼54%−56% 54\%-56\% range is consistent with this prediction and confirms that RPC operates in the intended ≈50%≈ 50\% token-budget regime while avoiding pathological ultra-short prefixes. 6 Related Work Efficient RL Training. Recent profiling of GRPO training reveals that the rollout generation phase accounts for approximately 70% of total training time (Gao et al., 2025; Zhou et al., 2025). This bottleneck is primarily driven by “long-tail rollouts,” where high variance in sequence lengths leads to significant GPU under-utilization and pipeline bubbles. Systems such as RollPacker (Gao et al., 2025) and SortedRL (Zhang et al., 2025) mitigate this via length-aware scheduling, isolating straggling sequences to maximize throughput. April (Zhou et al., 2025) further reduces latency through over-provisioning rollout requests and terminating stragglers once a target sample quota is met. To accelerate the generation itself, Tlt (Hu et al., 2025) leverages idle GPU cycles to train on-the-fly drafter models for speculative decoding. Xu et al. (2025) proposed the down-sampling rollouts and focusing on subset of trajectories which can maintain performance while significantly reducing gradient computation. Our work is orthogonal to these rollout optimizations; by focusing on token-level efficiency during the forward and backward passes, our method can be seamlessly integrated with existing rollout acceleration frameworks. Partial Tokens Training. Beyond system-level scheduling, recent research has explored selective token-level updates to reduce the computational cost of the forward and backward passes. Wang et al. (2025) demonstrate that RL performance is driven by a “high-entropy minority” of tokens, suggesting that gradients from low-entropy tokens are largely redundant for reasoning tasks. Most relevantly, Lee & Tong (2025) propose token-efficient RL schemes that mask random tokens. Our work extends this line of research by introducing a framework using reweighting to offer a unified theory analysis. Moreover, random masking and entropy-based masking cannot reduce forward pass cost since activation memory and forward computation are still done using full tokens. Our approach ensures that we match the full-sequence statistics while strictly capping the maximum sequence length processed by the policy network, leading to predictable and substantial memory savings. 7 Limitation and Future Work Beyond the empirical gains, there are two conceptual reasons why RPC is effective. First, RPC can be interpreted as a structured regularizer analogous in spirit to dropout: instead of randomly dropping individual hidden units, it randomly drops suffix segments while preserving causal prefix structure. This introduces stochasticity in the optimization path and reduces over-reliance on any single long trajectory realization, while remaining compatible with autoregressive decoding constraints. In that sense, RPC acts like sequence-level, causality-preserving dropout (Srivastava et al., 2014) rather than token-independent masking. Second, RPC is aligned with long-horizon credit-assignment dynamics. In many reasoning trajectories, early decisions constrain downstream token distributions; thus, late-token gradients can be high-cost but lower-marginal-value or noisier for policy improvement. Randomized prefix training increases exposure to early causal decisions across updates, while avoiding the systematic bias of deterministic fixed-length truncation. This offers a practical bias–efficiency balance for long-CoT RL, where temporal credit assignment is known to be a central challenge Pignatelli et al. (2023). That said, our current study only instantiates two simple samplers (URS and RPC). Both are information-agnostic: they do not use token-level uncertainty, gradient magnitude, or verifier-derived saliency. A key future direction is to learn or adapt inclusion probabilities within the same Horvitz–Thompson framework so that compute is preferentially allocated to high-information tokens, potentially reducing estimator variance at fixed token budget. Another direction is to co-design sampling with systems kernels (e.g., blockwise/prefix-aware attention and memory scheduling) to further improve end-to-end efficiency. Finally, while our results are strong on math RLVR, broader validation is still needed on domains with different reward sparsity, trajectory lengths, and verifier reliability. These observations motivate a broader research agenda on information-aware token selection and systems co-design. We now summarize the main takeaways of this work. 8 Conclusion In this work, we challenged the conventional assumption that full-sequence backpropagation is necessary for effective reinforcement learning in long-form reasoning tasks. We introduced Nat (Not All Tokens are Needed), a unified framework that utilizes Horvitz-Thompson estimation to decouple sequence-level reward evaluation from token-level policy optimization. By leveraging unbiased gradient estimators, Nat enables selective updates on trajectory subsets—such as high-entropy clusters or randomized prefixes—without inducing the systematic bias inherent in naive deterministic truncation. Our empirical results across challenging mathematical benchmarks demonstrate that Nat achieves performance parity with full-token GRPO while utilizing as few as 50% of the trajectory tokens. In particular, the Random Prefix Cutting (RPC) strategy provides a scalable solution to the memory and compute bottlenecks of long chain-of-thought training by significantly reducing activation memory and increasing training throughput. As reasoning trajectories continue to grow in complexity, Nat offers a theoretically grounded and orthogonal approach to existing rollout optimizations, providing a vital tool for the efficient scaling of frontier AI systems. References Gao et al. (2025) Wei Gao, Yuheng Zhao, Dakai An, Tianyuan Wu, Lunxi Cao, Shaopan Xiong, Ju Huang, Weixun Wang, Siran Yang, Wenbo Su, et al. Rollpacker: Mitigating long-tail rollouts for fast, synchronous rl post-training. arXiv preprint arXiv:2509.21009, 2025. Guo et al. (2025) Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025. Hendrycks et al. (2021) Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874, 2021. Hu et al. (2025) Qinghao Hu, Shang Yang, Junxian Guo, Xiaozhe Yao, Yujun Lin, Yuxian Gu, Han Cai, Chuang Gan, Ana Klimovic, and Song Han. Taming the long-tail: Efficient reasoning rl training with adaptive drafter. arXiv preprint arXiv:2511.16665, 2025. Jaech et al. (2024) Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai o1 system card. arXiv preprint arXiv:2412.16720, 2024. Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles, p. 611–626, 2023. Lee & Tong (2025) Alan Lee and Harry Tong. Token-efficient rl for llm reasoning. arXiv preprint arXiv:2504.20834, 2025. Leviathan et al. (2023) Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, p. 19274–19286. PMLR, 2023. Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022. Pan et al. (2024) Jiayi Pan, Xingyao Wang, Graham Neubig, Navdeep Jaitly, Heng Ji, Alane Suhr, and Yizhe Zhang. Training software engineering agents and verifiers with swe-gym. arXiv preprint arXiv:2412.21139, 2024. Pignatelli et al. (2023) Eduardo Pignatelli, Johan Ferret, Matthieu Geist, Thomas Mesnard, Hado van Hasselt, Olivier Pietquin, and Laura Toni. A survey of temporal credit assignment in deep reinforcement learning. arXiv preprint arXiv:2312.01072, 2023. Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Shao et al. (2024) Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024. Sheng et al. (2025) Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. In Proceedings of the Twentieth European Conference on Computer Systems, p. 1279–1297, 2025. Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929–1958, 2014. Stiennon et al. (2020) Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in neural information processing systems, 33:3008–3021, 2020. Team et al. (2025) Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al. Kimi k1. 5: Scaling reinforcement learning with llms. arXiv preprint arXiv:2501.12599, 2025. Wang et al. (2025) Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xionghui Chen, Jianxin Yang, Zhenru Zhang, et al. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for llm reasoning. arXiv preprint arXiv:2506.01939, 2025. Xi et al. (2026) Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, and Ligeng Zhu. Jet-rl: Enabling on-policy fp8 reinforcement learning with unified training and rollout precision flow. arXiv preprint arXiv:2601.14243, 2026. Xu et al. (2025) Yixuan Even Xu, Yash Savani, Fei Fang, and J Zico Kolter. Not all rollouts are useful: Down-sampling rollouts in llm reinforcement learning. arXiv preprint arXiv:2504.13818, 2025. Yang et al. (2024) An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, et al. Qwen2. 5-math technical report: Toward mathematical expert model via self-improvement. arXiv preprint arXiv:2409.12122, 2024. Yang et al. (2025) An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Yeo et al. (2025) Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long chain-of-thought reasoning in llms. arXiv preprint arXiv:2502.03373, 2025. Yu et al. (2025) Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476, 2025. Yue et al. (2025) Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model? arXiv preprint arXiv:2504.13837, 2025. Zhang et al. (2025) Yiqi Zhang, Huiqiang Jiang, Xufang Luo, Zhihe Yang, Chengruidong Zhang, Yifei Shen, Dongsheng Li, Yuqing Yang, Lili Qiu, and Yang You. Sortedrl: Accelerating rl training for llms through online length-aware scheduling. In ES-FoMo I: 3rd Workshop on Efficient Systems for Foundation Models, 2025. Zheng et al. (2024) Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Livia Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Sglang: Efficient execution of structured language model programs. Advances in neural information processing systems, 37:62557–62583, 2024. Zhou et al. (2025) Yuzhen Zhou, Jiajun Li, Yusheng Su, Gowtham Ramesh, Zilin Zhu, Xiang Long, Chenyang Zhao, Jin Pan, Xiaodong Yu, Ze Wang, et al. April: Active partial rollouts in reinforcement learning to tame long-tail generation. arXiv preprint arXiv:2509.18521, 2025. Ziegler et al. (2019) Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593, 2019. Appendix A Gradient Unbiasedness Proof of Gradient Unbiasedness. Given the Horvitz–Thompson estimator μ^iHT(θ) μ_i^HT(θ), its gradient is: ∇θμ^iHT(θ)=1Ti∑t=1Timi,tpi,t∇θLi,tGRPO(θ). _θ μ_i^HT(θ)= 1T_i _t=1^T_i m_i,tp_i,t _θL_i,t^GRPO(θ). Taking expectation over the mask distribution: m[∇θμ^iHT(θ)] _m [ _θ μ_i^HT(θ) ] =1Ti∑t=1Ti[mi,t]pi,t∇θLi,tGRPO(θ) = 1T_i _t=1^T_i E[m_i,t]p_i,t _θL_i,t^GRPO(θ) =1Ti∑t=1Ti∇θLi,tGRPO(θ) = 1T_i _t=1^T_i _θL_i,t^GRPO(θ) =∇θμi(θ). = _θ _i(θ). Thus, the gradient estimator is unbiased for the full-token gradient. ∎ Appendix B Appendix: Detailed Derivations and Extensions for Random Prefix Cutting B.1 Variance Derivation We analyze the variance of the RPC Horvitz–Thompson estimator by explicitly accounting for the stopping-time structure induced by prefix truncation. Let μ^iRPC-HT=1Ti∑t=1Timi,tpi,tLi,tGRPO(θ), μ_i^RPC-HT= 1T_i _t=1^T_i m_i,tp_i,tL_i,t^GRPO(θ), where the binary mask mi,tm_i,t indicates whether the rollout survives up to token t. Stopping-time formulation. Define the random stopping time τi≜max⁡t:mi,t=1, _i\; \; \t:m_i,t=1\, with the convention that τi=0 _i=0 if the rollout terminates immediately. By construction, mi,t=t≤τi,ℙ(τi≥t)=pi,t,m_i,t=1\t≤ _i\, ( _i≥ t)=p_i,t, and the estimator can be rewritten as μ^iRPC-HT=1Ti∑t=1τiLi,tGRPO(θ)pi,t. μ_i^RPC-HT= 1T_i _t=1 _i L_i,t^GRPO(θ)p_i,t. (10) Variance expression. Using equation 10, the variance admits the exact decomposition Var[μ^iRPC-HT]=1Ti2([(∑t=1τiLi,tGRPO(θ)pi,t)2]−(∑t=1TiLi,tGRPO(θ))2).Var\! [ μ_i^RPC-HT ]= 1T_i^2 (E [ ( _t=1 _i L_i,t^GRPO(θ)p_i,t )^2 ]- ( _t=1^T_iL_i,t^GRPO(θ) )^2 ). (11) This form is exact and does not rely on token-level independence assumptions. Discussion. Unlike standard Horvitz–Thompson estimators with fixed sampling designs, the random variables mi,tt=1Ti\m_i,t\_t=1^T_i are deterministically coupled through the stopping time τi _i: an early termination event sets all subsequent masks to zero. As a consequence, the variance of the RPC estimator cannot be decomposed into a sum of independent or pairwise token-level contributions. Instead, variance is dominated by early stopping events and scales inversely with the smallest survival probabilities pi,tp_i,t, particularly for late tokens. This stopping-time coupling explains the elevated variance observed when early prefixes are dropped with high probability, and motivates conservative scheduling of survival probabilities in practice. B.2 Optional Mandatory Prefix In practice, to reduce variance and guarantee that initial reasoning steps are always observed, we can enforce that the first C tokens are always included. This is done by restricting Li∈C,C+1,…,TiL_i∈\C,C+1,…,T_i\ and setting pi,t=1p_i,t=1 for t≤Ct≤ C. The survival function for t>Ct>C remains pi,t=Pr⁡(Li≥t)p_i,t= (L_i≥ t). The unbiased estimator equation 6 still applies. For a uniform cutoff over C,…,Ti\C,…,T_i\, we have [Li]=C+Ti2,[Li]Ti=12+C2Ti.E[L_i]= C+T_i2, E[L_i]T_i= 12+ C2T_i. (12) B.3 Optimal Survival Schedule under a Budget Given a compute budget B=∑t=1Tipi,tB= _t=1^T_ip_i,t, we might seek the survival schedule pi,t\p_i,t\ that minimizes the variance of the estimator. Ignoring the off-diagonal correlation terms (or assuming token losses are uncorrelated), the variance is proportional to ∑t=1TiL~i,t2pi,t _t=1^T_i L_i,t^2p_i,t. Using Lagrange multipliers to minimize this sum subject to ∑tpi,t=B _tp_i,t=B and monotonicity (1=pi,1≥⋯≥pi,Ti>01=p_i,1≥…≥ p_i,T_i>0) yields the intuitive solution pi,t∗∝|L~i,t|p_i,t^* | L_i,t|. However, the losses L~i,t L_i,t are unknown before the forward pass. Without prior knowledge, the uniform cutoff (linear pi,tp_i,t) emerges as a robust max-entropy choice, spreading probability mass evenly across feasible prefix lengths and minimizing worst-case variance. B.4 Comparison with Independent Token Masking For an independent masking scheme where each token t is kept independently with probability pi,tp_i,t, the covariance is Cov(mi,s,mi,t)=0Cov(m_i,s,m_i,t)=0 for s≠ts≠ t, and Var(mi,t)=pi,t(1−pi,t)Var(m_i,t)=p_i,t(1-p_i,t). The variance of the Horvitz–Thompson estimator then becomes Varindep=1Ti2∑t=1TiL~i,t21−pi,tpi,t.Var_indep= 1T_i^2 _t=1^T_i L_i,t^2 1-p_i,tp_i,t. (13) Comparing with Eq. equation 13, the RPC variance includes an extra positive term due to the covariances. Nevertheless, as argued in the main text, independent masking often leads to higher effective variance in practice because of unstable importance weights and the necessity to run the full forward pass (no computational saving). B.5 Bias-Variance Trade-off Formalization The mean squared error (MSE) of a gradient estimator μ μ relative to the true full-sequence gradient μ is MSE(μ^)=‖μ^−μ‖2=Var(μ^)+‖[μ^]−μ‖2.MSE( μ)=E\| μ-μ\|^2=Var( μ)+\|E[ μ]-μ\|^2. Deterministic truncation has low variance but high bias squared ‖[μ^det]−μ‖2\|E[ μ_det]-μ\|^2, which does not vanish with more steps. RPC has zero bias ([μ^RPC]=μE[ μ_RPC]=μ) but higher variance. Stochastic optimization algorithms (e.g., SGD) are robust to zero-mean noise (variance) but can diverge or converge to wrong solutions under persistent bias. Hence, the unbiasedness of RPC is critical for correct convergence. Appendix C Training Recipe C.1 Training Setup We employ Group Relative Policy Optimization (GRPO) to fine-tune mathematical reasoning models using reinforcement learning from outcome-based rewards. Our training infrastructure utilizes 16 H100 GPUs with distributed training via Fully Sharded Data Parallel (FSDP) and asynchronous rollout generation powered by SGLang. C.2 Model Configuration • Gradient checkpointing: Enabled to reduce activation memory at the cost of recomputation during backward passes. • Parameter offloading: Model parameters are offloaded to CPU memory when not actively used. • Optimizer offloading: Optimizer states (Adam moments) are offloaded to CPU to further reduce GPU memory pressure. • Fused kernels: Optimized CUDA kernels for improved throughput. C.3 Training Hyperparameters Our training configuration uses the following key hyperparameters for model Qwen 2.5 Math 7B: • Total epochs: 10 • Training batch size: 2048 samples • Mini-batch size: 1024 samples (for advantage computation) • Micro-batch size per GPU: 16 sequences • Learning rate: 1×10−51× 10^-5 with AdamW optimizer • Temperature: 1.0 for both training and validation sampling • Maximum prompt length: 1024 tokens • Maximum response length: 3000 tokens (total allowed context length is 4096) • Rollout samples per prompt: 8 responses generated per training iteration • Validation samples: 16 responses per prompt for evaluation • GPU memory utilization: 0.75 for rollout generation • KL penalty: Disabled (λKL=0 _KL=0) to focus purely on outcome rewards For Qwen3-8B model, we use: • Enable thinking: False • Total epochs: 2 • Training batch size: 1024 samples • Mini-batch size: 1024 samples (for advantage computation) • Micro-batch size per GPU: 16 sequences • Learning rate: 5×10−75× 10^-7 with AdamW optimizer • Temperature: 1.0 for both training and validation sampling • Maximum prompt length: 1024 tokens • Maximum response length: 8192 tokens • Rollout samples per prompt: 8 responses generated per training iteration • Validation samples: 16 responses per prompt for evaluation • GPU memory utilization: 0.75 for rollout generation • KL penalty: Disabled (λKL=0 _KL=0) to focus purely on outcome rewards Appendix D More Experiment Results (a) Qwen2.5-Math-7B (b) Qwen3-8B Figure 4: Gradient norm curves with 95% confidence interval across 5 runs for GRPO (vanilla GRPO), URS (GRPO with random sampling p=0.5p=0.5), Det. Trunc. (GRPO with deterministic prefix truncation of 50% of trajectory tokens) and RPC (GRPO with uniform random prefix cutting). (a) Qwen2.5-Math-7B (b) Qwen3-8B Figure 5: Time per step (excluding inference time) with 95% confidence interval across 5 runs for RPC (GRPO with uniform random prefix cutting). (a) Qwen2.5-Math-7B (b) Qwen3-8B 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.5p=0.5), Det. Trunc. (GRPO with deterministic prefix truncation of 50% of trajectory tokens) and RPC (GRPO with uniform random prefix cutting).

Instant research discovery