跳到主要内容
返回时间线
arXiv来源发表:

长上下文RL训练不再等自回归:稀疏自草稿加全注意力验证,吞吐提升1.39倍

核心概要

长上下文RLVR的rollout生成改用稀疏注意力起草、全注意力验证,在Qwen3-1.7B与Qwen3-8B上把端到端训练吞吐提高到自回归的1.39倍和1.35倍,同时保持任务表现相当。

Source-provided article image: S^3R: Fast and Faithful RL Rollouts via Self-Sparse Speculative Decoding
Figure 2 · arXiv

深度剖析

长上下文可验证奖励强化学习的rollout生成不再必须逐token跑完整注意力。S3R把稀疏注意力限制在候选token的起草阶段,最终响应仍由全注意力验证和精确投机采样产生,因此用于策略优化的响应组分布与稠密自回归一致。 此前降低rollout成本的做法是直接对稀疏策略采样,再用过滤或重要性重加权在训练侧补偿;在GRPO这类组相对优化里,每条响应的优势取决于同组其他响应的奖励,逐响应重加权无法一般地恢复稠密策略的期望更新。 在匹配的消融中,同样激进的稀疏起草预算下,S3R训练保持稳定,而带训练侧修正的直接稀疏采样出现梯度范数尖峰并在约第70步奖励恶化。

起草与验证共用同一份策略参数,不需要单独训练或更新一个草稿模型。起草阶段用注意力分数从历史KV中挑选少量条目,验证阶段用完整历史KV重算候选块,再按精确投机采样接受或从残差分布补采。 已有的投机解码方案依赖检索历史响应、外部草稿网络或层跳过等来源;外部草稿在策略持续更新后会变旧,接受率下降,而检索方法依赖历史轨迹覆盖当前策略的续写。 在Qwen3-1.7B和Qwen3-8B上、DAPO-MATH-17k数据、单节点八张H100的veRL训练中,S3R的块效率为5.59和5.47,高于N-gram Lookup的1.31和1.29、History Retrieval的2.03和2.04以及冻结EAGLE3辅助提议器的1.89和1.86。

稠密验证需要保留完整历史KV缓存,长响应下这部分状态会超出GPU容量,因此运行时用KV感知的AIMD调度器调节请求准入,在内存压力下限制KV换入换出的反复颠簸,同时维持rollout并行度。 固定较低的并发上限可以减少颠簸,但在显存充足时牺牲了批处理;直接重算被换出的KV状态则增加生成开销。 在Qwen3-1.7B的十步训练窗口里,AIMD调度器把生成吞吐从自回归的3989.0 tokens/s/GPU提高到6362.1,比固定并发256的卸载方案再高5.8%,每千生成token的传输页数从18.0降到1.3。

启示与展望

这套方法面向使用组相对优化的长上下文可验证奖励强化学习,训练在Qwen3-1.7B与Qwen3-8B、DAPO-MATH-17k数据、单节点八张H100的veRL环境中完成,评估用AIME26与ArXiv-Math的pass@1。它需要保留完整历史KV用于稠密验证,因此适合显存压力可被KV感知调度管理的场景。后续方向包括自适应稀疏预算与上下文选择,以及向MoE和线性注意力模型扩展。

稀疏预算与接受率之间存在权衡:预算从0.05升到0.50时接受率从80.4%升到96.7%,但加速比从1.370降到1.058;降到0.03则接受率与加速比同时下降,因此默认0.05是这一负载下的经验平衡,而非唯一最优点。接受率随起草位置下降,Qwen3-8B上第六位接受率的降幅大于第一位。周期性在线草稿适配在报告的配置中未带来接受率的单调提升。