# Tree Training: Accelerating Agentic LLMs Training via Shared Prefix Reuse

arXiv: https://arxiv.org/abs/2511.00413 · HTML: https://mlsanity.com/paper/2511.00413/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in General LLM agent frameworks and benchmarks, one of 16 directions in LLM agents and tool use.
- Builds on 25 of its 36 references that are in the map.
- First posted on arXiv 1 Nov 2025.

## Placed in

- LLM agents and tool use / [General LLM agent frameworks and benchmarks](https://mlsanity.com/agents/general-llm-agent-frameworks-and-benchmarks/)

## Builds on (most cited first, 25 of 25)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [A-MEM: Agentic Memory for LLM Agents](https://mlsanity.com/paper/2502.12110/) (2025)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search](https://mlsanity.com/paper/2406.03816/) (2024)
- [SWE-smith: Scaling Data for Software Engineering Agents](https://mlsanity.com/paper/2504.21798/) (2025)
- [CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion](https://mlsanity.com/paper/2405.16444/) (2024)
- [CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving](https://mlsanity.com/paper/2310.07240/) (2023)
- [Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving](https://mlsanity.com/paper/2407.00079/) (2024)
- [An LLM Compiler for Parallel Function Calling](https://mlsanity.com/paper/2312.04511/) (2023)
- [LMCache: An Efficient KV Cache Layer for Enterprise-Scale LLM Inference](https://mlsanity.com/paper/2510.09665/) (2025)
- [SWEET-RL: Training Multi-Turn LLM Agents on Collaborative Reasoning Tasks](https://mlsanity.com/paper/2503.15478/) (2025)
- [Agent Lightning: Train ANY AI Agents with Reinforcement Learning](https://mlsanity.com/paper/2508.03680/) (2025)
- [Tree Search for LLM Agent Reinforcement Learning](https://mlsanity.com/paper/2509.21240/) (2025)
- [Granite-Function Calling Model: Introducing Function Calling Abilities via Multi-task Learning of Granular Tasks](https://mlsanity.com/paper/2407.00121/) (2024)
- [FlashMask: Efficient and Rich Mask Extension of FlashAttention](https://mlsanity.com/paper/2410.01359/) (2024)
- [Do Large Language Models Need a Content Delivery Network?](https://mlsanity.com/paper/2409.13761/) (2024)
- [Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward](https://mlsanity.com/paper/2506.05433/) (2025)
- [Accelerating Direct Preference Optimization with Prefix Sharing](https://mlsanity.com/paper/2410.20305/) (2024)
- [One-Pass to Reason: Token Duplication and Block-Sparse Mask for Efficient Fine-Tuning on Multi-Turn Reasoning](https://mlsanity.com/paper/2504.18246/) (2025)

## Built on by (most cited first, 5 of 5)

- [KAT-Coder-V2 Technical Report](https://mlsanity.com/paper/2603.27703/) (2026)
- [AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models](https://mlsanity.com/paper/2602.00482/) (2026)
- [Rollplex: Cross-Phase GPU Spatial Sharing for Vision Language Model Post-Training](https://mlsanity.com/paper/2608.14498/) (2026)
- [QwenGyre: An Elastic Reinforcement Learning Framework for Training xLong-Horizon Agents](https://mlsanity.com/paper/2609.33848/) (2026)
- [ThunderSyncRL: Lossless Acceleration of Agentic Reinforcement Learning](https://mlsanity.com/paper/2610.05935/) (2026)

## Abstract (arXiv)

Agentic large language model (LLM) training often involves multi-turn interaction trajectories that branch into multiple execution paths due to concurrent tool use, think-mode, sub-agent, context management and other runtime designs. As a result, the tokens produced by a single task naturally form a tree-structured token trajectory with shared prefixes, rather than a linear sequence. Existing training pipelines linearize such trajectories and treat each branch independently, leading to substantial redundant computation in both forward and backward passes. We derive that averaging the loss over all branches independently is algebraically identical to a per-token weighted loss, where each token's weight equals the fraction of branches passing through it. The problem therefore reduces to computing the log-probability of every token in the prefix tree exactly once, with no repeated computation across shared prefixes: we propose DFS serialization of the tree, which visits every token exactly once, and adapt full-attention and SSM layers to ensure the resulting log-probabilities match independent per-branch calculation exactly. In practice, a single trajectory tree can be too large to fit in GPU memory; we therefore propose Redundancy-Free Tree Partitioning, which handles memory-constrained settings with zero redundant computation and peak memory bounded by a single root-to-leaf path. Together, these contributions form Tree Training, an efficient framework for training LLMs on tree-structured trajectories, achieving up to 6.2x end-to-end training speedup on dense and MoE models for both supervised fine-tuning and reinforcement learning.