多头Transformer通过梯度下降自动学会分步推理,无需深层结构
Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent
- 用梯度下降训练单层多头注意力,实现路径搜索的分步推理
- 在未见过的树结构上仍能正确找到从根到目标的路径
- 不同注意力头自主分工协作,揭示了浅层模型的推理机制
Transformer在多步推理任务中表现出色,但其训练过程中如何获得此类能力的理论理解仍有限。本文研究了Transformer如何通过链式思维解决符号化多步推理问题,聚焦于树结构中的路径查找。分析了两个相互关联的任务:反向推理(从目标节点返回根节点)和更复杂的正向推理(先识别目标到根的路径,再反转生成根到目标路径)。基于梯度下降的动力学理论分析表明,经过训练的一层Transformer可证明性地解决这两类任务,并对未见过的树具备泛化能力。特别是正向推理的多阶段训练动态揭示了不同注意力头如何自主分工与协作,在单个自回归序列中完成两个子任务。结果为训练后的Transformer实现序列算法过程提供了机制解释,也揭示了推理能力的涌现——当任务设计包含中间链式思维步骤时,即使浅层多头Transformer也能有效解决原本需要深层架构的问题。
原文摘要 · Abstract (English)
Transformers have demonstrated remarkable capabilities in multi-step reasoning tasks. However, understandings of the underlying mechanisms by which they acquire these abilities through training remain limited, particularly from a theoretical standpoint. This work investigates how transformers learn to solve symbolic multi-step reasoning problems through chain-of-thought processes, focusing on path-finding in trees. We analyze two intertwined tasks: a backward reasoning task, where the model outputs a path from a goal node to the root, and a more complex forward reasoning task, where the model implements two-stage reasoning by first identifying the goal-to-root path and then reversing it to produce the root-to-goal path. Our theoretical analysis, grounded in the dynamics of gradient descent, shows that trained one-layer transformers can provably solve both tasks with generalization guarantees to unseen trees. In particular, our multi-phase training dynamics for forward reasoning elucidate how different attention heads learn to specialize and coordinate autonomously to solve the two subtasks in a single autoregressive path. These results provide a mechanistic explanation of how trained transformers can implement sequential algorithmic procedures. Moreover, they offer insights into the emergence of reasoning abilities, suggesting that when tasks are structured to take intermediate chain-of-thought steps, even shallow multi-head transformers can effectively solve problems that would otherwise require deeper architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。