Transformer können Arithmetik mit den richtigen Einbettungen durchführen.

Transformers Can Do Arithmetic with the Right Embeddings

May 27, 2024

Autoren: 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

Zusammenfassung

Die schlechte Leistung von Transformatoren bei arithmetischen Aufgaben scheint zum großen Teil daher zu rühren, dass sie nicht in der Lage sind, die genaue Position jedes einzelnen Digits innerhalb eines großen Zahlenbereichs nachzuverfolgen. Wir beheben dieses Problem, indem wir jedem Digit eine Einbettung hinzufügen, die seine Position relativ zum Beginn der Zahl codiert. Neben dem Schub, den diese Einbettungen allein bieten, zeigen wir, dass diese Korrektur architektonische Modifikationen wie Eingabeinjektion und rekurrente Schichten ermöglicht, um die Leistung noch weiter zu verbessern. Mit gelösten Positionen können wir die logische Extrapolationsfähigkeit von Transformatoren untersuchen. Können sie arithmetische Probleme lösen, die größer und komplexer sind als die in ihren Trainingsdaten? Wir stellen fest, dass wir durch das Training mit nur 20-stelligen Zahlen mit einer einzigen GPU für einen Tag eine Spitzenleistung erreichen können, indem wir bis zu 99% Genauigkeit bei 100-stelligen Additionsaufgaben erzielen. Schließlich zeigen wir, dass diese Gewinne in der Numerik auch Verbesserungen bei anderen mehrstufigen Denkaufgaben wie Sortieren und Multiplikation freisetzen.

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.