Key and Value Weights Are Probably All You Need: On the Necessity of the Query, Key, Value weight Triplet in Self-Attention Transformers

AI summary

Overview

Research area: Transformer architecture theory and efficiency — specifically, parameter redundancy inside the multi-head self-attention mechanism.

Technical level: Advanced. The paper is theory-first, relying on linear algebra (change of basis, group actions, semi-conjugacy) and functional-equation analysis, though the empirical section is readable for anyone familiar with language model pretraining.

Scope: A combined theoretical and empirical study showing that one of the three attention projection weights (W_Q, W_K, W_V) is redundant and can be replaced with the identity matrix, plus GPT-style pretraining experiments validating the removal of W_Q.

Publication details as reported: arXiv:2510.23912v7 [cs.LG], dated 23 Apr 2026, licensed CC BY 4.0, by Marko Karbevski (In Simplicity Technologies) and Antonij Mijoski (Institut de Recherche Mathématique Avancée, Université de Strasbourg). The paper states it is to appear at the ICLR 2026 Workshop on Deep Generative Models (DeLTa). Code and checkpoints are released at a linked GitHub repository.

What This Paper Is About

Standard self-attention uses three learned weight matrices per layer — Query (W_Q), Key (W_K), and Value (W_V) — plus an output projection (W_O). The authors ask whether all three are actually necessary, motivated by the observation that attention only ever sees the input through the products XW_Q, XW_K, and XW_V, never through the raw input. Under mild assumptions they prove one of the three weights can be replaced by the identity matrix, cutting attention parameters by 25%, and they test the Query-weight version of this claim by pretraining GPT-style models from scratch.

Key Contributions

1. A redundancy proof covering three settings. The authors prove in the single-head case that the four attention matrices (W_Q, W_K, W_V, W_O) can be reduced to two (Proposition 4.1); prove single-layer Query weight elimination in any transformer without normalization (Theorem 4.1); and prove all-layer Query elimination under either skip connections placed exclusively around attention (Theorem 4.2) or weight sharing across layers (Theorem 4.3). They also identify a block-diagonal invariance of multi-head attention (Proposition 4.2), a redundancy parametrized by an (h · d_k²)-dimensional manifold present even in the full model with no simplifying assumptions.

2. A structural result about skip connections and ReLU MLPs. In the mathematically tractable ReLU setting, they explicitly solve the functional equation W₂ ReLU(W₁x) + x = V₂ ReLU(V₁x) for unknown V₁, V₂ given fixed W₁, W₂, characterizing when residual connections can be absorbed into MLPs. The abstract states that skip connections push MLPs into a generically disjoint function class at fixed width.

3. Analysis of how normalization obstructs basis changes. Section 8.3 derives sufficient conditions for basis transformations to commute with layer normalization (Lemma 8.2, Theorem 8.3), and the paper notes this obstruction is strictly milder than the per-head modulation introduced by QK-normalization.

4. Empirical validation on GPT-style pretraining. Models with W_Q = I_d are trained from scratch on OpenWebText and compared against parameter-matched baselines on fully converged runs, with reduced models matching or outperforming baselines.

The paper also introduces an index-free notation for multi-head attention (Section 8.1), building on variants of Block Hadamard Products, which the authors say makes the motivating observation mathematically trivial.

Main Findings

• 25% attention parameter reduction. One of W_Q, W_K, or W_V can be replaced with the identity matrix, reducing attention parameters by 25% per layer. For the empirically tested Query removal, this amounts to 8% of transformer block parameters.

• Simpler optimization when removing Query or Key. Because attention logits then depend on a single learned weight matrix rather than a product of two, the reduced architecture admits higher learning rates. The authors use 1.6×10⁻³ for the base reduced model and 2.2×10⁻³ for the larger-MLP reduced variant, versus the standard GPT-2/GPT-3 value of 6×10⁻⁴.

• Reduced 117M model matches the full 124M baseline. With an adjusted attention scaling factor, the reduced model (77.88M non-embedding parameters) matches the full baseline (84.95M non-embedding parameters) despite 8% fewer non-embedding parameters, while parameter-matched standard baselines at the same size (the 3.5×d-MLP baseline and the smaller-d baseline at d=744) perform measurably worse.

• Reallocating saved parameters gives the best result. Moving the saved parameters into the MLP (4.5×d hidden dimension) yields the best validation loss across all five configurations tested.

• A measured attention-scaling correction is required. With W_Q = I_d, queries become coordinate slices of the input rather than learned projections, producing attention scores with approximately 1.8× larger standard deviation at initialization. The paper adopts a scaling factor of 1/(2√d_k) instead of the standard 1/√d_k to prevent early softmax saturation.

• Single-layer elimination applies post-training. Theorem 4.1 needs only weight reparametrization, so for existing pretrained models whose LayerNorm has been removed via the techniques of Heimersheim and Baroni et al. (reported there as working up to GPT-2 XL), a single layer's W_Q can be eliminated after training.

• The motivating insight is simple, the propagation is not. Attention depends on X only through XW_Q, XW_K, and XW_V, which enables a telescoping construction across layers. The MLP skip connection forces a single basis to be used both in and out of each block, which is why full multi-layer elimination requires either attention-only skip connections or weight sharing.

Methodology in Plain English

The authors start from a reparametrization lemma: if a function only ever touches the products XW_Q, XW_K, and XW_V, then you can multiply X by any invertible matrix Θ while multiplying each of the weight matrices by Θ⁻¹ on the left, and nothing changes. Choosing Θ = W_Q turns the Query weight into the identity and pushes the change into W_K and W_V.

The hard part is doing this consistently across many layers, since each layer would prefer a different basis change. The authors analyze what breaks — specifically the MLP residual connection, which forces the same basis in and out of a block — and then prove elimination works when either the MLP skip is removed (attention-only skip connections) or weights are shared across layers so every layer wants the same basis.

They illustrate the resulting geometric obstruction by solving the functional equation MLP = Id + MLP in the ReLU case, and separately derive conditions under which a basis change commutes with LayerNorm (Section 8.3).

Empirically, they use Karpathy's NanoGPT implementation of GPT-2/GPT-3-small: 12 layers, 12 heads, d = 768, MLP hidden dimension 4d = 3072, GELU activations, LayerNorm, sequence length 1024, and the GPT-2 BPE tokenizer. Bias parameters are omitted. All models train for 60k gradient steps on OpenWebText with mixed precision, AdamW (β₁ = 0.9, β₂ = 0.95), gradient clipping at 1.0, weight decay 0.1, and roughly 490k tokens per gradient step (about 29.5B tokens total, approximately 12× Chinchilla-optimal). The cosine schedule completes fully at 60k steps after 2k warmup steps, so all models are evaluated at convergence. Training and evaluation batch indices are pre-generated from a fixed seed and reused across variants so every configuration sees identical data in identical order. Validation loss is estimated every 1000 steps by averaging over 2400 sequences (about 2.5M tokens). Training runs on a single NVIDIA RTX 5090 GPU with FlashAttention, and all configurations use tied embedding/LM-head weights.

Five configurations are compared, with non-embedding parameters reported following scaling-law convention: Baseline (124M) with 4×d MLP at 84.95M non-embedding parameters; Baseline (117M, Smaller MLP) at 3.5×d and 77.88M; Baseline (118M, Smaller d) at d = 744 and 79.73M; Reduced (117M) with W_Q = I_d at 77.88M; and Reduced (124M, Larger MLP) at 4.5×d and 84.95M. Total parameter counts are 117.30M, 117.92M, 117.30M, 124.37M, and 124.37M respectively in the order listed in Table 2.

Why This Matters

Impact on research. The paper adds to a growing line of work suggesting modern transformer architectures are overparameterized, and it takes a step beyond prior theory: Graef proved that in skipless transformers without normalization both W_Q and W_O can be eliminated simultaneously, whereas this work handles transformers with skip connections by retaining W_O to absorb basis changes, characterizes when ReLU MLPs can absorb skip connections, and validates the result through GPT-style pretraining. It also provides a canonical form that may simplify analysis of weight-shared, fixed-point, and continuous-time limit models. The departure point from He & Hofmann is notable: they simplified via block restructuring (parallel attention-MLP blocks), while this work eliminates weights inside the original architecture.

Real-world applications (as implications of the findings, not claims tested in the paper):

• Cheaper LLM pretraining and inference — removing 25% of attention parameters per layer reduces parameter count, memory, and compute.
• Compatibility with KV-cache and grouped-query attention (GQA) workflows — the authors deliberately choose to eliminate W_Q rather than W_K or W_V specifically for consistency with the KV cache and GQA optimization of LLMs.
• Post-training compression of existing models — once LayerNorm is removed from a pretrained model, a single layer's W_Q can be eliminated by reparametrization alone.
• Weight-shared and recursive architectures — the paper argues recursive models such as Tiny Recursion Models are natural candidates for Query weight elimination under the weight-sharing theorem, and notes that the TRM authors' choice of untied embedding and head weights satisfies the untied requirement.

Industry relevance. The paper emphasizes training and deploying transformer-based language models is computationally expensive, and positions its method as orthogonal to quantization, efficient attention, weight sharing, LoRA-style adaptation, pruning, and FlashAttention — meaning the savings could compound with those techniques. The authors estimate savings are multiplicative with GQA and MQA, since Query weight elimination applies to standard, GQA, and MQA architectures alike.

Future Directions

• Scaling the empirical validation. The current experiments cover 117M to 124M parameters. The authors explicitly frame these as establishing a baseline for systematic scaling, multi-seed validation, and downstream benchmarking.
• Testing elimination of W_V or W_K empirically. The theory treats the three symmetrically, but only Query weight removal was tested; the paper leaves empirical testing of W_V and W_K simplification for future work, noting the theoretical approach is identical.
• Extending through normalization layers. Section 8.3 derives sufficient conditions for basis transformations to commute with LayerNorm but the resulting obstruction is not fully resolved; extending elimination to normalized architectures in full generality remains open.
• Cross-modality and cross-architecture investigation. The abstract motivates investigation across modalities and at scale, where the authors expect the stability and efficiency gains to be most consequential.
• Weight decay tuning. The authors tested weight decay in [0.025, 0.1] and note that values above 0.1 may yield further improvements; this was not explored.

Target Audience

This paper is most valuable to transformer architecture researchers and theory-oriented machine learning scientists interested in parameter redundancy and expressivity boundaries, and to efficiency-focused practitioners — pretraining engineers, inference-optimization teams, and anyone working on compression, weight sharing, or attention redesign. Readers studying weight-shared and recursive architectures (including fixed-point and continuous-time limit analyses) will find the weight-sharing theorem directly relevant. A working comfort with linear algebra and the standard attention formulation is assumed; the empirical portion is accessible to anyone who has trained or fine-tuned a GPT-style model.

Authors’ abstract

We theoretically investigate whether the Query, Key, Value weight triplet can be reduced in encoder-only and decoder-only transformers. Under mild assumptions, we prove that one of the Query, Key or Value weights are redundant and can be replaced with the identity matrix, reducing attention parameters by 25\%. If applied to the Query or Key weights, this also simplifies optimization: attention logits depend on a single learned weight matrix rather than on a product of two. Validating the Query weight removal on decoder-only GPT-style small models trained from scratch, we find that reduced models match baseline performance despite fewer parameters, and outperform baselines when saved parameters are reallocated. Our analysis has also led us to a structural expressivity boundary: in the mathematically tractable ReLU setting, skip connections push MLPs into a generically disjoint function class at fixed width. These findings motivate investigation across modalities and at scale, where the observed stability and efficiency gains may prove most consequential.