Beyond Query Linearity: The Case for Nonlinear Queries

In self-attention transformers the query projection WQ can be set to the identity without noticeable performance deterioration, because attention depends on X only through the products XWQ, XWK, XWV, so a change of basis can be approximately absorbed by adjacent layers and propagated through the network. A learned linear query is therefore a choice of coordinates rather than of function, and capacity spent there is at least in part redundant. Adding capacity to the query pathway means breaking that symmetry, which forces nonlinearity. We replace WQ in R^(d x d) with a nonlinear residual Q(X) = (1/2)(X + ftheta(X)), where f_theta is a bottleneck MLP with d^2 + O(d) parameters. Empirically the modification lowers validation log-loss by 2.40% at parameter parity, and the gains transfer within the scales tested. The purpose of this paper is to motivate testing it across modalities and scales, to report the experiments we have under way, and to further that discussion with the community.