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

Each language version is independently generated for its own context, not a direct translation.

这篇论文提出了一种让大型人工智能（AI）模型“学得更聪明、更省钱”的新方法。我们可以把它想象成是在教一个学生做极其复杂的数学题。

1. 背景：为什么现在的 AI 训练太“烧钱”？

想象一下，你正在教一个超级聪明的学生（AI 模型）做奥数题。

传统的做法（全量训练）： 学生每写一步解题过程，老师都要停下来，从头到尾把整个解题过程（包括那些最基础的“因为 A 所以 B"、“把 X 移到等号右边”这种机械步骤）全部检查一遍，然后告诉学生哪里做得好，哪里需要改进。
问题所在： 现在的 AI 解题思路（Chain-of-Thought）越来越长，有时候要写几千个字。如果老师对每一个字都进行详细的检查和反馈，不仅老师累得半死（计算资源消耗巨大），而且学生学得很慢。这就好比为了教学生解一道题，老师把学生写的几千字草稿纸上的每一个标点符号都重新读了一遍，效率极低。

2. 核心创新：NAT（“并非所有字都需要”）

这篇论文的作者提出了一个叫 NAT (Not All Tokens are Needed) 的框架。它的核心理念是：并不是解题过程中的每一个字，都需要老师亲自检查一遍才能学会。

作者发现，在长长的解题过程中：

有些字是关键决策点（比如“这里应该用勾股定理”），这些字很重要，必须学。
有些字是机械性填充（比如“移项”、“合并同类项”），这些字只是把前面的逻辑顺下来，重复性很高，不需要每次都重新“背”一遍。

NAT 的做法是：
老师不再检查每一个字，而是随机挑选一部分关键的字进行检查和反馈。但是，为了保证学生学得不走样，老师用了一种特殊的“加权算法”（论文里叫 Horvitz-Thompson 估计），确保虽然只检查了一部分，但学到的知识总量和检查全部字是一样的，只是速度快了一倍，内存占用少了一半。

3. 两种“挑选”策略：谁更聪明？

论文里提出了两种挑选字的方法，我们可以用两个生动的比喻来理解：

方法 A：随机撒网 (URS - 均匀随机采样)

比喻： 老师手里有一把筛子，把学生写的几千字草稿纸倒进去，随机筛掉一半的字，只检查剩下的。
结果： 虽然检查的字少了，但阅读速度没变快。因为学生写的时候，是顺着逻辑写的，后面的字依赖于前面的字。老师为了检查第 1000 个字，还是得先读完前 999 个字（在计算机里叫“前向传播”）。所以，虽然省了点检查的力气，但阅读和记忆的负担（显存）并没有明显减少。

方法 B：随机截断 (RPC - 随机前缀切割) —— 这是本文的明星

比喻： 老师决定，每次只让学生写到一半就停下来。比如，学生本来要写 1000 字，老师随机决定：“今天只看到第 400 字，后面的不用看了，我们直接根据这 400 字来总结教训。”
关键点： 这种方法不仅检查的字少了，而且阅读和记忆的负担也大幅降低了！因为老师根本不需要去读后面那 600 个字。
为什么不会学偏？ 这里有个魔法：老师会告诉学生，“虽然我只看了前 400 字，但我给你的反馈权重会加倍（比如乘以 2.5），这样你学到的教训总量和看完全文是一样的。”
优势： 这种方法既省了时间，又省了内存（显存），而且因为每次截断的位置是随机的，学生不会只学会“只写前一半”，而是能学会处理各种长度的题目。

4. 实验结果：真的有效吗？

作者用最新的 AI 模型（Qwen3-8B）在数学题上做了测试，结果非常惊人：

成绩一样好： 使用“随机截断”（RPC）方法训练的 AI，做数学题的准确率，和那种“字字检查”的传统方法完全一样。
省了一半的力气：
- 显存（内存）： 峰值内存占用减少了约 18%。这意味着你可以用更便宜的显卡，或者在同样的显卡上训练更大的模型。
- 时间： 训练速度提升了约 29%。以前跑完一个训练步骤要 5 分钟，现在只要 3 分半。
对比“硬截断”： 如果老师只是机械地规定“永远只看前 50% 的字”（确定性截断），学生就会学偏，成绩会大幅下降。但 NAT 的“随机截断”因为加入了数学上的修正，完美避开了这个问题。

5. 总结：这对我们意味着什么？

这篇论文就像给 AI 训练界发了一张“节能通行证”。

以前，为了让 AI 变得更聪明（能处理更长的推理），我们需要更贵的显卡、更长的等待时间，因为我们要处理海量的文字数据。
现在，NAT 告诉我们：只要用对方法，我们可以只处理一半的数据，就能达到同样的效果。

对普通人： 未来的 AI 可能会更便宜、反应更快，因为它们训练起来不再那么“烧钱”了。
对开发者： 这是一个“即插即用”的工具，不需要改变 AI 的核心逻辑，就能让训练过程跑得更快、更稳。

简单来说，这就好比以前我们为了学会做一道菜，必须把厨师切菜、炒菜、装盘的每一个动作都看一遍；现在 NAT 告诉我们，只要随机看几个关键动作，再配合一点“脑补”技巧，就能学会同样的手艺，而且省了一半的精力！

Each language version is independently generated for its own context, not a direct translation.

这篇论文提出了一种名为 NAT (Not All Tokens are Needed) 的统一框架，旨在解决大语言模型（LLM）在长思维链（Chain-of-Thought, CoT）强化学习（RL）训练中的效率瓶颈问题。

以下是该论文的详细技术总结：

1. 问题背景 (Problem)

长思维链的扩展瓶颈： 尽管强化学习（特别是基于可验证奖励的 RLVR，如 GRPO）在提升模型推理能力方面效果显著，但随着思维链长度的增加，训练成本急剧上升。
全 Token 更新的代价： 传统的 RL 训练流程（如 GRPO）会对生成的每一个 Token 进行前向传播和反向传播。这导致：
- 显存爆炸： 长序列需要存储大量的激活值（Activations）用于反向传播，容易触发 OOM（显存溢出）。
- 计算冗余： 许多生成的 Token 只是机械性的延续或低熵的模板，对策略梯度的贡献较小，但计算成本相同。
- 效率不匹配： 虽然推理生成阶段（Rollout）已通过高性能引擎优化，但随后的学习阶段（前向/反向传播）仍然是内存和计算密集型瓶颈，限制了 RL 的扩展性。
核心问题： 是否真的需要所有生成的 Token 来训练一个强大的 RL 推理模型？

2. 方法论 (Methodology)

NAT 框架的核心思想是：在保持奖励信号完整性的前提下，仅使用生成 Token 的子集进行策略更新。

2.1 核心机制：Horvitz-Thompson (HT) 重加权

为了在仅使用部分 Token 进行反向传播时保持梯度的无偏性，NAT 引入了统计学中的 Horvitz-Thompson 估计量：

随机掩码 (Random Masking)： 对于每个生成的 Token $t$ ，定义一个包含概率 $p_{i,t}$ 。如果 Token 被选中（ $m_{i,t}=1$ ），则参与梯度计算；否则被忽略。
无偏估计： 通过将选中的 Token 的梯度贡献除以包含概率 $p_{i,t}$ $p_{i, t}$ （即 $1/p_{i,t}$ 重加权），可以证明该估计量是原始全序列梯度的无偏估计量。
- 公式： $\hat{\mu}^{HT}_i(\theta) = \frac{1}{T_i} \sum_{t=1}^{T_i} \frac{m_{i,t}}{p_{i,t}} L^{GRPO}_{i,t}(\theta)$
理论保证： 只要 $p_{i,t} > 0$ ，该方法的期望梯度与全 Token 训练完全一致，不会引入系统性偏差。

2.2 两种具体的 Token 选择方案

NAT 框架支持多种选择策略，论文重点实现了两种：

均匀随机采样 (URS, Uniform Random Sampling)：
- 以固定概率 $p$ （如 0.5）独立地随机选择 Token。
- 优点： 简单，减少反向传播的计算量。
- 缺点： 由于因果注意力机制（Causal Attention），前向传播仍需处理所有前置 Token，因此无法减少前向计算和激活显存。且当 $p$ 较小时，梯度方差会增大（$1/p$ 倍）。
随机前缀截断 (RPC, Random Prefix Cutting)：
- 机制： 不是独立采样，而是为每个轨迹随机选择一个截断长度 $L_i$ ，仅保留前 $L_i$ 个 Token 作为前缀。
- 优势：
  - 前向/反向双重节省： 由于只处理前缀，模型在前向传播时只需计算到 $L_i$ ，显著降低了激活显存和计算量（从 $O(T^2)$ 降至 $O(L^2)$ ）。
  - 无偏性： 通过 HT 重加权，即使只训练前缀，也能在统计上等价于训练全序列。
  - 避免偏差： 不同于确定性截断（Deterministic Truncation，即总是切掉后 50%），RPC 是随机截断，确保每个位置的 Token 都有非零概率被包含，避免了系统性忽略尾部关键信息（如验证步骤）的问题。

3. 主要贡献 (Key Contributions)

统一框架 (NAT)： 提出了首个将 Token 预算作为一级优化原语的 RLVR 框架，允许在保持全序列奖励评估的同时，仅使用 Token 子集进行策略优化。
理论无偏性证明： 证明了基于 HT 校正的 Token 掩码策略可以产生原始全序列 GRPO 梯度的无偏估计，为部分 Token 更新提供了严格的数学基础。
高效的 RPC 策略： 设计了随机前缀截断（RPC），这是唯一一种能同时减少前向和反向计算成本，且保持统计无偏性的方法。
实证性能： 在数学推理基准测试中，NAT（特别是 RPC）在仅使用 50% Token 的情况下，性能与全 Token GRPO 持平，同时显著降低了资源消耗。

4. 实验结果 (Results)

实验基于 Qwen2.5-Math-7B 和 Qwen3-8B 模型，在 MATH、AIME24、AIME25 等数学基准上进行测试。

推理性能 (Accuracy)：
- RPC 和 URS 的表现与全 Token GRPO 基本持平（Acc@16 和 Pass@16 在 95% 置信区间内重叠）。
- 确定性截断 (Det. Trunc.) 表现显著较差，因为系统性丢弃尾部 Token 破坏了学习信号。
显存效率 (GPU Memory)：
- RPC 实现了最佳平衡。在 Qwen3-8B 上，峰值显存从 47.72 GB 降至 39.23 GB（节省约 18%）。
- URS 由于无法减少前向计算，显存节省微乎其微。
训练时间 (Training Time)：
- RPC 显著加速了训练。在 Qwen3-8B 上，单步训练时间（不含推理）减少了 29%，总步长时间减少了 36%。
- URS 在 Qwen3-8B 上几乎没有加速效果（因为前向计算未减少）。
熵值分析： RPC 和 URS 的熵值曲线收敛至与 GRPO 相同的水平，而确定性截断导致熵值异常升高，表明优化不稳定。

5. 意义与影响 (Significance)

突破扩展瓶颈： NAT 提供了一种与现有推理引擎优化（如 vLLM、Speculative Decoding）正交的优化路径。它不改变生成过程，而是优化“如何消费”生成的轨迹，直接解决了长 CoT 训练中的显存和计算瓶颈。
理论指导实践： 证明了在 RL 中，并非所有 Token 都同等重要，通过统计校正的随机采样可以高效地提取学习信号。
未来方向： 为未来的信息感知 Token 选择（如基于梯度幅度或熵的动态选择）以及系统级内核协同设计（如块级注意力优化）奠定了基础。
实际应用价值： 对于需要长思维链推理的前沿 AI 系统（如数学解题、代码生成），NAT/RPC 提供了一种低成本、高效率的扩展方案，使得在有限硬件资源下训练更复杂的模型成为可能。

总结： 论文通过引入 Horvitz-Thompson 估计量，成功地将“全序列奖励”与“部分 Token 更新”解耦，提出的 RPC (随机前缀截断) 策略在保持模型推理能力不变的前提下，显著降低了长思维链强化学习的显存占用和训练时间，是解决 LLM 长序列训练效率问题的关键突破。