单GPU训练一天，Transformer在100位数字加法上就达能到99%准确率

乘法和排序也有效。

自 2017 年被提出以来，Transformer 已成为 AI 大模型的主流架构，一直稳站 C 位。

但所有研究者都不得不承认的是，Transformer 在算数任务中表现非常糟糕，尤其是加法，这一缺陷在很大程度上源于 Transformer 无法跟踪大范围数字中每个数字的确切位置。

为了解决这个问题，来自马里兰大学、CMU 等机构的研究者向这一问题发起了挑战，他们通过在每个数字中添加一个嵌入来解决这个问题，该嵌入编码数字相对于开头的位置。该研究发现，只用一天时间在单个 GPU 上训练 20 位数字，就可以达到最新的性能水平，100 位数字加法问题高达 99% 的准确率。

论文地址：https://arxiv.org/pdf/2405.17399

项目地址：https://github.com/mcleish7/arithmetic

标题：Transformers Can Do Arithmetic with the Right Embeddings

具体而言，研究者建议对数据表示进行一个简单的修改，就能解决这个缺点。他们提出了 Abacus 嵌入用于编码每个数字符号 token 范围内的位置。将 Abacus 嵌入与标准位置嵌入结合使用后，该研究观察到 Transformer 在算数任务上的准确率有显著提高，以至于最多只训练了 20 位数操作数的模型可以泛化到 120 位数操作数的问题。这一数字代表了 6 倍的 SOTA 泛化因子，而以前的最先进的泛化因子也只有 2.5 倍。据了解，这是迄今为止被证明的最长的学习加法序列。

此外，本文还研究了几种其他方法来改善 transformer 在算术和泛化方面的性能，他们发现结合输入注入（input injection），即在输入层和每个解码器层之间插入跳跃连接，可以在 Abacus 嵌入基线上减少 50% 的泛化误差。本文还发现，与嵌入结合使用的 looped transformer 架构可以在加法问题上实现几乎完美的泛化。

本文的贡献可以总结如下：

• 本文提出了一种新的位置嵌入，称为 Abacus 嵌入，以更好地捕获每个数字的重要性，从而实现近乎完美的分布内泛化；

• 研究表明，当将 Abacus 嵌入与输入注入和 looped transformer 相结合时，性能会进一步提高，分布外准确率从 92.9% 提高到 99.1%，与单独使用标准架构的嵌入相比，误差降低了 87%；

• 研究者将这些发现扩展到更复杂的问题，包括乘法和排序，在这些领域也展现出了长度泛化。

实现加法的长度泛化

作者研究了一系列方法，旨在提高从头开始训练的语言模型在算术能力上的表现。他们主要关注两个假设：1）数字内各个位数的位置信息正在丢失；2）循环可以提高 Transformer 架构在多步算术推理问题上的推理能力。在详细描述每项改进之前，作者简要讨论了训练和评估设置。

实验设置

作者训练了仅包含解码器的因果语言模型来解决加法问题。

他们考虑了两种标准 transformer 架构。首先，他们使用一个标准的自回归 transformer 模型，多个解码器层以前馈方式堆叠。其次，他们通过输入注入（input injection）增强了这一标准 transformer 模型，即把嵌入的输入添加到每个解码器层的输入中。作者在图 20 中直观地描述了这些架构。

Abacus 嵌入帮助对齐数字

通过之前的研究和初步实验，作者发现，即使输入的数字是先显示最不重要的数字，训练数据是分层的、丰富的（几百万个例子），标准 transformer 也很难学习多位数加法。他们还观察到，人类在进行长加法运算时，会先将数位相同的数字排列成列。因此，作者的第一个假设是，对于 transformer 来说，每个数字的数位并不容易表示，而且这个子问题比实际加法本身带来的障碍更大。

为了解决 transformer 在表示位置信息方面的局限性，作者设计了一种特殊的位置嵌入，它可以编码每个数字相对于当前数字起始位置的位置。作者将其称之为 Abacus 嵌入。他们将相同的位置嵌入应用于所有具有相同数位的数字，从而提供一个显式的信号，供模型用于对齐数字，如图 2 所示。

Abacus 嵌入解决加法问题

对于标准 transformer 架构，Abacus 嵌入可将泛化性能提高到 100 位及以上。在图 3（左）中，作者强调了 Abacus 嵌入与标准 transformer 架构和嵌入相比，在进行加法运算时所具有的比较优势，取三种模型在所有情况下的平均准确度。

图 1 还显示了使用 FIRE 和 Abacus 训练的标准 transformer 模型的准确度结果，这些模型经过了域内 (ID) 和域外 (OOD) 测试。

Transformer 中的循环提高了性能

在解决位置嵌入问题后，接下来作者探讨了循环架构能否进一步提高 transformer 执行多位数加法的能力。他们使用「循环块（recurrent block）」一词来指一组具有不同权重的解码器层，而「循环（recurrence）」则指循环块的重复次数。作者使用有效深度（effective depth）一词来指 transformer 中使用的层数，无论其权重是否唯一。除非另有说明，否则他们使用的是最大循环架构，即只循环一个唯一层来达到有效深度。他们还采用了输入注入、 残差连接的方式，将输入的副本传播到网络中的每一层。

循环的优势

在图 3（右）中，作者比较了使用 FIRE 和 NoPE 嵌入对操作数多达 40 位的加法进行训练的所有架构变体。尽管参数数量仅相当于其他模型的 1/10，但可以看到，looped transformer（循环的、有输入注入和渐进损失）在使用任何一种位置嵌入时都取得了最佳的分布外性能。在图 8 中，作者展示了这一结果在多种训练数据规模下的稳健性。

对于循环模型，可以选择在训练时改变每次前向传递的循环次数。这往往会提高模型测试时对较难任务的泛化能力，这也被称为渐进损失计算（progressive loss computation）。这个损失函数是两个前向传递的损失值的凸组合，一个使用字面上的循环数（1 × 16 模型为 16），另一个使用随机的较小循环数。

接下来，作者探讨了在保持有效深度固定的同时改变循环块大小的效果。他们将循环块中的层数减半，循环次数增加一倍，从块中有 16 层、循环次数只有一次（16 × 1，即标准 transformer）的模型，过渡到块中只有一层、循环次数有 16 次（1 × 16）的模型。

通过图 4 分析这些结果，作者发现在某些情况下，结合循环和 Abacus 嵌入可以进一步提高性能。具体来说，在 OOD 问题上，有两个循环的模型（8 × 2）产生的误差是纯非循环模型（16 × 1）的一半，而在 100 + 的 OOD 问题上，其准确率也有所提高。

最后，在附录 A.7.3 中，作者改变了模型的有效深度，以分析参数数量对这项任务的影响，包括 Abacus、FIRE 和 NoPE 嵌入。虽然图 4 中的实验是对不同深度的公平比较，但纯粹的标准 transformer 模型比相应的循环模型拥有更多的参数。在附录的表 3 中，作者记录了最接近百万的参数量。

实验

研究者不仅对加法问题进行了探讨，还对乘法和排序进行了研究。

整数乘法

图 5 展示了 Abacus 嵌入模型在 15 位数乘法的分布内准确率超过了之前的工作，且不需要用零将每个操作数填充到相同长度。特别地，该研究强调，与仅使用 FIRE 的基线相比，将 Abacus 嵌入与 FIRE 相结合也提高了分布问题中最难的分布准确率 (右下)。

数组排序

表 1 展示了使用不同嵌入 ——FIRE、Abacus 及其组合 —— 训练的标准 transformer（八层）的性能。结果显示，组合嵌入方法增强了模型的泛化能力。

如表 2 所示，研究者观察到在将 Abacus+FIRE 嵌入组合与不同的模型架构（有效深度为 8）配对时，结果表现出混合性。

Abacus 和相关嵌入

图 6 展示了将 Abacus 嵌入整合到更通用系统中的真正潜力，显示出 Abacus 嵌入与 FIRE 结合可以解锁远超 FIRE 嵌入解决问题的能力。 

更多研究细节，请参考原论文。