AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866910007310680064 |
|---|---|
| author | Zhang, Jiarui Yang, Yuchen Yan, Ran Mei, Zhiyu Zhang, Liyuan Li, Daifeng Fu, Wei Gao, Jiaxuan Xu, Shusheng Wu, Yi Yuan, Binhang |
| author_facet | Zhang, Jiarui Yang, Yuchen Yan, Ran Mei, Zhiyu Zhang, Liyuan Li, Daifeng Fu, Wei Gao, Jiaxuan Xu, Shusheng Wu, Yi Yuan, Binhang |
| contents | Reinforcement learning (RL) based post-training for large language models (LLMs) is computationally expensive, as it generates many rollout sequences that could frequently share long token prefixes. Existing RL frameworks usually process these sequences independently, repeatedly recomputing identical prefixes during forward and backward passes during policy model training, leading to substantial inefficiencies in computation and memory usage. Although prefix sharing naturally induces a tree structure over rollouts, prior tree-attention-based solutions rely on fully materialized attention masks and scale poorly in RL settings. In this paper, we introduce AREAL-DTA to efficiently exploit 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. Across the popular RL post-training workload, AREAL-DTA achieves up to $8.31\times$ in $τ^2$-bench higher training throughput. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_00482 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models Zhang, Jiarui Yang, Yuchen Yan, Ran Mei, Zhiyu Zhang, Liyuan Li, Daifeng Fu, Wei Gao, Jiaxuan Xu, Shusheng Wu, Yi Yuan, Binhang Machine Learning Reinforcement learning (RL) based post-training for large language models (LLMs) is computationally expensive, as it generates many rollout sequences that could frequently share long token prefixes. Existing RL frameworks usually process these sequences independently, repeatedly recomputing identical prefixes during forward and backward passes during policy model training, leading to substantial inefficiencies in computation and memory usage. Although prefix sharing naturally induces a tree structure over rollouts, prior tree-attention-based solutions rely on fully materialized attention masks and scale poorly in RL settings. In this paper, we introduce AREAL-DTA to efficiently exploit 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. Across the popular RL post-training workload, AREAL-DTA achieves up to $8.31\times$ in $τ^2$-bench higher training throughput. |
| title | AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2602.00482 |