Attention bottleneck explains one third power law in language model training
Emergent One-Third Scaling Law as Attention Tries to Concentrate
Machine LearningArtificial Intelligence
Summary
Training large language models usually gets better with more training, following a certain power law, but why was unclear. The authors show that when parts of the model called softmax heads focus sharply on specific information, their internal numbers grow following a power law with exponent 1/3. This growth creates a bottleneck that slows learning and makes the model's loss decrease follow the one-third power law. They find many of these softmax parts in language models behave this way, especially attention heads, suggesting the core method of attention may explain the observed training behavior.
What this means in practice
- •For machine learning engineers: Identify and address bottlenecks in transformer attention mechanisms to improve training efficiency and model performance scaling.
- •For ai hardware designers: Design hardware optimized for managing growth in attention head computations that follow one-third power scaling in large language models.
Authors
Yizhou Liu, Sara Kangaslahti, Jeff Gore
Abstract
The neural scaling law relating longer training to better performance through a power law is central to today's large language models (LLMs), yet its origin remains debated. One recent proposal is that power laws can emerge from the strong non-linearity of a single softmax head learning peaked distributions. What happens with multiple softmax functions, as in LLMs, is unclear. Here, we show through toy models that any softmax learning peaked distributions, regardless of its position in the model, can develop logit magnitudes that grow in a power law with exponent $1/3$, becoming a training bottleneck whose loss contribution decays as a power law with the same exponent $1/3$. The overall loss therefore obeys $1/3$ scaling whenever at least one softmax learns peaked distributions. We confirm that many softmax functions in LLMs learn peaked distributions and that LLM loss scaling matches this $1/3$ prediction. Moreover, logit growth dynamics reveal that attention heads, rather than the language modeling head, are the bottleneck likely driving the $1/3$ loss scaling in LLMs. Attention trying to concentrate on specific information, which is the heart of Transformers, may therefore also be the heart of the neural scaling law of training.