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.
Toy Models and Power Law Scaling
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.
Confirmation in LLMs
We confirm that many softmax functions in LLMs learn peaked distributions and that LLM loss scaling matches this $1/3$ prediction.
Attention Heads as Bottleneck
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.