[论文解读] Transformers Can Do Arithmetic with the Right Embeddings

Sean McLeish, Arpit Bansal+9 more|arXiv (Cornell University)|May 27, 2024

Computability, Logic, AI Algorithms被引用 6

一句话总结

本文提出 Abacus Embeddings，用于在数值中编码数字的位置，以使 Transformer 能执行长距离算术并外推到更大位数长度，借助输入注入和循环式 Transformer 架构。

ABSTRACT

The poor performance of transformers on arithmetic tasks seems to stem in large part from their inability to keep track of the exact position of each digit inside of a large span of digits. We mend this problem by adding an embedding to each digit that encodes its position relative to the start of the number. In addition to the boost these embeddings provide on their own, we show that this fix enables architectural modifications such as input injection and recurrent layers to improve performance even further. With positions resolved, we can study the logical extrapolation ability of transformers. Can they solve arithmetic problems that are larger and more complex than those in their training data? We find that training on only 20 digit numbers with a single GPU for one day, we can reach state-of-the-art performance, achieving up to 99% accuracy on 100 digit addition problems. Finally, we show that these gains in numeracy also unlock improvements on other multi-step reasoning tasks including sorting and multiplication.

研究动机与目标

• 识别限制 Transformer 算术能力的架构与数据表示瓶颈。
• 提出 Abacus Embeddings 以编码数字的数字位的重要性并改进长度泛化。
• 评估递归和输入注入如何与 Abacus Embeddings 交互以提升加法、乘法和排序任务的性能。
• 展示在 100+ 位数算术上的最先进外推能力以及向其他算法任务的可迁移性。

提出的方法

• 引入 Abacus Embeddings：对一个数字中同一数位的所有数位应用的学习位置嵌入。
• 在 20-million-sample 加法数据上训练仅解码器的因果 Transformer，采用 least-significant-digit-first Formatting 且不填充。
• 在绝对与相对嵌入方案下，对标准 Transformer、输入注入变体和循环 Transformer 进行比较。
• 评估分布内、分布外和极端 OOD 性能，包括 100+ digit 加法。
• 将实验扩展到乘法与排序，以测试对其他算法任务的泛化能力。
• 检查 Abacus Embeddings 与 FIRE 和 RoPE 相对位置嵌入的兼容性。

Figure 1: Zero shot exact match accuracy on addition using depth sixteen transformer (decoder only) models trained on operands of up to 20 digits. Compared to state-of-the-art embeddings (left), our new Abacus Embeddings (right) dramatically improve generalization to unseen digit lengths. The interi

实验结果

研究问题

• RQ1 Abacus Embeddings 能否实现对多位数加法在训练长度之外的长度泛化和零-shot 外推？
• RQ2 当与 Abacus Embeddings 搭配使用时，递归和输入注入是否能进一步降低泛化误差？
• RQ3 这些方法在更大规模的算术运算（乘法）以及非算术的算法任务（排序）中的迁移能力如何？
• RQ4 Abacus Embeddings 是否与现有的相对位置嵌入（如 FIRE 与 RoPE）兼容并互补？

主要发现

• Abacus Embeddings 在加法方面显著提升泛化能力，在 100-digit 加法上达到最高 99.1% 的准确率并外推到 120-digit 的问题。
• 将 Abacus Embeddings 与输入注入和循环 Transformer 相结合，达到近乎完美的泛化并实现相对于训练长度最多 6× 的长度扩展。
• 带递归的循环 Transformer 相对于无递归基线，在 out-of-distribution 加法上可将错误率降低多达 ~50%。
• Abacus Embeddings 与 FIRE 结合时对乘法的分布内性能也有所提升，并在分布外场景下提升排序的准确性。
• Abacus Embeddings 与 FIRE 和 RoPE 兼容，并且与 FIRE 一起能够实现超越 FIRE 单独所能达到的泛化。

Figure 2: Visualization of data formats and positional embeddings. Abacus Embeddings give the same positional embeddings to all digits of the same significance.