Abstract
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.
Abstract (translated)
Transformer模型在多步推理任务中表现出色，但它们通过训练获得这些能力的内在机制，尤其是在理论层面上的理解仍然有限。本研究探讨了Transformer是如何通过链式思维过程学习解决符号化多步推理问题的，特别是关注于树形结构中的路径寻找。我们分析了两个相互交织的任务：一个是反向推理任务，在此任务中模型输出从目标节点到根节点的路径；另一个是更复杂的正向推理任务，即模型首先识别出从目标节点到根节点的路径，然后反转这条路径以生成从根节点到目标节点的路径。基于梯度下降动力学的理论分析表明，经过训练的一层Transformer可以证明解决这两个任务，并且对未见过的树形结构具有泛化保证。 特别地，我们为正向推理开发出一个多阶段训练动态过程，阐明了不同的注意力头如何自主学习专业化并协作以在单一自回归路径中解决两个子任务。这些结果提供了关于经过训练的Transformer如何实现顺序算法程序的具体机制解释。此外，它们还为推理能力的出现提供了见解：当任务结构被设计成需要中间链式思维步骤时，即使是浅层多头Transformer也能有效解决问题，而这些问题在其他情况下可能需要更深层的架构来解决。
URL
https://arxiv.org/abs/2508.08222
