Généralisation de la longueur dans les Transformers arithmétiques

Length Generalization in Arithmetic Transformers

June 27, 2023

Auteurs: Samy Jelassi, Stéphane d'Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, François Charton

cs.AI

Résumé

Nous examinons comment les transformateurs font face à deux défis : l'apprentissage de l'arithmétique de base sur les nombres entiers, et la généralisation à des séquences plus longues que celles rencontrées pendant l'entraînement. Nous constatons que les embeddings de position relative permettent la généralisation en longueur pour des tâches simples, comme l'addition : des modèles entraînés sur des nombres à 5 chiffres peuvent effectuer des sommes à 15 chiffres. Cependant, cette méthode échoue pour la multiplication, et nous proposons le priming de l'ensemble d'entraînement : l'ajout de quelques (10 à 50) longues séquences à l'ensemble d'entraînement. Nous montrons que le priming permet à des modèles entraînés sur des multiplications de nombres à 5 chiffres par 3 chiffres de généraliser à des exemples de 35 chiffres par 3 chiffres. Nous montrons également que les modèles peuvent être primés pour différentes longueurs de généralisation, et que la taille de l'échantillon de priming évolue comme le logarithme de la taille de l'ensemble d'entraînement. Enfin, nous discutons des applications potentielles du priming au-delà de l'arithmétique.

English

We examine how transformers cope with two challenges: learning basic integer arithmetic, and generalizing to longer sequences than seen during training. We find that relative position embeddings enable length generalization for simple tasks, such as addition: models trained on 5-digit numbers can perform 15-digit sums. However, this method fails for multiplication, and we propose train set priming: adding a few (10 to 50) long sequences to the training set. We show that priming allows models trained on 5-digit times 3-digit multiplications to generalize to 35times 3 examples. We also show that models can be primed for different generalization lengths, and that the priming sample size scales as the logarithm of the training set size. Finally, we discuss potential applications of priming beyond arithmetic.