Half the Nonlinearity Is Wasted: Measuring and Reallocating the Transformer's MLP Budget

2026-03-03

AI Summary

This work investigates whether the nonlinear computation in the MLP modules of Transformers is redundant, revealing that such nonlinearity is highly context-dependent rather than token-specific. To address this, the authors propose a lightweight gating mechanism—requiring only \(d+1\) parameters—to dynamically determine whether to replace the MLP with a linear mapping. Combined with a two-stage training strategy and evaluated across models and corpora, the approach demonstrates that 25–56% of MLP layers in GPT-2 can be linearized with less than 1% perplexity degradation. Remarkably, under certain configurations, the method even improves performance by up to 17.3%, outperforming standard fine-tuning and achieving, for the first time, efficient and low-loss large-scale MLP linearization.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Transformer
• MLP
• nonlinearity
• linear surrogate
• routing

Innovation

• linear routing
• MLP sparsity
• contextual gating
• Transformer efficiency
• nonlinearity reallocation