2601.11471

# Low-Rank Key Value Attention **Authors**: James O’Neill, Robert Clancy, Mariia Matskevichus, Fergal Reid ## Abstract The key-value (KV) cache is a primary memory bottleneck in Transformers. We propose Low-Rank Key-Value (LRKV) attention, which reduces KV cache memory by exploiting redundancy across attention heads, while being compute efficient. Each layer uses a shared full-rank KV projection augmented with low-rank, head-specific residuals, providing a continuous trade-off between complete sharing and full independence. After pretraining models of size 128M to 6.3B parameters, LRKV consistently achieves the lowest test loss among standard MHA, MQA/GQA, and MLA while using only 45-53% of MHA’s KV cache. LRKV reaches equivalent baseline quality 18-25% faster (measured in training steps). After supervised midtraining, LRKV achieves the highest downstream task performance across ARC-Easy, ARC-Challenge, MMLU, GSM8K, and HumanEval benchmarks. Machine Learning, ICML ## 1 Introduction Transformers are the dominant architecture for large-scale sequence modeling in language, vision, and multimodal domains (Vaswani et al., 2017; OpenAI, 2023), but as their size, sequence length, and context window grow, so does, rapidly, their computational and memory costs. KV-caching, which stores the attention key and value representations, is a primary contributor to this overhead, as it spans every attention layer, and scales linearly with sequence length and count of heads. The cumulative KV footprint of modern models with tens of billion parameters can exceed the parameter memory itself, especially for long-context inference (Dao et al., 2024; Child et al., 2019). To illustrate the scale of this challenge, consider a 2.5B parameter model (18 layers, 18 heads, $d_h{=}128$ ) serving 8K-token contexts: the KV cache alone requires $2× 18× 18× 128× 8192× 2 bytes=648 MB$ per request. At batch size 4, KV cache memory (2.6 GB) exceeds the model’s parameter memory (5 GB in bfloat16), consuming more memory than the model itself. This forces practitioners into an impossible trilemma: use smaller batches (sacrificing throughput), shorter contexts (limiting capability), or frequent cache eviction (adding latency)—each option directly degrading production utility. Reducing KV cache size thus directly translates to 2 $×$ longer contexts or 2 $×$ larger batches at fixed memory budgets. <details> <summary>2601.11471v3/figures/fig1_intro_diagram.png Details</summary>  ### Visual Description ## Diagram: Comparison of Attention Mechanisms in Transformer Architectures ### Overview The diagram compares three attention mechanisms: **Standard MHA**, **MQA/GQA**, and **LRKV (Ours)**. Each section includes a visual workflow, labels, and textual descriptions of computational costs and design principles. --- ### Components/Axes 1. **Standard MHA** - **Visual Components**: - Three parallel "Head" blocks (Head 1, Head 2, Head 3) with dashed bidirectional arrows connecting to a central "X" node. - Text: "H independent projections. Cache Cost: High (2LHd_h)". - **Key Labels**: - "Head 1", "Head 2", "Head 3" (blue boxes). - "X" (central node). 2. **MQA/GQA** - **Visual Components**: - A central "Shared Projection" block (green) with dashed arrows connecting to multiple "X" nodes. - Text: "1 Shared projection. Cache Cost: Low, but loses detail." - **Key Labels**: - "Shared Projection" (green rectangle). 3. **LRKV (Ours)** - **Visual Components**: - Two stacked blocks: - **Shared Parameters (Full-Rank)** (green): Contains matrices `W_K_shared`, `W_V_shared`, and text "Captures global relational structures." - **Head-Specific Residuals (Low-Rank)** (orange): Contains matrices `U_h`, `B_h^T`, and text "Captures unique deviations and specialization." - Text: "1 Shared + H Residuals. Cache Cost: Low (2LL(d_h + Hr)) Preserves full token resolution." - **Key Labels**: - "Shared Parameters (Full-Rank)", "Head-Specific Residuals (Low-Rank)". --- ### Detailed Analysis 1. **Standard MHA** - **Structure**: Three independent attention heads process input `X` and aggregate results via bidirectional connections. - **Cache Cost**: Explicitly labeled as "High (2LHd_h)", where `L` and `H` likely represent sequence length and number of heads, respectively. 2. **MQA/GQA** - **Structure**: A single shared projection layer processes all heads, reducing redundancy. - **Cache Cost**: Labeled "Low", but notes a trade-off: "loses detail" due to shared parameters. 3. **LRKV (Ours)** - **Structure**: Combines a shared full-rank parameter layer (global structures) with low-rank head-specific residuals (local deviations). - **Cache Cost**: Labeled "Low (2LL(d_h + Hr))", where `d_h` is head dimension and `Hr` is residual rank. - **Key Advantage**: "Preserves full token resolution" by balancing shared and residual components. --- ### Key Observations 1. **Cache Cost Trade-offs**: - Standard MHA has the highest cache cost due to independent projections. - MQA/GQA reduces cost via shared projections but sacrifices detail. - LRKV optimizes cost further while retaining resolution through residual learning. 2. **Design Philosophy**: - LRKV explicitly separates global (shared) and local (residual) features, unlike MQA/GQA’s monolithic shared layer. 3. **Mathematical Notation**: - LRKV’s cache cost formula `2LL(d_h + Hr)` suggests linear scaling with sequence length (`L`) and residual rank (`Hr`), contrasting with Standard MHA’s quadratic dependency on heads (`H`). --- ### Interpretation - **Efficiency vs. Fidelity**: The diagram highlights a spectrum from computationally expensive but detailed (Standard MHA) to efficient but coarse (MQA/GQA). LRKV bridges this gap by decomposing attention into shared and residual components. - **Practical Implications**: LRKV’s design suggests suitability for tasks requiring both global context (e.g., long-range dependencies) and fine-grained local patterns (e.g., token-level nuances). - **Unresolved Questions**: The diagram does not quantify performance metrics (e.g., accuracy, latency) or specify hardware constraints, leaving implementation trade-offs ambiguous. This structured breakdown captures all textual and visual elements, emphasizing technical trade-offs and design principles. The LRKV method’s innovation lies in its hybrid approach, balancing efficiency and resolution—a critical insight for transformer architecture optimization. </details> Figure 1: Comparison of attention mechanisms. Standard MHA: Uses $H$ independent projections per head, requiring high cache cost ( $2LHd_h$ ). MQA/GQA: Shares projections across all heads, reducing cache but losing head-specific detail. LRKV (Ours): Combines shared full-rank parameters with head-specific low-rank residuals ( $U_hB_h^T$ , rank $r\ll d_h$ ), achieving low cache cost $2L(d_h+Hr)$ while preserving head diversity. A range of approaches have been proposed to alleviate the growing KV cache cost. Multi-Query Attention (Shazeer, 2019, MQA) and Grouped-Query Attention (Ainslie et al., 2023, GQA) reduce memory and latency by sharing key-value (KV) projections across heads or groups of heads, and are now standard in large-scale models such as PaLM and LLaMA (Touvron et al., 2023, 2024). However, aggressive KV sharing introduces a fundamental trade-off: while memory is reduced, head-level representational diversity is constrained, even though distinct attention heads are known to encode complementary syntactic and semantic patterns (Clark et al., 2019; Michel et al., 2019). At the same time, empirical studies and recent spectral analyses show that attention heads are not fully independent: head-specific KV projections are highly correlated and occupy overlapping subspaces, indicating substantial redundancy (Yunis et al., 2024). Crucially, this redundancy is structured rather than uniform—small, head-specific variations remain important for capturing nuanced dependencies. This raises a natural question: can KV memory be reduced by exploiting redundancy across heads, without collapsing the specialization that makes multi-head attention effective? Beyond KV sharing, efficiency research has pursued complementary directions such as sparse or kernelized attention (Wang et al., 2020; Beltagy et al., 2020; Child et al., 2019), architectural optimizations (Dao et al., 2024), and latent compression methods including Multi-Latent Attention (Liu et al., 2024, MLA). However, none of these approaches explicitly resolve the duplication of per-head key and value representations that dominates the KV memory footprint in large models. In this work, we propose Low-Rank Key-Value Attention (LRKV), a simple modification of multi-head attention that exploits structured redundancy across heads while preserving head specialization. Each Transformer layer maintains a shared, full-rank key and value projection that serves as a global basis, while each head learns a compact, trainable low-rank residual that captures head-specific deviations. The shared component encodes global relational structure, and the residuals restore the localized diversity otherwise lost under aggressive KV sharing. This factorization substantially reduces KV memory while retaining the expressivity of multi-head attention. LRKV differs fundamentally from prior KV-efficient designs in where compression is applied. Unlike LoRA (Hu et al., 2022), which introduces low-rank structure post hoc during fine-tuning, LRKV learns head-specific low-rank residuals jointly during pretraining. Unlike factorized QKV methods that reduce rank along the token or hidden dimension (Wang et al., 2020; Saxena et al., 2024; Chang et al., 2025), LRKV preserves full token resolution and compresses only across attention heads. We now discuss the main contributions of this paper. Contributions. (1) LRKV reduces KV cache by 45-53% while achieving lowest test loss across 128M-6.3B models, reaching baseline performance 18-25% faster in training steps. (2) LRKV achieves highest downstream performance across ARC, MMLU, GSM8K, demonstrating better pretraining translates to improved capabilities. (3) Gauge-invariant analysis shows LRKV effectively uses low-rank structure by preserving 93.5% head diversity vs 94% MHA. Code and trained models will be released upon acceptance. ## 2 Related Work KV-Sharing Mechanisms. MQA and GQA reduce KV cache by sharing projections across heads, but sacrifice head-level expressivity. MQA (Shazeer, 2019) uses a single shared KV projection for all heads, achieving maximal cache reduction but constraining diversity. GQA (Ainslie et al., 2023) groups heads to share KV projections within groups, interpolating between MQA and full MHA. While both methods reduce memory, they fundamentally limit the representational capacity available to each head. LRKV preserves shared projection efficiency while restoring diversity through structured low-rank head residuals, achieving better quality at comparable cache sizes. Latent Compression Approaches. MLA (Liu et al., 2024) takes a complementary approach: rather than sharing across heads, it compresses across the token dimension by projecting inputs into a low-dimensional latent space ( $d_c\ll d$ ) before caching. While effective for extreme compression, this bottleneck constrains all heads to operate through the same latent representation and requires T-dependent reconstruction overhead during generation (see Appendix F). LRKV instead preserves full token-level resolution and compresses along the head dimension via additive factorization ( $W_h=W_shared+U_hB_h^⊤$ ), avoiding latent bottlenecks while maintaining per-head specialization. These represent orthogonal design choices: MLA optimizes for minimal cache size via aggressive token compression, while LRKV balances memory efficiency with head-level expressivity. Low-Rank Parameterization. LoRA (Hu et al., 2022) introduces low-rank updates for parameter-efficient fine-tuning, intervening on the optimization pathway rather than the cached representations. LRKV applies the low-rank principle directly to KV projections during pretraining, shaping the structure of cached features from the start. Recent work explores low-rank QKV factorizations (Xie and others, 2023; Khalaf et al., 2025; Lv et al., 2024) to reduce parameters or computation; LRKV differs by preserving a full-rank shared base (capturing global structure) while using low-rank residuals only for head-specific deviations. Concurrent work on factorized KV mechanisms (MFA/MFA-KR (Hu et al., 2025), TPA (Zhang et al., 2025)) explores related structured compression ideas. Unlike joint-head low-rank decompositions such as J-LRD/Palu, which compress projections post hoc, LRKV is a pretraining-time architectural reparameterization that preserves a full-rank shared base while learning head-specific low-rank residuals. This enables exact associative decoding and reduces KV cache without post-hoc approximation. Complementary Efficiency Methods. LRKV is orthogonal to other efficiency techniques: sparse/kernelized attention (Child et al., 2019; Beltagy et al., 2020; Wang et al., 2020) reduce computational complexity but not KV cache size; architectural optimizations like FlashAttention (Dao et al., 2024) improve throughput through better memory access patterns; quantization and pruning reduce precision or parameters. For head diversity analysis, we extend prior metrics (Michel et al., 2019; Kornblith et al., 2019; Wang and Wang, 2025) using gauge-invariant bilinear forms with centered Gram matrix analysis (see Appendix A). ## 3 Methodology Transformers represent each token in a sequence through queries, keys, and values that interact via scaled dot-product attention. For an input matrix $X∈ℝ^L× d$ , the $h$ -th attention head computes $Q_h=XW_h^Q$ , $K_h=XW_h^K$ , and $V_h=XW_h^V$ , where $W_h^Q,K,V∈ℝ^d× d_h$ and $d_h=d/H$ for $H$ heads. The head output is $$ O_h=softmax≤ft(\frac{Q_hK_h^⊤}{√{d_h}}\right)V_h, \tag{1} $$ and outputs from all heads are concatenated and projected to form the layer output. KV caching. During autoregressive decoding, the KV cache stores $K_h$ and $V_h$ for all previous tokens and heads, incurring per-layer memory $M_standard=2LHd_h=2Ld$ . Equivalently, per token the cache stores $2Hd_h$ floating-point values (keys and values across $H$ heads). Over $N$ layers in bfloat16 precision, this is $2N·(2LHd_h)$ bytes of KV storage. Motivation. Empirical studies show attention heads within a layer are often correlated (Clark et al., 2019; Michel et al., 2019), suggesting that per-head key/value features contain substantial redundancy. Existing KV-sharing methods such as MQA and GQA reduce cache size by sharing K/V across heads, but can reduce head diversity and modeling capacity. Our goal is to reduce redundant KV storage while preserving head-specific flexibility. Low-Rank KV Attention. We parameterize each head’s key/value projection as a shared dense (unconstrained) base plus a head-specific low-rank residual, while keeping the resulting projection shape $d× d_h$ : $$ \displaystyleW_h^K \displaystyle=W_shared^K+U_h^K{B_h^K}^⊤, \displaystyleW_h^V \displaystyle=W_shared^V+U_h^V{B_h^V}^⊤, \tag{2} $$ where $W_shared^K,V∈ℝ^d× d_h$ are dense projection matrices with no rank constraint, $U_h^K,V∈ℝ^d× r$ , $B_h^K,V∈ℝ^d_h× r$ , and $r\ll d_h$ . The effective keys and values are $K_h=XW_h^K$ and $V_h=XW_h^V$ . When $r=0$ , LRKV reduces to complete KV sharing (MQA-style) within a layer; increasing $r$ interpolates continuously toward standard MHA. Training and initialization. All parameters in Equation 3 are optimized jointly during pretraining. Gradients from the shared and residual paths are additive, allowing the model to learn how much per-head variation is required. We initialize $W_shared^K,V$ using standard Kaiming initialization for attention projections. The per-head factors $U_h^K,V$ and $B_h^K,V$ are initialized to small random values scaled by $1/√{r}$ such that initial residuals $U_h{B_h}^⊤$ have magnitude $≈ 0.1×\|W_shared\|_F$ , ensuring the model starts close to the shared baseline and gradually learns head specialization. During training, we materialize full $K_h$ and $V_h$ matrices before attention computation for simplicity; the factored form is used only during inference to realize memory savings. This is an exact reparameterization, so the factorized and unfactorized forms are mathematically equivalent. We apply low-rank factorization only to key and value projections, as these are cached during inference; query projections remain independent per head at full rank. LRKV caching scheme. During decoding, the bottleneck is storing $K_h,V_h∈ℝ^L× d_h$ for every head. LRKV instead caches shared features once per layer, $$ K_shared=XW_shared^K, V_shared=XW_shared^V∈ℝ^L× d_h, $$ and compact per-head latents $$ R_h^K=XU_h^K, R_h^V=XU_h^V∈ℝ^L× r. $$ The full per-head features implied by Equation 3 are $$ \displaystyleK_h \displaystyle=K_shared+R_h^K{B_h^K}^⊤, \displaystyleV_h \displaystyle=V_shared+R_h^V{B_h^V}^⊤. \tag{4} $$ Importantly, the matrices $B_h^K,V$ are model parameters and are not cached per token. Attention computation without explicit reconstruction. Naively materializing $K_h,V_h$ for all cached tokens would cost $O(Lrd_h)$ memory and compute. Instead, we exploit associativity to compute attention exactly without forming full per-head KV tensors. For a decoding step with query $q_h∈ℝ^d_h$ , $$ q_hK_h^⊤=q_hK_shared^⊤+(q_hB_h^K) (R_h^K)^⊤, \tag{6} $$ and for attention weights $a_h∈ℝ^1× L$ , $$ a_hV_h=a_hV_shared+(a_hR_h^V) {B_h^V}^⊤. \tag{7} $$ Equations 6 and 7 compute the exact attention logits and outputs implied by Equation 4 – Equation 5, but avoid explicit reconstruction of $K_h,V_h$ . This form can be implemented inside fused attention kernels. See Table 4 for a detailed mathematical comparison to existing mechanisms. #### Compatibility with Positional Embeddings. Positional embeddings, such as RoPE, are applied after projection and distribute linearly over the LRKV decomposition: $$ RoPE(K_h)=RoPE(K_shared)+RoPE(R_h)B_h^⊤. $$ Both components are rotated once at caching time, as in standard KV caching, and reused during decoding. This introduces no additional memory or sequence-length-dependent computation. In contrast to latent-attention approaches such as MLA, which require partial RoPE to preserve projection absorption, LRKV supports full-dimension RoPE without modification. KV cache memory complexity (during decoding). Standard attention stores per-head keys and values: $$ M_standard=2LHd_h. $$ LRKV stores shared features plus per-head latents: $$ \displaystyle M_LRKV \displaystyle=\underbrace{2Ld_h}_shared (K,V)+\underbrace{2LHr}_per-head latents \displaystyle=2L(d_h+Hr). \tag{8} $$ Thus the memory ratio is $$ \frac{M_LRKV}{M_standard}=\frac{d_h+Hr}{Hd_h}=\frac{1}{H}+\frac{r}{d_h}. \tag{9} $$ Per token, LRKV stores $2d_h+2Hr$ values versus $2Hd_h$ for standard attention. Table 1: Pretraining performance across model scales. LRKV achieves competitive test loss across all model scales (128M, 1.2B, 2.5B, 6.3B) while maintaining efficient KV cache usage. | Standard MHA | 6/12/18/32 | 100 | 100 | 100 | 100 | 2.903 | 0.878 | 2.530 | 0.765 | 2.389 | 0.723 | 2.319 | 0.701 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | | GQA | 3/4/6/2 | 50 | 33 | 33 | 6 | 2.918 | 0.883 | 2.536 | 0.767 | 2.397 | 0.725 | 2.354 | 0.712 | | MQA | 1 | 17 | 8 | 6 | 3 | 2.929 | 0.886 | 2.573 | 0.778 | 2.408 | 0.729 | 2.304 | 0.697 | | MLA | – | 8.3 | 8.3 | 8.3 | 12.5 | 2.901 | 0.876 | 2.564 | 0.775 | 2.392 | 0.724 | 2.377 | 0.719 | | LRKV | – | 53 | 48 | 48 | 45 | 2.893 | 0.872 | 2.509 | 0.758 | 2.376 | 0.719 | 2.288 | 0.692 | KV Heads show values for 128M/1.2B/2.5B/6.3B scales. KV Cache % indicates memory relative to Standard MHA baseline (see Appendix G for profiling details). Test metrics on held-out FineWeb-Edu (100B tokens for 128M, 50B for 1.2B/2.5B/6.3B). Compute overhead during decoding (FLOPs). Per decoding step, standard attention has dominant per-head computational cost $O(Ld_h)$ from query–key dot products and value aggregation. Using the distributed forms in Equation 6 and Equation 7, LRKV adds an extra $ΔFLOPs=O(Lr+rd_h)$ per head per decoding step. For long contexts ( $L\gg 1$ ), the dominant additional term scales as $O(Lr)$ , yielding a relative overhead of approximately $r/d_h$ compared to standard attention. Memory bandwidth. Modern inference is often bandwidth-bound rather than compute-bound (Dao et al., 2024). Per decoding step, standard MHA reads $2H$ cached tensors ( $K_h$ and $V_h$ for each head). LRKV reads two shared tensors plus $2H$ per-head latent tensors. Since the shared tensors ( $L× d_h$ ) are reused across heads and the per-head latents ( $L× r$ ) are substantially smaller ( $r\ll d_h$ ), the total bytes transferred is $2L(d_h+Hr)$ for LRKV versus $2LHd_h$ for standard MHA, matching the cache-size reduction and directly translating memory savings into reduced bandwidth pressure. Parameter complexity. Standard per-layer K/V parameters scale as $P_MHA,KV=2Hdd_h$ . LRKV uses a shared base plus low-rank factors: $$ P_LRKV=\underbrace{2dd_h}_shared (K,V)+\underbrace{2Hr(d+d_h)}_per-head low-rank factors, \tag{10} $$ which may be lower or higher than standard depending on $(H,r,d_h)$ . In practice, we choose $r$ to prioritize KV-cache reduction with minimal quality impact; in our experiments, LRKV remains competitive or superior even when the total K/V parameter count is below that of standard MHA. Rank as a control for diversity. The residual rank $r$ controls a continuous spectrum between fully shared KVs (MQA-style, $r=0$ ) and fully independent per-head projections (standard MHA, large $r$ ). The shared base $W_shared^K,V$ captures globally useful features, while low-rank residuals enable head specialization within a constrained budget. Empirically, we find $r≈ 0.36$ - $0.43× d_h$ provides a critical threshold: LRKV achieves 93.5% PCA-based effective rank versus 94.0% for standard MHA at 2.5B scale (subsection 4.3), preserving nearly all head diversity while achieving substantial cache reduction. See Appendix B for detailed discussion of spectral structure and low-rank factorization. <details> <summary>2601.11471v3/x1.png Details</summary>  ### Visual Description ## Line Graphs: Test Cross-Entropy Loss vs. Training Tokens Across Model Sizes ### Overview The image contains four line graphs (a-d) comparing the performance of different training methods (Standard MHA, LRKV, MQA, GQA, MLA) across four model sizes (128M, 1B, 2.5B, 6.3B parameters). Each graph plots **Test Cross-Entropy Loss** (y-axis) against **Training Tokens (Billions)** (x-axis), with distinct line styles/colors for each method. All graphs show a general downward trend in loss as training tokens increase, with convergence patterns varying by model size. --- ### Components/Axes #### Common Elements Across All Graphs: - **X-axis**: Training Tokens (Billions), ranging from 0 to 100 (a), 0–50 (b), 0–70 (c), and 0–1750 (d). - **Y-axis**: Test Cross-Entropy Loss, ranging from ~2.85–3.15 (a), ~2.5–2.9 (b), ~2.35–2.65 (c), and ~2.3–2.8 (d). - **Legend**: Located in the top-left corner of each graph, with five methods: - **Standard MHA**: Solid blue line - **LRKV**: Solid red line - **MQA**: Dashed green line - **GQA**: Dashed orange line - **MLA**: Dotted purple line #### Graph-Specific Details: - **(a) 128M Parameters**: X-axis spans 0–100B tokens; y-axis loss starts near 3.15 and drops to ~2.9. - **(b) 1B Parameters**: X-axis spans 0–50B tokens; y-axis loss starts near 2.9 and drops to ~2.55. - **(c) 2.5B Parameters**: X-axis spans 0–70B tokens; y-axis loss starts near 2.65 and drops to ~2.4. - **(d) 6.3B Parameters**: X-axis spans 0–1750B tokens; y-axis loss starts near 2.8 and drops to ~2.3. --- ### Detailed Analysis #### Graph (a): 128M Parameters - **Trend**: All methods show steep initial declines, with lines converging toward ~2.9 loss at 100B tokens. - **Key Data Points**: - **Standard MHA** (blue): Starts at ~3.15, ends at ~2.9. - **LRKV** (red): Starts at ~3.1, ends at ~2.9. - **MQA** (green): Starts at ~3.1, ends at ~2.9. - **GQA** (orange): Starts at ~3.05, ends at ~2.9. - **MLA** (purple): Starts at ~3.0, ends at ~2.9. #### Graph (b): 1B Parameters - **Trend**: Faster initial decline than (a), with lines converging at ~2.55 loss by 50B tokens. - **Key Data Points**: - **Standard MHA**: Starts at ~2.9, ends at ~2.55. - **LRKV**: Starts at ~2.85, ends at ~2.55. - **MQA**: Starts at ~2.85, ends at ~2.55. - **GQA**: Starts at ~2.8, ends at ~2.55. - **MLA**: Starts at ~2.75, ends at ~2.55. #### Graph (c): 2.5B Parameters - **Trend**: Steeper initial drop than (b), with convergence at ~2.4 loss by 70B tokens. - **Key Data Points**: - **Standard MHA**: Starts at ~2.65, ends at ~2.4. - **LRKV**: Starts at ~2.6, ends at ~2.4. - **MQA**: Starts at ~2.6, ends at ~2.4. - **GQA**: Starts at ~2.55, ends at ~2.4. - **MLA**: Starts at ~2.5, ends at ~2.4. #### Graph (d): 6.3B Parameters - **Trend**: Most pronounced initial decline, with lines converging at ~2.3 loss by 1750B tokens. - **Key Data Points**: - **Standard MHA**: Starts at ~2.8, ends at ~2.3. - **LRKV**: Starts at ~2.75, ends at ~2.3. - **MQA**: Starts at ~2.7, ends at ~2.3. - **GQA**: Starts at ~2.65, ends at ~2.3. - **MLA**: Starts at ~2.6, ends at ~2.3. --- ### Key Observations 1. **Convergence**: All methods converge to similar loss values at maximum training tokens, suggesting diminishing returns with scale. 2. **MLA Advantage**: MLA (dotted purple) consistently starts with the lowest loss across all model sizes but converges to match others. 3. **Model Size Impact**: Larger models (6.3B) achieve lower final loss than smaller models (128M) despite requiring vastly more training tokens. 4. **Efficiency**: Smaller models (128M) reach near-optimal performance faster (100B tokens) than larger models (6.3B requires 1750B tokens). --- ### Interpretation The data demonstrates that: - **Larger models** achieve lower final loss but require disproportionately more compute (e.g., 6.3B needs 17.5x more tokens than 128M). - **Training methods** like MLA and GQA show early efficiency gains but face diminishing returns, while Standard MHA and LRKV perform comparably across scales. - The steepest declines in loss occur early in training, suggesting that most learning happens in the initial token ranges. For example, 6.3B models drop ~0.5 loss units in the first 100B tokens but only ~0.1 units in the final 1650B tokens. This highlights a trade-off between model size and training efficiency, with larger models offering better performance at the cost of significantly higher computational requirements. </details> Figure 2: Cross-scale pretraining curves with dual-axis compute metrics. Test cross-entropy loss across four model scales (128M, 1.2B, 2.5B, 6.3B) plotted against training tokens (bottom x-axis) and cumulative training compute in ExaFLOPs (top x-axis). LRKV demonstrates competitive performance across all scales, achieving the lowest test loss at 128M and 2.5B scales. Final results in Table 1. ## 4 Experiments We evaluate LRKV through large-scale pretraining experiments across four model sizes (128M, 1.2B, 2.5B, 6.3B parameters), comparing against standard MHA and state-of-the-art KV-efficient baselines. We measure pretraining loss, training efficiency, downstream task performance, and provide detailed analysis of why LRKV preserves modeling quality despite substantial cache reduction. Experimental setup. We pretrain decoder-only Transformers at four scales (128M, 1.2B, 2.5B, 6.3B) on FineWeb-Edu (Penedo et al., 2024) for 100B tokens (128M) and 50B tokens (others), comparing LRKV against Standard MHA, GQA, MQA, and MLA. Models use 2048-token context, Muon+AdamW optimizers (Sardana and Havens, 2024), and are pretrained on 8 $×$ H200 GPUs in bfloat16. After pretraining, we perform supervised midtraining on 568K examples from SmolTalk (Allal et al., 2024), MMLU (Hendrycks et al., 2021), and GSM8K (Cobbe et al., 2021), then evaluate on ARC-Easy, ARC-Challenge, MMLU, GSM8K, and HumanEval. LRKV uses 45-53% of Standard MHA’s KV cache (see Appendix D for details). ### 4.1 Pretraining Results #### Final pretraining performance. Table 1 shows LRKV achieves competitive test loss across all scales (128M, 1.2B, 2.5B, 6.3B), with the best performance at 128M and 2.5B scales. At 6.3B scale, LRKV reaches 2.288 CE (0.692 BPB), outperforming MHA (2.319 CE), MQA (2.304 CE), GQA (2.354 CE), and MLA (2.377 CE). At 1.2B scale, LRKV achieves 2.509 CE (0.758 BPB) with only 48% of MHA’s cache—outperforming all baselines while using half the memory. At 2.5B, LRKV achieves 0.719 BPB with only 48.4% of MHA’s cache—a strictly better accuracy-memory tradeoff. Cache efficiency improves at larger scales (52.6% → 48% → 48.4% → 45.1%), making LRKV increasingly attractive for large models. Table 2: Downstream task performance after midtraining (128M, 2.5B, and 6.3B scales). LRKV achieves the highest combined accuracy across all three scales (18.9%, 37.9%, 40.2%) on five diverse benchmarks, demonstrating that superior pretraining performance translates to stronger downstream capabilities. Combined score (Comb.) is the average across all five benchmarks. HE = HumanEval. | Standard MHA GQA MQA | 26.4 28.0 28.0 | 26.0 25.4 27.1 | 27.2 25.7 27.8 | 1.0 0.5 1.3 | 8.8 4.6 2.4 | 17.9 16.8 17.3 | 66.6 65.4 65.2 | 47.1 49.6 47.3 | 39.3 40.5 40.2 | 10.2 10.0 10.3 | 13.4 13.4 3.7 | 35.3 35.8 33.3 | 72.7 69.0 67.2 | 53.6 48.2 47.0 | 42.7 40.9 40.5 | 10.6 9.9 8.2 | 14.6 13.4 4.3 | 38.8 36.3 33.4 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | | MLA | 27.3 | 27.8 | 26.3 | 1.1 | 7.7 | 18.0 | 67.5 | 51.5 | 41.9 | 10.8 | 12.8 | 36.9 | 69.4 | 51.3 | 40.9 | 9.5 | 13.4 | 36.9 | | LRKV | 26.7 | 30.2 | 28.4 | 0.3 | 9.1 | 18.9 | 70.7 | 53.8 | 42.2 | 11.3 | 11.7 | 37.9 | 75.0 | 58.0 | 44.5 | 11.5 | 12.8 | 40.2 | Cross-scale training efficiency and compute normalization. Figure 2 presents test performance across all four model scales, measured in cross-entropy loss with dual-axis compute metrics (training tokens and cumulative FLOPs). LRKV demonstrates competitive performance across scales, achieving the lowest test loss at 128M and 2.5B scales while remaining highly competitive at 1.2B and 6.3B, showing consistent advantages across three orders of magnitude in model size. The top x-axis shows cumulative training compute (ExaFLOPs), accounting for mechanism-specific per-token costs (MQA -6%, GQA -3.2%, MLA -3.7%, LRKV +0.8% at 2.5B scale). The reported per-token costs reflect inference FLOPs. During training, LRKV materializes full K/V matrices (as do all methods), so training cost per step is comparable to MHA. The 18-25% “training compute savings” refer to sample efficiency (fewer steps to reach target loss), not per-step speedup. When normalized by total compute rather than token count, LRKV maintains its performance advantage: at every ExaFLOP milestone, LRKV achieves lower test loss than all baselines. At 2.5B scale, LRKV reaches equivalent baseline performance 18–30% faster (23.6% average across all methods), while no baseline reaches LRKV’s final performance even after exhausting the full compute budget. This establishes LRKV as dominant in both sample efficiency and final quality. At the largest 6.3B scale, the performance gap widens: LRKV achieves 2.288 CE versus Standard MHA’s 2.319 CE (1.3% improvement), demonstrating that LRKV’s architectural advantages strengthen rather than diminish with scale. <details> <summary>2601.11471v3/x2.png Details</summary>  ### Visual Description ## Chart/Diagram Type: Dual Axes Comparison Charts ### Overview The image contains two side-by-side charts comparing memory usage, performance, and training efficiency of different key-value (KV) cache methods relative to a standard multi-head attention (MHA) baseline. The left chart (a) focuses on memory vs. validation performance, while the right chart (b) highlights training compute savings. ### Components/Axes #### Chart (a): Memory vs Performance - **X-axis**: KV Cache Size (% of Standard MHA) - Scale: 0% to 100% in 20% increments - **Y-axis**: Final Validation BPB (lower is better) - Scale: 0.718 to 0.730 in 0.002 increments - **Legend**: - Green triangle: MHA - Red diamond: MLA - Orange square: GQA - Purple star: LRKV - Blue circle: Standard MHA - **Embedded Text**: - "LRKV (Best performance, moderate memory)" #### Chart (b): Training Efficiency Advantage - **X-axis**: KV Cache Size (% of Standard MHA) - Scale: 0% to 100% in 20% increments - **Y-axis**: Training Compute Savings (% to reach baseline performance) - Scale: 0% to 30% in 5% increments - **Legend**: - Same symbols/colors as Chart (a) - **Embedded Text**: - "LRKV reaches each baseline's final performance 18-30% faster" ### Detailed Analysis #### Chart (a) Data Points 1. **MHA** (Green triangle): - X: 0% cache size - Y: 0.728 BPB 2. **MLA** (Red diamond): - X: 10% cache size - Y: 0.724 BPB 3. **GQA** (Orange square): - X: 30% cache size - Y: 0.726 BPB 4. **LRKV** (Purple star): - X: 50% cache size - Y: 0.719 BPB 5. **Standard MHA** (Blue circle): - X: 100% cache size - Y: 0.723 BPB #### Chart (b) Data Points 1. **MHA** (Green triangle): - X: 0% cache size - Y: 29.7% compute savings 2. **MLA** (Red diamond): - X: 10% cache size - Y: 21.4% compute savings 3. **GQA** (Orange square): - X: 30% cache size - Y: 24.5% compute savings 4. **LRKV** (Purple star): - X: 50% cache size - Y: 18.7% compute savings 5. **Standard MHA** (Blue circle): - X: 100% cache size - Y: Not explicitly labeled (implied baseline) ### Key Observations 1. **Memory vs Performance (Chart a)**: - LRKV achieves the lowest BPB (0.719) at 50% cache size, outperforming all methods despite using moderate memory. - Standard MHA (100% cache size) has higher BPB (0.723) than LRKV but lower than MHA (0.728). - MHA (0% cache) has the highest BPB, indicating poor performance without cache optimization. 2. **Training Efficiency (Chart b)**: - MHA requires the most compute savings (29.7%) to match baseline performance. - LRKV achieves the lowest compute savings (18.7%) while still reaching baseline performance faster (18-30% speedup). - GQA and MLA show intermediate efficiency (24.5% and 21.4%, respectively). ### Interpretation - **LRKV Trade-offs**: LRKV demonstrates superior validation performance (lowest BPB) and training efficiency (fastest convergence) but uses moderate memory (50% cache size). This positions it as a balanced solution between memory and speed. - **Standard MHA Limitations**: Despite using full cache size (100%), Standard MHA underperforms LRKV in both validation BPB and training efficiency, suggesting LRKV’s architectural improvements (e.g., caching strategy) are critical. - **Memory-Performance Trade-off**: Increasing cache size generally reduces BPB (e.g., MHA at 0% → 0.728 BPB vs. Standard MHA at 100% → 0.723 BPB), but LRKV breaks this trend by achieving lower BPB at moderate cache size. - **Compute Savings Paradox**: LRKV’s lower compute savings (18.7%) correlate with its faster convergence, implying architectural optimizations reduce redundant computations during training. This analysis highlights LRKV as a Pareto-optimal method, balancing memory efficiency, validation performance, and training speed compared to baseline and alternative methods. </details> Figure 3: LRKV achieves superior training efficiency alongside best performance (2.5B scale). (a) Memory vs Performance: 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). (b) Training Efficiency Advantage: LRKV reaches each baseline’s final test loss, quantifying training compute savings. LRKV reaches all baselines’ performance earlier. Training efficiency analysis. Beyond superior converged performance, LRKV demonstrates remarkable sample efficiency at 2.5B scale (Figure 3). 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 50B token budget. Long Context Pretraining. We extend evaluation to 8192 token sequences using 512M parameter models trained for 50B tokens (153.6 ExaFLOPs). Figure 4 shows that LRKV achieves 2.67 test loss, outperforming MHA (2.74) by 2.7% while using only 48.2% of its KV cache. LRKV also surpasses GQA (2.70, -1.3%), MLA (2.71, -0.9%), and MQA (2.73, -0.4%). Notably, all KV-efficient methods outperform Standard MHA at long context, suggesting compression provides implicit regularization benefits, though LRKV’s architectural advantages remain pronounced. These results validate LRKV’s effectiveness at extended sequence lengths. <details> <summary>2601.11471v3/x3.png Details</summary>  ### Visual Description ## Line Chart: Training Compute (ExaFLOPs) vs. Test Cross-Entropy Loss ### Overview The chart visualizes the relationship between training compute (measured in ExaFLOPs) and test cross-entropy loss for five different model architectures. The x-axis represents training tokens (in billions), while the y-axis shows test cross-entropy loss values. Five distinct lines represent different methods: Standard MHA (blue), LRKV (red), MLA (purple), GQA (orange), and MQA (green). All lines exhibit a downward trend, indicating improved performance (lower loss) with increased training compute. ### Components/Axes - **X-axis**: "Training Tokens (Billions)" with markers at 5, 10, 15, 20, 25, 30, 35, 40, 45, and 50. - **Y-axis**: "Test Cross-Entropy Loss" with markers from 2.65 to 3.00 in increments of 0.05. - **Legend**: Located in the top-right corner, mapping colors to methods: - Blue: Standard MHA - Red: LRKV - Purple: MLA - Orange: GQA - Green: MQA - **Line Styles**: All lines are solid except for MLA (dotted) and GQA (dashed). ### Detailed Analysis 1. **Standard MHA (Blue)**: - Starts at ~3.00 loss at 5B tokens. - Gradually declines to ~2.74 loss at 50B tokens. - Steepest decline occurs between 5B and 20B tokens. 2. **LRKV (Red)**: - Begins at ~3.00 loss at 5B tokens. - Drops sharply to ~2.66 loss at 50B tokens. - Most rapid improvement observed between 5B and 15B tokens. 3. **MLA (Purple)**: - Starts at ~3.00 loss at 5B tokens. - Declines to ~2.72 loss at 50B tokens. - Moderate slope with consistent improvement. 4. **GQA (Orange)**: - Begins at ~3.00 loss at 5B tokens. - Reduces to ~2.71 loss at 50B tokens. - Slightly flatter slope compared to MLA. 5. **MQA (Green)**: - Starts at ~3.00 loss at 5B tokens. - Decreases to ~2.73 loss at 50B tokens. - Intermediate performance between MLA and GQA. ### Key Observations - **LRKV** demonstrates the most significant improvement, achieving the lowest loss (~2.66) at 50B tokens. - **Standard MHA** shows the slowest rate of improvement, with loss remaining above 2.74 even at 50B tokens. - All methods exhibit diminishing returns as training tokens increase beyond 20B. - The gap between LRKV and other methods widens as compute scales, suggesting superior efficiency. ### Interpretation The chart highlights the trade-off between training compute and model performance. LRKV outperforms other methods by achieving the lowest cross-entropy loss with the same compute budget, indicating superior architectural efficiency. Standard MHA, while requiring the most compute to reach comparable performance, may reflect baseline limitations in its design. The diminishing returns observed across all methods suggest that beyond a certain compute threshold (likely ~20B tokens), additional resources yield smaller performance gains. This data could guide resource allocation decisions, favoring architectures like LRKV for compute-constrained scenarios or Standard MHA for applications where compute is abundant but performance requirements are less stringent. </details> Figure 4: Long context pretraining for 512M parameter models. Test cross-entropy loss curves for models trained with 8192-token sequence length on 50B tokens (153.6 ExaFLOPs). LRKV outperforms Standard MHA by 2.7% while using 48.2% cache. ### 4.2 Downstream Task Performance To evaluate whether LRKV’s pretraining advantages translate to practical capabilities, we perform supervised midtraining on a diverse instruction-following dataset (568K examples from SmolTalk (Allal et al., 2024), MMLU (Hendrycks et al., 2021), and GSM8K (Cobbe et al., 2021)) and evaluate on five standard benchmarks. Table 2 shows final downstream performance across three scales (128M, 2.5B, 6.3B). <details> <summary>2601.11471v3/x4.png Details</summary>  ### Visual Description ## Line Charts: Final Test CE Loss vs KV Cache Size and Training Tokens ### Overview The image contains two line charts comparing model performance metrics. Chart (a) shows final test cross-entropy (CE) loss against KV cache size relative to standard MHA, while chart (b) displays CE loss trends across training token scales for different LRKV configurations. ### Components/Axes **Chart (a):** - **X-axis**: KV Cache Size (% of Standard MHA) [0-100%] - **Y-axis**: Final Test CE Loss [2.87-2.95] - **Legend**: - Red circles: LRKV - Purple squares: MLA - Blue diamond: Standard MHA - Orange triangles: GQA (3 groups) - Green triangle: MQA **Chart (b):** - **X-axis**: Training Tokens (Billions) [0-100B] - **Y-axis**: Test Cross-Entropy Loss [2.90-3.15] - **Legend**: - Red: r=8 - Pink: r=16 - Purple: r=32 - Blue: r=64 - Dark blue: r=128 ### Detailed Analysis **Chart (a) Data Points:** 1. **MQA** (Green triangle): 2.93 CE loss at ~20% KV cache 2. **GQA (3 groups)** (Orange triangle): 2.92 CE loss at ~50% KV cache 3. **Standard MHA** (Blue diamond): 2.90 CE loss at 100% KV cache 4. **LRKV configurations**: - r=8: 2.91 CE loss at ~25% KV cache - r=16: 2.90 CE loss at ~45% KV cache - r=32: 2.89 CE loss at ~55% KV cache - r=64: 2.88 CE loss at ~70% KV cache 5. **MLA configurations**: - d=128: 2.94 CE loss at ~15% KV cache - d=192: 2.93 CE loss at ~30% KV cache - d=384: 2.92 CE loss at ~60% KV cache **Chart (b) Trends:** - All LRKV configurations show decreasing CE loss with increased training tokens - Higher 'r' values (larger models) consistently achieve lower final loss: - r=128: ~2.90 CE loss at 100B tokens - r=64: ~2.91 CE loss at 100B tokens - r=32: ~2.92 CE loss at 100B tokens - Convergence pattern: Lines flatten as training approaches 100B tokens ### Key Observations 1. **Cache Efficiency**: Standard MHA achieves lowest CE loss (2.90) at full cache size 2. **Scaling Benefits**: LRKV configurations show progressive improvement with larger 'r' values 3. **MLA Tradeoff**: Larger 'd' values in MLA increase CE loss despite cache size 4. **Training Impact**: All models improve CE performance with more training tokens, with diminishing returns after ~50B tokens ### Interpretation The data demonstrates that: 1. **Model Architecture Matters**: Standard MHA outperforms specialized configurations (MQA, GQA) in CE efficiency 2. **Parameter Scaling**: Larger 'r' values in LRKV configurations yield better performance, suggesting that model capacity improvements translate to efficiency gains 3. **Training Efficiency**: While all models benefit from more training, the rate of improvement slows significantly after 50B tokens, indicating potential diminishing returns in large-scale training 4. **Cache Utilization**: Optimal KV cache usage varies by architecture, with Standard MHA maximizing efficiency at full capacity The charts collectively suggest that while specialized architectures may offer theoretical advantages, practical implementation (Standard MHA) and parameter scaling (LRKV 'r' values) are critical factors in achieving optimal performance. The training curve patterns indicate that model capacity (via 'r' values) plays a more significant role in final performance than training scale alone. </details> Figure 5: LRKV rank ablation with performance-memory tradeoff analysis (128M, 100B tokens). (a) Final test cross-entropy loss versus KV cache size (relative to Standard MHA) for LRKV rank ablations ( $r∈\{8,16,32,64,128\}$ ), MLA latent dimension ablations ( $d∈\{128,192,384\}$ ), and baselines (Standard MHA, GQA with 3 groups, MQA). LRKV dominates the performance-memory tradeoff space. (b) LRKV training dynamics across ranks, showing consistent convergence and monotonic improvement with increasing rank. <details> <summary>2601.11471v3/x5.png Details</summary>  ### Visual Description ## Heatmaps: Attention Distribution Across MHA Variants ### Overview The image displays five heatmaps comparing attention distribution patterns across different Multi-Head Attention (MHA) variants. Each heatmap visualizes the correlation between attention heads (x-axis) and their corresponding output indices (y-axis), with color intensity representing correlation strength (-1.00 to 1.00). All heatmaps share identical axis ranges (0-16) and color scales. ### Components/Axes - **X-axis**: "Head Index" (0-16) - **Y-axis**: "Head Index" (0-16) - **Color Scale**: - Red: Positive correlation (1.00) - Blue: Negative correlation (-1.00) - Yellow: Neutral (0.00) - **Legend**: Positioned right of each heatmap, consistent across all variants - **Heatmap Titles** (top-center): 1. Standard MHA 2. LRKV (r=64) 3. GQA 4. MQA 5. MLA ### Detailed Analysis 1. **Standard MHA** - Diagonal dominance with dark red squares (0.8-1.0 correlation) - Sparse off-diagonal activity (≤0.2 correlation) - No significant outliers 2. **LRKV (r=64)** - Strong diagonal (0.9-1.0) - Distinct off-diagonal pattern: - Secondary diagonal (i+1,j) shows moderate correlation (0.3-0.6) - Checkerboard-like pattern in lower-left quadrant 3. **GQA** - Diagonal dominance (0.7-0.9) - Increased off-diagonal noise: - Random clusters of 0.2-0.5 correlation - Blue patches (-0.1 to -0.3) in upper-right quadrant 4. **MQA** - Weakest diagonal (0.5-0.7) - Mixed off-diagonal signals: - Strong negative correlation (-0.4 to -0.6) in lower-left - Yellow/white dominance (0.0-0.2) elsewhere 5. **MLA** - Strongest diagonal (0.9-1.0) - Notable off-diagonal: - Dark red secondary diagonal (i+2,j) - Blue clusters (-0.3 to -0.5) in upper-right ### Key Observations - All variants show diagonal dominance (self-attention) - LRKV and MLA exhibit structured off-diagonal patterns - GQA shows highest noise variance - MQA demonstrates weakest diagonal correlation - Color consistency across heatmaps confirms uniform correlation scale ### Interpretation The heatmaps reveal fundamental differences in attention mechanism behavior: 1. **Standard MHA** represents baseline self-attention with minimal cross-head interaction 2. **LRKV's** structured off-diagonal pattern suggests engineered attention routing for specific head interactions 3. **GQA's** noise indicates potential instability in attention distribution 4. **MQA's** weak diagonal correlation aligns with its multi-query design reducing head specialization 5. **MLA's** strong secondary diagonal implies enhanced cross-head coordination The variations highlight how architectural choices (e.g., key-value sharing in MQA, relative position encoding in LRKV) fundamentally alter attention distribution patterns. The consistent diagonal dominance across all variants confirms the preservation of self-attention as a core mechanism, while off-diagonal patterns reveal unique variant characteristics. </details> Figure 6: Gauge-invariant head similarity matrices (2.5B scale). Heatmaps show pairwise similarities $s_ij=tr((W_i^K)^⊤W_j^K(W_j^Q)^⊤W_i^Q)$ normalized by Frobenius norms for all 18 heads. Red indicates high similarity, blue indicates independence. LRKV exhibits similarity structure nearly identical to Standard MHA with predominantly dark off-diagonal regions. Scale-consistent superiority across task categories. LRKV achieves the highest combined accuracy at all three scales: 18.9% (128M), 37.9% (2.5B), and 40.2% (6.3B), demonstrating consistent advantages across three orders of magnitude in model size. At 2.5B scale, LRKV outperforms Standard MHA (35.3%), GQA (35.8%), MQA (33.3%), and MLA (36.9%), with particularly strong gains on knowledge-intensive and reasoning tasks: ARC-Easy (+4.1pp over MHA), ARC-Challenge (+6.7pp), MMLU (+2.9pp), and GSM8K (+1.1pp). At 6.3B scale, LRKV achieves 40.2% combined versus MHA’s 38.8%, with the largest improvements on reasoning benchmarks (ARC-Challenge: 58.0% vs 53.6%, MMLU: 44.5% vs 42.7%). One exception is HumanEval at 6.3B scale, where MHA is slightly better, possibly suggesting code generation may be more sensitive to head specialization. Notably, MQA shows catastrophic degradation on code generation (HumanEval: 2.4% at 128M, 3.7% at 2.5B, 4.3% at 6.3B versus 13.4%, 11.7-13.4%, and 12.8-14.6% for other methods), confirming that complete KV sharing particularly harms structured generation tasks requiring precise long-range dependencies. LRKV avoids this pathology through per-head residuals that preserve specialization. Pretraining quality predicts downstream performance. The strong correlation between pretraining BPB and downstream accuracy (R²=0.786 at 2.5B scale, see Appendix E) confirms that architectural improvements generalizing across pretraining data translate directly to task-specific capabilities. LRKV’s 2.6 percentage point downstream advantage over Standard MHA at 2.5B scale (37.9% vs 35.3%) stems directly from its superior pretraining performance (0.719 vs 0.723 BPB). This validates that LRKV’s low-rank factorization provides fundamental capacity gains that manifest across both language modeling and downstream evaluation, rather than overfitting to pretraining objectives. ### 4.3 Why Low-Rank Key-Value Attention Works This subsection provides an empirical analysis explaining why LRKV preserves or exceeds the modeling quality of standard attention despite substantially reducing KV-cache memory. LRKV’s design is motivated by a practical constraint: standard MHA duplicates K/V representations across $H$ heads, causing memory cost to scale linearly with head count. While prior analyses suggest attention heads exhibit substantial redundancy (Michel et al., 2019; Clark et al., 2019), heads also specialize for distinct syntactic and semantic patterns (Zhang et al., 2023), implying that complete KV sharing (MQA) may degrade quality. LRKV addresses this tension through additive factorization: $W_h=W_shared+U_hB_h^⊤$ , which separates a full-rank shared basis from compact per-head residuals. We now examine three questions: (1) whether appropriate rank selection is critical for quality, (2) whether LRKV preserves architectural head diversity, and (3) whether the learned factorization approaches mathematical optimality. <details> <summary>2601.11471v3/x6.png Details</summary>  ### Visual Description ## Line Graphs: Effective Rank Across Layers for Different Model Scales ### Overview The image contains two line graphs comparing the "Effective Rank" of various neural network models across different layers. The left graph represents a 128M-scale model, while the right graph represents a 2.5B-scale model. Each graph tracks seven distinct model configurations over their respective layer ranges. ### Components/Axes - **X-axis (Layer)**: - 128M Scale: 0 to 10 (integer steps) - 2.5B Scale: 0 to 35 (integer steps) - **Y-axis (Effective Rank)**: - 128M Scale: 1 to 6 (continuous scale) - 2.5B Scale: 2.5 to 17.5 (continuous scale) - **Legend**: Located in the bottom-left corner of each graph, with color-coded labels: - Blue: Standard MHA - Orange: LRKV (r=16) - Green: LRKV (r=64) - Red: MQA - Purple: MLA - Brown: GQA - Dashed Gray: Fully redundant ### Detailed Analysis #### 128M Scale Graph - **Standard MHA (Blue)**: Starts at ~4.2 (Layer 0), declines steadily to ~3.8 (Layer 10). - **LRKV (r=16) (Orange)**: Begins at ~2.0 (Layer 0), rises sharply to ~5.8 (Layer 2), then fluctuates between ~5.2–5.8. - **LRKV (r=64) (Green)**: Starts at ~3.6 (Layer 0), remains flat at ~6.0 across all layers. - **MQA (Red)**: Begins at ~2.0 (Layer 0), peaks at ~5.8 (Layer 2), then declines to ~3.0 (Layer 10). - **MLA (Purple)**: Starts at ~4.8 (Layer 0), dips to ~3.5 (Layer 4), then rises to ~5.2 (Layer 10). - **GQA (Brown)**: Begins at ~3.2 (Layer 0), fluctuates between ~5.0–5.8, ending at ~5.3 (Layer 10). - **Fully Redundant (Dashed Gray)**: Constant at ~1.0 across all layers. #### 2.5B Scale Graph - **Standard MHA (Blue)**: Starts at ~14.8 (Layer 0), declines to ~15.2 (Layer 35). - **LRKV (r=16) (Orange)**: Begins at ~7.6 (Layer 0), peaks at ~17.5 (Layer 1), then fluctuates between ~16.0–17.5. - **LRKV (r=64) (Green)**: Starts at ~11.8 (Layer 0), stabilizes at ~17.5 (Layer 1), with minor fluctuations. - **MQA (Red)**: Begins at ~5.8 (Layer 0), peaks at ~17.5 (Layer 1), then declines to ~10.0 (Layer 20), with sharp recoveries. - **MLA (Purple)**: Starts at ~9.6 (Layer 0), peaks at ~17.5 (Layer 1), then fluctuates between ~14.0–17.5. - **GQA (Brown)**: Begins at ~12.8 (Layer 0), stabilizes at ~17.5 (Layer 1), with minor dips. - **Fully Redundant (Dashed Gray)**: Constant at ~1.0 across all layers. ### Key Observations 1. **Scale Correlation**: The 2.5B-scale models consistently achieve higher effective ranks (14–17.5) compared to the 128M-scale models (2–6). 2. **Layer Stability**: - LRKV (r=64) and GQA models stabilize at maximum effective rank (~17.5) after Layer 1 in the 2.5B graph. - MQA and MLA models in the 128M graph show significant volatility, with MQA dropping to ~3.0 by Layer 10. 3. **Baseline Consistency**: The "Fully redundant" line remains at ~1.0 in both graphs, suggesting a control baseline. ### Interpretation The data demonstrates that larger-scale models (2.5B) achieve higher effective ranks, particularly in early layers, indicating more efficient information propagation. The 128M-scale models exhibit greater variability, with MQA and MLA underperforming in later layers. The "Fully redundant" line’s constancy implies it represents a theoretical minimum, possibly reflecting non-informative or baseline behavior. The sharp peaks in LRKV (r=16) and GQA models at Layer 1 across both scales suggest architectural optimizations for early-layer performance. The divergence between 128M and 2.5B scales highlights the impact of model size on effective rank dynamics. </details> Figure 7: LRKV preserves head diversity across scales. Gauge-invariant effective rank shows LRKV with sufficient rank matches Standard MHA at 128M, while at 2.5B LRKV achieves 98.3% vs 98.9% for MHA using 48.4% of KV cache. Analysis setup. For a given Transformer layer, we analyze head-specific projection matrices $\{W_h^Q,W_h^K,W_h^V\}_h=1^H$ through their gauge-invariant bilinear forms on the final pretrained checkpoint. Our novelty is not the bilinear form itself, but its use for quantifying head diversity and the query-compensation effect under KV sharing. For LRKV, we quantify geometric separation using cosine similarity between shared and residual projections, and measure subspace overlap via principal-angle-based metrics. To assess head diversity, we use gauge-invariant similarity metrics based on attention bilinear forms $A_h=W_h^Q(W_h^K)^⊤$ , which are invariant to per-head rotations. We extend this analysis using PCA in bilinear form space by centering the Gram matrix $G$ (where $G_ij=⟨ A_i,A_j⟩_F$ ) via the kernel PCA transformation (Smola and Schölkopf, 1998): $G_centered=G-G_row-G_col+G_mean$ . This removes the mean bilinear form and reveals the intrinsic dimensionality of head specialization. To assess optimality, we compare LRKV’s learned decomposition to the mathematically optimal rank- $r$ truncated SVD of standard MHA projections, computed post-hoc. Rank selection determines capacity and performance. We first examine how residual rank $r$ affects modeling quality through a systematic ablation study. Figure 5 (b) shows training curves for LRKV with ranks $r∈\{8,16,32,64,128\}$ on 128M parameter models with 100B tokens, demonstrating monotonic improvement as rank increases: $r=128$ achieves the best performance (CE=2.877, BPB=4.156), while $r=8$ shows the worst (CE=2.908, BPB=4.201). The performance gap of approximately 1.06% confirms that residual rank is a critical capacity control. Figure 5 (a) contextualizes this ablation by plotting final performance against KV cache size for all evaluated methods, revealing that LRKV achieves superior performance-memory tradeoffs across the rank spectrum compared to MLA latent dimension ablations and baselines. When comparing LRKV and MLA across full sweeps of $r$ and $d_c$ , respectively, we see that LRKV dominates the memory-performance frontier even against MLA settings with larger latent dimensions. The constrained $r=16$ configuration underperforms standard MHA (0.881 vs 0.878 BPB at 128M), lacking capacity to capture head-specific variation. In contrast, sufficient rank ( $r≥ 64$ ) outperforms MHA (0.875 BPB at 128M, 0.719 at 2.5B vs 0.878 and 0.723), demonstrating that representational capacity is the limiting factor. In our experiments, we find that $r≈ 0.36$ - $0.43× d_h$ provides the necessary capacity for LRKV to exceed standard attention performance and MQA, GQA and MLA baselines (see Appendix G for scale-specific rank values). This empirical regularity across scales suggests that this rank range aligns with the intrinsic rank of attention projections while optimizing the performance-memory tradeoff (see Appendix B for spectral analysis). <details> <summary>2601.11471v3/x7.png Details</summary>  ### Visual Description ## Line Graphs: Frobenius Norm Comparison Across Layers ### Overview The image contains two line graphs comparing the Frobenius norm (in keys) of different attention mechanisms across neural network layers. Graph (a) represents a 128M parameter model, while graph (b) represents a 2.5B parameter model. Both graphs compare **Standard MHA** (Multi-Head Attention) with **LRKV** (Layer-wise Residual Key-Value) variants: total, shared, and residual components. --- ### Components/Axes #### Graph (a) 128M - **X-axis (Layer)**: 0 to 10 layers (discrete intervals). - **Y-axis (Frobenius norm)**: 0 to 3000 keys. - **Legend**: Top-right corner, color-coded: - Gray: Standard MHA - Red: LRKV (total) - Green: LRKV (shared) - Orange: LRKV (residual) - **LRKV/Standard Ratio**: 4.2× (annotated near bottom-left). #### Graph (b) 2.5B - **X-axis (Layer)**: 0 to 35 layers (discrete intervals). - **Y-axis (Frobenius norm)**: 0 to 1400 keys. - **Legend**: Top-right corner, identical color coding as graph (a). - **LRKV/Standard Ratio**: 5.1× (annotated near bottom-left). --- ### Detailed Analysis #### Graph (a) 128M 1. **Standard MHA**: Flat line at ~300 keys across all layers. 2. **LRKV (total)**: - Starts at ~3000 keys (layer 0), drops sharply to ~800 keys at layer 1. - Fluctuates between ~600–1000 keys for layers 2–10. 3. **LRKV (shared)**: - Starts at ~500 keys (layer 0), stabilizes at ~600–700 keys for layers 2–10. 4. **LRKV (residual)**: - Starts at ~400 keys (layer 0), stabilizes at ~500–600 keys for layers 2–10. 5. **Trend**: LRKV (total) dominates, with shared/residual components contributing smaller residuals. #### Graph (b) 2.5B 1. **Standard MHA**: Flat line at ~200 keys across all layers. 2. **LRKV (total)**: - Starts at ~1200 keys (layer 0), drops to ~800 keys by layer 5. - Stabilizes at ~1000–1200 keys for layers 10–35. 3. **LRKV (shared)**: - Starts at ~400 keys (layer 0), stabilizes at ~500–600 keys for layers 5–35. 4. **LRKV (residual)**: - Starts at ~300 keys (layer 0), stabilizes at ~400–500 keys for layers 5–35. 5. **Trend**: LRKV (total) remains significantly higher than Standard MHA, with shared/residual components contributing smaller residuals. --- ### Key Observations 1. **LRKV Efficiency**: - LRKV (total) consistently exceeds Standard MHA by 4.2× (128M) and 5.1× (2.5B), indicating higher parameter efficiency. - Shared/residual components account for ~20–30% of LRKV (total) in both models. 2. **Layer-Specific Behavior**: - Sharp initial drop in LRKV (total) at layer 1 (128M) and layer 5 (2.5B), suggesting architectural optimizations in early layers. - Stabilization of residuals (shared/residual) after early layers implies layer-wise parameter sharing. 3. **Scalability**: - Higher LRKV/Standard ratio in the 2.5B model (5.1× vs. 4.2×) suggests improved efficiency at scale. --- ### Interpretation The data demonstrates that **LRKV outperforms Standard MHA in parameter efficiency**, with the gap widening in larger models (2.5B). The **shared and residual components** of LRKV reduce redundancy, enabling better scalability. The sharp initial drop in LRKV (total) may reflect architectural choices to minimize early-layer complexity, while stabilization in later layers indicates optimized parameter reuse. These trends align with efforts to reduce computational overhead in large language models while maintaining performance. </details> Figure 8: Magnitude scaling in LRKV is absorbed by post-projection normalization. Frobenius norms of key projections comparing standard MHA with LRKV’s shared, residual, and total (shared + residual) components for (a) 128M and (b) 2.5B models. LRKV projections operate in a higher-magnitude regime than standard MHA, with both shared and residual components individually exceeding standard MHA magnitudes. LRKV preserves functional head diversity. Figure 6 confirms LRKV exhibits nearly identical similarity patterns to Standard MHA. Quantitatively, we measure head diversity using gauge-invariant metrics based on attention bilinear forms $A_h=W_h^Q(W_h^K)^⊤$ , computing effective rank via eigenvalue entropy ( Appendix B). LRKV (r=64) achieves 98.3% effective rank at 2.5B scale versus 98.9% for Standard MHA—a negligible 0.6pp difference (Figure 7). In contrast, MQA achieves only 86.2% and GQA 95.4%. Interpreting uncentered vs. PCA-based effective rank. The distinction between uncentered (98.3%) and PCA-based (93.5%) effective rank reveals LRKV’s factorization structure. Uncentered analysis measures total variance including the shared mean direction—the global structure captured by $W_shared$ . PCA-based analysis centers the Gram matrix, isolating variance around this mean and measuring true head independence. The modest 4.8pp 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 (see subsection C.4 for detailed analysis). PCA-based analysis reveals compensation mechanisms. Applying PCA in bilinear form space ( subsection C.4), LRKV achieves 93.5% PCA-based effective rank at 2.5B versus 94.0% for Standard MHA—within 0.5% despite rank-64 factorization. Remarkably, PCA reveals a compensation effect in MQA: its centered effective rank (91.0%) increases versus uncentered (86.2%), opposite to other architectures. This indicates MQA heads compensate for forced KV sharing by diversifying query projections more aggressively. Magnitude scaling and post-projection normalization. Figure 8 reveals LRKV projections operate in a higher-magnitude regime than standard MHA ( $3.1$ – $6.7×$ ), with both shared and residual components exceeding MHA norms. This arises naturally from the additive structure, allowing independent growth during optimization. Crucially, this scaling does not affect attention patterns because RMSNorm is applied after projection, making attention effectively use cosine similarity. Why magnitude scaling doesn’t degrade quality. The higher-magnitude regime emerges from unconstrained optimization of the additive factorization: gradients can independently scale $W_shared$ and residuals $U_hB_h^⊤$ without affecting their sum’s direction. RMSNorm applied after projection normalizes representations before attention computation, making the attention mechanism operate on directional information (cosine similarity) rather than absolute magnitudes. This architectural property ensures that magnitude scaling affects only the internal parameterization, not the functional behavior. The magnitude difference serves as a diagnostic: it indicates that optimization successfully allocated capacity between shared and residual pathways, with both contributing substantively to the final projection rather than one dominating. ## 5 Discussion The empirical analyses collectively reveal why LRKV achieves superior performance despite KV cache reduction: (1) Appropriate rank selection is critical. The range $r≈ 0.36$ - $0.43× d_h$ (46-55 for $d_h=128$ ) provides sufficient degrees of freedom for per-head specialization while constraining the factorization to exploit structured redundancy. Below this threshold (e.g., $r=16$ , 0.881 BPB), capacity becomes limiting; above it ( $r≥ 64$ , 0.875 BPB), performance exceeds MHA, confirming that representational capacity (not geometric properties) determines quality. (2) LRKV preserves architectural head diversity. PCA-based analysis shows LRKV achieves 93.5% effective rank versus 94.0% for standard MHA at 2.5B scale (Figure 7)—within 0.5% despite rank-64 factorization and 48.4% cache size. This near-perfect preservation validates that low-rank residuals provide sufficient capacity for heads to occupy independent dimensions in bilinear form space. The consistency across all 36 layers demonstrates depth-invariant factorization quality. (3) Compensation mechanisms explain baseline behavior. The PCA-based methodology reveals phenomena invisible to prior metrics: MQA’s 4.8pp improvement from uncentered to centered effective rank quantifies query compensation - heads recover diversity by specializing queries around forced shared KVs. This explains why MQA remains viable despite complete KV sharing, and positions LRKV’s approach (preserving KV and query diversity) as superior. (4) Emergent properties validate factorization quality. Geometric properties like moderate shared-residual orthogonality (cosine similarity 0.2-0.4) and magnitude scaling (4-5 $×$ MHA) emerge as consequences of effective optimization, not design constraints. These indicators confirm that end-to-end training discovers factorizations that efficiently allocate capacity between global structure (shared base) and local specialization (residuals). These properties collectively enable LRKV to achieve strong pretraining performance (0.692 BPB at 6.3B) and downstream accuracy (40.2% combined at 6.3B) while reducing KV cache to 45-53% of standard attention. The analyses confirm LRKV exploits structured redundancy without sacrificing the head specialization that makes MHA effective. ## 6 Conclusion We introduced Low-Rank Key-Value attention, which decomposes key/value projections into a shared dense component and compact per-head low-rank residuals. LRKV achieves (1) 45-53% KV cache reduction while attaining the best pretraining loss (0.692 BPB at 6.3B) and 18-25% training compute savings; (2) the highest downstream performance across 5 benchmarks for varying model sizes; and (3) near-complete preservation of head diversity give sufficiently lower rank (93.5% vs. 94% for standard MHA), confirming that LRKV exploits structured redundancy without sacrificing specialization. Our gauge-invariant analysis and rank ablations show that residual rank $r≈ 0.36$ - $0.43× d_h$ provides a critical threshold for preserving head specialization while enabling substantial KV cache reduction, positioning LRKV as a practical drop-in replacement for standard attention under memory constraints. ## References - J. Ainslie, S. Ontañón, V. Saxena, et al. (2023) GQA: training generalized multi-query transformer models from multi-head checkpoints. In arXiv:2305.13245, Cited by: Appendix A, §1, §2. - L. B. Allal, A. Lozhkov, L. von Werra, and T. Wolf (2024) SmolTalk: a conversational dataset for instruction tuning. arXiv preprint arXiv:2408.00833. Cited by: item 1, §4.2, §4. - I. Beltagy, M. E. Peters, and A. Cohan (2020) Longformer: the long-document transformer. arXiv:2004.05150. Cited by: Appendix A, §1, §2. - S. Bhojanapalli, A. Chakrabarti, A. Veit, M. Lukasik, H. Jain, F. Liu, Y. Chang, and S. Kumar (2021) Leveraging redundancy in attention with reuse transformers. arXiv preprint arXiv:2110.06821. Cited by: Appendix B. - C. Chang, W. Lin, C. Lin, C. Chen, Y. Hu, P. Wang, N. Huang, L. Ceze, M. S. Abdelfattah, and K. Wu (2025) Palu: kv-cache compression with low-rank projection. In The Thirteenth International Conference on Learning Representations, Cited by: §1. - V. Chari, G. Qin, and B. Van Durme (2025) Kv-distill: nearly lossless learnable context compression for llms. arXiv preprint arXiv:2503.10337. Cited by: Appendix A. - B. Chen, F. Zhang, A. Nguyen, D. Zan, Z. Lin, J. Lou, and W. Chen (2022) Codet: code generation with generated tests. arXiv preprint arXiv:2207.10397. Cited by: 4th item. - R. Child, S. Gray, A. Radford, and I. Sutskever (2019) Generating long sequences with sparse transformers. In ICML, Cited by: Appendix A, §1, §1, §2. - K. Choromanski, V. Likhosherstov, et al. (2021) Rethinking attention with performers. In ICLR, Cited by: Appendix A. - K. Clark, U. Khandelwal, O. Levy, and C. D. Manning (2019) What does bert look at? an analysis of bert’s attention. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 276––286. Cited by: Appendix A, Appendix B, §1, §3, §4.3. - P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord (2018) Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457. Cited by: 1st item. - K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and J. Schulman (2021) Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168. Cited by: item 3, 3rd item, §4.2, §4. - T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Re (2024) FlashAttention-3: fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §1, §1, §2, §3. - R. Fang and Y. Xu (2024) Addressing spectral bias of deep neural networks by multi-grade deep learning. In Advances in Neural Information Processing Systems (NeurIPS) 2024, pp. –. Cited by: Appendix B. - S. Ge et al. (2023) LongNet: scaling transformers to 1,000,000 tokens. arXiv:2307.02486. Cited by: Appendix A. - D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt (2021) Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR). Cited by: item 2, 2nd item, §4.2, §4. - J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. (2022) Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. Cited by: Table 6. - E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al. (2022) Lora: low-rank adaptation of large language models.. ICLR 1 (2), pp. 3. Cited by: Appendix A, §1, §2. - J. Hu, H. Li, Y. Zhang, Z. Wang, S. Zhou, X. Zhang, H. Shum, and D. Jiang (2025) Multi-matrix factorization attention. arXiv preprint arXiv:2412.19255. Cited by: Appendix A, §2. - A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret (2020) Transformers are rnns: fast autoregressive transformers with linear attention. In ICML, Cited by: Appendix A. - M. Khalaf, Y. Shamshoum, N. Hodos, Y. Sieradzki, and A. Schuster (2025) QKV projections require a fraction of their memory. arXiv preprint arXiv:2506.02939. Cited by: Appendix A, §2. - S. Kornblith, M. Norouzi, H. Lee, and G. Hinton (2019) Similarity of neural network representations revisited. In International Conference on Machine Learning, pp. 3519–3529. Cited by: Appendix A, §C.10, §2. - Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen (2024) Snapkv: llm knows what you are looking for before generation. Advances in Neural Information Processing Systems 37, pp. 22947–22970. Cited by: Appendix A. - A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al. (2024) Deepseek-v2: a strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434. Cited by: Appendix A, §1, §2. - X. Lv, N. Ding, K. Zhang, E. Hua, G. Cui, and B. Zhou (2024) Scalable efficient training of large language models with low-dimensional projected attention. arXiv preprint arXiv:2411.02063. Cited by: Appendix A, §2. - P. Michel, O. Levy, and G. Neubig (2019) Are Sixteen Heads Really Better than One?. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: Appendix A, Appendix A, Appendix B, §C.10, §1, §2, §3, §4.3. - OpenAI (2023) GPT-4. Note: https://openai.com/research/gpt-4 Accessed: 2025-10-02 Cited by: §1. - G. Penedo, H. Kydlíček, A. Lozhkov, M. Mitchell, C. A. Raffel, L. Von Werra, T. Wolf, et al. (2024) The fineweb datasets: decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems 37, pp. 30811–30849. Cited by: §D.1, §4. - Y. Peng, Y. Wang, Z. Fang, L. Zhu, Y. Deng, and Y. Duan (2025) Revisiting lora: a smarter low-rank approach for efficient model adaptation. In 2025 5th International Conference on Artificial Intelligence and Industrial Technology Applications (AIITA), pp. 1248–1252. Cited by: Appendix A. - M. Raghu, J. Gilmer, J. Yosinski, and J. Sohl-Dickstein (2017) SVCCA: singular vector canonical correlation analysis for deep learning dynamics and interpretability. In Advances in Neural Information Processing Systems, Vol. 30. Cited by: Appendix A, §C.10. - N. Rahaman, A. Baratin, D. Arpit, F. Draxler, M. Lin, F. Hamprecht, Y. Bengio, and A. Courville (2019) On the spectral bias of neural networks. In International conference on machine learning, pp. 5301–5310. Cited by: Appendix B, Appendix B. - N. Sardana and Z. Havens (2024) Muon: an optimizer for hidden layers. Note: https://github.com/KellerJordan/modded-nanogpt Accessed: 2024-12-22 Cited by: §D.1, §4. - U. Saxena, G. Saha, S. Choudhary, and K. Roy (2024) Eigen attention: attention in low-rank space for kv cache compression. arXiv preprint arXiv:2408.05646. Cited by: §1. - B. Schölkopf, A. Smola, and K. Müller (1998) Nonlinear component analysis as a kernel eigenvalue problem. In Neural computation, Vol. 10, pp. 1299–1319. Cited by: §C.1. - N. Shazeer (2019) Fast transformer decoding: one write-head is all you need. In arXiv:1911.02150, Cited by: Appendix A, §1, §2. - A. J. Smola and B. Schölkopf (1998) On a kernel-based method for pattern recognition, regression, approximation, and operator inversion. Algorithmica 22 (1), pp. 211–231. Cited by: §4.3. - Z. Sun, J. Liu, L. Dong, S. Wang, S. Huang, X. Chen, Y. Zhou, X. Wang, and F. Zhang (2024) RetNet: retentive network for efficient sequence modeling. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: Appendix A. - H. Touvron, L. Martin, T. Lavril, P. Albert, and et al. (2024) The llama 3 herd of models. Note: Meta AI Technical Report Cited by: §1. - H. Touvron, L. Martin, K. Stone, P. Albert, and et al. (2023) LLaMA 2: open foundation and fine-tuned chat models. Note: arXiv preprint arXiv:2307.09288 Cited by: §1. - A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is All You Need. In Advances in Neural Information Processing Systems, Cited by: §1. - E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov (2019) Analyzing multi-head self-attention: specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.09418. Cited by: Appendix A, §C.10. - H. Wang and K. Wang (2025) Complete characterization of gauge symmetries in transformer architectures. In NeurIPS 2025 Workshop on Symmetry and Geometry in Neural Representations, Cited by: Appendix A, §2. - S. Wang, M. Zaheer, N. Katiyar, R. Salakhutdinov, and A. Ahmed (2020) Linformer: self-attention with linear complexity. In NeurIPS, Cited by: Appendix A, §1, §1, §2. - Y. Xie et al. (2023) GPTQ: accurate post-training quantization for generative pretrained transformers. arXiv:2303.09035. Cited by: Appendix A, §2. - D. Yao, B. Shen, Z. Lin, W. Liu, J. Luan, B. Wang, and W. Wang (2025) TailorKV: a hybrid framework for long-context inference via tailored kv cache optimization. arXiv preprint arXiv:2505.19586. Cited by: Appendix A. - D. Yunis, K. K. Patel, S. Wheeler, P. Savarese, G. Vardi, K. Livescu, M. Maire, and M. R. Walter (2024) Approaching deep learning through the spectral dynamics of weights. arXiv preprint arXiv:2408.11804. Cited by: §1. - X. Zhang, A. Ghosh, G. Liu, and R. Wang (2023) Improving generalization of complex models under unbounded loss using pac-bayes bounds. arXiv preprint arXiv:2305.19243. Cited by: §C.10, §4.3. - Y. Zhang, Y. Liu, H. Yuan, Z. Qin, Y. Yuan, Q. Gu, and A. C. Yao (2025) Tensor product attention is all you need. In ICML 2025 Workshop ES-FoMo-III (OpenReview), Note: OpenReview publication Cited by: Appendix A, §2. ## Appendix A Related Work (Full Version) The memory footprint of the key-value (KV) cache has become a central bottleneck in deploying large autoregressive Transformers. While early work improved the computational efficiency of attention via sparsity, kernelization, or low-rank approximations of the attention matrix (Child et al., 2019; Beltagy et al., 2020; Wang et al., 2020; Katharopoulos et al., 2020; Choromanski et al., 2021), these methods do not reduce KV-cache size and therefore provide limited benefit for decoding-heavy LLMs where memory, rather than compute, dominates inference cost. Sharing K/V Projections. A widely adopted strategy for reducing KV memory is to share key and value projections across attention heads. Multi-Query Attention (MQA) (Shazeer, 2019) uses a single shared set of K/V projections, and Grouped-Query Attention (GQA) (Ainslie et al., 2023) extends this by sharing within small head groups. These approaches achieve substantial memory savings but sacrifice head-level expressivity. Analyses of pretrained models (Clark et al., 2019; Michel et al., 2019) consistently show that attention heads specialize for different syntactic or semantic functions; collapsing their projections can therefore degrade modeling quality. LRKV is motivated by this tension: we preserve the efficiency of shared projections while reintroducing diversity through low-rank head-specific residuals. Low-Rank Parameterization of Attention Weights. Our formulation is related to low-rank reparameterization techniques such as LoRA (Hu et al., 2022), which add low-rank updates to weight matrices for parameter-efficient fine-tuning. Follow-up work extends this idea to training-time regularization or improved factorization schemes (Peng et al., 2025). However, these methods intervene on the optimization pathway, not on the structure of stored activations. LRKV applies the same low-rank principle directly to attention projections, introducing low-rank representational deviations that shape the K/V features encoded in the cache. Recent work has also explored low-rank or factored $Q/K/V$ projections to improve parameter efficiency or computational throughput (Xie and others, 2023; Khalaf et al., 2025). The concurrent work of Lv et al. (2024) investigates low-rank parameterizations of attention projections more broadly. These methods typically replace full-rank projections with low-rank ones to reduce computation. LRKV differs conceptually: we preserve a full-rank shared projection and add only low-rank head-specific residuals, enabling memory reduction while maintaining full-rank representational capacity. KV-Cache Compression and Long-Context Modeling. A complementary body of research aims to compress or restructure the KV cache itself. Approaches include clustering or distilling K/V states (Chari et al., 2025), exact hybrid and compressed buffers (Yao et al., 2025; Li et al., 2024), token selection or pooling (Ge and others, 2023), and architectural changes that replace attention with recurrent or state-space alternatives (Sun et al., 2024). Unlike these activation-level methods, LRKV reduces the amount of information that must be stored by changing how each head generates its K/V features. Our hybrid long-context cache builds on this: exact short-range KVs are retained, while long-range information is reconstructed efficiently via low-rank residuals. Comparison to Multi-Latent Attention. Multi-Latent Attention (MLA) (Liu et al., 2024) reduces KV cache memory by compressing token representations into a shared low-dimensional latent space before caching. During attention, interaction with the cache therefore requires projecting queries and/or values through this latent bottleneck. This achieves strong memory compression but constrains all heads to operate through the same latent representation, limiting per-head expressivity and restricting positional encoding choices (e.g., requiring partial RoPE). LRKV addresses a different source of redundancy. Rather than compressing information across tokens, LRKV preserves full token-level resolution and reduces redundancy across heads by factorizing each head’s KV projection into a shared full-rank base plus low-rank head-specific residuals. This additive structure retains the original feature dimensionality, supports arbitrary positional encodings, and preserves head specialization while substantially reducing KV memory. As a result, LRKV and MLA represent complementary approaches: MLA compresses token space via a shared latent bottleneck, whereas LRKV compresses head space via structured sharing. Factorized or Shared KV Mechanisms. Concurrent work explores structured KV-cache reduction via factorized representations (e.g., MFA/MFA-KR; TPA (Hu et al., 2025; Zhang et al., 2025)) explores structured compression across heads. LRKV provides a principled middle ground between fully shared (MQA/GQA) and fully independent projections: the full-rank shared base captures global structure, while low-rank residuals preserve head-level variability. This structured additive decomposition is key for achieving memory savings without collapsing representational diversity. #### Analyzing attention head diversity. Prior work has measured head redundancy using attention pattern similarity (Michel et al., 2019; Voita et al., 2019), activation-based metrics such as CKA (Kornblith et al., 2019), or SVCCA (Raghu et al., 2017). Recent work explicitly characterizes gauge symmetries in attention parameterizations (Wang and Wang, 2025), recognizing that per-head rotations preserve attention function. However, these approaches either compare raw weights (not gauge-invariant) or analyze activation space rather than the functional operators $W^Q(W^K)^⊤$ that determine attention. We introduce a principled method combining gauge-invariant bilinear forms with centered Gram matrix analysis (kernel PCA) to measure intrinsic head diversity independent of parameterization choices. ## Appendix B Shared Subspaces, Spectral Bias, and Information Structure Prior empirical analyses of pretrained Transformers have shown that attention heads exhibit substantial redundancy and occupy overlapping subspaces (Michel et al., 2019; Clark et al., 2019; Bhojanapalli et al., 2021; Rahaman et al., 2019). To assess this structure in a gauge-invariant manner, we analyze attention bilinear forms $A_h=W_h^Q(W_h^K)^⊤$ and measure head diversity via PCA on their centered Gram matrix (subsection 4.3). We find that LRKV achieves 93.5% effective rank versus 94.0% for standard MHA at 2.5B scale, confirming that heads occupy nearly independent dimensions in bilinear form space despite using shared key-value bases. This validates that the functional behavior of attention projections lies near a shared low-dimensional manifold, with LRKV’s additive structure explicitly parameterizing this geometry. Importantly, while heads remain functionally diverse, the individual weight matrices $W_h^K$ exhibit low intrinsic rank—they concentrate energy in a small number of dominant singular directions. Such spectral concentration is consistent with the spectral bias of neural networks (Rahaman et al., 2019; Fang and Xu, 2024), whereby gradient-based optimization preferentially amplifies smooth, high-variance modes before fitting higher-frequency residual structure. As a result, attention projections tend to concentrate energy in a small set of globally shared directions. LRKV makes this implicit organization explicit: the shared projections $W_shared^K,V$ capture dominant spectral components, while per-head low-rank residuals model localized refinements aligned with lower-variance directions. Each head therefore operates within an affine subspace centered on $W_shared^K,V$ , with a tangent space spanned by a compact set of learned low-rank factors. #### Residual rank as a control knob for diversity. In LRKV, the residual rank $r$ controls a continuous spectrum between fully shared keys/values and fully independent per-head projections. At a high level, the shared base $W_shared^K,V$ serves as a common coordinate system capturing features that are broadly useful across heads, while the low-rank residuals provide a budgeted mechanism for head specialization. When $r$ is small, heads are strongly coupled through the shared representation; as $r$ grows, the model can allocate additional degrees of freedom to capture head-specific patterns. This view motivates treating $r$ as an architectural hyperparameter with a clear interpretation: it directly controls the efficiency-diversity trade-off in the KV cache and per-head specialization. Empirically, we find that $r≈ 0.36$ - $0.43× d_h$ provides the critical threshold: LRKV achieves 93.5% PCA-based effective rank versus 94.0% for standard MHA at 2.5B scale (subsection 4.3), confirming that moderate rank suffices to preserve nearly all head diversity while achieving substantial cache reduction. #### Decomposition geometry and (non-)orthogonality. The decomposition $W_h=W_shared+U_hB_h^⊤$ is not identifiable: for any $Δ$ , we can shift $W_shared←W_shared+Δ$ and $U_hB_h^⊤←U_hB_h^⊤-Δ$ without changing $W_h$ . In unconstrained Euclidean parameter space there is therefore no a priori reason to expect strict orthogonality between the shared and residual terms. Instead, we use overlap (e.g., cosine similarity or subspace angles) as a diagnostic: low overlap indicates that the residual contributes directions not already captured by the shared component, reducing redundant parameterization. In practice, any tendency toward reduced overlap is an emergent outcome of end-to-end training and the model’s incentives to allocate capacity efficiently, rather than an explicit constraint or training objective. #### Connection to low-rank approximation. LRKV can also be interpreted as learning a structured low-rank approximation of the family of per-head projections. In classical matrix approximation, the optimal rank- $r$ representation is given by truncated SVD with respect to a chosen error metric. LRKV differs in two important ways: (i) it learns the shared and residual factors jointly with the rest of the network under the task loss, and (ii) the residual factors are head-specific, enabling specialization without requiring each head to store full-rank keys/values. This perspective suggests that end-to-end training can discover factorizations that are close to classical optima while remaining aligned with the functional requirements of attention. #### Scaling, normalization, and stability. Because LRKV is additive, the norms of the shared and residual components can evolve independently during training. However, attention depends primarily on directional structure and relative alignment, and common transformer normalization (e.g., RMSNorm) reduces sensitivity to absolute scale. From this perspective, changes in parameter magnitudes are best interpreted as a redistribution of representational capacity between shared and residual pathways rather than as a direct indicator of instability. This motivates analyzing LRKV through geometry (alignment, overlap) and function (loss/perplexity), rather than through norms alone. ## Appendix C Principled Head Diversity Analysis via PCA in Bilinear Form Space This section presents our gauge-invariant PCA-based methodology for analyzing attention head diversity and provides comprehensive results across model scales. To our knowledge, this is the first work to combine (i) gauge-invariant bilinear form comparison, (ii) centered Gram matrix analysis (kernel PCA), and (iii) variance-explained interpretation for measuring head independence in transformers. ### C.1 Methodology: PCA in the Space of Bilinear Forms #### Gauge invariance motivation. Comparing attention heads via raw weight matrices $W_h^K$ or $W_h^Q$ is not meaningful because attention is invariant to coupled per-head rotations: for any orthogonal $R_h$ , the transformations $W_h^Q←W_h^QR_h$ and $W_h^K←W_h^KR_h$ leave attention outputs unchanged. The gauge-invariant object is the bilinear form $A_h=W_h^Q(W_h^K)^⊤$ , which determines attention logits. #### Inner product on bilinear forms. We define similarity between heads $i$ and $j$ using the Frobenius inner product on their bilinear forms: $$ ⟨A_i,A_j⟩_F=tr(A_i^⊤A_j)=tr((W_i^K)^⊤W_j^K(W_j^Q)^⊤W_i^Q). \tag{11} $$ Normalized by individual norms, this yields a similarity score $s_ij∈[-1,1]$ forming the Gram matrix $G$ . #### Centering for proper PCA. The key methodological contribution is recognizing that the Gram matrix $G$ can be centered without materializing the mean bilinear form: $$ G_centered[i,j]=G[i,j]-\frac{1}{H}∑_kG[i,k]-\frac{1}{H}∑_kG[k,j]+\frac{1}{H^2}∑_k,\ellG[k,\ell]. \tag{12} $$ This is the kernel PCA centering trick (Schölkopf et al., 1998), enabling PCA in the abstract space of bilinear forms using only their inner products. #### Interpretation as variance decomposition. The eigenvalues $\{λ_i\}$ of $G_centered$ represent variance explained by each principal component in bilinear form space. We compute: - Variance explained: $v_i=λ_i/∑_jλ_j$ (fraction of total variance in $i$ -th PC) - Cumulative variance: $∑_j=1^kv_j$ (variance captured by first $k$ PCs) - Effective rank (PCA-based): $\exp(-∑_iv_i\log v_i)$ via Shannon entropy High effective rank indicates heads occupy many independent dimensions (low sharing); low effective rank indicates clustering in few PCs (high sharing). #### Comparison to uncentered analysis. Prior work typically uses uncentered Gram matrices, whose eigenvalues include both the ”mean direction” (shared structure) and variance around the mean. The centered analysis isolates the latter, revealing head independence independent of shared baselines. This distinction is critical: MQA’s complete KV sharing creates a dominant mean direction (lowering uncentered effective rank) but heads compensate via query specialization (maintaining centered effective rank). The PCA-based metric reveals this compensation effect. ### C.2 Complete Results Across Scales Table 3: PCA-based effective rank comparison. Centered Gram matrix analysis reveals LRKV preserves head diversity within 1% of Standard MHA at both scales, while MQA shows surprising resilience through query compensation. | Standard MHA LRKV (r=64) GQA | 95.4% 96.2% 90.5%* | 82.4% 83.0% 79.6%* | 98.9% 98.3% 95.4% | 94.0% 93.5% 91.7% | | --- | --- | --- | --- | --- | | LRKV (r=16) | 88.8% | 81.5% | – | – | | MQA | 72.6% | 78.4% | 86.2% | 91.0% | | MLA | 83.9% | 78.5% | 92.7% | 91.6% | *GQA at 128M estimated from 2.5B patterns. PCA-based percentages are relative to maximum possible effective rank (number of heads). Note MQA’s 4.8pp improvement from uncentered to PCA-based at 2.5B scale, revealing query compensation. ### C.3 Effective Rank at 128M Scale <details> <summary>2601.11471v3/x8.png Details</summary>  ### Visual Description ## Line Chart: Head Diversity via Effective Rank (128M) ### Overview The chart visualizes the effective rank of different attention mechanisms across 11 layers (0–10) in a neural network model. Effective rank measures diversity, with lower values indicating higher redundancy. The y-axis ranges from 1 (fully redundant) to 6, while the x-axis represents layers. Six methods are compared: Standard MHA, LRKV (r=16), LRKV (r=64), MQA, MLA, and GQA. ### Components/Axes - **X-axis (Layer)**: Discrete values from 0 to 10. - **Y-axis (Effective Rank)**: Continuous scale from 1 (fully redundant) to 6. - **Legend**: Located in the bottom-right corner, mapping colors to methods: - Blue: Standard MHA - Orange: LRKV (r=16) - Green: LRKV (r=64) - Red: MQA - Purple: MLA - Brown: GQA - **Dashed Line**: Horizontal line at y=1 labeled "Fully redundant." ### Detailed Analysis 1. **Standard MHA (Blue)**: - Starts at 4.2 (layer 0), peaks at 6.0 (layers 1–3), then declines to 5.5 (layer 10). - Maintains high effective rank (5.5–6.0) across most layers. 2. **LRKV (r=16) (Orange)**: - Begins at 2.0 (layer 0), rises sharply to 6.0 (layers 1–3), stabilizes at ~5.8 (layers 4–10). - Shows the steepest initial increase. 3. **LRKV (r=64) (Green)**: - Starts at 3.5 (layer 0), increases to 6.0 (layers 1–3), remains flat at 6.0 (layers 4–10). - Highest effective rank for most layers. 4. **MQA (Red)**: - Begins at 2.0 (layer 0), peaks at 6.0 (layers 1–3), then declines to 3.2 (layer 10). - Most pronounced drop after layer 3. 5. **MLA (Purple)**: - Starts at 4.8 (layer 0), fluctuates between 4.5–5.8 (layers 1–9), drops to 3.3 (layer 10). - Sharp decline at layer 10. 6. **GQA (Brown)**: - Begins at 3.2 (layer 0), rises to 5.8 (layers 1–3), stabilizes at ~5.5 (layers 4–10). - Gradual increase followed by stabilization. 7. **Fully Redundant Line (Dashed)**: - Horizontal line at y=1, never intersected by any method. ### Key Observations - **No method reaches full redundancy** (y=1), indicating all maintain some diversity. - **LRKV (r=64)** and **Standard MHA** consistently exhibit the highest effective ranks (5.5–6.0), suggesting lower diversity. - **MQA** and **GQA** show significant variability, with MQA dropping sharply after layer 3. - **MLA** has the most erratic trend, including a steep drop at layer 10. - **LRKV (r=16)** and **GQA** demonstrate the most stable performance after initial increases. ### Interpretation The chart highlights trade-offs between attention mechanisms in balancing diversity and redundancy. LRKV (r=64) and Standard MHA prioritize higher effective ranks (lower diversity), while MQA and MLA show greater variability, potentially indicating adaptive behavior. The absence of any method reaching full redundancy suggests all mechanisms retain some level of unique information propagation. The sharp drop in MLA at layer 10 may reflect architectural constraints or optimization challenges in later layers. These trends could inform decisions about method selection based on desired diversity profiles for specific tasks. </details> Figure 9: Head diversity at 128M scale (PCA-based analysis). Effective rank computed from centered Gram matrices shows consistent patterns: LRKV (r=64) achieves 83.0% effective rank versus 82.4% for Standard MHA, demonstrating that sufficient residual capacity ( $r≈ 0.36$ - $0.43× d_h$ in deployed models) preserves head specialization. LRKV (r=16) achieves 81.5%, while MQA shows 78.4%. The PCA-based metric reveals that even aggressive compression methods maintain substantial head diversity through compensation mechanisms. At 128M scale with 6 attention heads, LRKV (r=64) achieves effective rank (83.0%) nearly matching Standard MHA (82.4%), demonstrating that the low-rank factorization does not inherently degrade head diversity when rank is appropriately chosen. The moderately constrained LRKV (r=16) configuration shows only slightly reduced diversity (81.5%), while MQA maintains 78.4% effective rank despite complete KV sharing—revealing the query compensation effect at smaller scale. ### C.4 PCA Eigenvalue Spectra <details> <summary>2601.11471v3/x9.png Details</summary>  ### Visual Description ## Line Graphs: PCA Eigenvalue Spectra in Bilinear Form Space (2.5B) ### Overview The image contains three line graphs comparing PCA eigenvalue spectra across three transformer layers (Early: L=0, Middle: L=18, Late: L=35). Each graph plots PCA index (x-axis) against variance (y-axis) for five methods: Standard MHA, LRKV (r=64), GQA, MQA, and MLA. The graphs reveal how variance distribution across principal components (PCs) differs by layer and method. ### Components/Axes - **X-axis**: "PC Index" (2.5 to 17.5, increments of 2.5) - **Y-axis**: "Variance (PCA Eigenvalue)" (0 to 7.5 for Early Layer, 0 to 2.5 for Middle Layer, 0 to 1.4 for Late Layer) - **Legends**: - Standard MHA (blue) - LRKV (r=64) (orange) - GQA (green) - MQA (red) - MLA (purple) ### Detailed Analysis #### Early Layer (L=0) - **Standard MHA**: Starts at ~7.5 variance at PC Index 2.5, drops sharply to ~0.5 by PC Index 17.5. - **LRKV (r=64)**: Begins at ~3.5, declines to ~0.2. - **GQA**: Starts at ~1.5, decreases to ~0.3. - **MQA**: Begins at ~0.8, drops to ~0.1. - **MLA**: Starts at ~0.5, declines to ~0.05. - **Trend**: All methods show rapid variance decay after PC Index 5. Standard MHA retains the highest early variance. #### Middle Layer (L=18) - **Standard MHA**: Starts at ~2.5, declines to ~1.0. - **LRKV (r=64)**: Begins at ~1.2, decreases to ~0.8. - **GQA**: Starts at ~1.0, drops to ~0.6. - **MQA**: Begins at ~0.9, declines to ~0.4. - **MLA**: Starts at ~0.8, decreases to ~0.3. - **Trend**: Variance decay slows after PC Index 10. Standard MHA maintains the highest values. #### Late Layer (L=35) - **Standard MHA**: Starts at ~1.4, drops to ~0.8. - **LRKV (r=64)**: Begins at ~1.0, decreases to ~0.6. - **GQA**: Starts at ~1.2, drops to ~0.5. - **MQA**: Begins at ~1.3, declines to ~0.4. - **MLA**: Starts at ~1.1, decreases to ~0.3. - **Trend**: Sharp variance drop after PC Index 15. GQA and MQA show steeper late-stage declines. ### Key Observations 1. **Layer-Specific Variance**: Early Layer retains the highest variance overall, while Late Layer shows the most pronounced decay. 2. **Method Differences**: - Standard MHA consistently retains the highest early variance. - MLA exhibits the lowest variance across all layers. - LRKV (r=64) and GQA show intermediate decay rates. 3. **Anomalies**: - Late Layer graphs exhibit abrupt drops near PC Index 17.5, suggesting potential numerical instability or sparsity in later PCs. - MQA and GQA show similar decay patterns but with higher initial variance in Late Layer. ### Interpretation The data suggests that PCA eigenvalue spectra vary significantly by transformer layer and attention mechanism. Standard MHA retains more early-stage variance, indicating stronger feature representation in initial PCs. The sharp late-layer drops (especially in GQA and MQA) may reflect overfitting or reduced information content in deeper layers. MLA's consistent low variance suggests a more regularized or efficient design. These patterns align with expectations for attention mechanisms balancing capacity and generalization across transformer depths. </details> Figure 10: PCA eigenvalue spectra reveal variance structure in bilinear form space (2.5B). Eigenvalues of centered Gram matrices at early, middle, and late layers show how variance is distributed across principal components. LRKV’s spectra closely match Standard MHA across all depth regimes, with similar leading eigenvalues and comparable tail decay, indicating nearly identical head correlation structure. MQA shows slightly elevated later eigenvalues, consistent with query compensation: while the first few PCs capture shared KV structure, remaining PCs capture query-driven diversity. Early layers show slightly higher leading eigenvalues across methods, suggesting a stronger shared component (greater alignment in bilinear form space) in lower layers; deeper layers exhibit more uniform spectra, consistent with variance being spread across more modes. <details> <summary>2601.11471v3/x10.png Details</summary>  ### Visual Description ## Line Graphs: PCA Eigenvalue Spectra in Bilinear Form Space (128M) ### Overview The image contains three line graphs comparing PCA eigenvalue spectra across three transformer layers (Early Layer L=0, Middle Layer L=6, Late Layer L=11). Each graph shows variance values (0-1.2) across principal component (PC) indices (1-6) for six methods: Standard MHA, LRKV (r=16), LRKV (r=64), MQA, MLA, and GQA. Lines are color-coded per legend. ### Components/Axes - **X-axis**: PC Index (1-6), labeled "PC Index" - **Y-axis**: Variance (PCA Eigenvalue), labeled "Variance (PCA Eigenvalue)" - **Legends**: - Standard MHA (blue) - LRKV (r=16) (orange) - LRKV (r=64) (green) - MQA (red) - MLA (purple) - GQA (brown) - **Graph Titles**: - Early Layer (L=0) - Middle Layer (L=6) - Late Layer (L=11) ### Detailed Analysis #### Early Layer (L=0) - **Standard MHA**: Starts at ~0.9 (PC1), declines sharply to ~0.3 (PC6) - **LRKV (r=16)**: Starts at ~0.3, declines to ~0.1 - **LRKV (r=64)**: Starts at ~0.7, declines to ~0.4 - **MQA**: Starts at ~0.4, declines to ~0.1 - **MLA**: Starts at ~1.2, declines to ~0.4 - **GQA**: Starts at ~0.5, declines to ~0.1 #### Middle Layer (L=6) - **Standard MHA**: Starts at ~0.9, declines to ~0.6 - **LRKV (r=16)**: Starts at ~0.9, declines to ~0.7 - **LRKV (r=64)**: Starts at ~1.0, declines to ~0.9 - **MQA**: Starts at ~0.9, declines to ~0.7 - **MLA**: Starts at ~0.9, declines to ~0.6 - **GQA**: Starts at ~0.9, declines to ~0.7 #### Late Layer (L=11) - **Standard MHA**: Starts at ~0.8, declines to ~0.4 - **LRKV (r=16)**: Starts at ~1.2, declines to ~0.3 - **LRKV (r=64)**: Starts at ~1.0, declines to ~0.8 - **MQA**: Starts at ~0.5, declines to ~0.2 - **MLA**: Starts at ~0.8, declines to ~0.4 - **GQA**: Starts at ~0.7, declines to ~0.4 ### Key Observations 1. **Variance Reduction**: All methods show decreasing variance with higher PC indices, indicating diminishing returns in eigenvalue capture. 2. **Layer-Specific Patterns**: - Early Layer (L=0): MLA shows highest initial variance (~1.2), while LRKV (r=16) has lowest (~0.3). - Middle Layer (L=6): LRKV (r=64) maintains highest variance (~1.0), while MLA drops to ~0.6. - Late Layer (L=11): LRKV (r=16) exhibits sharpest decline (~1.2→0.3), while MQA shows most gradual drop (~0.5→0.2). 3. **Method Performance**: - LRKV (r=64) consistently retains higher variance across layers compared to other methods. - MLA demonstrates strong early-layer performance but declines significantly in later layers. - GQA shows relatively stable variance across layers (~0.5→0.4 in Late Layer). ### Interpretation The data suggests that PCA eigenvalue spectra vary significantly across transformer layers and methods. Early layers (L=0) exhibit higher variance overall, with MLA capturing the most initial variance. As layers progress (L=6→L=11), variance generally decreases, but LRKV (r=64) maintains higher values, indicating better preservation of information. The sharp decline in LRKV (r=16) in Late Layer (L=11) suggests it may prioritize early principal components more aggressively. These patterns imply trade-offs between methods in balancing early vs. late-layer feature representation, with LRKV (r=64) showing the most consistent performance across layers. </details> Figure 11: PCA eigenvalue spectra at 128M scale. With 6 heads, the eigenvalue structure is more pronounced. LRKV (r=64) tracks Standard MHA closely across all layers, while LRKV (r=16) shows similar patterns with slightly elevated leading eigenvalues in later layers. MQA exhibits comparable spectra to other methods, with its eigenvalue distribution revealing that query compensation maintains diversity despite complete KV sharing. The consistency of spectra shapes across architectures suggests fundamental constraints on head organization. The eigenvalue spectra provide an alternative view of head diversity, complementing the aggregate effective rank metric. The PCA-based effective rank summarizes the entire spectrum via entropy, while these plots show the full eigenvalue distribution. Key observations: - LRKV matches Standard MHA spectra: At 2.5B scale (Figure 10), LRKV’s eigenvalue curves closely track Standard MHA across early, middle, and late layers. This indicates not just similar effective rank, but nearly identical head correlation structure in the space of bilinear forms. - MQA’s compensation appears in eigenvalue distribution: MQA shows slightly elevated later eigenvalues compared to its uncentered analysis would suggest, consistent with query compensation. While early PCs capture the shared KV structure (mean direction), later PCs reveal query-driven diversity that maintains head specialization. - Depth-dependent patterns: Early layers show slightly higher leading eigenvalues across all methods, suggesting that heads capture more structured relationships (greater specialization) for low-level features. Middle and late layers show more uniform spectra, consistent with increasingly abstract representations where heads become more similar. - Scale consistency: The 128M results (Figure 11) mirror the 2.5B patterns despite different head counts (6 vs 18), confirming that LRKV’s preservation of head diversity and MQA’s compensation mechanism are scale-invariant phenomena. ### C.5 Cumulative Variance Explained <details> <summary>2601.11471v3/x11.png Details</summary>  ### Visual Description ## Line Graphs: Cumulative Variance Explained by PCs (2.5B) ### Overview The image contains three line graphs comparing the cumulative variance explained by Principal Components (PCs) across three neural network layers: Early Layer (L=0), Middle Layer (L=18), and Late Layer (L=35). Each graph plots the cumulative variance explained (y-axis) against the number of PCs (x-axis), with five methods represented by distinct colored lines. A gray dashed line indicates the 90% variance threshold. --- ### Components/Axes - **Title**: "Cumulative Variance Explained by PCs (2.5B)" - **X-axis**: "Number of PCs" (ranges from 2.5 to 17.5 in increments of 2.5). - **Y-axis**: "Cumulative Variance Explained" (ranges from 0.0 to 1.0 in increments of 0.2). - **Legend**: - **Standard MHA** (blue) - **LRKV (r=64)** (orange) - **GQA** (green) - **MQA** (red) - **MLA** (purple) - **90% variance** (gray dashed line). --- ### Detailed Analysis #### Early Layer (L=0) - **Standard MHA**: Starts at ~0.15 (2.5 PCs) and rises to ~0.95 (17.5 PCs). - **LRKV (r=64)**: Starts at ~0.1 (2.5 PCs) and rises to ~0.9 (17.5 PCs). - **GQA**: Starts at ~0.12 (2.5 PCs) and rises to ~0.92 (17.5 PCs). - **MQA**: Starts at ~0.1 (2.5 PCs) and rises to ~0.95 (17.5 PCs). - **MLA**: Starts at ~0.15 (2.5 PCs) and rises to ~0.98 (17.5 PCs). - **90% threshold**: Reached at ~17.5 PCs for all methods. #### Middle Layer (L=18) - **Standard MHA**: Starts at ~0.2 (2.5 PCs) and rises to ~0.98 (17.5 PCs). - **LRKV (r=64)**: Starts at ~0.15 (2.5 PCs) and rises to ~0.95 (17.5 PCs). - **GQA**: Starts at ~0.18 (2.5 PCs) and rises to ~0.93 (17.5 PCs). - **MQA**: Starts at ~0.2 (2.5 PCs) and rises to ~0.97 (17.5 PCs). - **MLA**: Starts at ~0.25 (2.5 PCs) and rises to ~0.99 (17.5 PCs). - **90% threshold**: Reached at ~17.5 PCs for all methods. #### Late Layer (L=35) - **Standard MHA**: Starts at ~0.3 (2.5 PCs) and rises to ~0.99 (17.5 PCs). - **LRKV (r=64)**: Starts at ~0.25 (2.5 PCs) and rises to ~0.97 (17.5 PCs). - **GQA**: Starts at ~0.3 (2.5 PCs) and rises to ~0.95 (17.5 PCs). - **MQA**: Starts at ~0.3 (2.5 PCs) and rises to ~0.98 (17.5 PCs). - **MLA**: Starts at ~0.35 (2.5 PCs) and rises to ~1.0 (17.5 PCs). - **90% threshold**: Reached at ~17.5 PCs for all methods. --- ### Key Observations 1. **Consistent Trends**: All methods show upward-sloping lines, indicating increasing cumulative variance with more PCs. 2. **MLA Dominance**: MLA consistently achieves the highest cumulative variance across all layers, reaching ~1.0 in the Late Layer. 3. **Layer-Specific Differences**: - Early Layer: Lower starting points (~0.1–0.15) for all methods. - Middle Layer: Moderate starting points (~0.15–0.25). - Late Layer: Higher starting points (~0.25–0.35), suggesting greater initial variance capture. 4. **90% Threshold**: All methods reach the 90% variance mark at ~17.5 PCs, regardless of layer. --- ### Interpretation The graphs demonstrate that **MLA outperforms other methods** in capturing cumulative variance across all layers, particularly in the Late Layer (L=35), where it achieves near-complete variance explanation (~1.0). The **90% threshold** is consistently met at 17.5 PCs, suggesting a standardized benchmark for model performance. - **Early Layer (L=0)**: Methods require more PCs to explain variance, indicating less efficient feature extraction. - **Late Layer (L=35)**: Higher initial variance capture (e.g., MLA starts at ~0.35) suggests deeper layers encode more meaningful information. - **Method Comparisons**: - **Standard MHA** and **MQA** show similar performance, with MQA slightly outperforming in the Late Layer. - **LRKV (r=64)** and **GQA** lag behind MLA but remain competitive in the Middle and Late Layers. The data implies that **MLA’s architecture** is optimized for variance capture, making it suitable for tasks requiring high-dimensional data representation. The consistent 90% threshold across layers highlights the robustness of the PC-based analysis framework. </details> Figure 12: Cumulative variance explained quantifies dimensionality of head diversity (2.5B). Plots show how many principal components are needed to capture X% of total variance in bilinear form space at early, middle, and late layers. Standard MHA and LRKV require $∼$ 16-17 PCs for 90% variance (out of 18 total heads), demonstrating heads occupy nearly all available dimensions with minimal redundancy. GQA requires $∼$ 15-16 PCs, showing modest clustering from group-based sharing. MQA and MLA require $∼$ 14-15 PCs, indicating moderate concentration in fewer dimensions but still substantial spread. The 90% threshold (dashed line) serves as a practical measure of intrinsic dimensionality. All methods show relatively linear cumulative variance curves, indicating no single PC dominates—even MQA maintains distributed variance. <details> <summary>2601.11471v3/x12.png Details</summary>  ### Visual Description ## Line Graphs: Cumulative Variance Explained by PCs (128M) ### Overview The image contains three line graphs comparing the cumulative variance explained by principal components (PCs) across three transformer layers: Early Layer (L=0), Middle Layer (L=6), and Late Layer (L=11). Each graph plots the number of PCs (1–6) against cumulative variance explained (0–1.0), with multiple data series representing different methods. A dashed line at 0.9 (90% variance) serves as a reference threshold. ### Components/Axes - **X-axis**: "Number of PCs" (1–6, integer steps). - **Y-axis**: "Cumulative Variance Explained" (0.0–1.0, linear scale). - **Legend**: Positioned on the right of each subplot, with the following labels and colors: - Standard MHA (blue) - LRKV (r=16) (orange) - LRKV (r=64) (green) - MQA (red) - MLA (purple) - GQA (brown) - 90% variance (gray dashed line). ### Detailed Analysis #### Early Layer (L=0) - **Standard MHA**: Starts at ~0.3 (PC=1), rises steeply to ~1.0 by PC=6. - **LRKV (r=16)**: Begins at ~0.25, surpasses Standard MHA by PC=4 (~0.85). - **LRKV (r=64)**: Starts at ~0.2, reaches ~0.9 by PC=5. - **MQA**: Begins at ~0.35, plateaus at ~0.9 by PC=5. - **MLA**: Starts at ~0.3, reaches ~0.95 by PC=6. - **GQA**: Starts at ~0.3, plateaus at ~0.9 by PC=5. #### Middle Layer (L=6) - **Standard MHA**: Starts at ~0.4, reaches ~1.0 by PC=6. - **LRKV (r=16)**: Begins at ~0.35, surpasses Standard MHA by PC=4 (~0.85). - **LRKV (r=64)**: Starts at ~0.3, reaches ~0.9 by PC=5. - **MQA**: Begins at ~0.4, plateaus at ~0.9 by PC=5. - **MLA**: Starts at ~0.4, reaches ~0.95 by PC=6. - **GQA**: Starts at ~0.4, plateaus at ~0.9 by PC=5. #### Late Layer (L=11) - **Standard MHA**: Starts at ~0.5, reaches ~1.0 by PC=6. - **LRKV (r=16)**: Begins at ~0.45, surpasses Standard MHA by PC=4 (~0.85). - **LRKV (r=64)**: Starts at ~0.4, reaches ~0.9 by PC=5. - **MQA**: Begins at ~0.5, plateaus at ~0.9 by PC=5. - **MLA**: Starts at ~0.5, reaches ~0.95 by PC=6. - **GQA**: Starts at ~0.5, plateaus at ~0.9 by PC=5. ### Key Observations 1. **Consistency Across Layers**: All methods show similar trends, with variance explained increasing as more PCs are added. 2. **Performance Differences**: - **LRKV (r=16)** consistently outperforms other methods in later layers (L=6, L=11), reaching the 90% threshold earlier. - **Standard MHA** performs comparably to other methods in early layers but lags in later layers. - **MQA** and **GQA** plateau earlier than other methods, suggesting diminishing returns. 3. **90% Threshold**: All methods except **LRKV (r=16)** in late layers require 5–6 PCs to exceed 90% variance. ### Interpretation The data suggests that **LRKV (r=16)** is the most efficient method for capturing variance across all layers, particularly in later layers (L=6, L=11), where it achieves higher cumulative variance with fewer PCs. This implies LRKV (r=16) may be better suited for tasks requiring compact representations. In contrast, **Standard MHA** and other methods require more PCs to achieve similar performance, indicating potential inefficiencies in later layers. The consistent plateauing of **MQA** and **GQA** at ~0.9 variance suggests these methods may not scale as effectively for high-dimensional data. The 90% threshold acts as a practical benchmark, highlighting trade-offs between model complexity and performance. </details> Figure 13: Cumulative variance explained at 128M scale. With 6 heads, the variance structure is more visible. Standard MHA and LRKV (r=64) require $∼$ 5 PCs for 90% variance, demonstrating that heads remain largely independent. LRKV (r=16) shows similar patterns despite constrained rank. MQA requires $∼$ 4-5 PCs, confirming that query compensation maintains distributed variance even with complete KV sharing. The steeper curves at 128M compared to 2.5B suggest slightly more variance concentration at smaller scale. The cumulative variance plots directly answer the question: ”How many principal components capture most of the head diversity?” This provides an intuitive measure of intrinsic dimensionality that complements the entropy-based effective rank. Key findings: - LRKV preserves full dimensionality: At both scales, LRKV requires nearly all available PCs to capture 90% variance, matching Standard MHA. This confirms that low-rank residuals provide sufficient degrees of freedom for heads to occupy independent dimensions in bilinear form space. - No dominant principal component: The relatively linear cumulative variance curves indicate that no single PC captures disproportionate variance. Even the first PC accounts for only 15-20% of total variance, confirming that heads do not collapse to a single dominant direction. - MQA’s distributed compensation: Despite complete KV sharing, MQA’s cumulative variance curve shows substantial spread across PCs rather than concentration in the first few components. This quantifies the compensation effect: heads diversify their queries across many dimensions rather than collapsing to a low-dimensional subspace. - Scale-dependent behavior: The gap between methods narrows from 128M to 2.5B scale. At 128M, methods are more separated in the cumulative variance plot; at 2.5B, curves converge more closely, suggesting larger models develop more sophisticated compensation mechanisms. ### C.6 Comparison: Uncentered vs PCA-based Effective Rank <details> <summary>2601.11471v3/x13.png Details</summary>  ### Visual Description ## Line Graphs: Effective Rank Comparison (2.5B) ### Overview The image contains two side-by-side line graphs comparing the effective rank of different attention mechanisms across layers. The left graph uses an "Original (Uncentered Gram Matrix)" and the right graph uses a "PCA-based (Centered Gram Matrix)". Both graphs show the effective rank (y-axis) as a function of layer (x-axis), with five distinct data series representing different methods. ### Components/Axes - **X-axis (Layer)**: Ranges from 0 to 35, labeled "Layer". - **Y-axis (Effective Rank)**: Ranges from 6 to 18, labeled "Effective Rank (PCA)" for the right graph and "Effective Rank" for the left graph. - **Legends**: - **Left Graph (Original)**: - Standard MHA: Blue - LRKV (r=64): Orange - GQA: Green - MQA: Red - MLA: Purple - **Right Graph (PCA-based)**: - Standard MHA: Blue - LRKV (r=64): Orange - GQA: Green - MQA: Red - MLA: Purple - **Graph Titles**: - Left: "Original (Uncentered Gram Matrix)" - Right: "PCA-based (Centered Gram Matrix)" ### Detailed Analysis #### Left Graph (Original) - **Standard MHA (Blue)**: Starts at ~14 (layer 0), rises sharply to ~18 by layer 1, then fluctuates slightly around 18. - **LRKV (r=64) (Orange)**: Starts at ~7 (layer 0), rises sharply to ~18 by layer 1, then stabilizes. - **GQA (Green)**: Starts at ~11 (layer 0), rises sharply to ~18 by layer 1, then fluctuates slightly. - **MQA (Red)**: Starts at ~6 (layer 0), rises sharply to ~18 by layer 1, then shows significant dips (e.g., ~12 at layer 15, ~14 at layer 25). - **MLA (Purple)**: Starts at ~9 (layer 0), rises sharply to ~18 by layer 1, then fluctuates slightly. #### Right Graph (PCA-based) - **Standard MHA (Blue)**: Starts at ~16 (layer 0), rises sharply to ~17 by layer 1, then stabilizes. - **LRKV (r=64) (Orange)**: Starts at ~11 (layer 0), rises sharply to ~17 by layer 1, then stabilizes. - **GQA (Green)**: Starts at ~13 (layer 0), rises sharply to ~17 by layer 1, then fluctuates slightly. - **MQA (Red)**: Starts at ~9 (layer 0), rises sharply to ~17 by layer 1, then shows significant dips (e.g., ~15 at layer 15, ~16 at layer 25). - **MLA (Purple)**: Starts at ~9 (layer 0), rises sharply to ~17 by layer 1, then stabilizes. ### Key Observations 1. **Convergence**: All methods converge to similar effective ranks (~18 in the original graph, ~17 in the PCA-based graph) after the first layer (layer 1). 2. **Initial Variability**: The original graph shows greater variability in starting effective ranks (6–14) compared to the PCA-based graph (9–16). 3. **MQA Anomalies**: MQA (red) exhibits the most pronounced dips in both graphs, suggesting potential instability or sensitivity to layer-specific factors. 4. **PCA Impact**: The PCA-based graph shows slightly lower effective ranks overall, indicating possible differences in information retention or dimensionality reduction effects. ### Interpretation The data suggests that the effective rank of attention mechanisms stabilizes after the first layer, regardless of the Gram matrix type. The PCA-based approach results in marginally lower effective ranks, which could imply more efficient or constrained information encoding. The MQA method's sharp initial rise and subsequent dips highlight its sensitivity to layer-specific dynamics, potentially indicating overfitting or suboptimal performance in certain layers. The convergence of all methods after layer 1 underscores the robustness of these attention mechanisms in capturing information, while the PCA-based approach may offer a trade-off between efficiency and rank preservation. </details> Figure 14: Comparing uncentered vs PCA-based effective rank reveals compensation mechanisms (2.5B). Layer-by-layer comparison shows consistent patterns across all 36 layers. Left panel (uncentered): Standard MHA and LRKV maintain high effective rank (17-18), while MQA shows substantial degradation (15-16), suggesting significant diversity loss from complete sharing. Right panel (PCA-based): After centering, all methods converge to narrower range (16-17), with MQA showing dramatic improvement. This reveals the compensation effect: MQA’s shared KV structure creates a strong mean direction (visible in uncentered analysis), but heads compensate by diversifying queries around this mean (revealed by centering). LRKV consistently tracks Standard MHA in both metrics, confirming true head independence. GQA and MLA show intermediate behavior with modest gaps between uncentered and centered metrics. <details> <summary>2601.11471v3/x14.png Details</summary>  ### Visual Description ## Line Graphs: Effective Rank Comparison (128M) ### Overview The image contains two side-by-side line graphs comparing the effective rank of different attention mechanisms across transformer layers. The left graph uses an "Original (Uncentered Gram Matrix)" framework, while the right graph uses a "PCA-based (Centered Gram Matrix)" approach. Both graphs show six distinct data series representing different methods, with layers (0-10) on the x-axis and effective rank values on the y-axis. ### Components/Axes **Left Graph (Original):** - **X-axis (Layer):** Integer values 0-10 - **Y-axis (Effective Rank):** Continuous scale 2.0-6.0 - **Legend:** Positioned bottom-left, containing: - Standard MHA (blue solid line) - LRKV (r=16) (orange dashed line) - LRKV (r=64) (green dotted line) - MQA (red dotted line) - MLA (purple dash-dot line) - GQA (brown dashed line) **Right Graph (PCA-based):** - **X-axis (Layer):** Integer values 0-10 - **Y-axis (Effective Rank (PCA)):** Continuous scale 4.0-5.0 - **Legend:** Identical to left graph, with same color coding ### Detailed Analysis **Left Graph Trends:** 1. **Standard MHA (blue):** Starts at 3.6 (layer 0), peaks at 5.9 (layer 1), then declines to 5.5 (layer 10) 2. **LRKV (r=16) (orange):** Begins at 2.0 (layer 0), rises sharply to 5.8 (layer 1), then fluctuates between 5.5-5.8 3. **LRKV (r=64) (green):** Maintains near-constant 6.0 across all layers 4. **MQA (red):** Shows volatile pattern: 2.0 → 5.7 → 4.5 → 3.0 → 3.2 5. **MLA (purple):** Starts at 4.8, peaks at 5.7 (layer 1), then declines to 4.3 (layer 10) 6. **GQA (brown):** Begins at 3.3, rises to 5.7 (layer 1), then fluctuates between 5.3-5.6 **Right Graph Trends:** 1. **Standard MHA (blue):** Starts at 4.7, peaks at 4.9 (layer 1), then declines to 4.8 (layer 10) 2. **LRKV (r=16) (orange):** Begins at 4.8, rises to 4.95 (layer 1), then fluctuates between 4.85-4.95 3. **LRKV (r=64) (green):** Maintains near-constant 4.98 across all layers 4. **MQA (red):** Shows dramatic drop: 4.7 → 4.9 → 4.8 → 4.1 → 4.3 5. **MLA (purple):** Starts at 4.7, peaks at 4.9 (layer 1), then declines to 4.6 (layer 10) 6. **GQA (brown):** Begins at 4.4, rises to 4.9 (layer 1), then fluctuates between 4.8-4.9 ### Key Observations 1. **Dimensionality Reduction Effect:** PCA-based graph shows consistently lower effective ranks (4.0-5.0 vs 2.0-6.0 in original) 2. **Method Stability:** LRKV (r=64) maintains highest rank in both graphs, suggesting robustness 3. **MQA Anomaly:** Shows significant drop at layer 4 in original graph (3.0 vs 4.1 in PCA) 4. **Layer Sensitivity:** All methods show most dramatic changes between layers 0-2 5. **Normalization Impact:** PCA-based graph compresses rank values by ~20% compared to original ### Interpretation The data demonstrates that PCA-based processing reduces effective rank across all methods while preserving relative performance patterns. LRKV (r=64) consistently achieves highest ranks, suggesting it maintains more information. The MQA method's anomalous drop at layer 4 in the original graph may indicate layer-specific vulnerabilities. The PCA transformation appears to normalize rank distributions, potentially improving interpretability while sacrificing some absolute rank values. These findings suggest PCA-based approaches could be valuable for analyzing attention mechanisms in transformer models, particularly when comparing different architectural choices. </details> Figure 15: Effective rank comparison at 128M scale. Similar patterns emerge with 6 heads: uncentered metric shows larger separation between methods, while PCA-based metric reveals closer clustering. MQA’s improvement from uncentered (4.36/6 = 72.6%) to PCA-based (4.70/6 = 78.4%) demonstrates compensation effect at smaller scale. The larger percentage gaps at 128M versus 2.5B suggest smaller models rely more heavily on shared structure with compensation, while larger models develop more independent specialization. The side-by-side comparison of uncentered versus PCA-based effective rank reveals why centering is essential for understanding head diversity: #### The MQA compensation effect (quantified). At 2.5B scale, MQA improves from 86.2% (uncentered) to 91.0% (PCA-based)—a 4.8 percentage point gain. This quantifies how much diversity MQA recovers through query specialization after accounting for its shared KV baseline. Without centering, we would conclude MQA loses 12.7pp of diversity versus Standard MHA (98.9% → 86.2%); with centering, the loss is only 3.0pp (94.0% → 91.0%), revealing MQA is much less constrained than raw metrics suggest. #### LRKV’s consistency across metrics. LRKV shows minimal gap between uncentered (98.3%) and PCA-based (93.5%) effective rank—a 4.8pp difference similar to Standard MHA’s 4.9pp gap. This indicates LRKV achieves diversity through genuinely independent head specialization rather than compensation around a shared baseline. The low-rank residuals provide real degrees of freedom, not just perturbations of a dominant mean structure. #### Scale-dependent centering effects. The gaps between uncentered and PCA-based metrics are larger at 128M scale (LRKV: 13.2pp gap) than at 2.5B (4.8pp gap). This suggests smaller models rely more on shared structure with local compensation, while larger models develop more truly independent head specializations. The PCA-based metric, by removing the mean direction, reveals this scale-dependent transition. ### C.7 Eigenvalue Spectra of Uncentered Similarity Matrices For completeness, we also show the eigenvalue spectra of uncentered head similarity matrices, which have been used in prior work for head diversity analysis. <details> <summary>2601.11471v3/x15.png Details</summary>  ### Visual Description ## Line Chart: Eigenvalue Spectra of Head Similarity Matrices (2.5B) ### Overview The image presents three line charts comparing the eigenvalue spectra of head similarity matrices across three transformer layers: Early Layer (L=0), Middle Layer (L=18), and Late Layer (L=35). Five methods are compared: Standard MHA, LRKV (r=64), GQA, MQA, and MLA. Each chart shows eigenvalues plotted against eigenvalue indices (2.5–17.5), with distinct color-coded lines for each method. --- ### Components/Axes - **X-axis**: Eigenvalue Index (2.5 to 17.5, increments of 2.5). - **Y-axis**: Eigenvalue (ranging from 0 to ~10 in Early Layer, 0 to ~3.5 in Middle Layer, and 0 to ~1.8 in Late Layer). - **Legends**: Positioned on the right of each subplot, with colors: - Blue: Standard MHA - Orange: LRKV (r=64) - Green: GQA - Red: MQA - Purple: MLA --- ### Detailed Analysis #### Early Layer (L=0) - **Standard MHA (blue)**: Starts at ~10, drops sharply to ~0.5 by index 5, then plateaus. - **LRKV (r=64) (orange)**: Begins at ~8.5, declines to ~0.5 by index 5, then stabilizes. - **GQA (green)**: Starts at ~6, decreases to ~0.5 by index 5, then flattens. - **MQA (red)**: Peaks at ~10, drops to ~0.5 by index 5, then remains flat. - **MLA (purple)**: Starts at ~8, declines to ~0.5 by index 5, then stabilizes. #### Middle Layer (L=18) - **Standard MHA (blue)**: Begins at ~3.5, drops to ~1.0 by index 5, then plateaus. - **LRKV (r=64) (orange)**: Starts at ~1.2, declines to ~0.8 by index 5, then stabilizes. - **GQA (green)**: Begins at ~1.5, decreases to ~0.8 by index 5, then flattens. - **MQA (red)**: Peaks at ~3.5, drops to ~0.8 by index 5, then remains flat. - **MLA (purple)**: Starts at ~2.5, declines to ~1.0 by index 5, then stabilizes. #### Late Layer (L=35) - **Standard MHA (blue)**: Begins at ~1.8, drops to ~0.8 by index 5, then plateaus. - **LRKV (r=64) (orange)**: Starts at ~1.0, declines to ~0.8 by index 5, then stabilizes. - **GQA (green)**: Begins at ~1.6, decreases to ~0.8 by index 5, then flattens. - **MQA (red)**: Peaks at ~1.6, drops to ~0.8 by index 5, then remains flat. - **MLA (purple)**: Starts at ~1.4, declines to ~0.8 by index 5, then stabilizes. --- ### Key Observations 1. **Sharp Initial Drop**: All methods exhibit a steep decline in eigenvalues at the first few indices (2.5–5), followed by gradual flattening. 2. **Layer-Specific Trends**: - **Early Layer**: Highest eigenvalues (~10), indicating greater similarity variability. - **Middle Layer**: Moderate eigenvalues (~3.5), showing reduced variability. - **Late Layer**: Lowest eigenvalues (~1.8), suggesting stabilized similarity. 3. **Method Comparisons**: - **Standard MHA** and **MQA** consistently start with the highest eigenvalues but drop sharply. - **LRKV (r=64)** and **MLA** show more gradual declines, maintaining higher eigenvalues longer. - **GQA** exhibits intermediate behavior, with eigenvalues decreasing steadily across layers. --- ### Interpretation The eigenvalue spectra reveal how head similarity matrices evolve across transformer layers. Higher eigenvalues in early layers suggest greater diversity in attention head behavior, which diminishes as layers progress. Methods like **Standard MHA** and **MQA** aggressively reduce similarity early, potentially prioritizing distinct feature extraction. In contrast, **LRKV (r=64)** and **MLA** retain higher eigenvalues longer, possibly preserving more nuanced relationships. The Late Layer’s convergence across methods implies stabilized attention patterns, critical for final representations. These trends highlight trade-offs between diversity and stability in attention mechanisms, with implications for model efficiency and performance. </details> Figure 16: Uncentered eigenvalue spectra at early, middle, and late layers (2.5B scale). The eigenvalue distribution of uncentered head similarity matrices $S$ shows broader separation between methods than PCA-based analysis. LRKV maintains eigenvalue spectra nearly identical to Standard MHA across all depth regimes. MQA shows elevated leading eigenvalues and faster tail decay, indicating stronger head correlation in the uncentered view—but PCA analysis reveals much of this is due to the shared mean direction rather than loss of diversity. Early layers show slightly higher leading eigenvalues across all methods, indicating a stronger shared component in bilinear form space at lower layers; deeper layers exhibit more uniform spectra, consistent with variance being distributed across more modes. <details> <summary>2601.11471v3/x16.png Details</summary>  ### Visual Description ## Line Graphs: Eigenvalue Spectra of Head Similarity Matrices (128M) ### Overview The image presents three line graphs comparing the eigenvalue spectra of head similarity matrices across three neural network layers: Early Layer (L=0), Middle Layer (L=6), and Late Layer (L=11). Each graph shows six methods (Standard MHA, LRKV (r=16), LRKV (r=64), MQA, MLA, GQA) with eigenvalues plotted against eigenvalue indices (1–6). The y-axis represents eigenvalue magnitude, while the x-axis represents the eigenvalue index. --- ### Components/Axes - **X-axis**: "Eigenvalue Index" (1–6), labeled uniformly across all subplots. - **Y-axis**: "Eigenvalue" (ranging from 0 to 5 in Early Layer, 0 to 2.5 in Middle Layer, and 0 to 4 in Late Layer). - **Legend**: Located on the right side of each subplot, with six color-coded methods: - **Blue**: Standard MHA - **Orange**: LRKV (r=16) - **Green**: LRKV (r=64) - **Red**: MQA - **Purple**: MLA - **Brown**: GQA - **Subplot Titles**: - Top-left: "Early Layer (L=0)" - Center: "Middle Layer (L=6)" - Bottom-right: "Late Layer (L=11)" --- ### Detailed Analysis #### Early Layer (L=0) - **Standard MHA (Blue)**: Starts at ~5.0 (index 1), drops sharply to ~0.8 (index 2), then gradually declines to ~0.2 (index 6). - **LRKV (r=16) (Orange)**: Begins at ~4.5 (index 1), declines to ~0.3 (index 2), then stabilizes near ~0.1 (index 6). - **LRKV (r=64) (Green)**: Starts at ~3.8 (index 1), drops to ~0.5 (index 2), then remains flat (~0.3–0.4). - **MQA (Red)**: Peaks at ~4.0 (index 1), drops to ~0.4 (index 2), then declines to ~0.1 (index 6). - **MLA (Purple)**: Starts at ~2.5 (index 1), drops to ~0.6 (index 2), then stabilizes (~0.3–0.4). - **GQA (Brown)**: Begins at ~3.0 (index 1), declines to ~0.7 (index 2), then remains flat (~0.4–0.5). #### Middle Layer (L=6) - **Standard MHA (Blue)**: Starts at ~2.5 (index 1), drops to ~1.0 (index 2), then declines to ~0.8 (index 6). - **LRKV (r=16) (Orange)**: Begins at ~1.8 (index 1), drops to ~0.9 (index 2), then stabilizes (~0.7–0.8). - **LRKV (r=64) (Green)**: Starts at ~1.2 (index 1), drops to ~0.9 (index 2), then remains flat (~0.8–0.9). - **MQA (Red)**: Peaks at ~2.0 (index 1), drops to ~0.7 (index 2), then declines to ~0.5 (index 6). - **MLA (Purple)**: Starts at ~2.2 (index 1), drops to ~0.8 (index 2), then stabilizes (~0.6–0.7). - **GQA (Brown)**: Begins at ~1.5 (index 1), drops to ~0.9 (index 2), then remains flat (~0.7–0.8). #### Late Layer (L=11) - **Standard MHA (Blue)**: Starts at ~3.0 (index 1), drops to ~1.0 (index 2), then declines to ~0.5 (index 6). - **LRKV (r=16) (Orange)**: Begins at ~2.8 (index 1), drops to ~0.9 (index 2), then stabilizes (~0.7–0.8). - **LRKV (r=64) (Green)**: Starts at ~1.3 (index 1), drops to ~0.9 (index 2), then remains flat (~0.8–0.9). - **MQA (Red)**: Peaks at ~3.5 (index 1), drops to ~0.6 (index 2), then declines to ~0.4 (index 6). - **MLA (Purple)**: Starts at ~2.8 (index 1), drops to ~0.7 (index 2), then stabilizes (~0.5–0.6). - **GQA (Brown)**: Begins at ~2.0 (index 1), drops to ~0.8 (index 2), then remains flat (~0.6–0.7). --- ### Key Observations 1. **Eigenvalue Decay**: All methods show a sharp decline in eigenvalues from index 1 to 2, followed by gradual stabilization. This suggests that the first two eigenvalues capture the majority of the variance in the similarity matrices. 2. **Method-Specific Trends**: - **MLA (Purple)**: Consistently retains higher eigenvalues in early and middle layers compared to other methods, indicating stronger preservation of high-variance features. - **LRKV (r=64) (Green)**: Shows the most stable eigenvalues across layers, suggesting robustness in capturing consistent patterns. - **MQA (Red)**: Exhibits the steepest initial drop in all layers, implying a focus on early-layer features. 3. **Layer-Specific Patterns**: - **Early Layer (L=0)**: Higher eigenvalues overall, reflecting greater importance of early-layer features. - **Late Layer (L=11)**: Lower eigenvalues, indicating reduced significance of late-layer features in the similarity matrices. --- ### Interpretation The eigenvalue spectra reveal how different attention mechanisms prioritize and retain information across layers. - **MLA's performance** in early and middle layers suggests it effectively preserves critical features, potentially improving model interpretability or robustness. - **LRKV (r=64)** demonstrates stability, which may indicate its suitability for tasks requiring consistent feature representation. - The **sharp decay** in eigenvalues across all methods highlights the dominance of low-dimensional patterns in the similarity matrices, a common characteristic in attention mechanisms. These trends could inform the design of more efficient or interpretable attention mechanisms by emphasizing methods that balance feature retention and computational efficiency. </details> Figure 17: Uncentered eigenvalue spectra at 128M scale. With 6 heads, the eigenvalue structure is more pronounced. LRKV (r=64) tracks Standard MHA closely, while LRKV (r=16) shows slightly elevated leading eigenvalues in later layers, consistent with reduced effective rank. MQA exhibits the most concentrated spectra in the uncentered view, with a dominant leading eigenvalue suggesting strong head redundancy—but this is largely due to the shared KV mean direction rather than true loss of diversity, as PCA analysis reveals. ### C.8 Head Similarity Heatmaps <details> <summary>2601.11471v3/x17.png Details</summary>  ### Visual Description ## Heatmaps: Attention Head Correlation Across Mechanisms ### Overview The image presents five heatmaps comparing the correlation between attention heads in different neural network mechanisms. Each heatmap visualizes the relationship between head indices (0–5) using a color gradient from -1.00 (blue) to 1.00 (red), with red indicating strong positive correlation and blue indicating strong negative correlation. ### Components/Axes - **X-axis**: "Head Index" (0–5), labeled "Head Index" - **Y-axis**: "Head Index" (0–5), labeled "Head Index" - **Color Scale**: Ranges from -1.00 (blue) to 1.00 (red), with a legend on the right - **Heatmap Labels**: - Standard MHA - LRKV (r=16) - LRKV (r=64) - MQA - MLA ### Detailed Analysis 1. **Standard MHA** - Diagonal red squares dominate (e.g., (0,0), (1,1), ..., (5,5)), indicating strong self-correlation. - Off-diagonal values are lighter (yellow/orange), suggesting weaker cross-head correlations (~0.50–0.75). 2. **LRKV (r=16)** - Diagonal red squares persist but with slightly lighter shading (~0.60–0.80). - Off-diagonal red squares appear (e.g., (0,1), (1,2)), indicating moderate cross-head correlations (~0.40–0.60). 3. **LRKV (r=64)** - Diagonal red squares are darker (~0.70–0.90), showing stronger self-correlation. - Off-diagonal red squares are more frequent and darker (~0.50–0.70), suggesting increased cross-head correlation. 4. **MQA** - Diagonal red squares are less pronounced (~0.50–0.70). - Off-diagonal values are mixed (yellow/orange), with some red squares (e.g., (0,2), (3,4)), indicating sporadic cross-head correlations. 5. **MLA** - Diagonal red squares are faint (~0.40–0.60). - Off-diagonal red squares are scattered (e.g., (0,3), (2,5)), with weaker correlations (~0.30–0.50). ### Key Observations - **Standard MHA** exhibits the strongest diagonal dominance, typical of self-attention mechanisms. - **LRKV** variants show increased off-diagonal correlations, with **r=64** having the most uniform distribution. - **MQA** and **MLA** display less structured patterns, suggesting alternative attention distribution strategies. - All mechanisms show negative correlations (blue) in sparse regions, but these are minimal compared to positive values. ### Interpretation The data suggests that attention mechanisms vary in how they distribute correlations between heads: - **Standard MHA** prioritizes self-correlation, limiting cross-head interaction. - **LRKV** (especially with higher `r`) enables broader cross-head communication, potentially improving model flexibility. - **MQA** and **MLA** may sacrifice some self-correlation for distributed attention, which could enhance robustness but reduce direct head-to-head alignment. - The uniform red distribution in **LRKV (r=64)** implies that increasing `r` amplifies head interdependence, possibly at the cost of computational efficiency. This analysis highlights trade-offs between attention mechanism design and head correlation patterns, with implications for model interpretability and performance. </details> Figure 18: Head similarity matrices at 128M scale. With fewer heads (6 vs 18), individual head relationships are more visible. LRKV (r=64) shows similarity structure nearly identical to Standard MHA, with moderate positive correlations but no dominant clusters. LRKV (r=16) exhibits slightly stronger correlations, consistent with its reduced effective rank. The mean off-diagonal similarity for LRKV (r=64) is within 0.02–0.04 of Standard MHA across layers, providing quantitative confirmation of the visual similarity. These uncentered heatmaps complement the PCA-based analysis by showing the raw pairwise similarities before mean removal. The heatmaps provide direct visualization of head relationships in the uncentered similarity space, complementing the PCA-based aggregate metrics. Key observations: - LRKV preserves correlation structure: At both scales, LRKV’s similarity matrices closely match Standard MHA’s patterns. Off-diagonal entries (head-to-head similarities) show comparable magnitudes and spatial distribution, indicating that low-rank factorization does not artificially increase or decrease head correlations in the uncentered space. - No emergent clustering: Standard MHA shows diffuse correlations without strong block structure (aside from diagonal dominance). LRKV maintains this property, suggesting heads remain independently specialized rather than forming redundant groups, even before centering removes shared mean structure. - GQA shows group structure: At 2.5B scale with 6 KV groups of 3 heads each, GQA’s heatmap reveals visible within-group structure—heads sharing KV projections exhibit stronger mutual similarity. This validates that partial KV sharing induces measurable correlation visible in raw similarity space. - Quantitative confirmation: The mean off-diagonal similarity for LRKV (r=64) is within 0.02–0.04 of Standard MHA across layers, confirming the visual similarity is not merely qualitative. This tight correspondence holds in both uncentered and PCA-based analysis. ### C.9 Key Findings from PCA Analysis #### 1. LRKV preserves head independence within 1% of standard attention. At 2.5B scale, LRKV achieves 93.5% PCA-based effective rank versus 94.0% for Standard MHA (0.5pp gap). This near-perfect preservation occurs despite 52.6% KV cache size, validating that rank-64 residuals provide sufficient capacity for head specialization. The consistency across all 36 layers demonstrates depth-invariant factorization quality. #### 2. The MQA compensation effect. MQA’s 4.8pp improvement from uncentered (86.2%) to PCA-based (91.0%) effective rank at 2.5B scale reveals a previously unknown mechanism: forced KV sharing creates a strong mean bilinear form, but heads compensate by diversifying query projections around this mean. The centered analysis isolates this true independence, showing MQA is less constrained than prior metrics suggest. This effect is consistent at 128M scale (5.8pp improvement), confirming it is a fundamental property of MQA rather than a scale-specific artifact. #### 3. MLA sits between GQA and MQA in diversity space. MLA achieves 91.6% PCA-based effective rank at 2.5B, placing it between GQA (91.7%) and MQA (91.0%). This makes architectural sense: MLA’s latent bottleneck constrains heads more than GQA’s group sharing but less than MQA’s complete sharing, with per-head decompression allowing specialization after the bottleneck. At 128M scale, MLA shows 78.5% effective rank, slightly below MQA (78.4%), suggesting the latent bottleneck is more constraining at smaller scales where bottleneck capacity is limited. #### 4. Scale-dependent behavior. The gap between uncentered and PCA-based metrics shrinks from 128M to 2.5B (e.g., LRKV: 13.2pp gap at 128M vs 4.8pp at 2.5B), suggesting larger models develop stronger head specialization around shared structure rather than purely independent representations. This transition indicates that model scale affects not just capacity but the fundamental organization principle of attention heads. #### 5. Rank selection determines capacity. Comparing LRKV configurations reveals that appropriate rank is critical. At 128M scale, r=16 (12.5% of head dimension) achieves 81.5% PCA-based effective rank, while r=64 (50%) achieves 83.0%—a modest 1.5pp improvement in diversity but substantial performance difference (0.881 vs 0.875 BPB). This confirms that representational capacity, not geometric properties, determines quality: even moderate diversity can yield strong performance if the rank provides sufficient degrees of freedom. ### C.10 Comparison to Prior Diversity Metrics Our PCA-based approach differs from prior work in several key respects: #### CKA/SVCCA (Kornblith et al., 2019; Raghu et al., 2017) . These methods compare activation similarities (representations), not functional operators. While CKA uses centered Gram matrices (providing the ”C” in ”Centered Kernel Alignment”), it analyzes activation space rather than weight-space bilinear forms. Our contribution is applying the centering principle to the space of attention operators themselves, revealing compensation mechanisms invisible when analyzing activations. #### Attention pattern similarity (Michel et al., 2019; Voita et al., 2019) . These behavioral metrics measure which tokens heads attend to on specific inputs. They depend on input data distribution and don’t directly measure the representational capacity encoded in the parameterization. Our approach analyzes the parameterization itself, providing a data-independent measure of potential diversity. #### Raw weight comparison. Directly comparing $W_h^K$ or $W_h^Q$ matrices is not gauge-invariant: arbitrary per-head rotations that preserve attention function can make identical heads appear different or vice versa. The bilinear form $W^Q(W^K)^⊤$ is the minimal gauge-invariant object for comparison. #### Uncentered Gram matrices. Several prior analyses (Zhang et al., 2023) form head similarity matrices (often using uncentered correlations or kernel similarities) and interpret their spectra directly; our contribution is to apply explicit kernel centering in the space of gauge-invariant attention operators, which separates shared mean structure from variance around the mean. This conflates mean structure (shared baselines) with spread (true diversity), masking compensation effects like MQA’s query specialization. Our centered analysis isolates intrinsic dimensionality independent of mean direction. The combination of (i) gauge-invariant bilinear forms, (ii) centered Gram matrix (kernel PCA), and (iii) PCA interpretation as variance explained provides, to our knowledge, the first principled operator-space analysis of transformer head diversity. This methodology reveals phenomena invisible to prior approaches, such as the MQA compensation effect and the precise quantification of LRKV’s near-perfect diversity preservation. ### C.11 Methodological Limitations and Extensions #### Limitations. Our analysis focuses on final trained checkpoints and does not track how the factorization evolves during training. The PCA-based metric measures potential diversity encoded in the parameterization but not behavioral diversity on specific inputs. The gauge-invariant similarity metric is one of many possible inner products on bilinear forms; alternative metrics might reveal additional structure. #### Future extensions. Several directions merit exploration: (1) Analyzing PCA eigenvalue dynamics during pretraining to understand how shared and head-specific structure co-evolve. (2) Comparing parameterization-based diversity (this work) with input-dependent behavioral diversity to understand their relationship. (3) Extending the framework to cross-attention in encoder-decoder models, where query and key-value heads may have different sizes. (4) Investigating whether PCA-guided initialization or regularization can improve training dynamics by explicitly encouraging high-variance factorizations. ## Appendix D Experimental Setup (Full Details) Table 4: Mathematical comparison of attention mechanisms. We compare how different mechanisms parameterize key/value projections and where compression is applied. Here $X∈ℝ^T× d$ is the token sequence, $d_h=d/H$ is the head dimension, $H$ is the number of heads, and $r,d_c\ll d_h$ . | Standard MHA MQA GQA | $W_h^K,V∈ℝ^d× d_h$ (independent per head) $W_h^K,V=W_shared^K,V ∀ h$ $W_h^K,V=W_g(h)^K,V$ for group $g(h)$ | None (full rank) $rank(W_h^K,V)=rank(W_shared)$ Full rank within group | None Across heads (complete sharing) Across heads (group-wise sharing) | | --- | --- | --- | --- | | MLA | $W_h^K,V=W_downW_up,h^K,V$ | $rank(W_h^K,V)≤ d_c$ | Across tokens (latent bottleneck) | | LRKV (ours) | $W_h^K,V=W_shared^K,V+U_h^K,V(B_h^K,V)^⊤$ | $rank(W_h^K,V-W_shared^K,V)≤ r$ | Across heads (additive low-rank deviations) | ### D.1 Pretraining Configuration We pretrain from scratch a family of decoder-only Transformer language models with parameter counts ranging from 128M to 6.3B on the FineWeb-Edu dataset (Penedo et al., 2024). Unless otherwise stated, model scale (e.g., 128M, 2.5B, 6.3B) refers to the parameter count of the base architecture with standard multi-head attention (MHA). Alternative attention mechanisms (LRKV, MQA, GQA, MLA) replace only the attention module while keeping all other architectural components fixed, which will result in reduced total parameter count due to reparameterized key-value projections. #### Model Architecture Summary. Table 5 provides complete architectural specifications for all model scales evaluated in this work. Table 5: Complete model architectural specifications. All models use $d_head=128$ and FFN expansion ratio of 4 (i.e., $d_FFN=4× d_model$ ). | 128M 512M 1.2B | 12 24 24 | 6 12 12 | 768 1536 1536 | 3072 6144 6144 | 50304 50304 50304 | | --- | --- | --- | --- | --- | --- | | 2.5B | 18 | 18 | 2304 | 9216 | 50304 | | 6.3B | 32 | 32 | 4096 | 16384 | 50304 | #### Training Configuration Summary. Table 6 provides complete training configuration details across all model scales. Table 6: Training configuration across model scales. Batch size refers to number of sequences; total tokens per batch = batch size $×$ context length. All models trained on FineWeb-Edu with Muon+AdamW optimization. | 128M 512M 1.2B | 2048 8192 2048 | 16 4 8 (GA=2) | 32,768 32,768 32,768 | 100B 50B 50B | 31.0 153.6 304.7 | Ablations Long context Scaling study | | --- | --- | --- | --- | --- | --- | --- | | 2.5B | 2048 | 8 (GA=2) | 32,768 | 50B | 635.2 | Main experiments | | 6.3B | 2048 | 8 (GA=2) | 32,768 | 50B | 1603.3 | Scaling study | GA=gradient accumulation steps. EF=ExaFLOPs. 128M overtrained (100B tokens) for low-variance ablations; others use compute-near-optimal budgets (Hoffmann et al., 2022). #### Optimization Details. We use a hybrid optimizer strategy: the Muon optimizer (Sardana and Havens, 2024) for Transformer weight matrices and AdamW for embeddings, with component-specific learning rates (matrix: 0.02, embedding: 0.2, unembedding: 0.004). All models are optimized using cosine learning rate decay with linear warmup (2000 steps), weight decay 0.1, and gradient clipping at 1.0. #### Hardware and Precision. All experiments are run on a single 8 $×$ H200 GPU node in mixed-precision (bfloat16) mode. Full architectural and optimization hyperparameters, including learning rate schedules and detailed optimizer settings, are provided in Appendix F. #### Evaluation Metrics. Training progress is monitored using the cross-entropy loss and bits-per-byte (BPB) on held-out validation split of FineWeb-Edu. These metrics allow consistent comparison across model scales and serve as the primary indicators of data efficiency and convergence quality. ### D.2 Mid-training Configuration After pretraining, we perform supervised fine-tuning (mid-training) on a curated mixture of instructional data to improve downstream task performance. The mid-training dataset consists of three components: 1. SmolTalk (Allal et al., 2024) with 460K conversational examples for general instruction-following 1. MMLU auxiliary train (Hendrycks et al., 2021) with 100K multiple-choice questions spanning diverse academic subjects 1. GSM8K (Cobbe et al., 2021) with 8K grade-school math problems including calculator tool use This yields a total training mixture of approximately 568K examples. Validation is performed on held-out test splits using proportional sampling (24K SmolTalk, 5.2K MMLU, 420 GSM8K examples). We train for one full epoch over the mid-training mixture with a batch size of 524,288 tokens and sequence length of 2048. We use a decayed linear schedule for the learning rate, using AdamW for the embeddings and Muon for the remaining parameters. Models are evaluated every 150 steps using bits-per-byte on the validation set. All mid-training runs use the same hyperparameters across different attention mechanisms to ensure fair comparison. ### D.3 Attention Mechanism Configurations To evaluate the effectiveness of low-rank KV factorization, we compare LRKV against four baseline attention mechanisms across all model scales. Table 7 provides complete configuration details for all methods and scales. Table 7: Attention mechanism configurations across model scales. KV cache percentages are relative to Standard MHA. All models use $d_head=128$ and 2048-token context (except 512M long-context with 8192 tokens). | 128M 512M 1.2B | 6 12 12 | 6 12 12 | 100% 100% 100% | 1 (16.7%) 1 (8.3%) 1 (8.3%) | 3 (50%) 4 (33.3%) 4 (33.3%) | 50% 33.3% 33.3% | 128 (8.3%) 256 (8.3%) 256 (8.3%) | 46 51 51 | 52.6% 48.2% 48.2% | | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | | 2.5B | 18 | 18 | 100% | 1 (5.6%) | 6 (33.3%) | 33.3% | 384 (8.3%) | 55 | 48.5% | | 6.3B | 32 | 32 | 100% | 1 (3.1%) | 2 (6.3%) | 6.3% | 1024 (12.5%) | 54 | 45.3% | | 512M-8K † | 12 | 12 | 100% | 1 (8.3%) | 4 (33.3%) | 33.3% | 256 (8.3%) | 51 | 48.2% | † 512M-8K denotes the long-context configuration with 8192-token sequences. All other models use 2048-token context. #### KV Cache Calculation. All KV cache percentages are reported relative to Standard MHA’s memory usage. For LRKV, the reported cache size assumes an optimized implementation that caches the shared projection $W_shared$ and per-head low-rank latents $R_h^K,V$ rather than the fully reconstructed key-value matrices, following the theoretical analysis in Section 3. The cache ratio follows: $\frac{d_h+Hr}{Hd_h}=\frac{1}{H}+\frac{r}{d_h}$ . This configuration balances memory efficiency with model expressiveness: LRKV uses approximately half the KV cache of Standard MHA while maintaining full representational capacity through low-rank adaptation, outperforming both memory-minimal approaches (MQA, MLA) and full attention in downstream task performance. ### D.4 Rank Ablation Experiments To systematically evaluate how residual rank affects LRKV’s performance-memory tradeoff, we conduct comprehensive ablation studies at the 128M scale. These experiments isolate the effect of rank selection while controlling for all other architectural and training factors. #### LRKV Rank Ablation. We train five 128M parameter models with LRKV using ranks $r∈\{8,16,32,64,128\}$ on 100B tokens of FineWeb-Edu. All models use identical architecture (6 heads, $d_model=768$ , $d_head=128$ ), optimization hyperparameters, and training schedule. Table 8 presents complete results. Table 8: LRKV rank ablation and MLA latent dimension comparison (128M scale, 100B tokens). All models use 6 heads with $d_h=128$ . Cache percentages relative to Standard MHA baseline. | Configuration LRKV Rank Ablation LRKV-8 | Rank/Latent $r=8$ | KV Cache 22.9% | CE Loss $↓$ 2.908 | BPB $↓$ 0.881 | Notes Insufficient capacity | | --- | --- | --- | --- | --- | --- | | LRKV-16 | $r=16$ | 29.2% | 2.905 | 0.880 | Below MHA | | LRKV-32 | $r=32$ | 41.7% | 2.896 | 0.877 | Approaching MHA | | LRKV-64 | $r=64$ | 66.7% | 2.888 | 0.875 | Exceeds MHA | | LRKV-128 | $r=128$ | 116.7% | 2.877 | 0.873 | Best, but $>$ 100% cache | | MLA Latent Dimension Ablation | | | | | | | MLA-128 | $d_c=128$ | 8.3% | 2.901 | 0.876 | Minimal memory | | MLA-192 | $d_c=192$ | 12-5% | 2.898 | 0.876 | – | | MLA-384 | $d_c=384$ | 25.0% | 2.894 | 0.875 | Comparable to LRKV-64 | | Baselines | | | | | | | Standard MHA | 6 KV heads | 100% | 2.903 | 0.878 | Baseline | | GQA | 3 groups | 50.0% | 2.918 | 0.883 | – | | MQA | 1 KV head | 16.7% | 2.929 | 0.886 | Severe quality loss | #### Key Findings. The ablation reveals that performance improves monotonically with rank, confirming that representational capacity (not geometric properties of the factorization) determines modeling quality. Ranks around $r≈ 0.36$ – $0.43× d_h$ (46–55 for $d_h=128$ ) provide the optimal tradeoff: sufficient capacity to exceed Standard MHA performance while achieving substantial cache reduction. Comparing methods at matched cache budgets reveals LRKV’s architectural advantage: LRKV at 41.7% cache ( $r=32$ , BPB=0.877) slightly underperforms MLA at 50.0% cache ( $d_c=384$ , BPB=0.875), but LRKV at 66.7% cache ( $r=64$ , BPB=0.875) matches MLA-384 performance. However, LRKV avoids the sequence-length-dependent reconstruction overhead inherent to MLA’s latent compression (see Appendix F), providing better inference latency characteristics at long context lengths. All ablation models are trained for 100B tokens (rather than compute-optimal allocation) to reduce variance and enable high-confidence comparison of architectural capacity at fixed training budget. ### D.5 Long Context Experiments To evaluate LRKV’s effectiveness at extended sequence lengths, we conduct 8192-token context experiments using 512M parameter models trained on 50B tokens (153.6 ExaFLOPs). #### Configuration. The 512M models use 12 attention heads with $d_model=1536$ and $d_head=128$ , with 24 layers matching the 1.2B scale architecture. We compare LRKV ( $r=51$ , 48.2% cache), GQA (4 groups, 33.3% cache), MLA ( $d_c=256$ , 8.3% cache), and MQA (8.3% cache) against Standard MHA baseline. #### Batching Strategy. To maintain computational throughput while accommodating 4 $×$ longer contexts, we adjust batch size to keep total tokens per batch constant. Standard 2048-context runs use batch size 16 (32,768 tokens per batch); 8192-context runs use batch size 4 (32,768 tokens per batch). This ensures comparable gradient noise and optimizer dynamics while scaling to long contexts. #### Results. Table 9 presents validation performance at 8192-token context length after 50B tokens (153.6 ExaFLOPs). Table 9: Long context performance (512M models, 8K context, 50B tokens). Degradation computed relative to MHA baseline. | Standard MHA LRKV GQA | 100% 48.2% 33.3% | 2.740 2.665 2.704 | 0.494 0.481 0.488 | baseline -2.7% -1.3% | | --- | --- | --- | --- | --- | | MLA | 8.3% | 2.714 | 0.489 | -0.9% | | MQA | 8.3% | 2.728 | 0.492 | -0.4% | LRKV achieves the strongest performance, outperforming Standard MHA by 2.7% while using less than half the KV cache (48.2% vs 100%). All KV-efficient methods (LRKV, GQA, MLA, MQA) outperform the full Standard MHA baseline at 8K context, suggesting that some degree of KV compression may provide implicit regularization benefits at longer sequences. However, LRKV’s 2.7% advantage over MHA (compared to GQA’s 1.3%, MLA’s 0.9%, and MQA’s 0.4%) demonstrates that head-specific low-rank residuals are particularly effective for maintaining quality at extended context lengths. These results validate LRKV’s applicability to long-context scenarios where memory efficiency is critical, demonstrating that low-rank factorization preserves head diversity more effectively than aggressive sharing (MQA, GQA) or latent compression (MLA) approaches at extended sequence lengths. ### D.6 Downstream Evaluation Tasks After midtraining, we evaluate all models on five standard benchmarks: - ARC-Easy and ARC-Challenge (Clark et al., 2018): Question answering requiring scientific reasoning - MMLU (Hendrycks et al., 2021): Multi-task language understanding across 57 subjects - GSM8K (Cobbe et al., 2021): Grade-school math word problems - HumanEval (Chen et al., 2022): Python code generation from docstrings All evaluations follow standard zero-shot or few-shot protocols as specified in the original benchmark papers. ## Appendix E Pretraining-Downstream Correlation Analysis <details> <summary>2601.11471v3/x18.png Details</summary>  ### Visual Description ## Scatter Plot: Pretraining vs Downstream Performance ### Overview The image is a scatter plot comparing **Pretraining BPB (Bits Per Byte, lower is better)** on the x-axis and **Downstream Performance (%)** on the y-axis. Data points represent different models/methods, with a dashed trend line indicating a negative correlation (R²=0.786). The legend maps symbols/colors to specific models. --- ### Components/Axes - **Title**: "Pretraining vs Downstream Performance" (top center). - **X-axis**: "Pretraining BPB (lower is better)" with values ranging from ~0.720 to 0.728. - **Y-axis**: "Downstream Performance (%)" with values from ~33.5% to 38%. - **Legend**: Located in the bottom-right corner, with the following mappings: - **Red Diamond**: MLA - **Green Triangle**: MQA - **Orange Square**: GQA - **Blue Circle**: Standard MHA - **Purple Star**: LRKV - **Dashed Line**: Trend (R²=0.786). --- ### Detailed Analysis #### Data Points 1. **MQA (Green Triangle)**: - X: 0.728 (highest Pretraining BPB). - Y: ~33.5% (lowest Downstream Performance). 2. **GQA (Orange Square)**: - X: 0.726. - Y: ~35.8%. 3. **Standard MHA (Blue Circle)**: - X: 0.722. - Y: ~35.3%. 4. **MLA (Red Diamond)**: - X: 0.724. - Y: ~36.9%. 5. **LRKV (Purple Star)**: - X: 0.720 (lowest Pretraining BPB). - Y: ~38.0% (highest Downstream Performance). #### Trend Line - A **dashed line** slopes **upward** from left to right, indicating a **negative correlation** between Pretraining BPB and Downstream Performance (R²=0.786). This suggests that lower Pretraining BPB values are associated with higher downstream performance. --- ### Key Observations 1. **LRKV (Purple Star)** is the most efficient model, achieving the lowest Pretraining BPB (0.720) and highest Downstream Performance (38.0%). 2. **MQA (Green Triangle)** is the least efficient, with the highest Pretraining BPB (0.728) and lowest Downstream Performance (33.5%). 3. The trend line confirms a strong inverse relationship (R²=0.786), meaning Pretraining BPB is a significant predictor of downstream performance. 4. **Standard MHA (Blue Circle)** and **GQA (Orange Square)** fall in the middle, with moderate Pretraining BPB and performance. --- ### Interpretation The data demonstrates that **Pretraining BPB is a critical factor** in determining downstream performance. Models with lower Pretraining BPB (e.g., LRKV) outperform those with higher values (e.g., MQA). The trend line’s high R² value (0.786) underscores the robustness of this relationship. - **LRKV** stands out as the optimal model, suggesting it may employ more efficient pretraining strategies. - **MQA**’s poor performance highlights potential inefficiencies in its pretraining process. - The scatter plot implies that optimizing Pretraining BPB could significantly improve downstream outcomes, making it a key metric for model evaluation. This analysis aligns with the goal of identifying models that balance computational efficiency (low BPB) with high performance, guiding future research or deployment decisions. </details> Figure 19: 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\texttwosuperior=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. To validate the relationship between pretraining quality and downstream performance, we analyze the correlation between validation BPB and combined benchmark scores across all architectures after midtraining (Figure 19). The strong positive correlation (R²=0.786) demonstrates that models achieving lower pretraining loss consistently deliver superior downstream performance. LRKV achieves the lowest pretraining BPB (0.719) and highest downstream score (37.9%), while Standard MHA attains 0.723 BPB and 35.3% downstream. This correlation suggests that architectural improvements that enhance language modeling performance are not orthogonal to, but rather aligned with, practical task capabilities. The consistent 2.6 percentage point downstream advantage LRKV maintains over Standard MHA directly stems from its BPB pretraining improvement, confirming that low-rank KV factorization provides fundamental capacity gains that manifest across both pretraining and downstream evaluation regimes. This relationship validates that LRKV’s architectural improvements during pretraining directly translate to enhanced downstream capabilities, confirming the effectiveness of low-rank KV factorization for both language modeling and task-specific adaptation. ## Appendix F Additional Experimental Details and Computational Analysis ### F.1 Batching and Optimization Hyperparameters For models up to 256M parameters, we use a global batch size of 16 sequences (2048 tokens each). For the larger 1B–6.3B models, we reduce the physical batch size to 8 sequences and apply gradient accumulation of 2 steps to maintain an equivalent effective batch size. This setup ensures comparable gradient noise scale and optimizer dynamics across all model sizes. All runs employ weight decay of 0.1, Adam $β_1=0.9$ , $β_2=0.95$ , and gradient clipping at 1.0. ### F.2 Memory-Computation Trade-offs in Attention Mechanisms Modern attention mechanisms present a fundamental trade-off between cache memory (storage) and per-token computation (latency). We analyze this trade-off across standard attention, query-sharing variants, and compression-based approaches. #### The Compression Paradigm: Lossy vs. Lossless. Attention mechanisms can be understood through the lens of compression theory: - Multi-Latent Attention (MLA) implements lossy compression: it compresses K/V representations into a low-dimensional latent space $Z=XW_down$ , caching only $Z$ . During generation, full K/V need not be explicitly materialized; instead, interaction with the cached latent requires per-step projection overhead, e.g., $qK^⊤=(qW_up^K)Z^⊤$ and $aV=(aZ)W_up^V$ . This introduces additional per-token computation proportional to the latent dimension, and scales linearly with sequence length through dot products with $Z^⊤$ . - LRKV implements near-lossless compression with additive structure: it caches both shared full-rank features ( $\bar{K},\bar{V}$ ) and compact per-head latents ( $R_h^K,R_h^V$ ). It avoids explicitly materializing per-head $K_h,V_h$ tensors; instead, attention can be computed using associative forms (Eq. 6 – 7). Beyond the unavoidable $O(T)$ scan over cached tokens in attention, LRKV introduces only $O(Hrd_h)$ additional per-step projection work that does not grow with $T$ . This distinction is critical for autoregressive generation: both MLA and LRKV incur the unavoidable $O(T)$ cost to compare against cached tokens, but MLA introduces additional per-step projection work tied to the latent bottleneck, whereas LRKV’s extra work depends only on $r$ and does not increase with $T$ . #### Quantitative Analysis: Memory vs. Latency. We compare the memory footprint and inference-time computational characteristics of different attention mechanisms during autoregressive generation. Table 10 reports (1) the KV cache size stored per token and (2) the additional reconstruction overhead required per generation step, beyond the standard $O(T)$ attention computation inherent to cached autoregressive decoding. Table 10: Memory and inference-time characteristics of attention mechanisms. Cache size is measured per token. Reconstruction overhead denotes additional computation beyond standard $O(T)$ cached attention. $$ T 2Hd_h O(1) 2d_h O(1) G{=}3 2Gd_h O(1) d_c{=}128 d_c O(T d_cd_h) r{=}16 2(d_h{+}Hr) O(Hrd_h) \tag{1} $$ Extra cost refers to reconstruction needed to obtain per-head K/V representations. MLA requires latent-to-head expansion over the cached sequence unless expanded K/V are stored or fused kernels are used. Key observations. 1. Standard attention, MQA, and GQA incur no additional reconstruction cost at inference time, as per-head or shared key/value tensors are directly cached. 1. MLA achieves the smallest cache footprint but introduces an additional latent-to-head expansion step whose cost scales with the cached sequence length $T$ , unless expanded representations are stored (which increases memory) or specialized fused kernels are used. 1. LRKV trades a modest increase in cache size relative to MLA for sequence-length-independent reconstruction overhead. This avoids $T$ -dependent expansion during generation while preserving per-head expressivity through low-rank residuals. 1. As a result, LRKV occupies an intermediate point in the memory–latency design space: it substantially reduces KV memory relative to full attention while avoiding the sequence-dependent reconstruction overhead introduced by more aggressive compression schemes. ## Appendix G KV Cache Memory Measurements This appendix presents detailed measurements of KV cache memory usage during inference for all attention mechanisms evaluated in this work. ### G.1 Measurement Methodology All memory measurements are obtained by instantiating cache tensors with the exact shapes used during inference and measuring their memory footprint using PyTorch’s tensor.numel() and element_size() methods. Measurements use bfloat16 precision (2 bytes per element) with batch size 1 and sequence length 2048 tokens, matching our experimental configuration. ### G.2 Cache Structure by Method Table 11 summarizes the cache structure for each attention mechanism, showing tensor shapes and memory formulas for inference with batch size $B$ , sequence length $T$ , number of layers $L$ , number of heads $H$ , head dimension $d_h$ , group size $G$ , latent dimension $d_c$ , and rank $r$ . Table 11: KV cache structure by attention mechanism. All measurements use bfloat16 precision (2 bytes per element). | Standard MHA MQA GQA | Shape: $(L,B,H,T,d_h)$ Shape: $(L,B,1,T,d_h)$ Shape: $(L,B,H/G,T,d_h)$ | $2× L× B× H× T× d_h× 2$ $2× L× B× T× d_h× 2$ $2× L× B×(H/G)× T× d_h× 2$ | | --- | --- | --- | | MLA | Shape: $(L,B,T,d_c)$ (latents) | $L× B× T× d_c× 2$ | | LRKV (Optimized) | Shared: $(L,B,T,d_h)$ Latents: $(L,B,H,T,r)$ | $2× L× B× T×(d_h+H× r)× 2$ | ### G.3 Measured Results Table 12 presents measured memory usage across all three model scales. The optimized LRKV implementation (using LRKVCache from lrkv_attention_fused.py) achieves substantial memory savings compared to standard MHA while maintaining full representational capacity. Table 12: Measured KV cache memory during inference. All values measured for T=2048 tokens, batch size 1, bfloat16 precision, using rank $r=64$ for LRKV and latent dimensions as specified in main paper. | Standard MHA MQA GQA | 72.0 MB 12.0 MB (16.7%) 36.0 MB (50.0%) | 648.0 MB 36.0 MB (5.6%) 216.0 MB (33.3%) | 1152.0 MB 36.0 MB (3.1%) 72.0 MB (6.2%) | | --- | --- | --- | --- | | MLA | 24.0 MB (33.3%) | 216.0 MB (33.3%) | 288.0 MB (25.0%) | | LRKV | 48.0 MB (52.6%) | 360.0 MB (48.4%) | 612.0 MB (45.1%) | ### G.4 Scale-Dependent Efficiency LRKV’s cache percentage relative to standard MHA improves with model scale: | | Ratio | $\displaystyle=\frac{d_h+H· r}{H· d_h}=\frac{1}{H}+\frac{r}{d_h}$ | | | --- | --- | --- | --- | This demonstrates that LRKV’s efficiency increases with model scale: larger models with more heads benefit more from the shared component, achieving greater relative memory savings while maintaining fixed rank $r=64$ . ## Appendix H Limitations and Future Work ### H.1 Limitations Our experiments focus on decoder-only language models at scales up to 6.3B parameters. The design space for large-scale Transformer training is extremely large, spanning model size, training data, optimization schedules, attention mechanisms, and hardware/software configurations, many of which interact in nontrivial ways. As a result, it is infeasible to exhaustively benchmark all combinations or to claim that any single configuration is universally optimal across setups. Our empirical evaluation therefore samples this space using widely adopted architectures, training recipes, and system implementations on modern accelerator hardware. While this provides a representative and practically relevant assessment, different choices of model scale, data regime, training duration, or system optimizations may shift the precise efficiency–performance trade-offs. To facilitate independent validation and extension, we will release our training code and configuration details, enabling replication and exploration across alternative setups. We also note that the relative efficiency of different KV compression strategies may vary with hardware characteristics such as memory bandwidth, kernel fusion, and cache behavior.