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

2026-01-31

AI Summary

This work addresses the computational and memory inefficiencies in large language model reinforcement learning (RL) post-training, where redundant recomputation of shared prefixes across trajectories creates significant bottlenecks. To overcome this, the authors propose a dynamic prefix tree structure integrated with a depth-first search traversal strategy and a distributed load-balanced batching mechanism. This approach instantiates only a single root-to-leaf path during both forward and backward passes, drastically reducing redundant computations. Furthermore, they introduce a novel tree-based attention mechanism that eliminates the need for full attention masks, thereby enhancing both memory and computational efficiency. Empirical evaluations on standard RL post-training benchmarks demonstrate up to an 8.31× improvement in training throughput compared to baseline methods.

Technology Category

• Machine Learning: Reinforcement Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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 $\tau^2$-bench higher training throughput.

Problem

• reinforcement learning
• large language models
• prefix sharing
• computational inefficiency
• rollout sequences

Innovation

• Dynamic Tree Attention
• Prefix Sharing
• Depth-First Search (DFS)
• Distributed Batching
• Reinforcement Learning for LLMs