Voorbij schaalwetten: Transformer-prestaties begrijpen met associatief geheugen

Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

May 14, 2024

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

cs.AI

Samenvatting

Het vergroten van de omvang van een Transformer-model leidt niet altijd tot verbeterde prestaties. Dit fenomeen kan niet worden verklaard door de empirische schaalwetten. Bovendien treedt een verbeterd generalisatievermogen op naarmate het model de trainingsvoorbeelden uit het hoofd leert. We presenteren een theoretisch kader dat inzicht geeft in het memorisatieproces en de prestatie-dynamiek van transformer-gebaseerde taalmodellen. We modelleren het gedrag van Transformers met associatieve geheugens met behulp van Hopfield-netwerken, waarbij elke transformer-blok effectief een benaderende zoektocht naar de dichtstbijzijnde buur uitvoert. Op basis hiervan ontwerpen we een energie-functie die analoog is aan die in het moderne continue Hopfield-netwerk, wat een inzichtelijke verklaring biedt voor het aandacht-mechanisme. Met behulp van de majorization-minimization techniek construeren we een globale energie-functie die de gelaagde architectuur van de Transformer vastlegt. Onder specifieke voorwaarden laten we zien dat de minimaal haalbare kruis-entropie-verlies begrensd is van onderaf door een constante die ongeveer gelijk is aan 1. We onderbouwen onze theoretische resultaten door experimenten uit te voeren met GPT-2 op verschillende datagroottes, evenals het trainen van standaard Transformers op een dataset van 2M 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.