Transformer通过强化学习自动学会深度优先搜索,无需人工示范。
Agentic Transformers Provably Learn to Search via Reinforcement Learning

- 用双头Transformer实现随机深度优先搜索,头间分工协作。
- 仅训练深度1和2的树,模型就能在更深的树上成功找到目标。
- 在目标分布不均时,能自动优先探索高概率分支,适合智能体决策研究。
树搜索是语言智能体进行推理与决策的核心机制:智能体需探索动作、记忆失败并回溯至更有希望的路径。然而,基于Transformer的策略如何从强化学习(RL)的训练动态中获得这种搜索能力,仍缺乏理论理解。本文在随机k叉树环境中研究此问题,智能体仅通过交互历史观察环境,并在到达隐藏叶节点目标时获得最终奖励。首先构建一个双头Transformer,实现随机深度优先搜索(DFS):一个头追踪历史动作,另一个头检测失败并触发回溯。随后分析分层课程下的策略梯度训练动态,发现该DFS机制可从稀疏强化反馈中逐步涌现,无需专家示范。训练后的策略展现出深度泛化能力:仅在深度1和2的树上训练后,即可在更深层的完整树中成功寻得目标。此外,在目标分布不均时,采用回报折扣可生成按概率排序的DFS策略,优先探索高概率分支。总体结果揭示了基于Transformer的搜索机制的一种基本范式:注意力头分工合作,从上下文中提取决策相关轨迹,并通过强化学习转化为智能体动作选择。
原文摘要 · Abstract (English)
Tree search is a central abstraction behind many language-agent reasoning and decision-making tasks: agents must explore actions, remember failures, and backtrack toward promising alternatives. Yet, we lack a theoretical understanding of how transformer-based policies acquire such search capabilities from the training dynamics of reinforcement learning (RL). We study this question in a stochastic $k$-ary tree environment, where an agentic transformer observes only its trajectory history through interaction and receives a terminal reward for reaching a hidden leaf goal node. We first construct a two-head transformer that implements randomized depth-first search (DFS): one head tracks previous actions, while the other detects failure outcomes and triggers backtracking. We then analyze the training dynamics of policy gradient under a depth-wise curriculum, showing that this same DFS mechanism emerges in stages from sparse reinforcement feedback without expert demonstrations. The resulting policy exhibits depth generalization: after training only on depth-$1$ and depth-$2$ trees, it succeeds on deeper full trees. We further show that, under imbalanced goal distributions, discounting the return leads to a ranked DFS policy that prioritizes higher-probability branches. Overall, our results identify a mechanistic normal form for transformer-based search, in which attention heads specialize and cooperate to extract decision-relevant traces from context and convert them into agentic action selection via RL training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。