Transformers Can Do Arithmetic with the Right Embeddings

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: McLeish, Sean, Bansal, Arpit, Stein, Alex, Jain, Neel, Kirchenbauer, John, Bartoldson, Brian R., Kailkhura, Bhavya, Bhatele, Abhinav, Geiping, Jonas, Schwarzschild, Avi, Goldstein, Tom
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866909438759141376
author McLeish, Sean
Bansal, Arpit
Stein, Alex
Jain, Neel
Kirchenbauer, John
Bartoldson, Brian R.
Kailkhura, Bhavya
Bhatele, Abhinav
Geiping, Jonas
Schwarzschild, Avi
Goldstein, Tom
author_facet McLeish, Sean
Bansal, Arpit
Stein, Alex
Jain, Neel
Kirchenbauer, John
Bartoldson, Brian R.
Kailkhura, Bhavya
Bhatele, Abhinav
Geiping, Jonas
Schwarzschild, Avi
Goldstein, Tom
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2405_17399
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Transformers Can Do Arithmetic with the Right Embeddings
McLeish, Sean
Bansal, Arpit
Stein, Alex
Jain, Neel
Kirchenbauer, John
Bartoldson, Brian R.
Kailkhura, Bhavya
Bhatele, Abhinav
Geiping, Jonas
Schwarzschild, Avi
Goldstein, Tom
Machine Learning
Artificial Intelligence
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.
title Transformers Can Do Arithmetic with the Right Embeddings
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2405.17399