Low-Rank Key Value Attention

TL;DR

LRKV introduces a memory-efficient attention mechanism for Transformer pretraining by decomposing each head's key/value projections into a shared full-rank base plus head-specific low-rank residuals. This additive factorization reduces KV-cache memory by approximately 40-50% (up to 52.6% of the standard cache at 2.5B) while preserving near-full head diversity and representational capacity, enabling faster convergence and better downstream performance. Through gauge-invariant analyses (bilinear forms and kernel PCA), the authors show LRKV maintains functional head diversity comparable to full MHA, while also revealing compensation mechanisms in KV-sharing baselines like MQA. Empirically, LRKV delivers state-of-the-art pretraining results on FineWeb-Edu across 128M and 2.5B scale, and achieves superior downstream scores after midtraining, with notable improvements in memory/compute efficiency. The work demonstrates that structured low-rank factorization of KV projections is a principled route to scaling Transformer training under memory constraints without sacrificing quality, and outlines concrete directions for extending this approach to broader model classes and tasks.

Abstract

Transformer pretraining is increasingly constrained by memory and compute requirements, with the key-value (KV) cache emerging as a dominant bottleneck during training and autoregressive decoding. We propose \textit{low-rank KV adaptation} (LRKV), a simple modification of multi-head attention that reduces KV cache memory by exploiting redundancy across attention heads while preserving full token-level resolution. Each layer uses a shared full-rank KV projection augmented with low-rank, head-specific residuals, yielding a continuous trade-off between complete sharing and fully independent attention. LRKV is a drop-in replacement for standard multi-head attention and directly subsumes query-sharing approaches such as multi-query and grouped-query attention, while remaining distinct from latent-compression methods such as multi-latent attention (MLA). Across large-scale pretraining experiments, LRKV consistently achieves faster loss reduction, lower validation perplexity, and stronger downstream task performance than standard attention, MQA/GQA, and MLA. At the 2.5B scale, LRKV outperforms standard attention while using roughly half the KV cache, and reaches equivalent model quality with up to \textbf{20-25\% less training compute} when measured in cumulative FLOPs. To explain these gains, we analyze attention head structure in operator space and show that LRKV preserves nearly all functional head diversity relative to standard attention, whereas more aggressive KV-sharing mechanisms rely on compensatory query specialization. Together, these results establish LRKV as a practical and effective attention mechanism for scaling Transformer pretraining under memory- and compute-constrained regimes.

Figures (18)

• Figure 1: Comparison of attention mechanisms and their KV cache costs. (a) Multi-Head Attention (MHA) uses independent K/V projections per head with cache cost $2LHd_h$. (b) Multi-Latent Attention compresses inputs to latent dimension $d_c$ before per-head projections (cache: $Ld_c$ where $d_c \ll d$). (c) Multi-Query Attention (MQA) shares a single K/V across all heads (cache: $2Ld_h$) but suffers from low diversity. (d) Grouped-Query Attention (GQA) shares K/V within groups of heads (cache: $2L(H/G)d_h$), balancing efficiency and diversity. (e) Our proposed LR-KV method maintains shared full-rank $W_K, W_V$ projections while each head adds a low-rank residual $U_h (B_h)^\top$ (rank $r \ll d_h$). We cache both the shared features and per-head latents $R_h = XU_h$, achieving cache cost $2L(d_h + Hr)$. Dashed boxes indicate cached components. LR-KV provides head diversity while maintaining cache efficiency comparable to MLA.
• Figure 2: Pretraining test curves on FineWeb-Edu across model scales and metrics.(a,b) Test bits-per-byte (BPB) for 128M and 2.5B models respectively, measuring compression efficiency. (c,d) Test cross-entropy loss for 128M and 2.5B models respectively. LRKV (purple solid line) consistently achieves the lowest test loss across both scales and metrics, demonstrating superior language modeling performance. At 2.5B scale, LRKV achieves 0.719 BPB while using only 52.6% of Standard MHA's KV cache (with theoretical latent caching). The performance advantage is consistent throughout training, indicating that low-rank residuals effectively capture head-specific information while maintaining efficient memory usage. All models trained for 190K steps: 100B tokens (128M) and 50B tokens (2.5B).
• Figure 3: LRKV achieves superior training efficiency alongside best performance (2.5B scale).(a) Memory vs Performance: LRKV achieves the lowest validation BPB (0.719) across all methods, using 52.6% of Standard MHA's KV cache. While memory-minimal methods (MQA: 5.6%, MLA: 16.7%) use less cache, they show degraded performance (0.729 and 0.724 BPB respectively). (b) Training Efficiency Advantage: LRKV reaches each baseline's final performance 18--30% faster, quantifying training compute savings. To match MQA's final performance (0.729 BPB), LRKV requires 29.7% fewer training steps (14.9B tokens saved). For GQA: 24.5% savings (12.3B tokens), MLA: 21.4% savings (10.7B tokens), Standard MHA: 18.7% savings (9.4B tokens). Average savings: 23.6% across all baselines. This dual advantage—superior sample efficiency and better final performance—makes LRKV strictly preferable for large-scale pretraining.
• Figure 4: FLOPs-normalized pretraining curves (2.5B scale). Test BPB and cross-entropy plotted as a function of cumulative training compute (ExaFLOPs), rather than token count. LRKV (purple) achieves the lowest test loss throughout training despite using 0.8% more FLOPs per token than baselines. The x-axis accounts for mechanism-specific per-token costs: MQA -6%, GQA -3.2%, MLA -3.7%, LRKV +0.8%. Curves are truncated at 700 ExaFLOPs for equal-compute comparison.
• Figure 5: Pretraining quality predicts downstream performance. Scatter plot relating pretraining validation BPB (x-axis, inverted) to downstream combined performance (y-axis) across all attention mechanisms at 2.5B scale. A strong positive correlation emerges ($R²=0.786$), demonstrating that models with better pretraining performance consistently achieve higher downstream scores. LRKV occupies the optimal position with both the lowest pretraining BPB (0.719) and highest downstream performance (37.9%), while Standard MHA achieves 0.723 BPB and 35.3% downstream.