Source-linked AI summary
Emergent Inverse-Depth Scaling From Nonlinearity In Attention
Zirui Peng, Yizhou Liu, Ziming Liu, Jeff Gore
TL;DR
Scaling laws predict performance improvements with model size, but their mechanisms remain incompletely understood, including how attention nonlinearity affects depth scaling. The paper compares linear and nonlinear attention in controlled in-context-learning tasks and finds inverse-depth loss improvement beyond power-law data spectra. It links the gains to selective token focusing and aggregation across layers, while noting that the mechanism's broader applicability remains untested.
Problem
How attention nonlinearity affects depth scaling and its mechanisms remains an open question, while existing linear-attention theory ties power-law depth scaling to power-law data spectra.
Source: §1 · p. 2Method
The authors compare controlled Transformers with linear or nonlinear attention on in-context-learning tasks across independently varied data spectra.
Source: §2.1 · p. 3Results
Takeaways & Limitations
The findings suggest practical depth scaling can originate from attention nonlinearity rather than data structure alone.
Source: §6 · p. 9Experiments are restricted to in-context-learning tasks, leaving generalization to real-world datasets and broader task classes untested.
Source: §6 · p. 10
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper asks how attention nonlinearity affects depth scaling and its mechanisms, then studies the question in controlled in-context-learning tasks. It finds that nonlinear attention supports inverse-depth loss scaling across tested spectra, unlike the spectrum-dependent behavior of linear attention.
1 INTRODUCTION
- Linear-attention theory ties power-law depth scaling to power-law data spectra, while prior empirical inverse-depth accounts leave layerwise similarity unexplained.The paper asks whether attention nonlinearity accounts for depth scaling and its mechanism.Source: §1 · p. 2
- The study compares controlled Transformers with linear or nonlinear attention across independently varied data spectra in in-context learning.The tasks model retrieval of the query-matching token and copying its associated value.Source: §1 · p. 2
- Relaxed architectural assumptions retain similar inverse-depth signals, suggesting attention nonlinearity may contribute to depth scaling in real models and reduce dependence on global covariance structure.The paper frames these implications as suggestions, not universal claims.Source: §1 · p. 2
- Nonlinear attention enables inverse-depth loss scaling across all tested spectra, with strong and weak spectral directions learned in parallel.The paper connects depth gains to a loss decomposition and the central limit theorem.Source: §1 · p. 2
2 PROBLEM SETUP
The setup uses controlled in-context-learning tasks to compare retrieval and aggregation under varied data spectra. A simplified Transformer architecture isolates how linear and nonlinear attention process context inputs and labels.
2.1 TASK SETUP
- The controlled in-context-learning setting uses context-token interactions to study the origins of depth scaling.Each context contains inputs with labels generated by a teacher vector.Source: §2.1 · p. 3
- Input and teacher covariance spectra are varied independently, and context-specific rotations require the model to infer relevant directions from each context.The input and teacher eigenvalue spectra follow power laws with independently controlled exponents; zero exponents yield isotropic covariance.Source: §2.1 · p. 3
- The easy task copies the label of the closest context token, while the difficult task averages labels from the top-k tokens ranked by cosine similarity.The tasks model top-1 copying and top-k aggregation in the second half of an induction-head computation.Source: §2.1 · pp. 3–4, Figure 2 · p. 3
2.2 MODEL SETUP
- A simplified architecture retains context-dependent retrieval and label propagation to enable controlled comparison of linear and nonlinear attention.Input features determine token relevance, while attention passes label information to the test token.Source: §2.2 · p. 4
- Input features remain fixed across layers while layer-specific query and key matrices select tokens and a shared value matrix propagates labels.The setup also specifies attention masks, a linear readout, and mean-squared-error training over independently sampled contexts.Source: §2.2 · pp. 4–5
3 EASY TASK ANALYSIS
Linear attention’s easy-task loss follows spectrum-dependent power-law scaling, as stronger modes are learned before weaker ones. Nonlinear attention solves the easy task in one layer across tested spectra by selectively focusing on relevant tokens.
3.1 LINEAR-ATTENTION MODEL
- Linear-attention loss decreases with depth by a spectrum-dependent power law on the easy task.The exponent depends on the data spectra.Source: §3.1 · p. 5
- Greater depth lets linear attention use higher powers of the data covariance to better approximate the identity, producing spectrum-dependent convergence rates.Different spectra change the covariance in the polynomial approximation.Source: §3.1 · p. 6
- Linear attention learns stronger spectral modes within a few layers, while weaker modes require substantially greater depth.The combined mode dynamics yield spectrum-dependent power-law scaling.Source: §3.1 · p. 6
3.2 NONLINEAR-ATTENTION MODEL
- Across tested spectra, a single nonlinear-attention layer achieves zero loss on the easy task.Additional depth does not improve performance on this task.Source: §3.2 · p. 6
- Nonlinear attention focuses on the relevant token, assigning it nearly all attention mass and enabling its label to be copied in one layer.During training, query and key norms grow, amplifying logit differences and placing softmax in an effectively low-temperature regime.Source: §3.2 · p. 6
4 DIFFICULT TASK ANALYSIS
On the difficult task, linear attention quickly plateaus, whereas nonlinear attention continues improving with depth at approximately inverse-depth scaling. The nonlinear model learns spectral modes in parallel, and its depth gains are consistent with ensembling that cancels layer-specific error while leaving shared error as a floor.
4.1 LINEAR-ATTENTION MODEL
- Linear attention reaches a spectrum-dependent loss plateau within one or two layers because it cannot fully recover context-specific deviations in the target.Its leading predictable target is representable by a degree-one polynomial, leaving little further loss reduction from added depth.Source: §4.1 · p. 7
4.2 NONLINEAR-ATTENTION MODEL
- Nonlinear attention’s fitted loss exponents are close to one across tested spectra, indicating inverse-depth scaling as depth increases.Additional depth continues to reduce loss, unlike the rapid saturation of linear attention.Source: §4.2 · pp. 7–8
- Shared error across predictors sets the ensemble’s irreducible loss floor, while predictor-specific error can partially cancel through averaging and fall by a factor of 1/L.For the top-2 task, shared error reflects the mismatch between smooth softmax attention and the sharp target distribution.Source: §4.2 · p. 8
- Nonlinear attention learns strong and weak spectral modes at similar rates, with normalized mode-wise losses nearly collapsing onto a common curve.It can focus on relevant tokens without context updates that progressively expose finer-grained information.Source: §4.2 · p. 8, Figure 7 · p. 8
- On the difficult task, nonlinear attention achieves a substantially lower irreducible loss than linear attention.Source: §4.2 · p. 8, Table 1 · p. 9
- After adding shared input-label representations, untied value matrices, and context updates, the setup still shows similar signs of inverse-depth scaling.The authors suggest the mechanism may contribute to scaling in realistic nonlinear-attention models.Source: §4.2 · p. 9
5 RELATED WORK
Prior accounts explain depth through compositional learning, dynamical systems, or ensemble averaging, but do not fully explain the origins of depth scaling. This paper examines how attention nonlinearity shapes depth-scaling mechanisms and laws.
5 RELATED WORK
- Prior scaling work relates loss to model and dataset size, but much of the literature treats model size in aggregate rather than separating width and depth.Theoretical accounts include data geometry and covariance spectra, representation superposition, and learning dynamics of peaked distributions.Source: §5 · p. 9
- Compositional learning, dynamical-systems views, and ensemble averaging explain ways networks use depth, but not fully the origins of depth scaling.These perspectives respectively emphasize hierarchical feature combination, residual layers as ODE steps, and aggregation of related predictors.Source: §5 · p. 9
- Bordelon et al. (2026) derive spectrum-dependent depth scaling for deep linear attention, while mechanisms by which nonlinear attention affects scaling remain unresolved.This work aims to explain how nonlinear attention enables ensemble-like behavior across layers and inverse-depth scaling.Source: §5 · p. 9
6 DISCUSSION
Nonlinear attention is presented as a source of depth scaling beyond power-law spectra, with selective token focus and cross-layer aggregation explaining parallel feature learning and continued gains.
6 DISCUSSION
- Nonlinear attention is identified as a source of depth scaling beyond power-law spectra: selective token focus supports parallel learning, while layer aggregation reduces layer-specific error around a shared loss floor.Source: §6 · p. 9
- The discussion cautions that simplified architecture and ICL-only experiments leave full-LLM mechanisms and broader-task generalization unestablished; the mechanism may not be universal or dominant.Source: §6 · pp. 9–10
- Parameter reuse is raised as a possible efficiency route: looped Transformers may preserve aggregation benefits with fewer parameters, but this remains to be tested.Source: §6 · p. 10
A.1 COVARIANCE POLYNOMIALS IN THE LINEAR-ATTENTION MODEL
- For linear attention on the easy task, the output is an ordered matrix polynomial in data covariance of degree at most L; the isotropic analysis is a special case.Source: §A.1 · p. 13
A.2 CONSTRAINED AND FREE POLYNOMIAL DEPTH SCALING
- Measured depth-scaling exponents lie between repeated-update and free-polynomial references, whose constraints range from shared-layer conditions to independently optimized coefficients.Source: §A.2 · pp. 16–17, Table 2 · p. 17
- The polynomial-scaling analysis requires a population-covariance approximation S ≈ Σ (P → ∞), a dense spectrum with ζ > 0, and a sufficiently large spectral range λ1/λd.Source: §A.2 · p. 17
A.3 PREDICTING THE LINEAR-MODEL LOSS FLOOR
The linear-attention model’s predictable target component is proportional to Σq, but context-specific variation leaves an irreducible loss floor. A mode-wise heuristic shows how exponential feature convergence can aggregate into power-law depth scaling.
A.3 PREDICTING THE LINEAR-MODEL LOSS FLOOR
- The predictable component is proportional to Σq, so higher polynomial degrees cannot remove context-specific variation or the linear-attention model’s nonzero loss floor.The target’s context-dependent residual sets the irreducible component.Source: §A.3 · p. 18
- Monte Carlo predictions of the irreducible loss closely match the trained model’s actual loss plateau.Table 3 compares the prediction with the plateau.Source: §A.3 · p. 19, Table 3 · p. 19
A.4 POWER LAWS FROM MODE-WISE EXPONENTIAL DECAY
- Exponential convergence of individual features yields power-law depth scaling when their characteristic depths have a broad power-law distribution.Stronger features converge first while weaker features continue contributing at greater depth.Source: §A.4 · p. 20
B NONLINEAR-ATTENTION MODEL THEORY
The nonlinear-attention theory decomposes loss into shared error and layer disagreement, linking the plateau and depth-related gains to different components. The decomposition alone does not establish inverse-depth scaling, and its layerwise interpretation relies on a parallel controlled model.
B NONLINEAR-ATTENTION MODEL THEORY
- Shared error across layers sets the irreducible loss, while layer-specific disagreement can partially cancel through aggregation and produce reducible loss.Averaging reduces the disagreement contribution to (AL − CL)/L.Source: B NONLINEAR-ATTENTION MODEL THEORY · p. 22
- The decomposition identifies loss components but does not by itself imply inverse-depth scaling; experiments and power-law fitting establish that behavior first.The decomposition is then used to identify the components and mechanisms underlying the observed scaling.Source: B NONLINEAR-ATTENTION MODEL THEORY · p. 22
- The layerwise decomposition assumes independently parameterized predictors whose aggregate forms the output, unlike cascaded layers in shared-stream models.A mean-field treatment may be more appropriate for such generalized systems; a rigorous decomposition for complex models remains future work.Source: B NONLINEAR-ATTENTION MODEL THEORY · p. 22
C GENERALIZED MODEL EXPERIMENTS AND ANALYSIS
The generalized experiments relax controlled-model assumptions to bring the architecture closer to standard Transformers. In both shared-stream and layer-specific-value variants, approximate inverse-depth scaling remains visible beyond moderate depth.
C GENERALIZED MODEL EXPERIMENTS AND ANALYSIS
- The generalized experiments relax separate input and label streams, shared value matrices, and fixed context states one assumption at a time.These choices were used in the controlled setup to isolate retrieval and improve training stability.Source: C GENERALIZED MODEL EXPERIMENTS AND… · p. 22
C.1 ENCODING INPUT AND LABEL INFORMATION INTO ONE SPACE
- Combining token and label information in a shared, evolving state preserves an approximate inverse-depth loss trend beyond moderate depth.The shared-error and disagreement terms continue to approach finite constants.Source: Figure 9 · p. 23, §C.1 · p. 23
C.2 LAYER-DECOUPLED VALUE MATRICES
- Replacing the shared value matrix with independently parameterized layer-specific matrices preserves an approximate inverse-depth trend beyond moderate depth.The shared-error and disagreement terms still approach finite constants.Source: Figure 10 · p. 24, §C.2 · p. 24
C.3 ALLOWING CONTEXT UPDATES IN THE NONLINEAR-ATTENTION MODEL
With context updates enabled, inverse-depth scaling persists for top-2 targets but not top-16 targets. Additional experiments describe how update structure and layer aggregation relate to this task dependence, and show deeper models improve token-selection coverage across contexts.
C.3 ALLOWING CONTEXT UPDATES IN THE NONLINEAR-ATTENTION MODEL
- Top-2 loss continues decreasing with depth with an inverse-depth signal, whereas top-16 loss does not improve as depth increases.The learned update scalar varies with depth for top-2 but stays near −2.3 for top-16.Source: §C.3 · pp. 24–25, Figure 11 · p. 25
- The top-2 context-attention matrix is nearly identity, while top-16 attention is broader and flatter, reflecting their different target distributions.The sharp top-2 target favors preserving selected token identities; the flatter top-16 target favors stronger coupling across embeddings.Source: Figure 12 · p. 26, §C.3 · p. 26
- Context updates can add a uniform token-weighting direction; a negative update scalar subtracts the common attention floor and can yield low loss with few layers.The update does not introduce an independent token-weighting direction in the analyzed setting, and loss remains close to the fixed-context model with approximate inverse-depth scaling over the observed range.Source: §C.3 · p. 26
D.1 POWER-LAW FITTING PROCEDURES
- The power-law fit estimates the asymptotic floor and decay exponent, with the floor fitted jointly for the difficult task.The fitting form is RL = L0 + cL^-α, with L0 as the asymptotic loss floor.Source: §D.1 · p. 27
D.2 COMPLEMENTARY TOKEN SELECTION ACROSS CONTEXTS
- The 256-layer model markedly reduces the upper tail of token-selection loss, although the one-layer model is slightly more accurate in its lowest-error tail.The authors attribute broader coverage to aggregation of complementary selection rules across layers.Source: Figure 13 · p. 28, §D.2 · p. 28
- The one-layer model’s low-percentile advantage suggests a possible specialization–coverage tradeoff, with different layers potentially serving different context classes.The passage presents this as a testable hypothesis for finitely many or clustered context types.Source: §D.2 · p. 28
D.3 BROADER TARGETS AND DATA SPECTRA
The inverse-depth loss pattern extends to jointly structured spectra and top-16 aggregation, though top-16 exponents deviate more from one at the tested depths. For jointly nonzero top-2 spectra, decomposition terms approach spectrum-dependent finite values.
D.3 BROADER TARGETS AND DATA SPECTRA
- The extensions test jointly nonzero input and task spectra for top-2 targets and several spectral settings for flatter top-16 targets.The fits use the evaluation loss and log-space objective defined in Appendix D.1.Source: §D.3 · p. 28
- Top-16 exponents deviate relatively more from one at the tested depths, suggesting the task may reach the asymptotic regime only at greater depths.Source: §D.3 · p. 28
- Inverse-depth loss decay appears with jointly nonzero top-2 spectra and top-16 aggregation, extending beyond the main-text settings.The fitted top-2 exponents remain close to one; top-16 exponents deviate more at the tested depths.Source: §D.3 · pp. 28, 30, Figure 14 · p. 29, Table 4 · p. 29
- For jointly nonzero top-2 spectra, both decomposition terms approach spectrum-dependent finite values with depth.Source: §D.3 · p. 28, Figure 14 · p. 29
D.4 CONVERGENCE OF THE SHARED-ERROR TERM
The shared-error term’s gap from the fitted loss floor decays approximately inversely with depth over the fitted range. With the loss floor fixed from direct loss fits, the fitted exponents are close to one.
D.4 CONVERGENCE OF THE SHARED-ERROR TERM
- With L0 fixed from direct loss fits, least squares fits the shared-error term’s gap over tested depths.The fitting procedure follows the log-space method in Appendix D.1.Source: §D.4 · p. 30
- The gap L0 −CL decays approximately inversely with depth over the fitted range.The fitted decay exponents are close to one, consistent with the inverse-depth trend in the loss.Source: Table 5 · p. 30, §D.4 · p. 30
- High R2 values and small fitting uncertainties indicate that L0 −CL is well described by a power law over tested depths.Source: §D.4 · p. 30