Más allá de las leyes de escalamiento: Comprendiendo el rendimiento de los transformadores con memoria asociativa

Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

May 14, 2024

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

cs.AI

Resumen

Aumentar el tamaño de un modelo Transformer no siempre conduce a una mejora en el rendimiento. Este fenómeno no puede explicarse mediante las leyes de escalado empíricas. Además, la capacidad de generalización mejorada ocurre a medida que el modelo memoriza las muestras de entrenamiento. Presentamos un marco teórico que arroja luz sobre el proceso de memorización y la dinámica del rendimiento en los modelos de lenguaje basados en Transformers. Modelamos el comportamiento de los Transformers con memorias asociativas utilizando redes de Hopfield, de modo que cada bloque del Transformer efectivamente realiza una búsqueda aproximada del vecino más cercano. Basándonos en esto, diseñamos una función de energía análoga a la de la red de Hopfield continua moderna, lo que proporciona una explicación esclarecedora del mecanismo de atención. Utilizando la técnica de mayorización-minimización, construimos una función de energía global que captura la arquitectura en capas del Transformer. Bajo condiciones específicas, demostramos que la pérdida de entropía cruzada mínima alcanzable está acotada inferiormente por una constante aproximadamente igual a 1. Sustentamos nuestros resultados teóricos mediante experimentos con GPT-2 en varios tamaños de datos, así como entrenando Transformers básicos en un conjunto de datos de 2 millones 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.