Transformer学不会多位数乘法,因缺乏长程依赖的归纳偏置。
Why Can't Transformers Learn Multiplication? Reverse-Engineering Reveals Long-Range Dependency Pitfalls
- 用注意力构建有向无环图缓存和检索部分积。
- 成功模型通过傅里叶基表示数字并形成闵可夫斯基和。
- 标准微调易陷入局部最优,需引入辅助损失才能学会乘法。
语言模型能力不断增强,却仍无法完成看似简单的多位数乘法任务。本文通过逆向工程一个能成功学习乘法的隐式思维链模型,发现三个关键点:(1) 模型具备长程依赖结构,对数归因和线性探测显示其编码了乘法所需的长程依赖;(2) 模型利用注意力机制构建有向无环图,以缓存和检索成对的部分积;(3) 部分积在注意力头中通过数字对的闵可夫斯基和实现,数字采用傅里叶基表示,二者均为直观高效的表达,而标准微调模型缺乏此特性。我们重新审视标准微调的学习动态,发现模型收敛至缺乏所需长程依赖的局部最优解。进一步通过引入预测‘累加和’的辅助损失(线性回归探针),提供归纳偏置,使模型成功学习多位数乘法。综上,通过逆向工程揭示了Transformer在学习长程依赖中的陷阱,并展示了正确归纳偏置的作用。
原文摘要 · Abstract (English)
Language models are increasingly capable, yet still fail at a seemingly simple task of multi-digit multiplication. In this work, we study why, by reverse-engineering a model that successfully learns multiplication via \emph{implicit chain-of-thought}, and report three findings: (1) Evidence of long-range structure: Logit attributions and linear probes indicate that the model encodes the necessary long-range dependencies for multi-digit multiplication. (2) Mechanism: the model encodes long-range dependencies using attention to construct a directed acyclic graph to ``cache'' and ``retrieve'' pairwise partial products. (3) Geometry: the model implements partial products in attention heads by forming Minkowski sums between pairs of digits, and digits are represented using a Fourier basis, both of which are intuitive and efficient representations that the standard fine-tuning model lacks. With these insights, we revisit the learning dynamics of standard fine-tuning and find that the model converges to a local optimum that lacks the required long-range dependencies. We further validate this understanding by introducing an auxiliary loss that predicts the ``running sum'' via a linear regression probe, which provides an inductive bias that enables the model to successfully learn multi-digit multiplication. In summary, by reverse-engineering the mechanisms of an implicit chain-of-thought model we uncover a pitfall for learning long-range dependencies in Transformers and provide an example of how the correct inductive bias can address this issue.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。