Los Transformers pueden realizar operaciones aritméticas con los embeddings adecuados.

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

Resumen

El bajo rendimiento de los transformadores en tareas aritméticas parece deberse en gran parte a su incapacidad para rastrear la posición exacta de cada dígito dentro de un intervalo amplio de dígitos. Solucionamos este problema añadiendo una incrustación a cada dígito que codifica su posición relativa al inicio del número. Además del impulso que estas incrustaciones proporcionan por sí mismas, demostramos que esta corrección permite que modificaciones arquitectónicas, como la inyección de entrada y capas recurrentes, mejoren aún más el rendimiento. Con las posiciones resueltas, podemos estudiar la capacidad de extrapolación lógica de los transformadores. ¿Pueden resolver problemas aritméticos más grandes y complejos que los presentes en sus datos de entrenamiento? Descubrimos que, entrenando únicamente con números de 20 dígitos utilizando una sola GPU durante un día, podemos alcanzar un rendimiento de vanguardia, logrando hasta un 99% de precisión en problemas de suma de 100 dígitos. Finalmente, demostramos que estas mejoras en habilidades numéricas también desbloquean avances en otras tareas de razonamiento de múltiples pasos, como la ordenación y la multiplicación.

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.