Au-delà des lois d'échelle : Comprendre la performance des Transformers grâce à la mémoire associative

Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

May 14, 2024

Auteurs: Xueyan Niu, Bo Bai, Lei Deng, Wei Han

cs.AI

Résumé

L'augmentation de la taille d'un modèle Transformer ne conduit pas toujours à une amélioration des performances. Ce phénomène ne peut être expliqué par les lois empiriques de mise à l'échelle. Par ailleurs, une capacité de généralisation améliorée apparaît lorsque le modèle mémorise les échantillons d'entraînement. Nous présentons un cadre théorique qui éclaire le processus de mémorisation et la dynamique des performances des modèles de langage basés sur les Transformers. Nous modélisons le comportement des Transformers avec des mémoires associatives en utilisant des réseaux de Hopfield, de sorte que chaque bloc Transformer effectue efficacement une recherche approximative des plus proches voisins. Sur cette base, nous concevons une fonction d'énergie analogue à celle du réseau de Hopfield continu moderne, qui fournit une explication approfondie du mécanisme d'attention. En utilisant la technique de majorisation-minimisation, nous construisons une fonction d'énergie globale qui capture l'architecture en couches du Transformer. Sous certaines conditions, nous montrons que la perte d'entropie croisée minimale réalisable est bornée inférieurement par une constante approximativement égale à 1. Nous étayons nos résultats théoriques en menant des expériences avec GPT-2 sur différentes tailles de données, ainsi qu'en entraînant des Transformers standards sur un ensemble de données de 2 millions de tokens.

English

Increasing the size of a Transformer model does not always lead to enhanced performance. This phenomenon cannot be explained by the empirical scaling laws. Furthermore, improved generalization ability occurs as the model memorizes the training samples. We present a theoretical framework that sheds light on the memorization process and performance dynamics of transformer-based language models. We model the behavior of Transformers with associative memories using Hopfield networks, such that each transformer block effectively conducts an approximate nearest-neighbor search. Based on this, we design an energy function analogous to that in the modern continuous Hopfield network which provides an insightful explanation for the attention mechanism. Using the majorization-minimization technique, we construct a global energy function that captures the layered architecture of the Transformer. Under specific conditions, we show that the minimum achievable cross-entropy loss is bounded from below by a constant approximately equal to 1. We substantiate our theoretical results by conducting experiments with GPT-2 on various data sizes, as well as training vanilla Transformers on a dataset of 2M tokens.