单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 泛化因子…
实现加法的长度泛化
作者研究了一系列方法，旨在提高从头开始训练的语言模型在算术能力上的表现。他们主要关注两个假设：1）数字内各个位数的位置信息正在丢失；2）循环可以提高 Transformer 架构在多步算术推理问题上的推理能力。在详细描述每项改进之前，作者简要讨论了训练和评估设置。
实验设置
作者训练了仅包含解码器的因果语言模型来解决加法问题。
他们考虑了两种标准 transformer 架构。首先，他们使用一个标准的自回归 transformer 模型，多个解码器层以前馈方式堆叠。其次，他们通过输入注入（input injection）增强了这一标准 transformer 模型，即把嵌入的输入添加到每个解码器层的输入中。作者在图 20 中直观地描述了这些架构。
Abacus 嵌入帮助对齐数字
通过之前的研究和初步实验，作者发现…
Abacus 嵌入解决加法问题
对于标准 transformer 架构，Abacus 嵌入可将泛化性能提高到 100 位及以上。在图 3（左）中，作者强调了 Abacus 嵌入与标准 transformer 架构和嵌入相比，在进行加法运算时所具有的比较优势，取三种模型在所有情况下的平均准确度。
Transformer 中的循环提高了性能
在解决位置嵌入问题后，接下来作者探讨了循环架构能否进一步提高 transformer 执行多位数加法的能力。他们使用「循环块（recurrent block）」一词来指一组具有不同权重的解码器层，而「循环（recurrence）」则指循环块的重复次数。作者使用有效深度（effective depth）一词来指 transformer 中使用的层数，无论其权重是否唯一。除非另有说明，否则他们使用的是最大循环架构，即只循环一个唯一层来达到有效深度。他们还采用了输入注入、 残差连接的方式，将输入的副本传播到网络中的每一层。
循环的优势
在图 3（右）中，作者比较了使用 FIRE 和 NoPE 嵌入对操作数多达 40 位的加法进行训练的所有架构变体。尽管参数数量仅相当于其他模型的 1/10，但可以看到，looped transformer（循环的、有输入注入和渐进损失）在使用任何一种位置嵌入时都取得了最佳的分布外性能。在图 8 中，作者展示了这一结果在多种训练数据规模下的稳健性。
实验
研究者不仅对加法问题进行了探讨，还对乘法和排序进行了研究。
整数乘法
图 5 展示了 Abacus 嵌入模型在 15 位数乘法的分布内准确率超过了之前的工作，且不需要用零将每个操作数填充到相同长度。特别地，该研究强调，与仅使用 FIRE 的基线相比，将 Abacus 嵌入与 FIRE 相结合也提高了分布问题中最难的分布准确率 (右下)。
数组排序
表 1 展示了使用不同嵌入 ——FIRE、Abacus 及其组合 —— 训练的标准 transformer（八层）的性能。结果显示，组合嵌入方法增强了模型的泛化能力。
Abacus 和相关嵌入
图 6 展示了将 Abacus 嵌入整合到更通用系统中的真正潜力，显示出 Abacus 嵌入与 FIRE 结合可以解锁远超 FIRE 嵌入解决问题的能力。
更多研究细节，请参考原论文。
感谢阅读！如果您对AI的更多资讯感兴趣，可以查看更多AI文章：GPTNB。
单GPU训练一天，Transformer在100位数字加法上就达能到99%准确率
