Emergent Inverse-Depth Scaling From Nonlinearity In Attention
Authors: Zirui Peng, Yizhou Liu, Ziming Liu, Jeff Gore
Organizations: Department of Physics, Peking University · Department of Physics, Massachusetts Institute of Technology, Cambridge, MA · College of AI, Tsinghua University
Scaling laws describe power-law improvements in model performance with dataset size and parameter count, yet their underlying mechanisms are not fully understood. To explain the parameter count scaling, existing theory posits power-law scaling with model depth. In linear-attention models, this scaling is tied to a power-law data spectrum: unable to selectively attend to relevant tokens, these models learn according to global spectral strength, with stronger directions learned before weaker ones. Large language models, however, can be strongly nonlinear. Here, we show that nonlinear attention yields inverse-depth decay of loss across all tested data spectra. Nonlinearity enables attention to focus selectively on relevant tokens, allowing strong and weak spectral directions to be learned in parallel. Similar focusing across layers motivates a connection to the central limit theorem: shared error across layers sets the loss plateau, while aggregation turns layer-specific differences into continued gains with depth. Our findings suggest that depth scaling may arise from nonlinearity in attention, which allows large language models to focus locally and may make the global covariance structure less relevant.
Figures & tables
Figure 1: The nonlinear-attention model can yield inverse-depth scaling across all tested data spectra. This overview compares depth-dependent performance across models and tasks. Colors distinguish data spectra, and all panels are plotted on a log-log scale.
Figure 2: Both tasks model the second half of an induction head: the easy task copies the most relevant token’s label, while the difficult task averages the labels of the top- k relevant tokens.
Figure 3: On the easy task, the linear-attention model exhibits spectrum-dependent depth scaling and learns spectral modes sequentially. (a) Varying b at fixed a=2 . (b) Varying a at fixed b=2 . Dashed curves show power-law fits, with the corresponding exponents indicated. (c) Normalized loss of representative spectral modes. Values below 2% are displayed at the threshold marked by the dashed line.
Figure 4: A single nonlinear-attention layer solves the easy task by focusing on the matching token. (a) Varying a at fixed b=0 . (b) Varying b at fixed a=0 . (c) Top- 1 attention mass approaches one during training.
Figure 5: On the difficult task, linear attention rapidly reaches a spectrum-dependent loss plateau with little benefit from further depth. (a) Varying b at fixed a=0 . (b) Varying a at fixed b=0 .
Figure 6: On the difficult task, nonlinear attention yields inverse-depth decay across all tested data spectra. (a) Varying b at fixed a=0 . (b) Varying a at fixed b=0 . (c,d) CL and AL−CL converge to spectrum-dependent constants.
Figure 7: With nonlinear attention, strong and weak spectral modes are learned in parallel. (a) Different modes are learned at similar rates. (b) Normalized mode-wise reducible losses nearly collapse onto a common curve, indicating parallel learning across modes.
(a,b)
(0,0)
(0,0.25)
(0,0.5)
(0,1)
(0.25,0)
(0.5,0)
(1,0)
Nonlinear L0
2.007
1.268
0.836
0.424
1.255
0.798
0.341
Linear L0
59.853
23.055
9.766
2.504
22.962
9.501
2.120
Table 1: Nonlinear-attention models achieve substantially lower irreducible loss than linear-attention models on the difficult task.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
(a,b)
Constrained ζ
Measured α
Free 2ζ
(2,0.5)
0.750
1.121
1.500
(2,1)
1.000
1.485
2.000
(2,1.5)
1.250
1.684
2.500
(2,2)
1.500
1.987
3.000
(0.5,2)
3.000
3.310
6.000
(1,2)
2.000
2.552
4.000
Appendix
Table 2: Measured depth-scaling exponents of the linear-attention model on the easy task lie between two reference values.
(a,b)
Theory-predicted loss
Measured loss
Difference
k=2
(0,0)
59.3714±0.0213
59.8502
0.80%
(0.25,0.25)
9.7004±0.0049
9.7387
0.39%
(0.5,0.5)
2.2715±0.0061
2.3274
2.40%
(1,0.25)
1.2246±0.0106
1.2259
0.11%
(1,1)
0.4975±0.0102
0.5036
1.22%
Appendix
Table 3: The theoretically predicted irreducible loss closely matches the actual loss plateau of the linear-attention model.
Figure 8: The linear-attention model captures only the average predictable component of the difficult-task target. The green target is the average of the two selected inputs. Blue shows the best predictable component OC ; red shows the context-specific residual.
Figure 9: The single-stream nonlinear-attention model exhibits a similar inverse-depth scaling signal. (a) The model is evaluated under isotropic data spectra. (b,c) The shared-error term CL and disagreement term AL−CL both approach finite constants with depth.
Figure 10: The nonlinear-attention model with layer-specific value matrices exhibits a similar inverse-depth scaling signal. (a) The model is evaluated under isotropic data spectra. (b,c) The shared-error term CL and disagreement term AL−CL both approach finite constants with depth.
Figure 11: Enabling context updates preserves an inverse-depth scaling signal on the top- 2 task but not on the top- 16 task. Under isotropic data spectra, (a) the top- 2 loss exhibits an inverse-depth scaling signal, whereas (b) the top- 16 loss quickly saturates. (c,d) The learned context-update scalar m varies with depth for top- 2 but stabilizes for top- 16 .
Figure 12: The learned context-attention structure depends strongly on the aggregation target. The context-attention matrix A1 is (a) sharply concentrated on the diagonal for top- 2 task, but (b) distributed broadly across context tokens for top- 16 task.
Figure 13: Greater depth makes token selection more reliable across contexts. (a) Selection-loss percentiles for L=1 and L=256 . (b) Their difference ΔQ(q)=Q1(q)−Q256(q) ; positive values favor L=256 .
Figure 14: Emergent inverse-depth scaling in the deep regime generalizes across broader targets and data spectra. (a) Performance on the top- 2 task with both spectral parameters nonzero. (b) Performance on top- 16 aggregation across representative spectra. (c) CL and (d) AL−CL converge to finite constants with depth for the three top- 2 spectra.
Target
(a,b)
L0
α±SE
R2
Top- 2
(0.25,0.25)
0.818068
0.97±0.01
0.9998
(0.25,0.5)
0.559694
0.93±0.06
0.9984
(0.5,0.5)
0.382188
0.97±0.01
0.9999
Top- 16
(0,0.25)
0.173237
0.79±0.02
0.9999
(0,0.5)
0.114436
0.88±0.03
0.9998
(0.25,0.25)
0.115248
0.81±0.02
0.9998
Appendix
Table 4: Fitted depth-scaling exponents remain close to one across broader targets and data spectra.
(a,b)
α±SE
R2
(0,0)
0.980±0.007
0.99981
(0,0.25)
0.983±0.038
0.99396
(0,0.5)
0.977±0.014
0.99923
(0,1)
0.982±0.011
0.99946
(0.25,0)
0.990±0.015
0.99904
(0.5,0)
0.986±0.005
0.99988
Appendix
Table 5: The gap L0−CL decays approximately inversely with depth over the fitted range.
Neural scaling laws describe how pre-training loss decays as power laws with training time, model size, and compute. This position paper argues that the exponents of these power laws are fixed by generic mechanisms: a one-third time scaling due to the strong nonlinearity of Softmax, an inverse width scaling due to representational superposition, and an inverse depth scaling due to ensemble averaging of Transformer layers. These mechanisms are robust to a wide range of data structures and architectural details, placing current large language models in a universality class with fixed exponents. The coefficients, however, are expected to be sensitive to data and architecture details, and directly determine practical quantities such as the optimal model shape and the compute-optimal frontier. We therefore argue that understanding the coefficients is the key to near-term performance improvements, and that a closer examination of the current universality class may reveal pathways to better universality classes.
Yizhou Liu, Jeff Gore
Massachusetts Institute of Technology Cambridge, MA 02139
Learning-rate transfer can reduce the cost of training large language models: instead of sweeping learning rates at target scale, practitioners extrapolate from smaller runs. Existing approaches often assume that the optimal learning rate follows a log-linear scaling law in data scale and model size. We carefully examine and evaluate this scaling law. In our empirical study of GPT-2--style models from 22M to 707M parameters trained on 5B to 100B tokens, the optimal learning rate develops upward curvature at larger scales, leading to inaccurate extrapolation. We find that this curvature largely disappears when learning rates are replaced by effective learning rate (the step size in normalized weight space), and when data D extrapolation is used instead of model size N extrapolation. Next, we explain nonlinearity in scaling: weight-norm converges to equilibrium slower when optimal learning is small, requiring a larger step size to reduce the transient phase. Experiments with AdamH, which directly controls the effective learning rate, further support this explanation.
Zaiwen Yang, Huaqing Zhang, Jing Xu +1
Tsinghua University · Tsinghua University, Shanghai Qi Zhi Institute
Linear attention is a tractable model for understanding the mechanisms governing in-context learning in transformers. For linear regression tasks, recent asymptotic analyses have characterised its learning and generalisation behaviour. We extend this theory to nonlinear single-index targets, y=f(x⊤w)+ε. Our main result establishes a nonlinearity-noise equivalence: linear attention extracts only the linear Hermite component of f, while the remaining nonlinear structure contributes to the generalisation error as effective noise. This reduction allows results from the corresponding linear theory to be transferred to nonlinear tasks. We illustrate its implications for finite pretraining data and for the transition from task memorisation to task generalisation as task diversity increases. These results identify a limitation of the reduced linear-attention model and provide a tractable starting point for studying nonlinear in-context learning.
Mary Letey, Arman Rysmakhanov, Yue M. Lu +2
Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University