Transformer通过思维链实现多步梯度下降,精准恢复权重。
Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought
- 用思维链引导Transformer自回归完成多步梯度下降
- 可近乎精确恢复真实权重向量,优于单步梯度方法
- 适合研究模型推理机制与上下文学习的学者
思维链(CoT)提示已被证明能显著提升大语言模型在算术和推理任务中的表现,通过引导模型生成中间推理步骤。尽管CoT在实证上成效显著且理论上增强表达能力,其训练机制仍不明确。本文研究了Transformer在上下文权重预测任务中,针对线性回归使用CoT目标时的训练动态。理论证明:无CoT的一层线性Transformer仅能实现单步梯度下降(GD),无法恢复真实权重;而采用CoT的Transformer可自回归学习多步GD,实现近乎精确的权重恢复。此外,模型在未见数据上表现出良好泛化能力。我们还发现,循环式Transformer相比非循环结构,在线性回归的上下文学习中显著提升最终性能。实验进一步验证了CoT提示带来的显著性能提升。
原文摘要 · Abstract (English)
Chain of Thought (CoT) prompting has been shown to significantly improve the performance of large language models (LLMs), particularly in arithmetic and reasoning tasks, by instructing the model to produce intermediate reasoning steps. Despite the remarkable empirical success of CoT and its theoretical advantages in enhancing expressivity, the mechanisms underlying CoT training remain largely unexplored. In this paper, we study the training dynamics of transformers over a CoT objective on an in-context weight prediction task for linear regression. We prove that while a one-layer linear transformer without CoT can only implement a single step of gradient descent (GD) and fails to recover the ground-truth weight vector, a transformer with CoT prompting can learn to perform multi-step GD autoregressively, achieving near-exact recovery. Furthermore, we show that the trained transformer effectively generalizes on the unseen data. With our technique, we also show that looped transformers significantly improve final performance compared to transformers without looping in the in-context learning of linear regression. Empirically, we demonstrate that CoT prompting yields substantial performance improvements.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。