Beyond Linearity in Attention Projections: The Case for Nonlinear Queries

Marko Karbevski

cs.LGMar 11, 2026

Beyond Linearity in Attention Projections: The Case for Nonlinear Queries

Authors: Marko Karbevski

Organizations: In Simplicity Technologies

Abstract

Recent algebraic analysis shows that in decoder-only and encoder-only transformers, the Query projection $W_Q$ may be set to identity without noticeable performance deterioration. This is possible because attention depends on $X$ only through the products $XW_Q, XW_K, XW_V$ , allowing basis transformations to be absorbed by adjacent layers and propagated through the network. We replace $W_Q \in \R^{d \times d}$ with a nonlinear residual of the form $Q(X) = X + f_θ(X)$ , where $f_θ$ is a bottleneck MLP with $d^2 + O(d)$ parameters. The identity term anchors the nonlinearity to a known-good prior. Experiments on GPT-3 small style models show consistent improvement over the baseline ( $2.40\%$ lower validation log-loss, $6.81\%$ lower perplexity), comfortably outperforming a model with 12.5% more non-embedding parameters. These results motivate investigation at larger scales and across modalities.

Figures & tables

Config	$\Delta$ N.E. Params	$\lambda$ (WD)	$\eta_{\max}$	$\eta_{\min}$	Val loss	$\Delta$ LL	$\Delta$ PPX
Baseline	ref.	$0.10$	$6{\cdot}10^{-4}$	$6{\cdot}10^{-5}$	$2.9441$	ref.	ref.
MLP 4.75	$+12.5\%$	$0.10$	$6{\cdot}10^{-4}$	$6{\cdot}10^{-5}$	$2.9165$	$+0.94\%$	$+2.72\%$
MLP 4.75 (scaled)	$+12.5\%$	$0.10$	$5.66{\cdot}10^{-4}$	$5.66{\cdot}10^{-5}$	$2.9173$	$+0.91\%$	$+2.65\%$
Res. GELU	${\sim}0.03\%$	$0.05$	$2.4{\cdot}10^{-3}$	$3{\cdot}10^{-5}$	$2.9015$	$+1.45\%$	$+4.17\%$
Res. GELU	${\sim}0.03\%$	$0.10$	$4.0{\cdot}10^{-3}$	$1{\cdot}10^{-5}$	$2.8750$	$+2.35\%$	$+6.68\%$
Res. GELU (best)	${\sim}0.03\%$	$0.10$	$3.6{\cdot}10^{-3}$	$2{\cdot}10^{-5}$	$\mathbf{2.8736}$	$\mathbf{+2.40\%}$	$\mathbf{+6.81\%}$

Table 1: Configurations and validation loss at step 60k. MLP w denotes the baseline with MLP hidden dimension

w{\cdot}d

. The

\Delta

N.E. Params column (N.E. = non-embedding) reports the change in parameter count relative to baseline; the bottleneck MLP

f_{\theta}

is sized to match

W_{Q}

d^{2}

parameters, leaving only a small LayerNorm overhead. The

\Delta

LL and

\Delta

PPX columns report relative improvement over baseline:

\Delta\text{LL}=(\mathcal{L}_{\text{base}}-\mathcal{L})/\mathcal{L}_{\text{base}}

and

\Delta\text{PPX}=1-\exp(\mathcal{L}-\mathcal{L}_{\text{base}})

Figure 1: Training dynamics (steps 5k–60k). Solid: validation; dashed: training curve for the best configuration. Top: absolute validation log-loss for baseline, both MLP 4.75 variants, and the best Res. GELU. Bottom: relative log-loss improvement over baseline for all five non-baseline configurations.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Relative improvement over baseline (steps 1k to 59k). Nonlinear configurations: 84.97M parameters; MLP-widened controls: 89.6 to 95.6M.

Explore similar work

Half the Nonlinearity Is Wasted: Measuring and Reallocating the Transformer's MLP Budget
Mar 3, 2026Peter BaloghTransformer FFNs Conditional Computation
Peter Balogh
Position-Agnostic Pre-Projection for Transformer Attention: Nonlinear Feature Construction and Content Skip Before Q/K/V
Apr 12, 2026Chirag ShindeTransformer Attention Attention Residual Connections
Independent Researcher
A Controlled Study of Attention-Only Transformers
Jul 20, 2026Henry Ndubuaku, Karen Mosoyan, Jakub Mroz +5Transformer FFNs Transformer
Cactus Compute, Inc.

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

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

Peter Balogh

Transformer FFNs Conditional Computation

Open paper AskPDF 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.

Abstract

Explore similar work

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

Position-Agnostic Pre-Projection for Transformer Attention: Nonlinear Feature Construction and Content Skip Before Q/K/V

A Controlled Study of Attention-Only Transformers