Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

TL;DR

This paper presents a theoretical framework using associative memories and Hopfield networks to explain why increasing transformer size doesn't always improve performance, linking memorization to model behavior.

Contribution

It introduces a novel energy-based model of transformers with associative memory, providing insights into memorization and performance limits beyond empirical scaling laws.

Findings

Transformer performance is linked to memorization of training samples.

A global energy function captures the layered architecture of transformers.

There is a size-dependent optimal performance bound related to dataset size.

Abstract

Increasing the size of a Transformer does not always lead to enhanced performance. This phenomenon cannot be explained by the empirical scaling laws. Furthermore, the model's enhanced performance is closely associated with its memorization of the training samples. We present a theoretical framework that sheds light on the memorization during pre-training 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. In particular, the energy function in modern continuous Hopfield networks serves as an explanation for the attention mechanism, which we approximate with a distance-based energy function. By observing that the softmax function corresponds to the gradient of the LogSumExp function in the energy, and employing the majorization-minimization technique, we construct a global energy function designed to capture the layered architecture. We demonstrate a dependency between the model size and the dataset size for the model to achieve optimal performance, and we show that the achievable cross-entropy loss is bounded from below.