Transformadores Podem Realizar Operações Aritméticas com os Embeddings Corretos

Transformers Can Do Arithmetic with the Right Embeddings

May 27, 2024

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

Resumo

O desempenho insatisfatório dos transformadores em tarefas aritméticas parece decorrer, em grande parte, de sua incapacidade de rastrear a posição exata de cada dígito dentro de uma sequência extensa de dígitos. Corrigimos esse problema adicionando uma incorporação (embedding) a cada dígito que codifica sua posição relativa ao início do número. Além do benefício que essas incorporações proporcionam por si só, demonstramos que essa correção permite modificações arquitetônicas, como injeção de entrada e camadas recorrentes, para melhorar ainda mais o desempenho. Com as posições resolvidas, podemos estudar a capacidade de extrapolação lógica dos transformadores. Eles conseguem resolver problemas aritméticos maiores e mais complexos do que aqueles presentes em seus dados de treinamento? Descobrimos que, ao treinar apenas com números de 20 dígitos usando uma única GPU por um dia, é possível alcançar desempenho de ponta, atingindo até 99% de precisão em problemas de adição com 100 dígitos. Por fim, mostramos que esses ganhos em habilidades numéricas também desbloqueiam melhorias em outras tarefas de raciocínio multi-etapas, incluindo ordenação e multiplicação.

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.