If you've written, fine-tuned, or even just looked at the architecture diagram of a modern Large Language Model (LLM), you’ve likely accepted one fundamental truth as gospel: Self-Attention requires Queries (Q), Keys (K), and Values (V).
Since the groundbreaking 2017 paper "Attention Is All You Need", this tripartite projection step has been the bedrock of deep learning. Every token in your input sequence is multiplied by three distinct weight matrices ($W_Q$, $W_K$, $W_V$) to project it into three different vector spaces. We’ve memorized the formula, optimized the CUDA kernels for it, and designed entire clusters of GPUs to calculate it as fast as possible.
But a fascinating new systematic study has just hit the AI research community, and it asks a question that might make you question everything you know about transformer architecture: Do transformers actually need three projections?
As developers, we care deeply about this. Why? Because the computational bottleneck of running LLMs—especially at the edge or during high-throughput inference—is memory bandwidth and parameter count. If we can eliminate one, or even two, of these projections without sacrificing model quality, we can build smaller, faster, and cheaper models. Let’s dive deep into the math, the architectural variations, and what this means for the future of developer-facing AI tools.
The Status Quo: Why We Have Q, K, and V
To understand how we can simplify the transformer, we first need to recall why we have three projections in the first place. In standard scaled dot-product attention, we start with an input matrix $X$. We then project it into three representations:
Query (Q) = X * W_Q
Key (K) = X * W_K
Value (V) = X * W_V
The intuition behind this is analogous to a database lookup system:
- The Query represents what the current token is looking for.
- The Key represents what information other tokens contain.
- The Value is the actual content that is retrieved once we match the Query with the Key.
Mathematically, the attention matrix is computed as:
Attention(Q, K, V) = softmax( (Q * K^T) / sqrt(d_k) ) * V
While elegant, this requires three separate matrix multiplications ($W_Q$, $W_K$, $W_V$) for every single token at every single layer. This contributes heavily to the model's parameter footprint and consumes precious memory bandwidth during the autoregressive generation phase (a bottleneck known as the KV-cache bottleneck).
Enter the QKV Variants: What If We Simplify?
Researchers have recently systematically analyzed several alternative architectures that reduce the number of projections. What happens if we share weights? What if we completely eliminate the Key or Query projections?
Let's look at the three most promising variants evaluated in the study, ranging from moderate simplification to radical reduction.
1. The Shared Query-Key (QK-Shared) Variant
In this architecture, we assert that a token's search intent (Query) and its search index (Key) can exist in the same vector space. Instead of having $W_Q$ and $W_K$, we use a single projection matrix $W_{QK}$.
Q = X * W_QK
K = X * W_QK (or simply Q)
V = X * W_V
By forcing $Q$ and $K$ to share parameters, we immediately slice off a massive chunk of the model's parameter budget. The dot product is then calculated between a token's projection and other tokens' projections from the same space.
2. The "Value-Only" (VO) and "Query-Value" (QV) Architectures
Going even further, what if we eliminate the distinct Key and Query projections entirely? In a "Value-Only" setup, attention weights are computed directly using the raw residual stream representation ($X$), and only the Value ($V$) is projected.
Alternatively, in a "Single-Projection" model (sometimes called the Q-Only or V-Only variant), we only maintain one projection matrix. For instance, if we only project the Query/Key:
Q = X * W_P
K = X * W_P
V = X (No projection!)
In this case, the value vector is simply the unprojected representation of the token itself. This bypasses the $W_V$ multiplication entirely, sending the raw input vectors directly to be weighted by the attention map.
The Findings: Can Less Be More?
The systematic study evaluated these variants across various downstream tasks, measuring both pre-training perplexity and zero-shot downstream task performance. The results are incredibly encouraging for software engineers who build and deploy these models:
- Minimal Loss in Accuracy: Models with shared QK projections or single-projection configurations retained up to 95-98% of the performance of standard three-projection transformers, despite using significantly fewer parameters in their attention blocks.
- Massive Parameter Savings: In a standard transformer, attention projections account for a massive chunk of non-feed-forward network (FFN) parameters. By eliminating even one projection, we can reduce attention parameter counts by up to 33%.
- KV-Cache Reductions: If we can merge or eliminate $K$ or $V$ projections, the size of the KV-cache stored in GPU memory during inference drops dramatically. This means we can fit much larger batch sizes or longer context windows on the same hardware.
Implementing a Single-Projection Attention Layer in PyTorch
To see how simple this is to conceptualize, let's write a custom PyTorch module. We will implement a simplified Shared-QK Attention layer, where the Query and Key share the same projection matrix, leaving only the Value as a separate projection.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SharedQKAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
# Instead of three projection layers (Q, K, V), we only instantiate TWO!
self.q_k_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
self.out_proj = nn.Linear(embed_dim, embed_dim)
def forward(self, x):
batch_size, seq_len, embed_dim = x.size()
# 1. Project Q and K using the SAME weights
qk = self.q_k_proj(x) # Shape: [batch_size, seq_len, embed_dim]
# 2. Project V normally
v = self.v_proj(x) # Shape: [batch_size, seq_len, embed_dim]
# Reshape for multi-head attention: [batch, heads, seq_len, head_dim]
qk = qk.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# Queries and Keys are identical in representation here
q = qk
k = qk
# 3. Compute Attention Scores: (Q * K^T) / sqrt(d_k)
scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
# Apply causal mask if doing autoregressive generation (optional, omitted for simplicity)
attn_weights = F.softmax(scores, dim=-1)
# 4. Multiply by Value
context = torch.matmul(attn_weights, v) # Shape: [batch_size, num_heads, seq_len, head_dim]
# Concatenate heads back
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
# 5. Final output projection
return self.out_proj(context)
# Quick sanity check
if __name__ == "__main__":
dummy_input = torch.randn(2, 16, 256) # Batch size 2, Sequence length 16, Embedding dim 256
attention_layer = SharedQKAttention(embed_dim=256, num_heads=8)
output = attention_layer(dummy_input)
print("Output shape:", output.shape) # Expected: torch.Size([2, 16, 256])
As you can see, the code is remarkably clean. By reusing qk for both our queries and keys, we completely sidestep the generation of a third projection matrix. In production, this means fewer weights to load from memory, leading to faster execution times.
Why This Matters for Software Engineers
If you aren't training foundation models from scratch, you might wonder why you should care about architectural tweaks like this. The answer boils down to three words: cost, speed, and democratization.
1. Ultra-Low Latency Edge AI
Running LLMs on consumer devices (phones, laptops, IoT devices) is highly constrained by memory bandwidth. By utilizing models trained with single-projection architectures, we reduce the footprint of the model weights. This allows larger context sizes to be processed directly on-device without running out of RAM or relying on expensive cloud APIs.
2. Green Computing & Lower Cloud Bills
If you run high-throughput LLM pipelines in the cloud, inference costs can skyrocket. Architectures that require fewer projection steps translate directly to fewer FLOPS (Floating Point Operations) per token. When multiplied by millions of API requests, these minor efficiency gains scale into thousands of dollars saved on your monthly cloud bill.
3. Simpler Custom Architectures
For developers building domain-specific transformers (such as time-series forecasting, custom recommendation engines, or small code-generation assistants), using a simplified QKV structure makes training faster and less prone to overfitting, especially when working with smaller datasets.
Conclusion
The assumption that transformers strictly require distinct Query, Key, and Value projections is slowly being dismantled. As this latest systematic study demonstrates, we can achieve remarkably similar performance with simpler, sleeker architectures like Shared-QK or single-projection setups.
Just as Grouped-Query Attention (GQA) became the industry standard for optimizing LLM inference over the past year (used in Llama 3, Mistral, and others), we may soon see "Two-Projection" or "Single-Projection" attention variants become the default choice for the next generation of highly efficient open-source models.
What do you think? Are you ready to ditch the traditional QKV split for more streamlined architectures, or do you think the slight performance tradeoffs aren't worth the change? Let’s chat in the comments below!
Looking for more deep dives into cutting-edge developer tools and architecture? Don't forget to subscribe to the "Coding with Alex" newsletter at sysseder.com!