Nonlinear attention improves learning speed and depth gains in models
Emergent Inverse-Depth Scaling From Nonlinearity In Attention
Machine LearningArtificial Intelligence
Summary
Large language models get better as they get bigger, but it’s unclear why. The authors show that a special kind of attention in these models, called nonlinear attention, helps the model focus on important pieces of information. This lets the model learn both easy and hard parts of data at the same time, improving learning speed across layers. This discovery connects deep model improvements to how attention works, rather than just looking at the whole data at once.
What this means in practice
- •For machine learning engineers: Design deep learning models that use nonlinear attention to improve training speed and depth scaling efficiency.
- •For ai hardware architects: Optimize hardware for nonlinear attention mechanisms to enhance model performance and resource use based on improved scaling insights.
A theory result. No direct application yet.
Authors
Zirui Peng, Yizhou Liu, Ziming Liu, Jeff Gore
Abstract
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.