arXiv:2602.00482cs.LG2026-02中稿 · ICML

通过动态树注意力提升大模型强化学习效率,大幅减少重复计算

AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models

  • 用深度优先遍历树结构动态处理共享前缀,只保留当前路径
  • 在τ²-bench上训练吞吐量提升8.31倍(相比密集训练)
  • 适合大规模语言模型强化学习训练,尤其关注效率优化的研究者

基于强化学习的大型语言模型后训练计算成本高昂,因生成的多个回溯序列常共享长标记前缀。现有RL框架在策略训练中通常独立处理这些序列,即在策略梯度计算的前向和反向传播中反复重新计算相同的前缀,造成严重的计算资源与内存浪费。尽管前缀共享自然形成回溯序列的树结构,但现有的打包树掩码方法在强化学习场景下扩展性差。本文提出AReaL-DTA,高效利用强化学习训练中的前缀共享。AReaL-DTA采用基于深度优先搜索(DFS)的执行策略,在前向和反向计算中动态遍历回溯前缀树,每次仅实例化一条根到叶的路径。为进一步提升可扩展性,AReaL-DTA引入负载均衡的分布式批处理机制,动态在多张GPU间构建并处理前缀树。在τ²-bench上,AReaL-DTA相比密集训练提升最高达8.31倍的训练吞吐量,相比稀疏训练提升最高达1.70倍。代码已开源:https://github.com/areal-project/AReaL/tree/feat/dta。

原文摘要 · Abstract (English)

Reinforcement learning (RL)-based post-training for large language models (LLMs) is computationally expensive, as it generates many rollout sequences that frequently share long token prefixes. Existing RL frameworks usually process these sequences independently during policy training, i.e., repeatedly recomputing identical prefixes in both the forward and backward passes of policy gradient computation, leading to substantial inefficiencies in computation resources and memory usage. Although prefix sharing naturally induces a tree structure over rollouts, packed tree-mask approaches scale poorly in RL settings. In this paper, we introduce AReaL-DTA, which efficiently exploits prefix sharing in RL training. AReaL-DTA employs a depth-first search (DFS)-based execution strategy that dynamically traverses the rollout prefix tree during both forward and backward computation, materializing only a single root-to-leaf path at a time. To further improve scalability, AReaL-DTA incorporates a load-balanced distributed batching mechanism that dynamically constructs and processes prefix trees across multiple GPUs. On $τ^2$-bench, AReaL-DTA improves training throughput by up to $8.31\times$ over dense training and up to $1.70\times$ over sparse training. Our code is available at https://github.com/areal-project/AReaL/tree/feat/dta.

强化学习大模型训练效率优化

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