Beyond Linearity in Attention Projections: The Case for Nonlinear Queries

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 $X{W}_Q,X{W}_K,X{W}_V$, allowing basis transformations to be absorbed by adjacent layers and propagated through the network. We replace $W_Q\in {\mathbb{R}}^{d\times d}$ with a nonlinear residual of the form $Q(X)=X+f_{\hat{I}}¸(X)$, where $f_{\hat{I}}¸$ 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.

Accepted at the ICLR 2026 GRaM workshop: https://openreview.net/forum?id=pwdnneFiNZ#discussion