arXiv:2511.00413cs.LG2025-11被引 4

通过复用共享前缀,让智能体大模型训练快6倍以上

Tree Training: Accelerating Agentic LLMs Training via Shared Prefix Reuse

  • 把多分支对话轨迹转为树结构,用深度优先遍历避免重复计算
  • 在稠密和MoE模型上实现最高6.2倍的端到端训练加速
  • 适合做智能体系统、强化学习或高效微调的研究者

智能体大模型训练常因并发工具调用、思考模式、子代理等设计产生多轮交互的分支路径,导致生成的文本序列呈现树状结构且存在共享前缀。现有训练流程将这些轨迹线性化并独立处理每个分支,造成前向与反向传播中大量冗余计算。我们推导出:对所有分支独立平均损失,等价于按每个令牌所处分支比例加权的逐令牌损失。问题转化为仅需计算前缀树中每个令牌一次,不重复计算共享前缀。为此提出深度优先序列化方法,遍历每个令牌恰好一次,并适配全注意力与状态空间模型层,使最终概率结果与独立分支计算完全一致。由于单个轨迹树可能超出显存容量,我们进一步提出无冗余树划分策略,在内存受限场景下仍保证零冗余计算,峰值显存仅受单条根到叶路径限制。结合上述贡献,形成树训练框架,在密集与MoE模型上,监督微调与强化学习任务中均实现最高6.2倍的端到端训练加速。

原文摘要 · Abstract (English)

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.

大模型训练智能体加速树结构

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