Beyond Linearity in Attention Projections: The Case for Nonlinear Queries

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_{\theta }(X)$, where $f_{\theta }$ 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\mathrm{\%}$ lower validation log-loss, $6.81\mathrm{\%}$ lower perplexity), comfortably outperforming a model with 12.5% more non-embedding parameters. These results motivate investigation at larger scales and across modalities.