论文
HARTS:支持任意采样树的混合注意力智能体强化学习
HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees
摘要
智能体强化学习(RL)通常会产生具有共享历史的不规则的展开树。训练根到叶轨迹会独立地重新计算这些共享前缀。现有系统主要针对全注意力模型,缺乏与激活重计算兼容的密集、可微的混合注意力执行。我们提出了 HARTS(基于树结构的混合注意力强化学习)。 HARTS 在前缀压缩后使用非重播紧凑词元工作来联合规划微批次、数据并行 (DP) 副本分配和微批次时隙调度。对于分块线性注意力,线性时间算法协调块边界状态恢复和重放,并在我们的打包执行模型下产生最小数量的顺序线性注意力调用。 HARTS 保留了轨迹训练的块状状态划分:它不重复投影、MLP/MoE 计算或最终输出,并且仅执行数值对齐的有界状态重播。每一轮,HARTS 将所有分支批处理为一个打包调用,通过可微状态切换传播梯度,支持激活重计算,并恢复每个词元的对数概率。对于确定性、无词元丢弃的 top-$k$ MoE 路由,语义多重性可恢复 MoE 目标词元权重和负载统计数据。现有的 RL 目标保留其接口。据我们所知,HARTS 是第一个在真正的混合注意力模型上展示任意采样树前缀共享加速的系统。在 SWE 基准任务生成的 Agentic RL 工作负载上,HARTS 通过跨多个并行配置的激活重计算实现了 4.81$--$4.87\times$ 前向/后向/梯度加速。其数值差异与基线自我重新运行变化相当,其奖励趋势与 $τ^3$-Bench 训练的前 120 个步骤的基线相似。