Transformers Can Do Arithmetic with the Right Embeddings

Sean McLeish,Arpit Bansal,Alex Stein,Neel Jain,John Kirchenbauer,Brian R. Bartoldson,Bhavya Kailkhura,Abhinav Bhatele,Jonas Geiping,Avi Schwarzschild,Tom Goldstein

2024-05-28

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.

Machine Learning,Artificial Intelligence

What problem does this paper attempt to address?

This paper mainly discusses the problems encountered by Transformer models when performing arithmetic tasks, such as addition, that is, they have difficulty tracking the exact position of each number in long sequences of numbers. To address this issue, researchers propose a new embedding method called "Abacus Embeddings," which encodes the position of each number relative to the start of the number. Experiments show that this improvement not only improves the performance of arithmetic tasks but also allows the model to expand from 20-digit numbers in the training data to 120-digit numbers, achieving a 6-fold generalization factor, which is two and a half times better than the previous best result.

Furthermore, by combining Abacus Embeddings, input injection, and a recurrent transformer architecture, the model's accuracy on out-of-distribution addition problems increased from 92.9% to 99.1%, reducing the error rate by 87%. These improvements also have a positive impact on other multi-step reasoning tasks such as sorting and multiplication. The study also compared different position embedding methods, including NoPE and FIRE, and found that Abacus Embeddings performed best on addition tasks.

Overall, the paper aims to address the limitations of Transformer models in handling arithmetic tasks, particularly for large numbers and complex logical reasoning, and significantly improve the model's generalization and algorithmic reasoning capabilities by introducing Abacus Embeddings and other structural optimizations.