TLDR: A new research paper explores the redundancy of Query (WQ) weights in the Transformer architecture’s attention mechanism. Theoretically, it proves WQ can be eliminated through basis transformations. Empirically, GPT-3 small models trained without WQ (replaced by an identity matrix) achieved comparable performance to baselines, with adjustments to attention scaling and weight decay. This leads to an 8.3% reduction in transformer block parameters, suggesting significant efficiency gains for large language models.
In the rapidly evolving world of artificial intelligence, large language models (LLMs) built on the Transformer architecture have become foundational. These models, known for their impressive capabilities, are also notoriously expensive to train and deploy. This has led researchers to constantly seek architectural optimizations to make them more efficient without sacrificing performance.
A recent research paper titled “WK, WV is probably all you need: On the necessity of the Query, Key and Value weight triplet in the decoder-only Transformer architecture” by Marko Karbevski and Antonij Mijoski delves into one such optimization. The paper investigates a core component of the Transformer’s attention mechanism: the Query, Key, and Value (QKV) weight triplet. This triplet of weight matrices (WQ, WK, WV) is crucial for how a Transformer processes information, but the authors theorize that one part of it, the Query (WQ) weights, might be redundant.
The central hypothesis of the paper is that the Query weights (WQ) can be eliminated, meaning they can be replaced by an identity matrix (Id) without losing the model’s ability to perform its task. This would effectively reduce the number of parameters in the attention mechanism, leading to a more efficient model. The theoretical investigation in the paper, under certain simplifying assumptions, indeed proves that WQ can be made redundant through a process called basis transformations. This means that the information typically encoded in WQ can be implicitly handled by adjusting other parts of the network.
To validate their theory, the researchers conducted empirical experiments using GPT-3 small architectures, which include common features like layer normalization, skip connections, and weight decay. They trained these models from scratch, comparing a standard baseline model with a modified version where WQ was set to the identity matrix. The results were compelling: the reduced model achieved validation loss comparable to the standard baseline. This suggests that even in complex, realistic Transformer architectures, the Query weights might not be strictly necessary.
Achieving this comparable performance required some practical adjustments. The attention scaling, which controls the magnitude of pre-softmax activations, was modified from 1/√dk to 1/(2√dk). Additionally, the weight decay, a regularization technique, was significantly reduced from 0.1 to 2^-5 (approximately 0.03125) for all parameters in the reduced model. These adjustments were crucial to rebalance the model’s capacity and regularization, as removing WQ acts as a form of architectural regularization itself.
The implications of these findings are significant for the efficiency of LLMs. By eliminating the Query weights, the researchers demonstrated a reduction of 25% of attention parameters per layer, which translates to over 8% of the total non-embedding/LM-head parameters in a Transformer block. This parameter reduction can lead to lower computational costs for both training and inference. Furthermore, this approach is orthogonal to other efficiency techniques like Grouped-Query Attention or Mixture-of-Experts, meaning it can be combined with them for even greater savings.
While the theoretical proofs relied on simplified architectural assumptions (e.g., without layer normalization in some cases), the empirical results on full-complexity GPT-style models provide strong evidence that the redundancy persists in practical settings. The paper acknowledges a gap between theory and practice, particularly concerning the interaction with layer normalization, which often requires approximate solutions and hyperparameter tuning. However, the consistency between theoretical predictions and empirical observations is a strong indicator of the potential for Query weight redundancy.
Also Read:
- Dynamic Plasticity Enhances Transformer Learning and Adaptation
- Boosting Circuit Discovery in LLMs with Per-Attention Head Quantization
This research opens new avenues for exploring architectural simplifications in Transformers. It challenges the long-held assumption about the necessity of the full QKV triplet and motivates further investigation into parameter redundancy at a larger scale. The code and checkpoints for this research will be made available at https://github.com/MarkoKarbevski/Wqkv_necessity.
