This post introduces Low-Rank Key-Value (LRKV) attention<\/strong>, a drop-in modification to multi-head attention that reduces KV cache size by 45\u201353%<\/strong> vs standard MHA, while achieving lower test loss<\/strong> across model scales (128M \u2192 6.3B), faster convergence in training steps, and stronger downstream performance after supervised midtraining.<\/p>\n\n\n\n The key idea is that attention heads are not independent. There\u2019s structured redundancy across heads – yet fully sharing keys\/values (like in MQA\/GQA) can constrain expressivity. LRKV instead exploits redundancy using a shared full-rank KV basis<\/strong> plus head-specific low-rank residuals<\/strong>, yielding a continuous knob<\/strong> between complete sharing and full per-head independence<\/p>\n\n\n\n <\/p>\n\n\n\n In autoregressive decoding, each layer caches keys and values for all previously generated tokens.<\/p>\n\n\n\n For a sequence length L<\/strong>, number of heads H<\/strong>, per-head dimension dh<\/sub><\/strong>, and hidden dimension d<\/strong> =Hdh<\/sub><\/strong>, standard multi-head attention (MHA) caches, for each head, keys and values $$\\mathbf{K_h}, \\mathbf{V_h} \\in \\mathbb{R}^{L \\times d_h}$$<\/p>\n\n\n\n So the KV cache memory per layer scales as:<\/p>\n\n\n\n $$M_{\\text{standard}} = 2 L H d_h = 2 L d$$<\/p>\n\n\n\n Existing methods such as MQA (Multi-Query Attention) and GQA (Grouped-Query Attention) reduce cache size by sharing K\/V across heads (or groups). This often improves throughput, but it forces heads to look through the same K\/V representations, reducing representational diversity.<\/p>\n\n\n\n Empirically and theoretically, we know that heads specialize: different heads represent complementary syntactic and semantic patterns, and so fully sharing K\/V across attention heads can degrade capabilities such as code generation and structured reasoning.<\/p>\n\n\n\n At the same time, we know that head specialization is not fully independent: recent analyses show high correlation and overlapping subspaces across head projections, so the redundancy is structured rather than random.<\/p>\n\n\n\n This motivates the central question:<\/p>\n\n\n\n Can we compress the KV cache by exploiting cross-head redundancy without comprimising head specialization?<\/strong><\/p>\n<\/blockquote>\n\n\n\n This bring us to our main contribution, Low-Rank Key Value Attention (LRKV).<\/em><\/p>\n\n\n\n Before discussing our proposed LRKV attention mechanism, we give a concise refresher on the well-established baselines that have been the mainstay mechanisms used in Transformers. If you are already familiar with Multi-Head Attention (MHA), Multi-Query Attention (MQA), Grouped Query Attention (GQA) and Multi-Latent Attention (MLA), you can skip this section.<\/p>\n\n\n\n For Multi-Head Attention<\/strong>, each attention head maintains its own independent key and value projection matrices. Standard attention uses per-head key and value projections:<\/p>\n\n\n\n $$\\mathbf{K}_h = \\mathbf{X} \\mathbf{W}_h^K, \\quad \\mathbf{V}_h = \\mathbf{X} \\mathbf{W}_h^V \\quad \\text{where} \\quad \\mathbf{W}_h^{K,V} \\in \\mathbb{R}^{d \\times d_h}$$<\/p>\n\n\n\n The full projection matrix is formed by concatenating H independent weight blocks side by side, i.e. there is no parameter sharing whatsoever between heads. This gives each head complete freedom to learn specialised key\/value representations, making MHA the most expressive configuration available. However, this expressiveness comes at a steep cost: the KV cache scales linearly with the number of heads (\\(2Hd_h\\) values stored per token), making it the most memory-intensive option during long-context inference.<\/p>\n\n\n\n All other attention variants can be understood as different strategies for reducing this cost while preserving as much of MHA’s expressiveness as possible.<\/p>\n\n\n\n Multi-Query Attention<\/strong> takes the most aggressive approach to KV cache reduction: all heads share a single key and value projection. Each head still computes its own query, preserving some capacity for diverse attention patterns. However, every head attends over the exact same keys and values. The KV cache shrinks from \\(2Hd_h\\) to just \\(2d_h\\) per token – a 75% parameter reduction that is transformative for long-context inference where cache memory is the primary bottleneck. <\/p>\n\n\n\n But the trade-off is severe: because all heads attend over identical keys and values, the only way they can differentiate is through their query projections. In practice, heads tend to converge on similar attention patterns, reducing the model’s ability to capture diverse linguistic phenomena simultaneously. In fact, in our paper (see at the end), we quantitatively show that the query parameters are forced to diversify more compared to other less restrictive mechanisms because K and V are the same across heads. <\/p>\n\n\n\n Grouped-Query Attention<\/strong> offers a middle ground between MHA and MQA. Attention heads are divided into G groups, and all heads within the same group share a single key and value projection. In the example above, \\(G= 2\\): heads 1 & 2 share one projection while heads 3 & 4 share another. This allows groups to specialize in different roles. For instance, one group might focus on local syntax while another captures long-range dependencies. The KV cache reduces from \\(2Hd_h\\) to \\(2Gd_h\\) per token. In practice, G is typically set to H\/4 or H\/8, yielding a 4-8\u00d7 cache reduction.<\/p>\n\n\n\n The key limitation of this approach is that sharing boundaries are fixed at architecture design time and cannot adapt to the data. Heads assigned to the same group must share representations regardless of whether their learned roles are compatible. This is a coarse-grained constraint that limits flexibility compared to methods that can adapt per-head. <\/p>\n\n\n\n Introduced in DeepSeek-V2, Multi-Latent Attention<\/strong> takes a fundamentally different approach to KV cache compression. Rather than sharing projections between heads (like MQA\/GQA), MLA compresses all key\/value information into a single low-dimensional latent vector per token. During inference, only this compact latent is stored in the KV cache. When attention is computed, the full per-head keys and values are reconstructed on the fly via learned up-projection matrices. This decouples the cache cost from the number of heads entirely. <\/p>\n\n\n\n This architecture works in two stages. First, a down-projection compresses each token into a latent of dimension \\(d_c\\) (much smaller than \\(d\\)). Second, separate up-projections reconstruct per-head keys and values from this shared latent. The effective key projection is an 8\u00d78 matrix but has rank at most d_c. This resembles MHA’s full projection but lives in a lower-dimensional subspace.<\/p>\n\n\n\n There’s a fundamental trade-off here: heads can specialise (unlike MQA\/GQA) – but they must reconstruct their keys and values from a shared compressed representation. Any per-head information that does not survive this bottleneck is permanently lost at inference time. We can compare this with LRKV, where each head’s low-rank correction is baked directly into the weight matrix – getting around the information bottleneck during inference.<\/p>\n\n\n\n Given this background on related attention mechanisms, we now move to explaining LRKV. <\/p>\n\n\n\n LRKV factorizes each head\u2019s key and value projection into a shared full-rank base plus a low-rank residual:<\/p>\n\n\n\n $$\\mathbf{W}_h^K = \\mathbf{W}_{\\text{shared}}^K + \\mathbf{U}_h^K (\\mathbf{B}_h^K)^\\top, \\quad \\mathbf{W}_h^V = \\mathbf{W}_{\\text{shared}}^V + \\mathbf{U}_h^V (\\mathbf{B}_h^V)^\\top$$<\/p>\n\n\n\n Here Wshared<\/sub><\/strong> is dense, full-rank, and shared across all heads in the layer,<\/p>\n\n\n\n $$\\mathbf{U}_h^{K,V} \\in \\mathbb{R}^{d \\times r}, \\quad Finally, \\(r << d_h\\) is the residual rank. <\/p>\n\n\n\n Our interpretation is that the shared base learns a global KV basis<\/strong> for the layer, and each head then learns a small low-rank deviation from that basis. This gives a continuous interpolation<\/em>: \\(r = 0\\) reduces to full sharing of K\/V within a layer (MQA-style limit) and increasing \\(r\\) increases head-specific capacity, moving toward MHA.<\/p>\n\n\n\n A subtle but important design choice: LRKV applies the factorization only to K and V<\/strong> (the cached parts). Queries remain unconstrained (per-head), preserving attention expressivity while targeting the true inference bottleneck.<\/p>\n\n\n\n At inference time, we want to avoid caching Kh<\/sub><\/strong> and Vh<\/sub><\/strong> for every head and every token. LRKV caches shared features once per layer:<\/p>\n\n\n\n $$ \\mathbf{K}_{\\text{shared}} = \\mathbf{X} W_{\\text{shared}}^K \\in \\mathbb{R}^{L \\times d_h}, \\quad \\mathbf{V}_{\\text{shared}} = \\mathbf{X} W_{\\text{shared}}^V \\in \\mathbb{R}^{L \\times d_h} $$<\/p>\n\n\n\n Per-head latents:<\/p>\n\n\n\n $$\\mathbf{R}_h^K = \\mathbf{X} \\mathbf{U}_h^K \\in \\mathbb{R}^{L \\times r}, \\quad Then the implied per-head features are:<\/p>\n\n\n\n $$ \\mathbf{K}_h = \\mathbf{K}_{\\text{shared}} + \\mathbf{R}_h^K (\\mathbf{B}_h^K)^\\top, \\quad \\mathbf{V}_h = \\mathbf{V}_{\\text{shared}} + \\mathbf{R}_h^V (\\mathbf{B}_h^V)^\\top $$<\/p>\n\n\n\n <\/p>\n\n\n\n Crucially, B<\/strong>h<\/sub> are parameters, not cached per token. So the cache memory becomes:<\/p>\n\n\n\n $$M_{\\text{LRKV}} = 2 L d_h \\; (\\text{shared } K,V)\\quad 2 L H r \\; (\\text{per-head latents}) = 2L(d_h + Hr)$$<\/p>\n\n\n\n Relative to standard MHA:<\/p>\n\n\n\n $$\\frac{M_{\\text{LRKV}}}{M_{\\text{standard}}}\\frac{d_h + Hr}{H d_h}\\frac{1}{H} + \\frac{r}{d_h}.$$<\/p>\n\n\n\n This is the cleanest engineering knob LRKV provides: for fixed H and dh<\/sub>, the residual rank r trades cache size for per-head flexibility.<\/p>\n\n\n\n Naively reconstructing the full per-head keys and values for every cached token would erase the practical gain. Instead, LRKV computes logits and outputs exactly via associativity.<\/p>\n\n\n\n For a decoding step with query qh<\/sub><\/strong> the attention logits are:<\/p>\n\n\n\n $$\\mathbf{q}_h \\mathbf{K}_h^\\top=\\mathbf{q}_h\\mathbf{K}_{\\text{shared}}^\\top+(\\mathbf{q}_h \\mathbf{B}_h^K)(\\mathbf{R}_h^K)^\\top, \\quad \\mathbf{q}_h \\in \\mathbb{R}^{d_h}$$<\/p>\n\n\n\n For attention weights ah<\/sub><\/strong> the value aggregation is:<\/p>\n\n\n\n $$\\mathbf{a}_h \\mathbf{V}_h=\\mathbf{a}_h \\mathbf{V}_{\\text{shared}}+(\\mathbf{a}_h \\mathbf{R}_h^V)(\\mathbf{B}_h^V)^\\top, \\quad \\mathbf{a}_h \\in \\mathbb{R}^{1 \\times L}$$<\/p>\n\n\n\n These expressions compute the same outputs as full reconstruction, but operate on smaller cached tensors: That is precisely where the memory win is realized.<\/p>\n\n\n\n During decoding, standard MHA\u2019s dominant per-head cost scales as:<\/p>\n\n\n\n $$O(L d_h), \\quad \\text{while LRKV adds } O(L r + r d_h).$$<\/p>\n\n\n\n For long contexts: $$L \\gg 1 \\quad \\Rightarrow \\quad O(L r) \\text{ dominates, giving overhead } \\sim \\frac{r}{d_h}.$$<\/p>\n\n\n\n In modern inference, we\u2019re often memory bandwidth bound<\/strong>, so reducing the bytes moved can dominate a modest FLOP increase. LRKV reduces total bytes read from cache proportionally to its cache reduction: it reads two shared tensors<\/strong> plus small per-head latents<\/strong> instead of full per-head K\/V.<\/p>\n\n\n\n With the design space now fully mapped, from MHA’s full independence, through MQA and GQA’s discrete sharing strategies, to MLA’s latent bottleneck and our proposed LRKV’s continuous low-rank interpolation, the natural question is: how do these architectural choices actually play out in practice?<\/strong> We evaluate all five methods under identical pretraining and midtraining conditions across three model scales (128M, 2.5B, and 6.3B parameters), measuring both pretraining loss and downstream task performance on five diverse benchmarks. <\/p>\n\n\n\n Our results are very encouraging: LRKV reaches each baseline’s final validation performance 18-30% faster, averaging 23.6% training compute savings across all baselines while achieving better final performance. Critically, this reveals an asymmetric advantage: LRKV reaches any baseline’s performance target early in training, but no baseline reaches LRKV’s final performance (0.719 BPB) even after the full token budget.<\/p>\n\n\n\n <\/p>\n\n\n\n The residual rank r <\/em><\/strong>controls the tradeoff. The ablation shows monotonic improvement with larger rank, and a strong Pareto frontier relative to other KV-efficient methods.<\/p>\n\n\n\n LRKV appears to be not merely \u201ca compression trick\u201d, but a bias that improves optimization and\/or effective capacity under the same token budget. Empirically, we see a consistent \u201cuseful rank\u201d regime around: $$r \\approx 0.36\u20130.43 \\times d_h$$ as a threshold where LRKV matches\/exceeds MHA while still delivering ~50% cache reduction.<\/p>\n\n\n\n At longer sequence lengths, the benefits of KV-efficient attention become more pronounced. The below figure shows that LRKV not only maintains its advantage over standard MHA, but actually widens the gap in the long-context regime. At 8K context, LRKV achieves lower test loss while using roughly half the KV cache, outperforming both MHA and other KV-efficient baselines such as MQA, GQA, and MLA. This suggests that the low-rank decomposition is not merely compressing redundant structure, but acting as an effective inductive bias for long-range modeling. As context length increases and KV cache pressure becomes the dominant bottleneck, LRKV\u2019s combination of memory efficiency and preserved head diversity translates directly into improved modeling performance. A reasonable question is whether these gains turn into better downstream task performance. On a standardized evaluation after supervised mid-training, LRKV achieved the highest combined score across ARC, MMLU, GSM8K, and HumanEval, confirming that its improved pretraining efficiency translates into better downstream performance.<\/p>\n\n\n\n The below figure shows LRKV consistently achieves the highest combined accuracy at every scale – 18.9% (128M), 37.9% (2.5B), and 40.2% (6.3B) – demonstrating that the pretraining gains afforded by LRKV transfer reliably to downstream capabilities. At the 2.5B and 6.3B scales, LRKV leads on four of five benchmarks, with particularly strong margins on knowledge-intensive tasks: at 6.3B it surpasses the next-best method by +2.3% points on ARC-Easy, +4.4% on ARC-Challenge, and +1.8 on MMLU. Notably, MQA suffers a pronounced collapse on HumanEval at all three scales (2.4%, 3.7%, 4.3%). Crucially, the gap between LRKV and competing methods widens with scale, rising from +0.9 over MLA at 128M to +3.3 at 6.3B, reinforcing the conclusion that LRKV’s architectural advantages compound as model capacity increases.<\/p>\n\n\n\n A common failure mode of aggressive KV sharing is that heads lose the ability to represent distinct interactions. LRKV claims you can reduce KV cache without<\/em> degrading diversity by preserving a shared basis plus low-rank head-specific deviations. We measure head diversity using gauge-invariant similarity metrics derived from bilinear forms:<\/p>\n\n\n\n $$\\mathbf{A}_h = \\mathbf{W}_h^Q (\\mathbf{W}_h^K)^\\top$$<\/p>\n\n\n\n We then compare the architectures via similarity matrices and effective-rank via eigenvalue entropy.<\/p>\n\n\n\n We see LRKV exhibits very similar patterns to Standard MHA with sufficient rank of r=64 and achieves 98.3% effective rank at 2.5B scale versus 98.9% for Standard MHA. In contrast, MQA achieves only 86.2% and GQA 95.4%.<\/p>\n\n\n\n Interpreting uncentered vs. PCA-based effective rank. Cache layout $$\\mathbf{K}_{\\text{shared}}, \\mathbf{V}_{\\text{shared}} \\in \\mathbb{R}^{L \\times d_h}, \\quad\\mathbf{R}_h^K, \\mathbf{R}_h^V \\in \\mathbb{R}^{H \\times L \\times r}$$<\/p>\n\n\n\n while Bh<\/sub>K<\/sup><\/strong>, Bh<\/sub>V<\/sup><\/strong> are parameters stored in the weights, not the KV cache.<\/p>\n\n\n\n Kernel fusion $$\\mathbf{q}_h \\mathbf{K}_h^\\top = \\mathbf{q}_h \\mathbf{K}_{\\text{shared}}^\\top + (\\mathbf{q}_h \\mathbf{B}_h^K)(\\mathbf{R}_h^K)^\\top, \\quad These allow exact attention computation to be embedded in fused kernels (e.g. FlashAttention) without reconstructing full Kh<\/sub><\/strong> and Vh<\/sub><\/strong> tensors.<\/p>\n\n\n\n Choosing rank r<\/strong><\/p>\n\n\n\n Use the memory ratio:<\/p>\n\n\n\n $$\\text{ratio} = \\frac{1}{H} + \\frac{r}{d_h}, \\quad r \\approx (0.36\\text{\u2013}0.43)\\,d_h$$<\/p>\n\n\n\n Choose r to match your cache target, then validate quality.<\/p>\n\n\n\n LRKV is a structural change to attention that:<\/p>\n\n\n\n![\"\"](\"https:\/\/s47652.pcdn.co\/research\/wp-content\/uploads\/2026\/04\/Code_Generated_Image.jpeg\")
Why this matters: the KV cache bottleneck<\/strong><\/h2>\n\n\n\n
\n
Background<\/h2>\n\n\n\n
![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/mha_diagram-2-1024x976.png\")
<\/figure>\n\n\n\n![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/mqa_diagram-2-1024x981.png\")
<\/figure>\n\n\n\n![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/gqa_diagram-2-1024x981.png\")
<\/figure>\n\n\n\n![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/mla_diagram-11-577x1024.png\")
<\/figure>\n\n\n\nLRKV parameterization: shared full-rank base + per-head low-rank residual<\/strong><\/h2>\n\n\n\n
\\mathbf{B}_h^{K,V} \\in \\mathbb{R}^{d_h \\times r}$$<\/p>\n\n\n\n![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/lrkv_diagram-5-502x1024.png\")
<\/figure>\n\n\n\nKV caching in LRKV<\/strong><\/h2>\n\n\n\n
\\mathbf{R}_h^V = \\mathbf{X} \\mathbf{U}_h^V \\in \\mathbb{R}^{L \\times r}$$<\/p>\n\n\n\nExact attention without explicitly reconstructing full K\/V tensors<\/strong><\/h2>\n\n\n\n
$$\\mathbf{K}_{\\text{shared}}, \\mathbf{V}_{\\text{shared}} \\in \\mathbb{R}^{L \\times d_h} \\quad \\text{and} \\quad \\mathbf{R}_h^K, \\mathbf{R}_h^V \\in \\mathbb{R}^{L \\times r}.$$<\/p>\n\n\n\nWhat does LRKV cost in compute?<\/strong><\/h2>\n\n\n\n
Results<\/strong><\/h2>\n\n\n\n
Cross-scale pretraining curves<\/strong><\/h3>\n\n\n\n
![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/updated_cross_scale_ce_loss_beautiful-1024x233.jpg\")
![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/updated_training_efficiency_analysis-1-1024x433.jpg\")
Memory vs Performance<\/strong> (left): Test BPB versus KV cache percentage for all methods. LRKV achieves optimal trade-off with lowest BPB at 48.4% cache usage (2.5B scale). Training Efficiency Advantage<\/strong> (right): LRKV reaches each baseline’s final test loss, quantifying training compute savings. LRKV reaches all baselines’ performance earlier.<\/figcaption><\/figure>\n\n\n\n\n
Training efficiency & memory\/performance tradeoff <\/h3>\n\n\n\n
![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/updated_ablation_cache_and_curves-1024x431.jpg\")
Long-context pretraining<\/strong><\/h3>\n\n\n\n
<\/strong><\/p>\n\n\n\n![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/updated_512m_8k_context_ce_loss_beautiful-1024x671.jpg\")
Downstream Task Performance <\/strong><\/h3>\n\n\n\n
![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/midtraining_results_v2-1-1024x618.png\")
<\/figure>\n\n\n\nWhy does LRKV preserve head diversity (and why that\u2019s non-trivial)?<\/strong><\/h3>\n\n\n\n
![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/updated_similarity_heatmaps_2.5B-1024x186.jpg\")
Effective rank across scales<\/strong><\/h4>\n\n\n\n
<\/em>The distinction between uncentered and PCA-based effective rank reveals LRKV’s factorization structure. Uncentered analysis measures total variance including the shared mean direction, the global structure captured by W<\/strong>shared<\/sub>. PCA-based analysis centers the Gram matrix, isolating variance around this mean and measuring true head independence. The modest 4.8% gap indicates LRKV achieves diversity primarily through genuine per-head specialization rather than merely perturbing a dominant shared structure. For comparison, MQA shows dramatic improvement from uncentered (86.2%) to centered (91.0%), a compensation effect where forced KV sharing creates a strong mean direction, but heads recover diversity by aggressively diversifying query projections around this baseline.<\/p>\n\n\n\n![\"\"](\"https:\/\/fin.ai\/research\/wp-content\/uploads\/2026\/04\/updated_effective_rank_combined-1024x380.jpg\"){kind=icon}
\n\n\n\nImplementation notes<\/strong><\/h2>\n\n\n\n
<\/strong>Store:<\/p>\n\n\n\n
<\/strong>The associativity forms used in LRKV are:<\/p>\n\n\n\n
\\mathbf{a}_h \\mathbf{V}_h = \\mathbf{a}_h \\mathbf{V}_{\\text{shared}} + (\\mathbf{a}_h \\mathbf{R}_h^V)(\\mathbf{B}_h^V)^\\top$$<\/p>\n\n\n\n
\n\n\n\nTakeaway<\/strong>s<\/h2>\n\n\n\n