WK, WV is (Linearly) All You Need: On the Necessity of the QKV Weight Triplet in Self-Attention Transformers
Marko Karbevski ⋅ Antonij Mijoski
Abstract
Multi-head attention is invariant under the joint $GL(d)$ action $(X, W_Q, W_K, W_V) \mapsto (X\Theta,\Theta^{-1}W_Q,\Theta^{-1}W_K,\Theta^{-1}W_V)$ with two consequences for the QKV triplet. First, any one of $W_Q, W_K, W_V$ can be fixed to $I_d$ without loss of expressivity if $X W$ is precomputed; under mild structural conditions the precomputation folds into the preceding MLP (or, in the first layer, the embedding) at no parameter cost, removing $25\%$ of attention parameters per layer. We prove the multi-layer reduction and analyse the normalisation obstructions. Second, every linear $W_Q$ already lies on the orbit of $I_d$, so a learned linear query is redundant: expressive gains in the QKV pathway require at least one of the three to be nonlinear, a branch we realise with the residual query $Q(X)=\frac{1}{2}(X + f_\theta(X))$ at parity of parameters. This research also led us to examine residual skip connections: MLPs with and without a skip form generically disjoint function classes for modern activations. We validate both halves on GPT-style models trained from scratch under batch-matched comparisons: the reduced $117$M model matches the dense $124$M baseline; reallocating the saved parameters to the feed-forward sublayer strictly improves on it; and the residual nonlinear query outperforms a baseline with a wider feed-forward sublayer carrying $12.5$ % more non-embedding parameters.
Chat is not available.
Successful Page Load