Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory
This provides theoretical insights into Transformer memorization for researchers in language modeling, though it is incremental as it builds on existing associative memory concepts.
The paper tackles the problem that increasing Transformer size does not always improve performance, which scaling laws cannot explain, by developing a theoretical framework linking memorization to performance using associative memory models, and shows that optimal performance depends on model and dataset sizes with a lower bound on cross-entropy loss.
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.