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

Peter Balogh

cs.LGMar 3, 2026

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

Authors: Peter Balogh

Organizations: Peter Balogh

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.

Figures & tables

Name	Architecture	Parameters
linear	$\mathrm{sigmoid}(w^{T}x+b)$	$d+1$
b=1	$\mathrm{sigmoid}(v^{T}\cdot\mathrm{ReLU}(P_{1}^{T}x)+c)$	$\sim 2d$
b=3	PCA to 6 dims + logistic	$\sim 6d$
b=6	PCA to 12 dims + logistic	$\sim 12d$

Table 1: Gate architectures evaluated.

Model	Family	Params	Layers	$d$	Layers tested
GPT-2 Medium	GPT-2	345M	24	1024	1–23 (all non-embedding)
GPT-2 Large	GPT-2	774M	36	1280	0–35 (all)
Pythia-160M	GPT-NeoX	162M	12	768	0,2,5,8,11
Pythia-410M	GPT-NeoX	405M	24	1024	0,2,6,12,18,23
Pythia-1B	GPT-NeoX	1,012M	16	2048	0,2,6,8,10,15
Pythia-2.8B	GPT-NeoX	2,775M	32	2560	0–31 (all)

Table 2: Models evaluated.

Model	Params	Cheapest	$\Delta\%$	Most Expensive	$\Delta\%$
GPT-2 Medium	345M	L8	$+1.6\%$	L23	$+12.3\%$
GPT-2 Large	774M	L5	$+0.38\%$	L31	$+3.7\%$
Pythia-160M	162M	L5	$+8.6\%$	L2	$+117.9\%$
Pythia-410M	405M	L12	$+3.1\%$	L6	$+28.6\%$
Pythia-1B	1,012M	L6	$+7.4\%$	L2	$+31.7\%$
Pythia-2.8B	2,775M	L10	$+1.9\%$	L0	$+513.2\%$

Table 3: All-linear MLP replacement cost. Each row shows the cheapest and most expensive layers to linearize.

Figure 1: All-linear replacement cost vs. best-gate cost by layer (GPT-2 Medium). Red circles mark layers where gating improves perplexity over baseline. The gate reduces cost dramatically at every layer, and the middle layers (L8–L18) are nearly free to linearize.

Layer	Full	Token Only	Context Only	Token Type	Random
L1	0.536	0.533	0.538	0.545	0.500
L12	0.579	0.568	0.575	0.579	0.500
L23	0.609	0.567	0.605	0.607	0.500

Table 4: Gate AUC by input component (GPT-2 Medium, 50K tokens). Context alone matches the full gate; token identity adds nothing.

Figure 2: Mean MLP contribution ( $\delta$ ) by token type and layer depth (GPT-2 Medium). All categories show increasing nonlinearity demand at later layers, but within-category variance (not shown) is large, confirming that token type is a weak predictor.

Test Corpus	Found	Still No-Fly	Flipped ( $\delta<0$ )	Correlation
WikiText-103 train	149 / 221	47 (31.5%)	40 (26.8%)	$r=0.041$
LAMBADA	100 / 221	22 (22.0%)	—	$r=-0.062$

Table 5: Cross-corpus stability of No-Fly lists (L12, GPT-2 Medium). Lists built on WikiText-103 test do not transfer.

Layer	PCA Var	Cluster AUC	Clean Clusters	Ambiguous
L1	61%	0.519	0 / 20	100%
L12	37%	0.570	0 / 20	100%
L23	56%	0.611	0 / 20	100%

Table 6: Residual stream clustering results (GPT-2 Medium, 15K tokens). No cluster cleanly separates linearizable from non-linearizable tokens.

Metric	Value
Total gate parameters	104,659 (0.03% of model)
Average linear routing	39.9% of positions
Layers beating baseline	4 / 23
MLP FLOPs saved	${\sim}35\%$
Total forward-pass FLOPs saved	${\sim}21\%$
Best single layer	L9: $-0.10\%$ at 35% linear

Table 7: GPT-2 Medium compound gating summary.

Model	Params	Base PPL	Layers	Best ( $\Delta\%$ )	Gate	Avg % Lin
GPT-2 Med	345M	34.8	23	L9 ( $-0.10\%$ )	linear	39.9%
GPT-2 Large	774M	30.6	36	L5 ( $-0.25\%$ )	b=1	40.9%
Pythia-160M	162M	45.1	5	L0 ( $+1.57\%$ )	b=1	42.4%
Pythia-410M	405M	29.5	6	L12 ( $+1.25\%$ )	linear	43.5%
Pythia-1B	1,012M	23.2	6	L0 ( $+0.13\%$ )	b=1	26.9%
Pythia-2.8B	2,775M	17.7	32	L3 ( $-0.13\%$ )	b=1	37.1%

Table 8: Cross-model scale validation. GPT-2 models linearize cheaply; Pythia models show higher costs but improve at scale. Base PPL is evaluated on the gating evaluation split (tokens 30K–80K of WikiText-103 test); progressive linearization (§ 4.9 ) and fine-tuning experiments use different splits and sequence lengths, producing different baseline values.

Figure 3: Median all-linear perplexity cost vs. model size. Within Pythia, linearization cost decreases monotonically with scale. Both GPT-2 models (Medium with all 23 layers, Large with all 36 layers) are substantially cheaper than Pythia at comparable size, reflecting the architectural divide. Pythia-2.8B uses all 32 layers; other Pythia models use sampled layers.

Model	Params	All-Linear $\Delta\%$	Best Gate $\Delta\%$	% Linear
GPT-2 Large	774M	$+1.77\%$	$+0.43\%$ (b=1)	53.5%
Pythia-160M	162M	$+16.9\%$	$+1.57\%$ (b=1)	18.0%
Pythia-410M	405M	$+9.5\%$	$+1.29\%$ (b=3)	34.9%
Pythia-1B	1,012M	$+27.1\%$	$+0.13\%$ (b=1)	8.2%
Pythia-2.8B	2,775M	$+513.2\%$	—	—

Table 9: Layer 0 all-linear damage across scale. The Pythia scaling threshold does not replicate in GPT-2.

# Linearized	Layer Added	PPL (after FT)	$\Delta\%$	Note
0	(baseline)	27.17	—
1	L12	24.79	$-8.8\%$	$\star$ beats baseline
2	+L11	24.65	$-9.3\%$	$\star$ beats baseline
3	+L13	25.13	$-7.5\%$	$\star$ beats baseline
4	+L10	25.97	$-4.4\%$	$\star$ beats baseline
5	+L14	27.17	$+0.0\%$	$=$ matches baseline

Table 10: Progressive linearization of GPT-2 Medium. Layers are linearized in center-outward order with 50 fine-tuning steps after each. The first 4–5 layers can be linearized at zero cost or with improved perplexity.

Step	WikiText PPL	$\Delta\%$	LAMBADA $\Delta\%$
0 (baseline)	22.98	—	—
0 (post-lin)	27.09	$+17.9\%$	$+36.7\%$
200	22.76	$-1.0\%$	$+23.7\%$
400	22.27	$-3.1\%$	$+39.2\%$
800	21.89	$-4.8\%$	$+43.3\%$
1000	21.32	$-7.2\%$	$+40.8\%$

Table 11: Fine-tuning with 117.9M-token corpus after linearizing layers 10–13. The improvement is monotonic and never plateaus. WikiText-103 test evaluation; LAMBADA shown for cross-corpus reference.

Explore similar work

The Discrete Charm of the MLP: Binary Routing of Continuous Signals in Transformer Feed-Forward Layers
Mar 11, 2026Peter BaloghTransformer FFNs Transformer Interpretability
Beyond Linearity in Attention Projections: The Case for Nonlinear Queries
Mar 11, 2026Marko KarbevskiTransformer Attention Decoder-Only Language Models
In Simplicity Technologies
How Linear Is a Transformer Feed-Forward Block? Per-Block Linear Recoverability Is Learned, Not Architectural
Jun 12, 2026Stuart WhippTransformer FFNs Transformer Expressivity
Independent Research

Mar 11, 2026J/K move · Enter open · S save

The Discrete Charm of the MLP: Binary Routing of Continuous Signals in Transformer Feed-Forward Layers

Peter Balogh

Transformer FFNs Transformer Interpretability Circuits in Language Models

Open paper AskPDF arXiv

We show that MLP layers in transformer language models perform binary routing of continuous signals: the decision of whether a token needs nonlinear processing is well-captured by binary neuron activations, even though the signals being routed are continuous. In GPT-2 Small (124M parameters), we find that specific neurons implement a consensus architecture -- seven "default-ON" neurons and one exception handler (N2123 in Layer 11) that are 93-98% mutually exclusive -- creating a binary routing switch. A cross-layer analysis reveals a developmental arc: early layers (L1-3) use single gateway neurons to route exceptions without consensus quorums; middle layers (L4-6) show diffuse processing with neither gateway nor consensus; and late layers (L7-11) crystallize full consensus/exception architectures with increasing quorum size (1 to 3 to 7 consensus neurons). Causal validation confirms the routing is functional: removing the MLP at consensus breakdown costs 43.3% perplexity, while at full consensus removing it costs only 10.1% -- exceeding a 4x difference. Comparing binary vs. continuous features for the routing decision confirms that binarization loses essentially no information (79.2% vs. 78.8% accuracy), while continuous activations carry additional magnitude information (R^2 = 0.36 vs. 0.22). This binary routing structure explains why smooth polynomial approximation fails: cross-validated polynomial fits (degrees 2-7) never exceed R^2 = 0.06 for highly nonlinear layers. We propose that the well-established piecewise-affine characterization of deep networks can be complemented by a routing characterization: along the natural data manifold, the piecewise boundaries implement binary decisions about which tokens need nonlinear processing, routing continuous signals through qualitatively different computational paths.

Abstract

Explore similar work

The Discrete Charm of the MLP: Binary Routing of Continuous Signals in Transformer Feed-Forward Layers

Beyond Linearity in Attention Projections: The Case for Nonlinear Queries

How Linear Is a Transformer Feed-Forward Block? Per-Block Linear Recoverability Is Learned, Not Architectural