Les Transformers peuvent effectuer des opérations arithmétiques avec les bonnes représentations vectorielles.

Transformers Can Do Arithmetic with the Right Embeddings

May 27, 2024

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

cs.AI

Résumé

La faible performance des transformers sur les tâches arithmétiques semble provenir en grande partie de leur incapacité à suivre avec précision la position exacte de chaque chiffre au sein d'une longue séquence de chiffres. Nous corrigeons ce problème en ajoutant à chaque chiffre un embedding qui encode sa position relative par rapport au début du nombre. En plus de l'amélioration apportée par ces embeddings eux-mêmes, nous montrons que cette correction permet à des modifications architecturales telles que l'injection d'entrée et les couches récurrentes d'améliorer encore davantage les performances. Une fois les positions résolues, nous pouvons étudier la capacité d'extrapolation logique des transformers. Peuvent-ils résoudre des problèmes arithmétiques plus grands et plus complexes que ceux présents dans leurs données d'entraînement ? Nous constatons qu'en nous entraînant uniquement sur des nombres de 20 chiffres avec un seul GPU pendant une journée, nous pouvons atteindre des performances de pointe, obtenant jusqu'à 99 % de précision sur des problèmes d'addition de 100 chiffres. Enfin, nous montrons que ces gains en numératie débloquent également des améliorations sur d'autres tâches de raisonnement multi-étapes, y compris le tri et la multiplication.

English

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.