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

发布于 2024-6-3 09:56
浏览
0收藏

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


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


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


论文地址: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 中直观地描述了这些架构。


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


Abacus 嵌入帮助对齐数字


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


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


Abacus 嵌入解决加法问题


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



单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


Transformer 中的循环提高了性能


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


循环的优势


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


对于循环模型,可以选择在训练时改变每次前向传递的循环次数。这往往会提高模型测试时对较难任务的泛化能力,这也被称为渐进损失计算(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 中,作者记录了最接近百万的参数量。


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


实验


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


整数乘法


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


数组排序


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


Abacus 和相关嵌入


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


单GPU训练一天,Transformer在100位数字加法上就达能到99%准确率-AI.x社区


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


本文转自 机器之心 ,作者:机器之心


原文链接:​​https://mp.weixin.qq.com/s/zHx_pMk6sHmIm-AJCS_sRA​

标签
收藏
回复
举报
回复
相关推荐