只用部分生成的token就能高效训练大模型,省时省显存。
Not all tokens are needed(NAT): token efficient reinforcement learning
- 用重加权方法仅更新选定的生成token,保持学习信号不变。
- 在数学推理任务上仅用50%的token,性能接近全序列训练。
- 适合需要长推理链的大模型训练,尤其节省显存和时间。
强化学习(RL)已成为大语言模型发展的关键驱动力,但将RL扩展到长链式思维(CoT)轨迹时,受限于对每个生成token进行反向传播。即使使用优化的推演引擎,全量token更新仍占总训练成本的很大比例,使序列长度成为隐性负担。本文提出一种统一框架NAT(Not All Tokens Are Needed),将token预算作为首要优化目标。NAT通过选择性更新生成token,在不损失全序列学习信号的前提下减少计算开销。核心是基于霍维茨-汤普森重加权的无偏部分token策略梯度估计器,确保子采样后梯度统计正确。我们实现了两种即插即用的token选择策略:均匀随机采样(URS)与随机前缀截断(RPC),两者均降低前向与反向计算及内存占用,且无需改动奖励计算或推演流程。在数学推理基准测试中,NAT在仅使用50% token的情况下,性能与完整序列的GRPO相当,提供了一条高效、独立于现有架构的可扩展路径。实验表明,针对Qwen3-8B模型,RPC可节省18%峰值GPU内存和29%前向与反向训练时间。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。