cs.LGSep 26, 2026
SaveEmergent One-Third Scaling Law as Attention Tries to Concentrate
Organizations: Massachusetts Institute of Technology · Harvard University
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 , becoming a training bottleneck whose loss contribution decays as a power law with the same exponent . The overall loss therefore obeys 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 prediction. Moreover, logit growth dynamics reveal that attention heads, rather than the language modeling head, are the bottleneck likely driving the 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.
Figures & tables
Figure 1: In the gate-head (GH) model, any cold softmax, whether the head or the gate, leads to time scaling of loss. (a) The GH model has two parts: a gate layer and a head layer. The gate layer uses the key matrix and softmax to weight the value matrices. The head layer uses a matrix and softmax to project hidden states to distributions. (b) The GH model’s loss exhibits time scaling whenever the target coldness , , or both are large. Experiment details in Appendix A.1 .
Figure 2: When gate coldness is large but head coldness is low, the gate logit magnitude grows as a power law after the head saturates, sustaining loss scaling. (a) The late-time loss follows time scaling across different (represented by colors). Different loss curves enter their time scaling regions at different times, with the transitions marked by the vertical bands. (b) When loss enters the time scaling region, student head coldness saturates. (c) The scaling in loss also corresponds to student gate coldness growing as a power law. Details in Appendix A.1 .
Figure 3: Training only the gate with MSE between hidden states yields scaling in both loss and logit magnitude. (a) The loss enters the asymptotic time scaling regime, with the vertical line marking its onset. (b) The student coldness exhibits positive time scaling, beginning at the same time. (c) The value matrix displacement from initialization saturates as the loss and coldness enter their scaling regimes. Curves from different learning rates collapse when plotted against , supporting as the fundamental variable governing dynamics. See Appendix A.2 for details.
Figure 4: In Pythia LLMs, attention heads are more strongly peaked than the LM head, with logit magnitudes growing as power laws, identifying attention as the bottleneck for loss scaling. (a) Entropy suggests that later-layer attention heads and the LM head are in the low-temperature regime (below the horizontal dashed lines). (b) Probability densities at zero energy gap are non-zero. (c) LLMs have losses following the time scaling. (d) Logit coldness of attention heads rather than the LM head grows as the positive time scaling at late training. Details in Appendix C.1 .
Figure 5: Using MSE between final logits for training, attention logit magnitude grows like a positive scaling while the LM head saturates (panel a), ruling out the LM head as the only bottleneck for loss scaling. (b) MSE loss has a tail following the time scaling, as suggested by loss universality. Details in Appendix C.2 .
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Under Adam, the GH model’s loss exhibits approximate time scaling whenever the target coldness , , or both are large. Each panel shows the test loss against dynamic time on log-log axes, with a dashed reference line of slope . The layout arranges the four corners of the sweep, with gate coldness increasing rightward and head coldness increasing upward. When both coldnesses are small (bottom left), the loss initially descends rapidly as in the SGD case, but Adam becomes unstable at late times. In the remaining three panels, approximate loss scaling is observed, consistent with the SGD results in Figure 1 b.
Figure 7: Under Adam, when gate coldness is large but head coldness is small, the head saturates and the gate logit magnitude subsequently grows approximately as , sustaining approximate loss scaling. Three representative runs are shown, color-coded by run. (a) Test loss against ; the dashed line has slope . (b) Head logit standard deviation against ; grows initially then saturates. (c) Gate logit standard deviation against ; the dashed line has slope . The sequence of head saturation followed by gate growth and continued loss decay mirrors the SGD behavior in Figure 2 , though the alignment with the reference slopes is less tight.
Figure 8: Under Adam, training only the gate with MSE loss yields approximate scaling in both loss and gate logit magnitude, mirroring the SGD results in Figure 3 . We fix and show three representative learning rates, color-coded by from dark purple to light green. (a) Test loss against ; the dashed line has slope . (b) Gate logit standard deviation against ; the dashed line has slope . (c) Value tensor displacement against ; rises and then saturates. Curves from different learning rates collapse well when plotted against .
Figure 9: The GH-II model reproduces time scaling in loss and gate coldness for both sigmoid and tanh gate non-linearities ( , ). Each pair of panels shows (left) and (right) against on log-log axes. (a,b) Sigmoid variant. (c,d) Tanh variant. Dashed reference lines have slope (panels a,c) and (panels b,d).
Figure 10: The GH-III model with real attention as the gate reproduces time scaling in loss and gate coldness ( , ). (a) The loss follows scaling. (b) The gate coldness grows as . Dashed reference lines have slope (panel a) and (panel b).
Figure 11: Later-layer attention heads and the LM head are in the low-temperature regime for Pythia-410M (a) and Pythia-160M (b). Same format as Figure 4 a: black dots show per-head entropy at the final checkpoint, red dots show layer averages, and the red diamond shows the upper bound of the LM head target entropy ( nats). Dashed lines mark the entropy thresholds for attention heads (left) and the LM head (right). The pattern is consistent with Pythia-12B in the main text.
Figure 12: Energy gap density distributions for all attention heads across all 36 layers of Pythia-12B (final checkpoint), extending Figure 4 b. Each panel corresponds to one layer; curves are colored by head index. For most heads across all layers, the density is non-zero at , confirming the theoretical condition broadly. Heads with density approaching zero at small are in the high-temperature regime. Analysis follows Appendix C.1.2 .
Figure 13: Attention logit coldness versus dynamic time for all attention heads across all 36 layers of Pythia-12B, extending Figure 4 d. Each panel corresponds to one layer; curves are colored by head index. The dashed red line has slope for reference. Low-temperature heads broadly track the slope at late training, while high-temperature heads grow more slowly or saturate. Analysis follows Appendix C.1.2 .
Figure 14: Energy gap density (left column) and attention logit coldness versus (right column) for Pythia-410M (a–b) and Pythia-160M (c–d), replicating Figure 4 b and d. In each energy gap panel, black shows a representative attention head and gray shows the LM head; both have non-zero density at . In each coldness panel, circles show an attention head and diamonds show the LM head; the dashed line has slope for reference. Analysis follows Appendix C.1.2 .
Figure 15: The raw cross-entropy loss and the fitting procedure for all eight Pythia model sizes (70M to 12B), extending Figure 4 c. Left: raw loss versus with power-law-plus-constant fits overlaid (solid lines). The loss converges to a strictly positive floor that decreases with model size, requiring the constant to be included in the fit. Right: training-dependent component after subtracting the fitted irreducible loss. Curves from all model sizes collapse and follow a single power law (dashed reference line). Analysis follows Appendix C.1.2 .
Figure 16: Replication of the Figure 4 analysis for OLMo-2, extending the results to a different model family. (a) Per-head entropy of OLMo-2-13B at the final checkpoint. The red diamond shows the upper bound of the LM head target entropy. Although the mean entropy decreases in later layers, later-layer heads fall at most slightly below the entropy threshold (horizontal dashed line), and the heads deepest in the low-temperature regime appear in early layers. The layer-1 head highlighted by a blue square is in the low-temperature regime. (b) Energy gap density for the highlighted layer-1 head; the density is non-zero at . (c) Loss versus dynamic time for OLMo-2-1B, 7B, and 13B, with fitted exponents , , and . (d) Attention logit coldness of the highlighted layer-1 head grows as (dashed reference line), while the LM head saturates. Analysis follows Appendix C.1.2 .