Do Transformers Need Three Projections? A Systematic Study of QKV Variants

The Core of Attention: Re-evaluating the QKV Paradigm

Since Vaswani et al. introduced the Transformer architecture in 2017, the Query-Key-Value (QKV) projection paradigm has stood as an undisputed foundation of modern deep learning. From BERT and GPT-4 to LLaMA and Claude, virtually every state-of-the-art large language model relies on this tripartite projection scheme to calculate self-attention. But as we push the limits of scale, memory bandwidth, and compute efficiency, researchers are asking a fundamental question: Do transformers actually need three separate projection matrices?

This deep dive explores the mathematical underpinnings of QKV projections, examines modern research into alternative projection variants, and analyzes how reducing or sharing these projections impacts representational capacity, memory efficiency, and downstream task performance.

The Standard Transformer: Why Three Projections?

To understand why we might want to alter the QKV structure, we must first look at why it exists in its current form. In a standard scaled dot-product attention mechanism, an input sequence matrix $X \in \mathbb{R}^{N \times d_{model}}$ is projected into three distinct spaces:

Queries ($Q$): $Q = X W_Q$, where $W_Q \in \mathbb{R}^{d_{model} \times d_k}$
Keys ($K$): $K = X W_K$, where $W_K \in \mathbb{R}^{d_{model} \times d_k}$
Values ($V$): $V = X W_V$, where $W_V \in \mathbb{R}^{d_{model} \times d_v}$

The attention matrix is then computed as:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

Each projection serves a distinct conceptual purpose:

Queries and Keys define the "routing" or "relevance" mechanism, determining which tokens should pay attention to which other tokens.
Values represent the actual content being routed once the attention weights are established.

Having three separate projection matrices allows the model to decouple the token's identity as a searcher (Query), a target (Key), and a payload (Value). However, this decoupling comes at a heavy cost in terms of parameter budget, FLOPs, and—critically—KV cache memory during inference.

Mathematical Redundancies in Self-Attention

Is the high dimensionality of $W_Q$, $W_K$, and $W_V$ strictly necessary? Theoretical and empirical analyses suggest significant redundancy.

If we examine the rank of the attention matrix $A = \text{softmax}(Q K^T / \sqrt{d_k})$, we often find it is low-rank, meaning the effective capacity of the projection matrices is underutilized. Furthermore, because $Q$ and $K$ are immediately multiplied together ($Q K^T = X W_Q W_K^T X^T$), the model is effectively learning a bilinear form parameterized by the matrix product $M = W_Q W_K^T \in \mathbb{R}^{d_{model} \times d_{model}}$.

This mathematical formulation opens up several possibilities:

Symmetric Attention ($W_Q = W_K$): If we force $W_Q = W_K$, the bilinear product becomes $W_Q W_Q^T$, which is symmetric and positive semi-definite. While this reduces the parameter count of the attention mechanism by 33%, it severely restricts the model's ability to represent asymmetric relationships (e.g., token A attending to token B does not necessarily imply token B should attend to token A).
The Identity Projection: What if we eliminate $W_Q$ or $W_K$ completely and use the raw input $X$ directly? For instance, setting $Q = X$ and $K = X W_K$. The bilinear product becomes $X W_K^T X^T$. While this breaks symmetry, it forces the query space to align perfectly with the input representation space, reducing the model's capacity to abstract away query-specific features.

A Systematic Study of QKV Variants

Recent systematic studies have mapped out the performance of various "ablated" QKV configurations. Let's look at the primary variants tested in modern architecture search:

1. Dual-Projection (QV / KV / QK Sharing)

In these architectures, one of the three projections is eliminated or tied.

No-Key Projection (Q-I-V): Here, $K = X$ (Identity). The attention equation simplifies to $A = \text{softmax}(Q X^T / \sqrt{d_{model}}) V$. Empirical studies show that this model struggles with deep hierarchical representation because the keys cannot be transformed to match the abstraction level of the queries.
Tied Query-Key (Tied-QK): $W_Q$ and $W_K$ share weights but may use a small transformation or relative position encoding to break symmetry. This has shown surprisingly robust performance in encoder-only architectures (like BERT) but degrades in autoregressive, decoder-only models (like GPT) where causal masking already imposes directional constraints but representational asymmetry remains critical.

2. Multi-Query and Grouped-Query Attention (MQA / GQA)

While not eliminating the projections entirely, Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) drastically alter the projection landscape.

MQA: Uses a single Key and Value projection head shared across all Query heads.
GQA: Groups Query heads and assigns a single Key and Value head per group.

These variants target the memory bandwidth bottleneck during autoregressive decoding. By reducing the number of $K$ and $V$ projections, they shrink the KV cache size by up to 8x, allowing for massive context windows and higher batch sizes during inference, with negligible drops in perplexity.

Implementing a Custom QKV-Optimized Attention Layer

To understand how these variants are implemented, let's look at a PyTorch implementation of a modified Multi-Query Attention layer where the Value projection is shared, and the Key projection is replaced with a low-rank bottleneck to minimize parameter overhead.

import torch
import torch.nn as nn
import math

class OptimizedAttention(nn.Module):
    def __init__(self, d_model, n_heads, d_k, r_bottleneck=4):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_k
        self.d_model = d_model
        
        # Standard Query projection
        self.q_proj = nn.Linear(d_model, n_heads * d_k, bias=False)
        
        # Low-rank Key projection to save parameters
        self.k_down = nn.Linear(d_model, d_model // r_bottleneck, bias=False)
        self.k_up = nn.Linear(d_model // r_bottleneck, d_k, bias=False)
        
        # Single shared Value projection (Multi-Query style)
        self.v_proj = nn.Linear(d_model, d_k, bias=False)
        
        # Output projection
        self.out_proj = nn.Linear(n_heads * d_k, d_model, bias=False)
        
    def forward(self, x):
        B, N, d_model = x.shape
        
        # Query: [B, N, n_heads, d_k] -> transpose to [B, n_heads, N, d_k]
        q = self.q_proj(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        
        # Key: [B, N, d_k] -> project via low-rank bottleneck
        k = self.k_up(self.k_down(x)) # Shape: [B, N, d_k]
        # Unsqueeze to align with heads: [B, 1, N, d_k]
        k = k.unsqueeze(1)
        
        # Value: [B, N, d_k] -> Unsqueeze to [B, 1, N, d_k]
        v = self.v_proj(x).unsqueeze(1)
        
        # Compute attention scores: [B, n_heads, N, N]
        # k.transpose(-2, -1) shape is [B, 1, d_k, N] (broadcasts across heads)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn_weights = torch.softmax(scores, dim=-1)
        
        # Compute output: [B, n_heads, N, d_k]
        context = torch.matmul(attn_weights, v)
        
        # Reshape and project out
        context = context.transpose(1, 2).contiguous().view(B, N, self.n_heads * self.d_k)
        return self.out_proj(context)

Empirical Trade-offs: Parameter Efficiency vs. Expressivity

When we systematically strip away or compress these projections, how does it affect downstream tasks?

The empirical consensus is clear: We cannot easily eliminate the query projection, as queries must remain highly dynamic and head-specific to capture complex, multi-layered dependencies. However, the key and value projections are highly compressible. GQA represents the current optimal sweet spot between memory performance and expressivity.

Future Outlook: Beyond Linear Projections

As the deep learning community shifts from purely heuristic designs to mathematically principled architectures, the future of the QKV paradigm lies in dynamic, non-linear, or implicit projections.

Researchers are actively exploring Kernelized Attention and State Space Models (SSMs) like Mamba, which bypass the quadratic $QK^T$ calculation entirely by maintaining a continuous-time hidden state. In these architectures, the concept of separate "Query", "Key", and "Value" projections is completely redefined, merging state transitions with projection spaces.

For now, if you are designing or fine-tuning transformer architectures for edge deployment or high-throughput APIs, leveraging GQA or implementing low-rank bottleneck projections for keys and values offers the most reliable mechanism to slash latency without sacrificing the contextual reasoning capabilities that make transformers so powerful.