arXiv:2608.28158cs.LGcs.DC2026-08

HARTS让智能体强化学习在任意树状路径中高效复用计算,提速近5倍。

HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees

  • 通过压缩前缀+分块线性注意力,实现共享路径的高效复用
  • 在SWE-bench任务上达4.81–4.87倍正反向传播加速
  • 适合做复杂推理路径的强化学习系统开发者

智能体强化学习常生成具有共享历史的不规则回溯树。现有系统主要针对全注意力模型,缺乏与激活重计算兼容的密集可微分混合注意力执行。本文提出HARTS(面向树结构的混合注意力强化学习)。HARTS在前缀压缩后,联合规划微批次、数据并行副本分配与微批次槽调度,采用非重放紧凑令牌工作模式。对于分块线性注意力,提出线性时间算法,协调块边界状态恢复与重放,在打包执行模型下实现最少序列线性注意力调用次数。HARTS保持轨迹级训练的分块状态划分:不重复投影、MLP/MoE计算或最终输出,仅进行有限状态重放以保证数值对齐。每轮将所有分支合并为一次打包调用,通过可微状态传递传播梯度,支持激活重计算,并恢复每个标记的对数概率。对于确定性无丢弃的top-k MoE路由,语义多重性恢复MoE目标的标记权重与负载统计。现有强化学习目标接口不变。据我们所知,HARTS是首个在真实混合注意力模型上实现任意回溯树前缀共享加速的系统。在由SWE-bench任务生成的智能体强化学习工作负载上,结合激活重计算,HARTS在多种并行配置下实现4.81–4.87倍的前向/反向/梯度加速。其数值差异与基线自重跑波动相当,且在τ^3-Bench训练前120步内奖励趋势与基线相似。

原文摘要 · Abstract (English)

Agentic reinforcement learning (RL) often produces irregular rollout trees with shared histories. Training root-to-leaf trajectories independently recomputes these shared prefixes. Existing systems primarily target full-attention models and lack dense, differentiable hybrid-attention execution compatible with activation recomputation. We present HARTS (Hybrid-Attention RL over Tree Structures). HARTS jointly plans microbatches, data-parallel (DP) replica assignments, and microbatch-slot schedules using non-replay compact-token work after prefix compression. For chunkwise linear attention, a linear-time algorithm coordinates chunk-boundary state recovery and replay and produces the minimum number of sequential linear-attention calls under our packed execution model. HARTS preserves the chunkwise state partitioning of trajectory-wise training: it does not repeat projections, MLP/MoE computation, or final outputs, and performs only bounded state replay for numerical alignment. Per round, HARTS batches all branches into one packed call, propagates gradients through differentiable state handoffs, supports activation recomputation, and restores per-token log-probabilities. For deterministic, no-token-drop top-$k$ MoE routing, semantic multiplicities restore MoE-objective token weights and load statistics. Existing RL objectives retain their interface. To our knowledge, HARTS is the first system to demonstrate arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model. On an Agentic RL workload generated from SWE-bench tasks, HARTS achieves $4.81$--$4.87\times$ forward/backward/gradient speedup with activation recomputation across multiple parallel configurations. Its numerical differences are comparable to baseline self-rerun variation, and its reward trend is similar to the baseline over the first 120 steps of $τ^3$-Bench training.

强化学习混合注意力高效推理智能体系统

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。