Half the Nonlinearity Is Wasted: Measuring and Reallocating the Transformer's MLP Budget
Not on the map. Why?
Abstract (arXiv)
We investigate when transformer MLP nonlinearity is actually necessary. A gate with d+1 parameters decides when to replace the full MLP with a linear surrogate. Through systematic investigation across six models (162M-2.8B parameters), two architectures, and three corpora, we establish that nonlinearity need cannot be predicted from token identity: cross-corpus correlation is zero (r < 0.05). The routing decision is fully contextual. Despite weak per-instance predictability, the gate exploits a heavily skewed distribution where most MLP computations are near-linear, achieving 25-56% linear routing at <1% perplexity cost in GPT-2. In GPT-2 Large, 11 of 36 layers beat baseline with gating and no layer exceeds 3.7% all-linear cost. This success is architecture-dependent: Pythia models show higher costs, though Pythia-2.8B's full 32-layer sweep reveals one layer that narrowly beats baseline. As a proof of concept, we progressively replace middle-layer MLPs with frozen linear matrices: 5 of 24 layers linearize at zero cost. With a full training budget, 4 linearized layers yield a 10.2% perplexity improvement -- and a two-phase gated approach pushes this to 17.3%, beating a vanilla fine-tuning control and confirming that the nonlinear MLPs at these layers were actively harmful.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 10 of its 15 references that are in the map. First posted on arXiv 3 Mar 2026.
Builds on
10 of 15 references.
References · 10
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
- 2016Pointer Sentinel Mixture Modelscited by 2,992
- 2023Pythia: A Suite for Analyzing Large Language Models Across Training and Scalingcited by 1,734
- 2023Sparse Autoencoders Find Highly Interpretable Features in Language Modelscited by 1,359
- 2019BERT Rediscovers the Classical NLP Pipelinecited by 1,046
- 2016The LAMBADA dataset: Word prediction requiring a broad discourse contextcited by 960
- 2021MoEfication: Transformer Feed-forward Layers are Mixtures of Expertscited by 149
- 2024Neural Redshift: Random Networks are not Random Functionscited by 43
- 2024Bilinear MLPs enable weight-based mechanistic interpretabilitycited by 28
- 2025Do We Always Need the Simplicity Bias? Looking for Optimal Inductive Biases in the Wildcited by 13
Built on by
Not cited yet.