Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent

Transformer 在多步推理任务中展现出了卓越的能力。然而，对于它们通过训练获得这些能力的底层机制，尤其是从理论角度来看，目前的理解仍然有限。本文研究了 Transformer 如何通过思维链过程来学习解决符号多步推理问题，重点关注树结构中的路径寻找任务。我们分析了两个相互关联的任务：一个是反向推理任务，模型从目标节点输出到根节点的路径；另一个更为复杂的正向推理任务，模型首先识别目标节点到根节点的路径，然后将其反转以生成根节点到目标节点的路径，从而实现两阶段推理。 我们的理论分析基于梯度下降的动力学，表明经过训练的单层 Transformer 可以在理论上保证解决这两个任务，并且能够泛化到未见过的树结构上。特别是我们对正向推理任务的多阶段训练动力学分析，揭示了不同的注意力头如何自主地学习专业化并相互协作，在单一的自回归路径中完成两个子任务。这些结果为训练后的 Transformer 如何执行顺序性的算法过程提供了机制性的解释。此外，它们还为推理能力的产生提供了洞见，表明当任务被设计为需要中间思维步骤时，即使是浅层的多头 Transformer 也能有效解决那些原本需要更深架构才能处理的问题。

论文旨在探讨Transformer模型如何通过思维链（chain-of-thought）过程学习解决符号化的多步推理问题，特别是在树结构中的路径查找任务。该问题关注Transformer在多步推理任务中表现出色的能力，但其背后的机制，尤其是理论层面的理解仍有限。

论文提出了一种基于梯度下降动力学的理论分析框架，证明单层Transformer可以通过多阶段训练动态解决前向和后向推理任务。关键的新意在于揭示了不同注意力头如何自主专业化和协作，以在单一自回归路径中解决两个子任务，从而提供了一种机制解释Transformer如何实现顺序算法过程。

1. 提出了一个理论框架，用于分析Transformer在多步推理任务中的训练动态。 2. 通过设计的后向和更复杂的前向推理任务，验证了浅层多头Transformer的有效性。 3. 实验表明，即使浅层架构也能解决需要深度架构的问题，前提是任务结构化为中间思维链步骤。 4. 论文提供了对未见过的树结构的泛化保证，表明模型的鲁棒性。 5. 未来研究可探索更复杂的推理任务以及模型架构优化。

1. Chain-of-Thought Prompting Elicits Reasoning in Large Language Models 2. Language Models as Fast Program Interpreters 3. Theoretical Analysis of Transformers for Sequence Modeling 4. On the Expressive Power of Deep Learning: A Tensor Analysis 5. Emergent Multi-step Reasoning Capabilities in Pretrained Transformers