Ensinando Aritmética a Pequenos Transformers

Teaching Arithmetic to Small Transformers

July 7, 2023

Autores: Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, Dimitris Papailiopoulos

cs.AI

Resumo

Grandes modelos de linguagem como o GPT-4 exibem capacidades emergentes em tarefas de propósito geral, como aritmética básica, quando treinados em extensos conjuntos de dados textuais, mesmo que essas tarefas não sejam explicitamente codificadas pelo objetivo não supervisionado de previsão do próximo token. Este estudo investiga como pequenos transformadores, treinados a partir de inicialização aleatória, podem aprender eficientemente operações aritméticas como adição, multiplicação e funções elementares como raiz quadrada, utilizando o objetivo de previsão do próximo token. Primeiro, demonstramos que os dados de treinamento convencionais não são os mais eficazes para o aprendizado de aritmética, e que simples mudanças na formatação podem melhorar significativamente a precisão. Isso leva a transições de fase abruptas em função da escala dos dados de treinamento, que, em alguns casos, podem ser explicadas por conexões com a completação de matrizes de baixo posto. Baseando-se em trabalhos anteriores, em seguida treinamos com dados no estilo de cadeia de pensamento que incluem resultados de etapas intermediárias. Mesmo na completa ausência de pré-treinamento, essa abordagem melhora significativamente e simultaneamente a precisão, a complexidade amostral e a velocidade de convergência. Também estudamos a interação entre dados aritméticos e textuais durante o treinamento e examinamos os efeitos de prompts de poucos exemplos, pré-treinamento e escala do modelo. Além disso, discutimos os desafios de generalização de comprimento. Nosso trabalho destaca a importância de dados de alta qualidade e instrutivos que consideram as características particulares do objetivo de previsão da próxima palavra para elicitar rapidamente capacidades aritméticas.

English

Large language models like GPT-4 exhibit emergent capabilities across general-purpose tasks, such as basic arithmetic, when trained on extensive text data, even though these tasks are not explicitly encoded by the unsupervised, next-token prediction objective. This study investigates how small transformers, trained from random initialization, can efficiently learn arithmetic operations such as addition, multiplication, and elementary functions like square root, using the next-token prediction objective. We first demonstrate that conventional training data is not the most effective for arithmetic learning, and simple formatting changes can significantly improve accuracy. This leads to sharp phase transitions as a function of training data scale, which, in some cases, can be explained through connections to low-rank matrix completion. Building on prior work, we then train on chain-of-thought style data that includes intermediate step results. Even in the complete absence of pretraining, this approach significantly and simultaneously improves accuracy, sample complexity, and convergence speed. We also study the interplay between arithmetic and text data during training and examine the effects of few-shot prompting, pretraining, and model scale. Additionally, we discuss length generalization challenges. Our work highlights the importance of high-quality, instructive data that considers the particular characteristics of the next-word prediction objective for rapidly eliciting arithmetic capabilities.