AREAL-DTA：基于动态树注意力的高效强化学习大语言模型
机器学习
2026-02-03 v1
摘要
基于强化学习 (RL) 的大语言模型 (LLM) 后训练计算成本高昂，因为其生成的许多 rollout 序列可能频繁共享长 token 前缀。现有 RL 框架通常独立处理这些序列，在策略模型训练的前向与反向传递期间反复计算相同的前缀，导致计算和内存使用效率极低。尽管前缀共享自然诱导出 rollout 前缀上的树结构，但先前基于树注意力的解决方案依赖完全物化的注意力掩码，在 RL 场景中规模不佳。本文引入 AREAL-DTA 以高效利用前缀共享进行 RL 训练。AREAL-DTA 采用基于深度优先搜索 (DFS) 的执行策略，在前向与反向计算期间动态遍历 rollout 前缀树，仅物化单一根到叶子的路径。为进一步提升可扩展性，AREAL-DTA 集成负载均衡的分布式批处理机制，动态在多个 GPU 上构建和处理前缀树。在流行的 RL 后训练工作负载中，AREAL-DTA 实现了最高可达 的 -bench 训练吞吐量提升。
引用
@article{arxiv.2602.00482,
title = {AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models},
author = {Jiarui Zhang and Yuchen Yang and Ran Yan and Zhiyu Mei and Liyuan Zhang and Daifeng Li and Wei Fu and Jiaxuan Gao and Shusheng Xu and Yi Wu and Binhang Yuan},
journal= {arXiv preprint arXiv:2602.00482},
year = {2026}
}