Paper deep dive
Virtual Head Attention
Xibo Ding, Guoxia Wang, Jinle Zeng, Jiabin Yang, Sirui Tao, Chao Yang, Dianhai Yu
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The design of attention mechanisms in long-context large language models faces an ``impossible triangle'' among KV cache size, training and prefill FLOPs, and model quality. GQA achieves strong quality but incurs high KV cache overhead; alternative approaches such as MLA and MFA can compress the cache, but MLA suffers from quality loss while MFA substantially increases training and prefill compute. We propose Virtual Head Attention (VHA), which simultaneously improves all three vertices through two lightweight linear operations: Q Premix applies a near-identity transformation to queries within each KV group to recover virtual head diversity from halved physical query heads, and Linear Postmix fuses inter-head features via a low-rank $I + AB^\top$ residual structure. Since VHA introduces no nonlinear operations, Premix can be folded into the query projection and Postmix into the output projection at inference time, reducing the model to a standard GQA-2 architecture that directly reuses all existing GQA inference optimizations. With only two KV heads, the group size in large models readily reaches 64, fully exploiting Tensor Core parallelism during decoding. Across four scales with 4K-context 50B-token pretraining, VHA consistently surpasses the standard GQA baseline at every scale ($+$0.05\% to $+$0.81\%), while reducing KV cache by $4\times$ vs.\ GQA-8 at dense scales and $2\times$ vs.\ GQA-4 at the MoE scale, and lowering training and prefill FLOPs by 7.9\% (1.7B, $S{=}4096$). On an industrial 30B-A3B MoE model trained on 500B tokens at 8K context, VHA outperforms both MLA and GQA-4 by $+$0.70\,/\,$+$2.60\,\% on a 12-benchmark base-model evaluation suite; in a separate matched 10B-A1B MoE comparison after long-context mid-training, VHA attains an average RULER score of 84.4 over 8K--128K versus 73.0 for GQA, while using two rather than four KV heads in the full-attention layers.
Tags
Links
- Source: https://arxiv.org/abs/2610.17959v1
- Canonical: https://arxiv.org/abs/2610.17959v1
Trouble viewing inline? Open PDF directly →
Full Text
75,708 characters extracted from source content.
Expand or collapse full text
Virtual Head Attention Xibo Ding †thanks: Equal contribution. Work done at Baidu Inc. Affiliation: Baidu Inc. Peking University Email: dingxibo@baidu.com Guoxia Wang11footnotemark: 1 Email: wangguoxia@baidu.com Jinle Zeng Email: yudianhai@baidu.com Jiabin Yang Sirui Tao Affiliation: Baidu Inc. Peking University Chao Yang Affiliation: Baidu Inc. Peking University Dianhai Yu †thanks: Corresponding author. Affiliation: Baidu Inc. Peking University Abstract The design of attention mechanisms in long-context large language models faces an “impossible triangle” among KV cache size, training and prefill FLOPs, and model quality. Grouped-Query Attention (GQA) achieves strong quality but incurs high KV cache overhead; alternative approaches such as Multi-head Latent Attention (MLA) and Multi-matrix Factorization Attention (MFA) can compress the cache, but MLA suffers from quality loss while MFA substantially increases training and prefill compute. We propose Virtual Head Attention (VHA), which simultaneously improves all three vertices through two lightweight linear operations: Q Premix applies a near-identity transformation to queries within each KV group to recover virtual head diversity from halved physical query heads, and Linear Postmix fuses inter-head features via a low-rank I+AB⊤I+AB residual structure. Since VHA introduces no nonlinear operations, Premix can be folded into the query projection and Postmix into the output projection at inference time, reducing the model to a standard GQA-2 architecture that directly reuses all existing GQA inference optimizations. With only two KV heads, the group size in large models readily reaches 64, fully exploiting Tensor Core parallelism during decoding. Across four scales (0.6B, 1.7B, 4B dense, and 8B-A1B MoE) with 4K-context 50B-token pretraining, VHA consistently surpasses the standard GQA baseline at every scale (++0.05% to ++0.81%), while reducing KV cache by 4×4× vs. GQA-8 at dense scales and 2×2× vs. GQA-4 at the MoE scale, and lowering training and prefill FLOPs by 7.9% (1.7B, S=4096S=4096). On an industrial 30B-A3B MoE model trained on 500B tokens at 8K context, VHA outperforms both MLA and GQA-4 by ++0.70 / ++2.60 % on a 12-benchmark base-model evaluation suite; in a separate matched 10B-A1B MoE comparison after long-context mid-training, VHA attains an average RULER score of 84.4 over 8K–128K versus 73.0 for GQA, while using two rather than four KV heads in the full-attention layers. The source code is available at https://github.com/PaddlePaddle/VHA. 1 Introduction Large language models are rapidly scaling toward longer contexts, with context windows expanding from 2K tokens to 128K and beyond (Grattafiori et al., 2024; Yang et al., 2024). As sequence lengths grow, the attention mechanism (Vaswani et al., 2017) becomes a central efficiency bottleneck in both inference and training. During inference, the key-value (KV) cache grows linearly with sequence length, consuming substantial GPU memory and constraining both serving concurrency and generation latency. During training and prefill, the FLOPs of attention computation scale quadratically with sequence length. Since the design of the attention mechanism jointly determines KV cache size, computational cost, and model expressiveness, it is a critical factor governing the overall efficiency and quality of large language models. The design of existing attention mechanisms faces three mutually constraining objectives, forming an “impossible triangle”: KV cache size, training and prefill FLOPs, and model quality. Consider the widely adopted Grouped-Query Attention (GQA) (Ainslie et al., 2023): GQA with 8 KV heads (GQA-8) achieves strong quality, but incurs a KV cache of 4 KB per layer per token (in BF16). Directly reducing the number of KV heads to 2 (GQA-2) compresses the cache to 1 KB, but dramatically enlarges the group size so that far more query heads must share the same key-value pairs, causing attention diversity to drop sharply and resulting in a 1.00% quality degradation at the 1.7B scale. Multi-head Latent Attention (MLA) (Liu et al., 2024) applies low-rank compression to reduce KV cache to 1.12 KB, but yields the lowest quality across multiple model scales. Multi-matrix Factorization Attention (MFA) (Hu et al., 2025) reduces the cache to 1 KB by enlarging the head dimension, achieving quality close to GQA-8, but at the cost of a 49.8% increase in training and prefill FLOPs (1.7B, S=4096S=4096). No existing approach achieves a satisfactory balance across all three dimensions. We propose Virtual Head Attention (VHA), which simultaneously improves all three vertices of the impossible triangle through two lightweight linear operations. The core idea of VHA is to halve the number of physical query heads while setting the number of KV heads to only 2. Q Premix then applies a near-identity linear transformation to queries within each KV group, expanding the reduced physical query heads into virtual query heads to recover attention diversity. Linear Postmix applies a low-rank identity-residual transformation (I+AB⊤I+AB , r=Hq′/2r=H_q /2, e.g., r=4r=4 at 1.7B) on the attention output, fusing cross-group inter-head features. Since VHA introduces no nonlinear operations, Premix can be directly folded into the query projection matrix and Postmix into the output projection matrix at inference time, reducing the model to a standard GQA-2 architecture that reuses all existing GQA inference optimizations in mainstream serving frameworks such as vLLM (Kwon et al., 2023) and SGLang (Zheng et al., 2024). Moreover, with only 2 KV heads, the group size in large models readily reaches 64, fully exploiting Tensor Core parallelism during the decoding phase. Across four model scales (0.6B, 1.7B, 4B dense, and 8B-A1B MoE) with 4K-context 50B-token pretraining, VHA consistently surpasses the standard GQA baseline at every scale (++0.05% to ++0.81%), while reducing KV cache by 4×4× vs. GQA-8 at dense scales and 2×2× vs. GQA-4 at the MoE scale, and lowering training and prefill FLOPs by 7.9% (1.7B, S=4096S=4096). On an industrial 30B-A3B MoE model trained on 500B tokens at 8K context, VHA outperforms both MLA and GQA-4 by ++0.70 / ++2.60 % on a 12-benchmark base-model evaluation suite, demonstrating scalability to industrial model scale. A separate matched 10B-A1B MoE comparison further extends this advantage to 128K context under more aggressive KV compression. Our main contributions are as follows: 1. We propose VHA, employing Q Premix and Linear Postmix, two lightweight linear operations that simultaneously reduce KV cache and training/prefill FLOPs while improving model quality. 2. VHA uses only linear operations and can be fully absorbed into a standard GQA-2 model at inference time, reusing existing serving frameworks such as vLLM and SGLang with zero additional overhead. 3. Direct inspection of the trained Premix and Postmix matrices reveals two complementary structures (group-conditioned query specialization and structured low-rank cross-head communication) that jointly close the gap of the doubly-compressed configuration; component ablations confirm that neither module is sufficient on its own. 4. Comprehensive evaluations across five scales (0.6B to 30B-A3B MoE) against MHA, GQA, MFA, and MLA establish VHA as the best attention mechanism that is Pareto-optimal on KV cache, training/prefill FLOPs, and downstream quality simultaneously; matched long-context evaluations through 128K further validate this advantage under aggressive KV compression. 2 Related Work 2.1 KV Cache Efficient Attention Standard Multi-Head Attention (MHA) (Vaswani et al., 2017) maintains independent key and value projections per head, causing the KV cache to grow linearly with the number of heads. Multi-Query Attention (MQA) (Shazeer, 2019) shares a single KV head across all query heads, drastically reducing the cache at the cost of quality. Grouped-Query Attention (GQA) (Ainslie et al., 2023) interpolates between MHA and MQA by allowing multiple query heads to share a group of KV heads, and has been widely adopted in LLaMA 2/3 (Touvron et al., 2023; Grattafiori et al., 2024) and Qwen2 (Yang et al., 2024). Several recent designs further compress the KV cache. Multi-head Latent Attention (MLA) (Liu et al., 2024) projects KV into a shared low-rank latent space and absorbs the up-projection matrices into Q and O at inference, but requires custom kernels and a non-standard inference pipeline. Multi-matrix Factorization Attention (MFA) (Hu et al., 2025) enlarges the head dimension while reducing the number of KV heads, preserving model capacity under a tight cache budget but substantially increasing training and prefill FLOPs. Tensor Product Attention (TPA) (Zhang et al., 2025) decomposes QKV into context-dependent tensor products. Grouped-head Latent Attention (GTA) (Sun et al., 2025) extends MLA’s global latent space to a group-wise design. Simulated Attention Score (SAS) (Zheng et al., 2025) projects a small number of low-dimensional KV heads into higher-dimensional spaces to simulate more attention heads, an idea related to VHA’s virtual head expansion; however, SAS operates on the KV side whereas VHA virtualizes on the query side. Low-Rank Key-Value Attention (LRKV) (O’Neill et al., 2026) shares a full-rank KV basis across heads and restores head-specific flexibility through low-rank residuals, a form reminiscent of VHA’s Postmix (I+AB⊤I+AB ); yet LRKV acts in KV space and its additional structure remains at inference time. Cross-Layer Attention (CLA) (Brandon et al., 2024) and YOCO (Sun et al., 2024) reduce memory by sharing or reusing KV caches across layers. Unlike all of the above, VHA employs a richer training-time parameterization that merges exactly back into standard GQA at deployment, requiring no custom kernels or specialized inference procedures. 2.2 Inter-Head Communication and Virtual Heads Talking-Heads Attention (Shazeer et al., 2020) applies learned linear projections across heads before and after softmax, pioneering inter-head communication. Cordonnier et al. (2020) analyze multi-head attention from a collaborative perspective, revealing substantial redundancy among heads. More recently, Explicit Multi-head Attention (MEA) (Peng et al., 2026) introduces a head-level linear composition module for explicit cross-head interaction, and Interleaved Head Attention (Duvvuri et al., 2026) interleaves head representations to enable information exchange. In contrast to these methods, VHA does not mix across heads in the attention weight space; instead, it virtualizes in the query space via Premix and fuses inter-head features in the output space via a low-rank Postmix, with all operations being strictly linear and fully absorbable at inference. 2.3 Reparameterization RepVGG (Ding et al., 2021) uses a multi-branch training architecture (3×3+1×1+identity3×3+1×1+identity) that merges into a single convolution at inference, demonstrating that decoupled training-time parameterization provides implicit regularization at zero inference cost. LoRA (Hu et al., 2022) performs parameter-efficient fine-tuning via low-rank decomposition (W+BAW+BA). VHA extends this reparameterization paradigm from a single weight matrix to the head dimension of attention: Premix structurally factorizes the per-virtual-head query projections (W~Q(i,g)=WQ(i)Wgpremix W_Q^(i,g)=W_Q^(i)W_g^premix), while Postmix introduces a low-rank identity-residual mixing on the head axis (I+AB⊤I+AB ). Both modules are mathematically equivalent in the forward pass to a standard GQA configuration, yet the decoupled parameterization yields distinct optimization dynamics during training and merges back into a standard GQA-2 architecture at inference with zero overhead. 3 Method: Virtual Head Attention Figure 1: Overview of Virtual Head Attention. (a) VHA improves the trade-off among KV cache size, training/prefill FLOPs, and model quality by combining Q Premix and Linear Postmix during training, while absorbing both linear modules into standard GQA projections at inference. (b) Schematic of the KV cache–FLOPs–quality impossible triangle faced by existing attention designs. (c) Method comparison at 1.7B and 4B scales, plotting KV cache size against training FLOPs, with bubble area proportional to 8-task average accuracy. 3.1 Preliminaries Grouped-Query Attention (GQA) (Ainslie et al., 2023) partitions HqH_q query heads into HkvH_kv groups, where heads within each group share a common key-value pair. Given input X∈ℝS×hX ^S× h, GQA computes Q=XWQQ=XW_Q, K=XWKK=XW_K, V=XWVV=XW_V, where WQ∈ℝh×HqdW_Q ^h× H_qd, WK∈ℝh×HkvdW_K ^h× H_kvd, WV∈ℝh×HkvdW_V ^h× H_kvd. Each query head attends to its group’s key-value pair via scaled dot-product attention (Vaswani et al., 2017), and all head outputs are concatenated and projected by WO∈ℝHqd×hW_O ^H_qd× h. During inference, K and V are cached, consuming 4Hkvd4H_kvd bytes per layer per token in BF16. Reducing HkvH_kv compresses the cache but enlarges the group size G=Hq/HkvG=H_q/H_kv, forcing more query heads to share the same key-value pair and degrading attention diversity. 3.2 Overall Design VHA halves the number of physical query heads while setting Hkv=2H_kv=2, then recovers model expressiveness through two lightweight linear operations. Specifically, Hq′=Hq/2H_q =H_q/2 physical query heads and Hkv=2H_kv=2 KV heads yield Hv=Hkv×Hq′=HqH_v=H_kv× H_q =H_q virtual heads (matching the original GQA query head count), with WO∈ℝHvd×hW_O ^H_vd× h identical in size to the original GQA output projection. The resulting KV cache is 4×2×d4× 2× d bytes per layer per token, 4×4× smaller than GQA-8 at dense scales and 2×2× smaller than GQA-4 at MoE scales. On top of this configuration, VHA introduces Q Premix, which expands Hq′H_q physical query heads into HvH_v virtual query heads to restore attention diversity, and Linear Postmix, which applies a low-rank identity-residual transformation on the attention output to fuse inter-head features. Inference absorption principle. VHA uses only linear operations (no nonlinear activations), so that at inference time Premix can be folded into WQW_Q and Postmix into WOW_O, reducing the entire model to a standard GQA-2 architecture that directly reuses all existing GQA optimizations in frameworks such as vLLM (Kwon et al., 2023) and SGLang (Zheng et al., 2024). 3.3 Q Premix: Virtual Query Head Expansion Q Premix applies a near-identity linear transformation to physical query heads within each KV group, expanding them into virtual query heads. For KV group g∈1,…,Hkvg∈\1,…,H_kv\, each physical query head Qi∈ℝS×dQ_i ^S× d, i∈1,…,Hq′i∈\1,…,H_q \, is transformed as: Q~i,g=QiWgpremix,Wgpremix∈ℝd×d Q_i,g=Q_i\,W_g^premix, W_g^premix ^d× d (1) where WgpremixW_g^premix is a per-group transformation matrix. We initialize Wgpremix=I+α⋅ϵW_g^premix=I+α·ε, ϵ∼(0,σ2)ε (0,σ^2), with a small scaling factor α. This near-identity initialization ensures that virtual query heads start as slight perturbations of the physical query heads, preserving training stability at the early stage. As training progresses, each WgpremixW_g^premix gradually deviates from the identity to produce increasingly diverse attention patterns across virtual heads. Through this operation, Hq′H_q physical query heads are expanded into Hv=Hkv×Hq′H_v=H_kv× H_q virtual query heads. Structured factorization view. Premix is mathematically equivalent in the forward pass to a standard HvH_v-head GQA whose query projections are structurally factorized. Specifically, defining W~Q(i,g)=WQ(i)Wgpremix∈ℝh×d, W_Q^(i,g)\;=\;W_Q^(i)\,W_g^premix\;∈\;R^h× d, (2) the entire Premix-then-attend computation reduces, for any input X and any parameter values, to softmax(Q~i,gKg⊤/d)Vg=softmax(XW~Q(i,g)Kg⊤/d)Vg.softmax\! ( Q_i,gK_g / d )V_g\;=\;softmax\! (X W_Q^(i,g)K_g / d )V_g. (3) Thus VHA is a standard GQA with Hv=HkvHq′H_v=H_kvH_q virtual query heads and HkvH_kv KV heads, where the per-virtual-head query projections are constrained to lie on the structured low-dimensional manifold WQ(i)Wgpremix\W_Q^(i)W_g^premix\. Why this factorization is useful. The structured parameterization brings two benefits over an unconstrained HvH_v-head GQA. First, it is parameter- and compute-efficient: the Q path uses Hq′hd+Hkvd2H_q hd+H_kvd^2 parameters and a comparable number of FLOPs, versus HvhdH_vhd for a fully decoupled HvH_v-head Q projection. At the 1.7B scale, this yields a roughly 50% reduction in Q-projection parameters and FLOPs (the primary source of VHA’s training-FLOPs reduction). Second, although the forward pass is algebraically equivalent to training a single W~Q(i,g) W_Q^(i,g), the gradients of WQ(i)W_Q^(i) and WgpremixW_g^premix are coupled across all groups and physical heads, yielding a different optimization landscape and providing implicit regularization analogous to LoRA (Hu et al., 2022) and RepVGG (Ding et al., 2021) (see Appendix G for a formal analysis). Inference-time absorption. Once training is complete and parameters are frozen, for each virtual head (i,g)(i,g) we precompute W~Q(i,g)=WQ(i)⋅Wgpremix W_Q^(i,g)=W_Q^(i)· W_g^premix, after which VHA reduces to a standard GQA-2 with HvH_v independent query projections and Hkv=2H_kv=2 KV heads. 3.4 Linear Postmix: Low-Rank Identity-Residual Mixing While Premix restores virtual head diversity within each KV group, the outputs across different groups still lack cross-group information exchange. Linear Postmix addresses this through a low-rank identity-residual structure. Let Y∈ℝS×Hv×dY ^S× H_v× d denote the attention output. Postmix computes: Y′=Y+BA⊤Y=(I+BA⊤)Y =Y+BA Y= (I+BA )Y (4) where A,B∈ℝHv×rA,B ^H_v× r and r is the bottleneck rank (default r=Hq′/2r=H_q /2, e.g., r=4r=4 at 1.7B). The operation acts along the head dimension independently for each token and head-dimension channel, applying the transformation I+BA⊤∈ℝHv×HvI+BA ^H_v× H_v to the head axis. Initialization. A is initialized with small random values and B with zeros, ensuring BA⊤=0BA =0 at the start of training so that Postmix acts as the identity and does not disturb the initial training dynamics. As training progresses, B gradually learns effective cross-group mixing patterns. Linear redundancy and reparameterization benefit. Mathematically, Postmix can be absorbed into WOW_O. Using the row-major Kronecker identity vecrow(MY)⊤=vecrow(Y)⊤(M⊤⊗Id)vec_row(MY) =vec_row(Y) (M I_d) (Roger and Charles, 1994): vecrow(Y)⊤((IHv+AB⊤)⊗Id)WO=vecrow(Y)⊤WO′vec_row(Y) \! ((I_H_v+AB ) I_d )W_O=vec_row(Y) W_O (5) where WO′=((IHv+AB⊤)⊗Id)WO∈ℝHvd×hW_O = ((I_H_v+AB ) I_d )\,W_O ^H_vd× h. Thus, linear Postmix is theoretically redundant (see Appendix H for the complete derivation). However, ablation experiments (Section 4.4) show that adding Postmix on top of Premix yields a ++1.19 % improvement at 1.7B, analogous to the implicit regularization observed in RepVGG (Ding et al., 2021) and the low-rank decomposition in LoRA (Hu et al., 2022): the separated gradient paths of A and B provide a different optimization landscape, while the low-rank constraint (r≪Hvr H_v) acts as an additional regularizer preventing cross-head mixing from overfitting. Inference-time absorption. At inference, we precompute WO′=((IHv+AB⊤)⊗Id)WOW_O = ((I_H_v+AB ) I_d )W_O, after which VHA reduces to standard GQA-2. Algorithm 1 (Appendix F) summarizes the complete training forward pass and the one-time inference-absorption procedure. 3.5 Efficiency Analysis Training and prefill FLOPs. By halving the physical HqH_q, VHA directly reduces Q projection FLOPs. The additional cost of Premix (2HkvHq′d22H_kvH_q d^2) and Postmix (4Hvrd4H_vrd) is negligible compared to the FLOPs saved by halving the Q projection. Table 5 quantifies this trade-off in detail across all four model scales (S=4096S=4096); VHA emerges as the best Pareto-optimal point on KV cache and FLOPs at every scale (Section 4.5). Decoding inference. Autoregressive decoding is memory-bound (Pope et al., 2023), so KV cache size directly governs latency and serving concurrency. With KV cache reduced by 4×4× vs. GQA-8 (dense) or 2×2× vs. GQA-4 (MoE), VHA achieves a corresponding 2–4× theoretical speedup in memory-bound decoding. Since Premix and Postmix are fully absorbed at inference, VHA’s decoding cost is identical to standard GQA-2. Table 1: Pretraining accuracy across four model scales. We report zero-shot accuracy (%) on eight standard multiple-choice benchmarks. Cache is KV cache per layer per token in BF16 (KB). Bold marks the best per scale per column; ties are bolded jointly. VHA achieves the highest average accuracy at every scale while using the smallest KV cache. Per-method full configurations and benchmark citations are in Appendix D; per-method training-FLOPs comparison is in Table 5. Scale Attn Q-K-V Cache ARC-C BoolQ HSwag OBQA PIQA Wino SciQ SIQA Avg 0.6B MHA 12-12-12 6.00 29.10 60.40 44.88 33.20 68.66 51.85 74.50 39.15 50.22 GQA 16-8-8 4.00 29.86 58.01 44.77 35.80 68.17 52.33 75.70 38.69 50.42 MLA 14-14-14 1.12 29.27 55.96 43.20 33.80 68.23 50.67 76.40 38.59 49.52 MFA 17-1-1 1.00 30.80 60.18 45.47 32.80 69.04 54.70 73.20 39.05 50.66 VHA 8-2-2 1.00 31.91 61.19 45.26 35.60 68.44 50.99 75.60 39.36 51.04 1.7B MHA 12-12-12 6.00 35.24 58.53 51.86 35.20 71.33 54.06 80.30 39.36 53.24 GQA 16-8-8 4.00 35.15 62.11 51.98 35.60 71.71 55.25 81.10 39.82 54.09 MLA 12-12-12 1.12 31.57 61.31 49.22 36.20 70.51 54.22 77.60 39.30 52.49 MFA 19-1-1 1.00 34.47 59.45 52.16 37.80 71.82 54.85 81.50 40.43 54.06 VHA 8-2-2 1.00 35.92 60.64 51.71 39.00 71.38 54.70 81.70 40.94 54.50 4B MHA 20-20-20 10.00 38.31 60.21 57.52 39.60 73.88 56.99 85.40 41.56 56.68 GQA 32-8-8 4.00 39.16 62.32 58.39 39.40 73.83 57.06 83.80 41.20 56.90 MLA 28-28-28 1.12 37.20 61.65 56.55 38.80 73.29 56.75 84.90 41.86 56.38 MFA 33-1-1 1.00 38.40 62.20 57.67 39.80 73.88 57.70 83.60 42.27 56.94 VHA 16-2-2 1.00 38.91 64.74 58.67 39.20 74.43 58.80 84.60 42.34 57.71 8B-A1B MHA 20-20-20 10.00 35.67 59.63 53.98 37.40 72.25 53.04 82.70 42.12 54.60 GQA 32-4-4 2.00 34.73 61.53 54.03 38.60 71.55 55.09 82.60 40.74 54.86 MLA 28-28-28 1.12 33.11 60.86 52.29 36.60 72.63 53.28 82.00 39.56 53.79 MFA 33-1-1 1.00 35.07 62.48 54.08 36.60 72.80 55.17 81.80 40.48 54.81 VHA 16-2-2 1.00 36.69 62.17 53.78 37.00 71.71 54.06 82.40 41.45 54.91 4 Experiments 4.1 Experimental Setup Baselines. We compare VHA with MHA, GQA, MFA, and MLA across all scales. Following common practice, GQA uses 8 KV heads at dense scales (0.6B/1.7B/4B) and 4 KV heads at MoE scales (8B-A1B and 30B-A3B). VHA halves the physical query heads relative to GQA and uses Hkv=2H_kv=2, yielding 4×4× KV cache compression at dense scales and 2×2× at MoE scales. Small-scale pretraining. We pretrain Qwen3-architecture models at four scales (dense at 0.6B, 1.7B, and 4B parameters, plus an 8B-A1B MoE) on 50B tokens of FineWeb-Edu with BF16 mixed precision and a 4K context length. We report zero-shot accuracy on eight standard benchmarks. Industrial-scale pretraining. We pretrain an ERNIE-style 30B-A3B MoE on an internal multi-domain pretraining corpus (web text, books, code, and multilingual data) with FP8 mixed precision on 500B tokens at 8K context. We evaluate on 12 standard base-model benchmarks spanning general knowledge, math & STEM, coding, and multilingual capabilities. Full model configurations, benchmark details, and training hyperparameters are in Appendix D. Long-context mid-training. We additionally compare VHA and GQA on a matched MiMo-2.5-Flash-style 10B-A1B MoE backbone that interleaves full and sliding-window attention at a 1:3 ratio. The two models share the backbone, data, optimization, and sliding-window layers; only the full-attention layers differ, using two KV heads for VHA and four for GQA. We evaluate through 128K, with the stage-wise recipe and benchmark definitions reported in Appendix A. 4.2 Main Results Cross-scale consistency. VHA achieves the highest average accuracy at every scale, surpassing the GQA baseline by ++0.05% to ++0.81% while using the smallest KV cache: 1.00 KB per layer per token, tied with MFA, 4×4× smaller than GQA-8 at dense scales and 2×2× smaller than GQA-4 at the MoE scale. The advantage holds across both dense and MoE architectures. A separate matched-harness study at 0.6B confirms the same ordering across repeated runs, with VHA reaching 50.40±0.0750.40±0.07 (n=4n=4) versus GQA’s 50.00±0.3450.00±0.34 (n=3n=3); full results are in Appendix B. Trade-offs against each baseline. At equal KV cache (1.00 KB), VHA dominates MFA at every scale (++0.10% to ++0.77%) while reducing training and prefill FLOPs by 7.9%, compared with MFA’s ++49.8% increase at the 1.7B scale (S=4096S=4096; Table 5). Full-cache MHA falls below GQA-8 at every dense scale despite using a ≥1.5×≥ 1.5× larger cache, confirming that grouping the queries is essentially free. MLA underperforms at every scale, suggesting that latent KV compression is less effective than head-level virtualization at the 50B-token training budget used here. Per-method training loss curves at all four scales are reported in Appendix E. Table 2: Recovering from a doubly-compressed shape (1.7B). GQA 16-8-8 is the standard GQA-8 baseline; GQA 16-2-2 reduces only HkvH_kv to obtain a 1 KB KV cache; GQA 8-2-2 additionally halves the physical query heads to match VHA’s shape but lacks Premix and Postmix. Δ is the gap to the GQA 16-8-8 baseline. Method Cache ARC-C BoolQ HSwag OBQA PIQA Wino SciQ SIQA Avg Δ GQA 16-8-8 4.00 35.15 62.11 51.98 35.60 71.71 55.25 81.10 39.82 54.09 — GQA 16-2-2 1.00 33.19 58.96 50.92 36.00 71.27 54.22 79.60 40.53 53.09 −-1.00 GQA 8-2-2 1.00 33.02 59.11 49.98 34.00 71.44 53.28 77.80 40.43 52.38 −-1.71 VHA 8-2-2 1.00 35.92 60.64 51.71 39.00 71.38 54.70 81.70 40.94 54.50 ++0.41 Recovering from a doubly-compressed shape. Reducing GQA’s KV head count from 8 to 2 (GQA 16-2-2) costs −-1.00% in average accuracy at the 1.7B scale; further halving the physical query heads to match VHA’s shape (GQA 8-2-2) costs another −-0.71%, totaling −-1.71% vs. the GQA-8 baseline. With the same (8, 2, 2) shape, VHA recovers all of this loss and exceeds the original baseline by ++0.41%, despite using only 25% of GQA-8’s KV cache and 50% of its physical query heads. The doubly-compressed (8, 2, 2) shape that on its own forfeits −-1.71% thus becomes a winning configuration once Premix and Postmix are introduced. 4.3 Large-scale and long-context validation Table 3: Pretraining accuracy at the 30B-A3B MoE scale. We report accuracy (%) averaged within four task categories (General, Math & STEM, Coding, Multilingual) over 12 base-model benchmarks; Avg is the macro-average over the four task categories. Cache is KV cache per layer per token in BF16 (KB). Bold marks the best per column. Per-benchmark scores and benchmark citations are in Appendix C. Method Q-K-V Cache General Math & STEM Coding Multilingual Avg GQA 48-4-4 2.00 52.86 43.09 51.85 43.26 47.76 MLA 28-28-28 1.12 55.57 43.49 54.37 45.20 49.66 VHA 24-2-2 1.00 56.25 46.40 53.23 45.56 50.36 At the 30B-A3B MoE scale (500B tokens; FP8 mixed precision; full setup in Appendix D), Table 3 shows that VHA 24-2-2 attains the highest 12-benchmark average while using only 50% of GQA-4’s KV cache and 89% of MLA’s, confirming that the trends observed at 0.6B∼ 8B-A1B carry over to the production-scale MoE regime. The category-level breakdown shows that gains are broadly distributed: General ++3.39 / ++0.68%, Math & STEM ++3.31 / ++2.91%, Coding ++1.38 / −-1.14%, Multilingual ++2.30 / ++0.36% (VHA vs. GQA / MLA), indicating head-level virtualization scales gracefully across heterogeneous downstream tasks. Per-benchmark scores for all 12 evaluations are in Appendix C. Long-context validation. At the final 128K checkpoint, VHA outperforms the matched GQA baseline on RULER at every evaluated length from 8K to 128K, with a length-average of 84.4 versus 73.0 (++11.4 points). The advantage extends to complementary evaluations at 128K: VHA outperforms GQA on LongBench-v2 (24.2 vs. 18.9) and BABILong (38.2 vs. 31.1), while remaining on par on InfiniteBench (31.5 vs. 31.8); averaged across lengths, VHA wins, ties, and loses 3/1/1 HELMET categories. Together, these results demonstrate that VHA preserves its quality advantage through 128K while halving the KV heads and KV cache of the full-attention layers relative to GQA-4; complete results are in Appendix A. 4.4 Ablations We ablate VHA’s design choices at the 1.7B scale (additionally at 4B for the component contribution), isolating (i) the contribution of Premix and Postmix, (i) the Postmix rank, and (i) whether non-linearity helps either module. Reading the Δ column of Table 4 (left and middle), the two modules address complementary deficits of the doubly-compressed shape: only their joint VHA configuration consistently achieves the largest gains at both scales at 25%25\% of GQA-8’s KV cache; we provide a mechanistic interpretation of this complementarity in Section 5. Table 4: VHA ablations. Left, middle: component contribution at 1.7B and 4B (each row uses VHA’s shape, 8-2-2 or 16-2-2, with KV cache 1.00 KB / layer / token vs 4.00 KB for the GQA-8 baseline; Premix and Postmix are toggled individually). Right: activation ablation at 1.7B (replacing the linear formulation with SiLU on either or both modules). Avg (%) is over the eight benchmarks of Table 1; Δ is the average delta vs the GQA-8 baseline (component) or the linear default (activation); the linear VHA default is bold. Each module on its own only partially recovers the doubly-compressed loss from Table 2; the full VHA configuration exceeds the GQA-8 baseline at both scales, and SiLU consistently underperforms the linear default. Component @ 1.7B Variant Avg Δ GQA 54.09 — Premix only 53.31 −-0.78 Postmix only 53.95 −-0.14 VHA 54.50 ++0.41 Component @ 4B Variant Avg Δ GQA 56.90 — Premix only 56.62 −-0.28 Postmix only 57.22 ++0.32 VHA 57.71 ++0.81 Activation @ 1.7B Variant Avg Δ + SiLU on Premix 53.83 −-0.67 + SiLU on Postmix 53.79 −-0.71 + SiLU on both 53.67 −-0.83 Linear (default) 54.50 — Postmix rank. With Premix fixed at the 1.7B scale, sweeping the Postmix rank gives Avg 54.50/54.24/54.1554.50/54.24/54.15 for r∈4,32,128r∈\4,32,128\: larger r monotonically lowers training loss but degrades downstream evaluation, indicating that the low-rank constraint provides useful implicit regularization. We therefore fix r=Hq′/2r=H_q /2 across all scales (Section 3.4 / Appendix D). Activation. Replacing the linear formulation with SiLU on either or both modules consistently degrades quality at 1.7B (Table 4, right), regardless of where the non-linearity is inserted. Beyond the empirical loss, every non-linear variant forfeits inference-time absorption: with both modules linear the entire VHA computation collapses to standard GQA-2 at serving time, whereas any non-linearity must remain in the runtime kernel. We therefore keep both modules strictly linear, simultaneously achieving the highest empirical quality and lossless absorption into standard GQA-2 at inference. 4.5 Computational Efficiency Table 5: Computational efficiency across four model scales (S=4096S=4096). For each scale we list per-method query/key/value head counts, KV cache per layer per token (BF16, KB), training FLOPs per token aggregated over all transformer layers (G), and the relative change vs. the per-scale GQA baseline (GQA-8 dense, GQA-4 MoE). Bold highlights VHA, the method that achieves the best joint KV cache++FLOPs Pareto optimum at every scale. MHA’s lower FLOPs at three scales reflect its smaller query head count under our parameter-matched configurations (Hq=12H_q=12 at 0.6B/1.7B, Hq=20H_q=20 at 4B/8B-A1B versus GQA’s Hq=16H_q=16 or Hq=32H_q=32), at the cost of a 66–10×10× larger KV cache. 0.6B 1.7B 4B 8B-A1B Method Q-K-V Cache FLOPs Δ % Q-K-V Cache FLOPs Δ % Q-K-V Cache FLOPs Δ % Q-K-V Cache FLOPs Δ % MHA 12-12-12 6.00 6.87 −-9.3 12-12-12 6.00 9.91 −-6.6 20-20-20 10.00 21.66 −-11.1 20-20-20 10.00 8.76 −-14.7 GQA 16-8-8 4.00 7.57 0 16-8-8 4.00 10.61 0 32-8-8 4.00 24.38 0 32-4-4 2.00 10.27 0 MLA 14-14-14 1.12 8.25 ++8.9 12-12-12 1.12 10.42 −-1.8 28-28-28 1.12 25.09 ++2.9 28-28-28 1.12 11.39 ++10.9 MFA 17-1-1 1.00 11.54 ++52.3 19-1-1 1.00 15.90 ++49.8 33-1-1 1.00 36.04 ++47.9 33-1-1 1.00 17.82 ++73.5 VHA 8-2-2 1.00 6.96 −-8.1 8-2-2 1.00 9.78 −-7.9 16-2-2 1.00 22.52 −-7.6 16-2-2 1.00 9.60 −-6.5 Table 5 quantifies the joint KV cache–FLOPs trade-off across the four model scales of Table 1. VHA is the only method that simultaneously matches the smallest KV cache (1.00 KB) and reduces training/prefill FLOPs below the per-scale GQA baseline (−-8.1 / −-7.9 / −-7.6 / −-6.5 % across scales). MFA matches the cache but pays a ++48–74 % FLOPs penalty; MLA’s latent compression yields a slightly larger cache (1.12 KB) and raises FLOPs at three of four scales (up to ++10.9 % at 8B-A1B). MHA appears competitive on FLOPs only because its smaller query head count under our parameter-matched configurations comes with a 66–10×10× larger KV cache. Combined with VHA’s complete inference absorption to standard GQA-2 (Section 3.4), this makes VHA the best Pareto-optimal point on the cache–FLOPs–quality triangle in our experiments. 5 Analysis: How VHA Reparameterization Works Beyond task accuracy, VHA should exploit the limited information available under aggressive KV sharing. We therefore inspect the Premix and Postmix matrices of trained VHA models (Premix on the 1.7B, Postmix on the 4B model), asking whether (i) Premix learns a group-conditioned query geometry, (i) Postmix learns structured low-rank cross-head communication, and (i) the two jointly close the doubly-compressed gap of Section 4.4. Figure 2: Q Premix weight analysis on a trained Qwen-1.7B VHA model. (a) Per-layer Frobenius deviation ‖Wgpremix−I‖F\|W_g^premix-I\|_F per KV group (group mean overlaid). (b) Per-layer mean pairwise distance ‖Wg1premix−Wg2premix‖F\|W_g_1^premix-W_g_2^premix\|_F between KV groups. Premix learns non-trivial, group-differentiated query transforms. Premix introduces a learnable per-group query transform WgpremixW_g^premix before RoPE, initialized near identity. Two trivial outcomes would falsify our hypothesis: (a) WgpremixW_g^premix remains near identity after training, rendering the module a no-op; or (b) all KV groups converge to the same matrix, degenerating Premix into a group-agnostic shared adapter without group-conditioned specialization. Figure 2 refutes both: panel (a) shows that ‖Wgpremix−I‖F\|W_g^premix-I\|_F is substantially nonzero and most pronounced in mid-to-late layers (peak 5.85), and panel (b) shows that the inter-group distance ‖Wg1premix−Wg2premix‖F\|W_g_1^premix-W_g_2^premix\|_F is of comparable magnitude (peak 6.13), yielding a ratio well below the 2 2 value expected under fully independent groups and far above the zero expected under a shared deviation. This quantifies that the two KV groups learn substantially but not independently different query transforms. Premix therefore learns differentiated query probing directions per KV group under a fixed minimal KV cache, realizing genuine group-conditioned virtual-head expansion rather than mere parameter augmentation. Figure 3: Postmix head-mixing analysis on a trained Qwen-4B VHA model. Sankey diagrams of |AB⊤||AB | across five representative layers, with ribbons flowing from source heads (left, colored by KV group) to recombined output heads (right); width and color saturation proportional to |AB⊤|j,i|AB |_j,i. Postmix learns structured low-rank cross-head communication. Postmix recombines the HvH_v attention output heads via the low-rank residual I+BA⊤I+BA , where entry (j,i)(j,i) of |AB⊤||AB | encodes the contribution of source head j to recombined head i. The central question is whether AB⊤AB merely performs a per-head rescaling (a strictly diagonal pattern) or genuinely redistributes information across heads (off-diagonal flow). Figure 3 visualizes |AB⊤||AB | for the trained 4B model and reveals pronounced off-diagonal flow, ruling out the trivial per-head-rescaling outcome. The dominant pattern is within-group redistribution: Postmix primarily rebalances information among heads sharing the same KV cache, with a sparse but consistent set of cross-group bridges providing inter-group connectivity without additional KV heads. Joint reparameterization closes the doubly-compressed gap. Tying these analyses back to the component ablation in Table 4 clarifies the three numbers reported there. Premix-only retains the group-conditioned query geometry of Figure 2 but lacks Postmix’s cross-group repair pathway, hence the −-0.78% gap at 1.7B. Postmix-only owns the cross-head mixing channels of Figure 3 but operates over collapsed, unspecialized virtual heads, hence the −-0.14% gap. Only the joint configuration realizes both structures simultaneously, yielding a ++0.41% advantage over GQA-8 at 25%25\% of its KV cache. VHA’s effectiveness therefore arises from the functional complementarity of its two modules: Premix and Postmix are not interchangeable, and neither suffices alone. 6 Conclusion We have presented Virtual Head Attention (VHA), which mitigates the tension among KV cache size, training/prefill FLOPs, and model quality through two lightweight linear operations: Q Premix and Linear Postmix. By halving physical query heads and recovering expressiveness via virtual head expansion and low-rank cross-group mixing, VHA reduces the KV cache to its architectural minimum while simultaneously lowering training FLOPs and improving model quality. Since all added operations are strictly linear, VHA fully absorbs into standard GQA-2 at inference, reusing existing serving optimizations with zero additional overhead. Analysis of trained models reveals that Premix and Postmix are functionally complementary and neither suffices alone, explaining why the joint configuration consistently surpasses the GQA baseline across all five evaluated scales (0.6B to 30B-A3B MoE). A matched 10B-A1B evaluation further shows that VHA maintains this advantage over GQA through 128K context. Acknowledgments We thank Wenbin Chang for valuable input and helpful discussions. We also thank the anonymous reviewers, area chairs, and senior area chairs for their insightful comments and constructive suggestions. This work was supported by the National Engineering Research Center of Deep Learning Technology and Application and the Fundamental and Interdisciplinary Disciplines Breakthrough Plan of the Ministry of Education of China (Grant No. JYB2025XDXM202). References Ainslie et al. (2023) J. Ainslie, J. Lee-Thorp, M. De Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai Gqa: training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, p. 4895–4901. Cited by: §1, §2.1, §3.1. Austin et al. (2021) J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732. Cited by: §D.2. Bisk et al. (2020) Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, Vol. 34, p. 7432–7439. Cited by: §D.2. Brandon et al. (2024) W. Brandon, M. Mishra, A. Nrusimha, R. Panda, and J. Ragan-Kelley Reducing transformer key-value cache size with cross-layer attention. Advances in Neural Information Processing Systems 37, p. 86927–86957. Cited by: §2.1. Clark et al. (2019) C. Clark, K. Lee, M. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova Boolq: exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 conference of the north American chapter of the association for computational linguistics: Human language technologies, volume 1 (long and short papers), p. 2924–2936. Cited by: §D.2. Clark et al. (2018) P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457. Cited by: §D.2. Cobbe et al. (2021) K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168. Cited by: §D.2. Cordonnier et al. (2020) J. Cordonnier, A. Loukas, and M. Jaggi Multi-head attention: collaborate instead of concatenate. arXiv preprint arXiv:2006.16362. Cited by: §2.2. Ding et al. (2021) X. Ding, X. Zhang, N. Ma, J. Han, G. Ding, and J. Sun Repvgg: making vgg-style convnets great again. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 13733–13742. Cited by: Appendix G, Appendix H, §2.3, §3.3, §3.4. Duvvuri et al. (2026) S. S. Duvvuri, C. Ekbote, R. Bansal, R. Tiwari, D. Khatri, D. Brandfonbrener, P. Liang, I. Dhillon, and M. Zaheer Interleaved head attention. arXiv preprint arXiv:2602.21371. Cited by: §2.2. Gema et al. (2025) A. P. Gema, J. O. J. Leang, G. Hong, A. Devoto, A. C. M. Mancino, R. Saxena, X. He, Y. Zhao, X. Du, M. R. G. Madani, et al. Are we done with mmlu?. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), p. 5069–5096. Cited by: §D.2. Grattafiori et al. (2024) A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783. Cited by: §1, §2.1. Hendrycks et al. (2020) D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300. Cited by: §D.2. Hendrycks et al. (2021) D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874. Cited by: §D.2. Hu et al. (2022) E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al. Lora: low-rank adaptation of large language models.. Iclr 1 (2), p. 3. Cited by: Appendix G, §2.3, §3.3, §3.4. Hu et al. (2025) J. Hu, H. Li, Y. Zhang, Z. Wang, S. Zhou, X. Zhang, and H. Y. Shum Multi-matrix factorization attention. In Findings of the Association for Computational Linguistics: ACL 2025, p. 25114–25126. Cited by: §1, §2.1. Kwon et al. (2023) W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. Gonzalez, H. Zhang, and I. Stoica Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles, p. 611–626. Cited by: §1, §3.2. Liu et al. (2024) A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al. Deepseek-v2: a strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434. Cited by: §1, §2.1. Liu et al. (2023) J. Liu, C. S. Xia, Y. Wang, and L. Zhang Is your code generated by chatgpt really correct? rigorous evaluation of large language models for code generation. Advances in neural information processing systems 36, p. 21558–21572. Cited by: §D.2. Loshchilov and Hutter (2017) I. Loshchilov and F. Hutter Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. Cited by: §D.1. Mihaylov et al. (2018) T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of the 2018 conference on empirical methods in natural language processing, p. 2381–2391. Cited by: §D.2. OpenAI (2024) OpenAI MMMLU: multilingual massive multitask language understanding. External Links: Link Cited by: §D.2. O’Neill et al. (2026) J. O’Neill, R. Clancy, M. Matskevichus, and F. Reid Low-rank key value attention. arXiv preprint arXiv:2601.11471. Cited by: §2.1. Peng et al. (2026) R. Peng, Y. Zhou, D. Song, K. Lv, B. Wang, Q. Guo, and X. Qiu Explicit multi-head attention for inter-head interaction in large language models. arXiv preprint arXiv:2601.19611. Cited by: §2.2. Pope et al. (2023) R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, J. Heek, K. Xiao, S. Agrawal, and J. Dean Efficiently scaling transformer inference. Proceedings of machine learning and systems 5, p. 606–624. Cited by: §3.5. Rein et al. (2023) D. Rein, B. L. Hou, A. C. Stickland, J. Petty, R. Y. Pang, J. Dirani, J. Michael, and S. R. Bowman Gpqa: a graduate-level google-proof q&a benchmark. arXiv preprint arXiv:2311.12022. Cited by: §D.2. Roger and Charles (1994) H. Roger and R. J. Charles Topics in matrix analysis. Cambridge University Press Cambridge, UK. Cited by: Appendix H, §3.4. Romanou et al. (2024) A. Romanou, N. Foroutan, A. Sotnikova, Z. Chen, S. H. Nelaturu, S. Singh, R. Maheshwary, M. Altomare, M. A. Haggag, A. Amayuelas, et al. Include: evaluating multilingual language understanding with regional knowledge. arXiv preprint arXiv:2411.19799. Cited by: §D.2. Sakaguchi et al. (2021) K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi Winogrande: an adversarial winograd schema challenge at scale. Communications of the ACM 64 (9), p. 99–106. Cited by: §D.2. Sap et al. (2019) M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi Social iqa: commonsense reasoning about social interactions. In Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing (EMNLP-IJCNLP), p. 4463–4473. Cited by: §D.2. Shazeer et al. (2020) N. Shazeer, Z. Lan, Y. Cheng, N. Ding, and L. Hou Talking-heads attention. arXiv preprint arXiv:2003.02436. Cited by: §2.2. Shazeer (2019) N. Shazeer Fast transformer decoding: one write-head is all you need. arXiv preprint arXiv:1911.02150. Cited by: §2.1. Sun et al. (2025) L. Sun, C. Deng, J. Jiang, X. Wu, H. Zhang, L. Chen, L. Ni, and J. Wang Gta: grouped-head latent attention. arXiv preprint arXiv:2506.17286. Cited by: §2.1. Sun et al. (2024) Y. Sun, L. Dong, Y. Zhu, S. Huang, W. Wang, S. Ma, Q. Zhang, J. Wang, and F. Wei You only cache once: decoder-decoder architectures for language models. Advances in Neural Information Processing Systems 37, p. 7339–7361. Cited by: §2.1. Suzgun et al. (2023) M. Suzgun, N. Scales, N. Schärli, S. Gehrmann, Y. Tay, H. W. Chung, A. Chowdhery, Q. Le, E. Chi, D. Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. In Findings of the Association for Computational Linguistics: ACL 2023, p. 13003–13051. Cited by: §D.2. Touvron et al. (2023) H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Cited by: §2.1. Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin Attention is all you need. Advances in neural information processing systems 30. Cited by: §1, §2.1, §3.1. Wang et al. (2024) Y. Wang, X. Ma, G. Zhang, Y. Ni, A. Chandra, S. Guo, W. Ren, A. Arulraj, X. He, Z. Jiang, et al. Mmlu-pro: a more robust and challenging multi-task language understanding benchmark. Advances in Neural Information Processing Systems 37, p. 95266–95290. Cited by: §D.2. Welbl et al. (2017) J. Welbl, N. F. Liu, and M. Gardner Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy User-generated Text, p. 94–106. Cited by: §D.2. Yang et al. (2024) A. Yang, B. Yang, B. Hui, B. Zheng, B. Yu, C. Zhou, C. Li, C. Li, D. Liu, F. Huang, G. Dong, H. Wei, H. Lin, J. Tang, J. Wang, J. Yang, J. Tu, J. Zhang, J. Ma, J. Yang, J. Xu, J. Zhou, J. Bai, J. He, J. Lin, K. Dang, K. Lu, K. Chen, K. Yang, M. Li, M. Xue, N. Ni, P. Zhang, P. Wang, R. Peng, R. Men, R. Gao, R. Lin, S. Wang, S. Bai, S. Tan, T. Zhu, T. Li, T. Liu, W. Ge, X. Deng, X. Zhou, X. Ren, X. Zhang, X. Wei, X. Ren, X. Liu, Y. Fan, Y. Yao, Y. Zhang, Y. Wan, Y. Chu, Y. Liu, Z. Cui, Z. Zhang, Z. Guo, and Z. Fan Qwen2 technical report. External Links: 2407.10671, Link Cited by: §1, §2.1. Zellers et al. (2019) R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi Hellaswag: can a machine really finish your sentence?. In Proceedings of the 57th annual meeting of the association for computational linguistics, p. 4791–4800. Cited by: §D.2. Zhang et al. (2025) Y. Zhang, Y. Liu, H. Yuan, Z. Qin, Y. Yuan, Q. Gu, and A. C. Yao Tensor product attention is all you need. arXiv preprint arXiv:2501.06425. Cited by: §2.1. Zheng et al. (2025) C. Zheng, J. Sun, Y. Gao, Y. Wang, P. Wang, J. Xiong, L. Ren, H. Cheng, J. Kulkarni, Y. Shen, et al. Sas: simulated attention score. arXiv preprint arXiv:2507.07694. Cited by: §2.1. Zheng et al. (2024) L. Zheng, L. Yin, Z. Xie, C. Sun, J. Huang, C. H. Yu, S. Cao, C. Kozyrakis, I. Stoica, J. E. Gonzalez, et al. Sglang: efficient execution of structured language model programs. Advances in neural information processing systems 37, p. 62557–62583. Cited by: §1, §3.2. Limitations and Future Work Our largest validation is at the 30B-A3B MoE scale trained on 500B tokens; scaling VHA to larger MoE models and substantially more training data (e.g., 20T tokens) remains an important direction. We further validate VHA through 128K on a 10B-A1B MoE, while broader long-context architectures and training recipes remain to be explored. VHA’s behavior under reinforcement learning training paradigms is also worth exploring. Finally, VHA is orthogonal to other KV cache compression techniques such as quantization and sparsification, and combining them may yield further inference savings. Appendix A Long-Context Evaluation A.1 Model and Stage-wise Training We conduct a controlled long-context comparison on a MiMo-2.5-Flash-style 10B-A1B MoE backbone that interleaves full and sliding-window attention at a 1:3 ratio. The GQA and VHA models share the same backbone, data, optimization settings, and sliding-window layers. They differ only in the full-attention layers: GQA uses four KV heads, whereas VHA uses two; all sliding-window layers use eight KV heads in both models. Thus, VHA halves the KV cache of the full-attention layers without changing the sliding-window component. The 8K pretraining stage uses 500B tokens and ends at 60K steps. Long-context mid-training then uses 100B tokens at 32K followed by another 100B tokens at 128K. These stages add 12K and 24K steps, yielding cumulative checkpoints at 72K and 96K steps, respectively. The RoPE base is increased from 10410^4 at 8K to 6.4×1056.4×10^5 at 32K and 5×1065×10^6 at 128K. We report intermediate results from the 32K checkpoint and final results from the 128K checkpoint. A.2 RULER and HELMET We evaluate RULER and HELMET from 8K to 128K. RULER Overall is the mean over 13 tasks: eight Retrieval, two Aggregation, two QA, and one Variable Tracking task. HELMET is reported by category because its categories measure distinct capabilities. Table 6: Long-context results on RULER and HELMET (%). RULER Overall averages all 13 tasks. Bold marks the better result within each matched checkpoint and evaluation length; exact ties are unbolded. RULER HELMET Checkpoint Length Model Overall Retr. Aggr. QA VT RAG LongQA Summ. ICL Rerank 32K 8K GQA 80.9 92.4 73.1 56.5 53.6 24.5 9.0 7.2 60.7 0.3 VHA 85.7 95.9 78.0 60.1 69.8 29.6 11.0 6.6 57.1 0.4 16K GQA 73.7 90.0 51.8 53.7 27.6 22.6 11.0 8.0 70.4 0.0 VHA 83.2 93.3 71.5 55.4 80.8 32.3 10.5 7.2 68.6 0.0 32K GQA 69.7 81.9 51.9 49.0 49.4 22.3 17.0 8.1 78.3 0.0 VHA 73.5 83.5 56.8 51.8 71.4 33.4 11.0 7.6 75.1 0.5 128K 8K GQA 79.1 89.5 87.9 46.0 44.6 27.4 12.5 5.9 58.7 0.6 VHA 90.1 96.7 92.8 59.4 94.0 34.2 12.5 5.6 46.1 0.1 16K GQA 78.2 88.1 95.1 46.7 28.4 28.4 13.0 6.4 66.4 0.5 VHA 88.5 93.5 98.9 55.9 93.4 34.3 14.5 6.3 61.0 0.7 32K GQA 72.6 83.3 93.6 43.2 3.2 26.7 21.0 6.7 75.6 0.4 VHA 85.3 89.6 96.0 53.7 92.6 36.9 18.5 6.9 72.1 0.4 64K GQA 69.5 81.0 86.4 39.5 4.2 27.8 19.0 7.9 80.1 0.7 VHA 81.3 85.6 90.2 52.2 87.2 39.5 20.0 7.7 79.6 0.4 128K GQA 65.7 75.3 92.5 31.5 3.0 30.4 13.0 7.8 82.6 2.0 VHA 76.8 79.9 92.5 45.5 83.0 42.5 15.0 9.5 82.2 2.6 At the final checkpoint, VHA leads RULER Overall at every evaluated length, with an average of 84.4 across lengths versus 73.0 for GQA. The largest gap occurs in Variable Tracking, where VHA scores 92.6/87.2/83.0 at 32K/64K/128K versus 3.2/4.2/3.0 for GQA. HELMET is more heterogeneous: averaged across lengths, VHA wins RAG, LongQA, and Summarization, ties Reranking, and trails on ICL. A.3 Complementary Long-Context Evaluations We additionally report LongBench-v2, InfiniteBench, BABILong, a short-context eight-task average, and language-modeling bits per byte on the four ProLong domains. At the final checkpoint, VHA wins seven of the eight reported metrics; InfiniteBench is nearly tied. Table 7: Complementary evaluation results. LB-v2, InfBench, and BABI denote LongBench-v2, InfiniteBench, and BABILong; Short is the short-context eight-task average. The remaining columns report bits per byte (BPB) on ProLong domains. Bold marks the better matched result. Checkpoint Model LB-v2 ↑ InfBench ↑ BABI ↑ Short ↑ arXiv ↓ Books ↓ Code ↓ Web ↓ 32K GQA 28.0 29.9 36.5 62.7 0.568 0.812 0.294 0.763 VHA 22.3 27.3 37.4 63.7 0.568 0.812 0.295 0.765 128K GQA 18.9 31.8 31.1 62.6 0.582 0.819 0.311 0.768 VHA 24.2 31.5 38.2 64.1 0.564 0.804 0.296 0.760 Appendix B Additional Robustness Experiments The following experiments use a more efficient training implementation than the main 50B-token runs. Absolute scores therefore need not match the main tables, but all methods within each comparison use the same implementation and training recipe. B.1 Repeated-run Robustness at 0.6B We repeat the 0.6B experiment four times for VHA and three times for GQA. We report every run using zero-shot BF16 evaluation, with normalized accuracy used where available. Avg is the unweighted mean over eight tasks, and ± denotes the sample standard deviation across run-level averages. VHA achieves a higher mean with lower observed dispersion. Table 8: Repeated-run results at 0.6B (%). Every run is reported. Bold in the summary rows marks the higher across-run mean. Model Run ARC-C BoolQ HSwag OBQA PIQA Wino SciQ SIQA Avg VHA 01 30.12 60.18 44.46 33.80 68.77 52.25 74.20 39.56 50.42 VHA 02 29.35 59.27 44.36 32.80 68.55 55.96 74.00 39.56 50.48 VHA 03 29.27 59.36 44.79 33.80 67.14 52.25 76.00 40.23 50.35 VHA 04 30.12 56.09 44.68 34.80 68.55 52.64 76.10 39.76 50.34 VHA mean 29.72 58.73 44.57 33.80 68.25 53.28 75.08 39.78 50.40±0.0750.40±0.07 GQA 01 29.44 58.59 44.49 32.40 68.28 51.14 75.50 38.79 49.83 GQA 02 30.12 57.86 44.55 32.60 68.34 50.91 74.10 39.87 49.79 GQA 03 30.20 60.55 44.61 35.00 68.99 51.93 72.50 39.30 50.39 GQA mean 29.92 59.00 44.55 33.33 68.54 51.33 74.03 39.32 50.00±0.3450.00±0.34 B.2 Sensitivity to Premix Initialization We train the 1.7B model from scratch with the same recipe while changing only the Premix initialization: the default near-identity initialization (α=0.1α=0.1), a larger identity-anchored perturbation (α=1.0α=1.0), and fully random initialization without the identity anchor. All three runs remain stable, with no gradient explosion, divergence, NaN, or Inf. Their early loss and pre-clipping gradient norms are similar (Table 9), while the default near-identity initialization gives the highest downstream average (Table 10). Table 9: Early training trajectory under different Premix initializations. Each entry is loss / pre-clipping gradient norm. Step α=0.1α=0.1 α=1.0α=1.0 Random 10 12.2976 / 14.74 12.3072 / 14.80 12.3132 / 14.98 20 11.1501 / 15.97 11.1581 / 16.26 11.1576 / 16.25 30 9.9346 / 4.98 9.9389 / 5.01 9.9400 / 5.07 50 9.1168 / 2.67 9.1243 / 3.76 9.1255 / 2.03 100 7.9670 / 2.96 7.9788 / 2.93 7.9669 / 4.02 200 6.4777 / 2.99 6.4878 / 2.21 6.4898 / 1.65 500 4.9822 / 1.70 4.9990 / 2.47 5.0198 / 2.85 1000 3.8647 / 1.17 3.8851 / 0.94 3.8888 / 0.89 2000 3.1699 / 0.27 3.1842 / 0.33 3.1967 / 0.54 Table 10: Downstream accuracy under different Premix initializations (%). Bold marks the best result per row. Task α=0.1α=0.1 α=1.0α=1.0 Random ARC-C 34.81 34.13 34.39 BoolQ 60.06 58.04 57.03 HellaSwag 51.80 51.72 51.76 OpenBookQA 37.60 37.00 36.00 PIQA 71.06 71.33 72.14 Winogrande 55.72 54.85 53.28 SciQ 80.70 80.30 80.30 Social IQA 40.58 40.23 39.97 Avg 54.04 53.45 53.11 Appendix C Per-benchmark Scores for the 30B-A3B MoE Evaluation Table 11 reports per-benchmark scores for all evaluated models on the 30B-A3B MoE suite. Category-averaged numbers are reproduced in Table 3 in the main text. Model configurations: GQA 48-4-4, MLA 28-28-28, VHA 24-2-2. Bold marks the best per row. Table 11: Per-benchmark scores on the 30B-A3B MoE evaluation suite (accuracy %). Category Benchmark Shot GQA MLA VHA General MMLU 5-shot 55.97 59.34 60.24 MMLU-Redux 5-shot 63.05 65.25 66.10 MMLU-Pro 5-shot 36.20 41.12 40.03 BBH 3-shot CoT 56.23 56.57 58.61 Math & STEM GPQA 10-shot CoT 34.62 31.62 34.11 GSM8K 4-shot CoT 56.48 55.42 56.25 MATH 4-shot CoT 38.18 43.42 48.84 Coding HumanEval+ 0-shot 50.00 52.44 47.56 MBPP 3-shot 45.00 48.60 49.20 MBPP+ 3-shot 60.54 62.07 62.93 Multilingual MMMLU 5-shot 39.82 42.69 42.95 INCLUDE 5-shot 46.69 47.71 48.16 Appendix D Training Configuration and Benchmark Citations D.1 Model Architecture and Training Hyperparameters Table 12 lists the backbone architecture and training hyperparameters for all five scales. Small-scale models follow the Qwen3 architecture; the 30B-A3B model uses an ERNIE-style MoE architecture with auxiliary-free routing (noaux_tc) and z-loss (λ=4×10−4λ=4×10^-4). We use AdamW [Loshchilov and Hutter, 2017] with β1=0.9 _1=0.9, β2=0.95 _2=0.95, weight decay 0.10.1, and gradient clipping at 1.01.0. Table 12: Model architecture and training hyperparameters across all five scales. “–” denotes not applicable. 0.6B 1.7B 4B 8B-A1B (MoE) 30B-A3B (MoE) Architecture Precision BF16 BF16 BF16 BF16 FP8 Num layers 28 28 36 24 46 Hidden size 1024 2048 2560 2048 1536 FFN dim 3072 6144 9728 768×64768×64 1024×1281024×128 Num experts – – – 64 128 (+1 shared) Experts per token – – – 4 8 Vocab size 151936 151936 151936 151936 102400 RoPE base 10610^6 10610^6 10610^6 10610^6 10610^6 Tied embeddings ✓ ✓ ✓ × × Training Training tokens 50B 50B 50B 50B 500B Context length 4096 4096 4096 4096 8192 Peak learning rate 2.5×10−42.5×10^-4 2×10−42×10^-4 2×10−42×10^-4 2×10−42×10^-4 1×10−41×10^-4 Min learning rate 1×10−51×10^-5 1×10−51×10^-5 1×10−51×10^-5 1×10−51×10^-5 1×10−51×10^-5 LR schedule cosine cosine cosine cosine wsd Warmup steps 2000 2000 2000 2000 2000 Total steps 24000 24000 24000 24000 40000 Batch size (tokens) 2M 2M 2M 2M 13.1M Adam β1 _1 0.9 0.9 0.9 0.9 0.9 Adam β2 _2 0.95 0.95 0.95 0.95 0.95 Weight decay 0.1 0.1 0.1 0.1 0.1 Gradient clipping 1.0 1.0 1.0 1.0 1.0 VHA-specific Premix init scale α 0.1 0.1 0.1 0.1 0.1 Premix init noise σ 1 1 1 1 1 Postmix rank r 4 4 8 8 12 Baseline-specific MLA latent dim ckvc_kv 512 512 512 512 512 MLA latent dim cqc_q 512 1536 1536 1536 1536 MFA head dim d′d 256 256 256 256 – D.2 Evaluation Benchmarks Small-scale benchmarks (8 tasks). We report zero-shot accuracy on the following benchmarks: ARC-Challenge [Clark et al., 2018], BoolQ [Clark et al., 2019], HellaSwag [Zellers et al., 2019], OpenBookQA [Mihaylov et al., 2018], PIQA [Bisk et al., 2020], Winogrande [Sakaguchi et al., 2021], SciQ [Welbl et al., 2017], Social IQA [Sap et al., 2019]. Large-scale benchmarks (12 tasks). We report accuracy with the shot counts listed in Table 11 on the following benchmarks: MMLU [Hendrycks et al., 2020], MMLU-Redux [Gema et al., 2025], MMLU-Pro [Wang et al., 2024], BBH [Suzgun et al., 2023], GPQA [Rein et al., 2023], GSM8K [Cobbe et al., 2021], MATH [Hendrycks et al., 2021], HumanEval+ [Liu et al., 2023], MBPP [Austin et al., 2021], MBPP+ [Liu et al., 2023], MMMLU [OpenAI, 2024], INCLUDE [Romanou et al., 2024]. Appendix E Training Loss Curves Figure 4 shows per-method training loss curves at four model scales (50B tokens, 4K context). VHA achieves the lowest convergence loss at the 4B scale and remains competitive at 0.6B, 1.7B, and 8B-A1B, where its gap from the best-converging method is within 0.01 nats. Figure 4: Training loss curves across four model scales. Each panel compares MHA, GQA, MLA, MFA, and VHA on FineWeb-Edu (50B tokens, 4K context). All curves are plotted from step 4K onward after the learning-rate warm-up; smoothing is applied with an exponential moving average (α=0.93α=0.93). VHA (blue, solid) achieves the lowest loss at convergence at the 4B scale. Appendix F VHA Algorithm Algorithm 1 provides a self-contained reference for VHA’s training forward pass and the one-time inference-absorption procedure that reduces the trained model to a standard GQA-2 deployment. Algorithm 1 Virtual Head Attention (VHA) Training Forward Pass 0: Input X∈ℝS×hX ^S× h; parameters WQ,WK,WV,WO,Wgpremixg=1Hkv,A,BW_Q,W_K,W_V,W_O,\W_g^premix\_g=1^H_kv,A,B 1: Q←XWQ∈ℝS×Hq′×dQ← XW_Q ^S× H_q × d ⊳ Physical query heads (Hq′=Hq/2H_q =H_q/2) 2: K←XWK∈ℝS×Hkv×dK← XW_K ^S× H_kv× d, V←XWV∈ℝS×Hkv×dV← XW_V ^S× H_kv× d 3: for each KV group g=1,…,Hkvg=1,…,H_kv do 4: Q~g←QWgpremix Q_g← Q\,W_g^premix ⊳ Q Premix: virtual head expansion 5: end for 6: Q~←concat(Q~1,…,Q~Hkv)∈ℝS×Hv×d Q ( Q_1,…, Q_H_kv) ^S× H_v× d ⊳ Hv=Hkv×Hq′H_v=H_kv×H_q virtual heads 7: Yg←softmax(Q~gKg⊤/d)VgY_g \! ( Q_gK_g / d )V_g for each group g ⊳ Grouped attention 8: Y←concat(Y1,…,YHkv)Y (Y_1,…,Y_H_kv) 9: Y′←(IHv+BA⊤)Y ←(I_H_v+BA )\,Y ⊳ Linear Postmix: cross-group mixing 10: output←reshape(Y′,[S,Hvd])WOoutput (Y ,[S,H_vd])\;W_O Inference Absorption (one-time precomputation) 0: WQ∈ℝh×Hq′dW_Q ^h× H_q d (sliced into per-head WQ(i)∈ℝh×dW_Q^(i) ^h× d); Wgpremixg=1Hkv\W_g^premix\_g=1^H_kv each ∈ℝd×d ^d× d; WO∈ℝHvd×hW_O ^H_vd× h; A,B∈ℝHv×rA,B ^H_v× r 1: for each (i,g)∈1,…,Hq′×1,…,Hkv(i,g)∈\1,…,H_q \×\1,…,H_kv\ do 2: W~Q(i,g)←WQ(i)Wgpremix∈ℝh×d W_Q^(i,g)← W_Q^(i)\,W_g^premix\;∈\;R^h× d ⊳ Fold Premix into the (i,g)(i,g)-th virtual query projection 3: end for 4: WQabs←concat(W~Q(i,g)i,g)∈ℝh×HvdW_Q^abs \! ( \ W_Q^(i,g) \_i,g )\;∈\;R^h× H_vd ⊳ Hv=Hkv×Hq′H_v=H_kv×H_q 5: WOabs←((IHv+AB⊤)⊗Id)WO∈ℝHvd×hW_O^abs← ((I_H_v+AB ) I_d )\,W_O\;∈\;R^H_vd× h ⊳ Fold Postmix into WOW_O 6: Discard WQ,Wgpremix,A,BW_Q,\W_g^premix\,A,B; run standard GQA-2 with WQabs,WK,WV,WOabsW_Q^abs,W_K,W_V,W_O^abs Appendix G Premix as Structured Factorization of Virtual-Head Query Projections We show that, in the forward pass, Q Premix is mathematically equivalent to a standard multi-head attention with Hv=HkvHq′H_v=H_kvH_q virtual query heads, whose per-virtual-head query projections are constrained to a structured low-dimensional manifold. The training-time benefit of the separated parameterization arises from the coupled gradient structure induced by this factorization. Setup. Let WQphys∈ℝh×Hq′dW_Q^phys ^h× H_q d be the physical query projection (with Hq′H_q heads, each of dimension d), and let Wgpremix∈ℝd×dW_g^premix ^d× d be the Premix matrix for KV group g∈1,…,Hkvg∈\1,…,H_kv\. For physical head i∈1,…,Hq′i∈\1,…,H_q \ within group g, the virtual query is Q~i,g=QiWgpremix=XWQ(i)Wgpremix, Q_i,g\;=\;Q_i\,W_g^premix\;=\;X\,W_Q^(i)\,W_g^premix, (6) where WQ(i)∈ℝh×dW_Q^(i) ^h× d is the i-th slice of WQphysW_Q^phys. The attention output for virtual head (i,g)(i,g) is headi,g=softmax(Q~i,gKg⊤d)Vg.head_i,g\;=\;softmax\! ( Q_i,gK_g d )V_g. (7) Forward equivalence to standard multi-head attention. Define W~Q(i,g)≜WQ(i)Wgpremix∈ℝh×d. W_Q^(i,g)\; \;W_Q^(i)\,W_g^premix\;∈\;R^h× d. (8) By the associativity of matrix multiplication, Q~i,g=XW~Q(i,g) Q_i,g=X W_Q^(i,g), and therefore headi,g=softmax(XW~Q(i,g)Kg⊤d)Vghead_i,g\;=\;softmax\! ( X W_Q^(i,g)K_g d )V_g (9) holds as a pointwise identity for every input X and every parameter assignment. Hence VHA is, at the level of input-output behavior, a standard multi-head attention with HvH_v query heads and HkvH_kv KV heads, in which the i-th virtual head of group g uses the projection matrix W~Q(i,g) W_Q^(i,g). The factorization defines a low-dimensional parameter manifold. The standard HvH_v-head Q projection has Hvhd=HkvHq′hdH_vhd=H_kvH_q hd free parameters. The Premix factorization uses Hq′hd+Hkvd2H_q hd\;+\;H_kvd^2 (10) parameters, parameterized by (WQphys,Wgpremixg=1Hkv)(W_Q^phys,\W_g^premix\_g=1^H_kv). The set of attainable W~Q(i,g)\ W_Q^(i,g)\ tuples is therefore a structured submanifold of the unconstrained parameter space: ℳ=WQ(i)Wgpremixi,g:WQ(i)∈ℝh×d,Wgpremix∈ℝd×d.M\;=\; \ \W_Q^(i)W_g^premix \_i,g\;:\;W_Q^(i) ^h× d,\;W_g^premix ^d× d \. (11) At the 1.7B scale (h=2048h=2048, d=128d=128, Hq′=8H_q =8, Hkv=2H_kv=2, Hv=16H_v=16), the factorized parameterization uses ∼ 2.13M parameters, compared to ∼ 4.19M for the unconstrained HvH_v-head Q projection, a ∼ 50% reduction. The corresponding Q-projection FLOPs are reduced by a similar factor, as the additional cost 2HkvHq′d22H_kvH_q d^2 of applying Premix is much smaller than the savings 2hHq′(Hkv−1)d2hH_q (H_kv-1)d achieved by halving the physical head count. This is the principal source of VHA’s reduction in training and prefill FLOPs. Reparameterization benefit via gradient coupling. Although Eq. 9 shows that the forward pass is algebraically equivalent to training a single matrix W~Q(i,g) W_Q^(i,g) per virtual head, the gradients under the factorized parameterization are coupled across groups and physical heads: ∂ℒ∂WQ(i) ∂ W_Q^(i) =∑g=1HkvX⊤∂ℒ∂Q~i,g(Wgpremix)⊤, = _g=1^H_kvX \, ∂ Q_i,g\, (W_g^premix ) , (12) ∂ℒ∂Wgpremix ∂ W_g^premix =∑i=1Hq′(XWQ(i))⊤∂ℒ∂Q~i,g. = _i=1^H_q (XW_Q^(i) ) \, ∂ Q_i,g. (13) Each WQ(i)W_Q^(i) aggregates contributions from all HkvH_kv groups, modulated by the corresponding Premix matrices, while each WgpremixW_g^premix receives gradient signal through all Hq′H_q physical heads. SGD on the factorized parameters therefore traverses the manifold ℳM along directions induced by this coupled gradient structure, which differs from the trajectories obtained by independently optimizing HvH_v unconstrained virtual-head matrices. This is the same source of implicit regularization that underlies LoRA [Hu et al., 2022] and RepVGG [Ding et al., 2021]: a forward-equivalent factorization of a large weight tensor yields different optimization dynamics and, empirically, different (and often better) final solutions. Inference-time absorption. Since the forward equivalence in Eq. 9 holds for any frozen parameter values, at inference we simply precompute W~Q(i,g)=WQ(i)Wgpremix W_Q^(i,g)=W_Q^(i)W_g^premix for all (i,g)(i,g) pairs and discard WQphysW_Q^phys and Wgpremix\W_g^premix\. The model then runs as a standard GQA-2 with HvH_v independent query projections, Hkv=2H_kv=2 KV heads, and the same output projection WOW_O, fully compatible with existing GQA serving kernels. Unified view with Postmix. Premix and Postmix are two instantiations of the same head-axis reparameterization principle. Premix factorizes the per-virtual-head query projection on the Q side (W~Q(i,g)=WQ(i)Wgpremix W_Q^(i,g)=W_Q^(i)W_g^premix); Postmix factorizes the merged output projection on the O side (WO′=((IHv+AB⊤)⊗Id)WOW_O = ((I_H_v+AB ) I_d )W_O, see Appendix H). Both are forward-equivalent to a standard GQA in their merged form, both gain training-time benefits from coupled gradients of the separated parameters, and both fold back into a standard GQA-2 architecture at inference with zero overhead. Appendix H Proof of Postmix Linear Redundancy We prove that Linear Postmix can be mathematically absorbed into the output projection WOW_O, and discuss why the separated parameterization nonetheless improves training. Setup. For a single token, let Y∈ℝHv×dY ^H_v× d be the concatenated attention output across all HvH_v virtual heads. Postmix computes: Y′=(IHv+BA⊤)Y =(I_H_v+BA )\,Y (14) where A,B∈ℝHv×rA,B ^H_v× r. The final output is obtained by applying the output projection: output⊤=vecrow(Y′)⊤WO∈ℝ1×houtput =vec_row(Y ) W_O ^1× h (15) where WO∈ℝHvd×hW_O ^H_vd× h and vecrow(⋅)vec_row(·) flattens a matrix in row-major order (i.e., Y.reshape(Hvd)Y.reshape(H_vd) in PyTorch, concatenating head outputs). Absorption via Kronecker product. Using the row-major identity: for M∈ℝHv×HvM ^H_v× H_v and Y∈ℝHv×dY ^H_v× d, vecrow(MY)⊤=vecrow(Y)⊤(M⊤⊗Id)vec_row(MY) =vec_row(Y) (M I_d ) (16) Setting M=IHv+BA⊤M=I_H_v+BA (so M⊤=IHv+AB⊤M =I_H_v+AB ) [Roger and Charles, 1994]: vecrow(Y′)⊤ _row(Y ) =vecrow((IHv+BA⊤)Y)⊤ =vec_row\! ((I_H_v+BA )\,Y ) =vecrow(Y)⊤((IHv+AB⊤)⊗Id) =vec_row(Y) ((I_H_v+AB ) I_d ) (17) Substituting into the output computation: output⊤ =vecrow(Y′)⊤WO =vec_row(Y ) W_O =vecrow(Y)⊤((IHv+AB⊤)⊗Id)WO =vec_row(Y) ((I_H_v+AB ) I_d )W_O =vecrow(Y)⊤WO′ =vec_row(Y) W_O (18) where WO′=((IHv+AB⊤)⊗Id)WO∈ℝHvd×hW_O = ((I_H_v+AB ) I_d )W_O ^H_vd× h. Since IHv+BA⊤I_H_v+BA is generically invertible (it equals the identity at initialization and remains near-identity during training), WO′W_O has the same number of free parameters as WOW_O. Therefore, for any choice of WO,A,BW_O,A,B, there exists a single matrix WO′W_O that produces the identical output, confirming that Postmix is theoretically redundant. Reparameterization benefit. Despite this redundancy, the separated parameterization WO,A,B\W_O,A,B\ provides training advantages analogous to those observed in RepVGG [Ding et al., 2021]: 1. Factored gradient paths. The gradients with respect to A, B, and WOW_O follow different paths through the computation graph. This factorization provides a form of implicit regularization, as the effective weight update to WO′W_O at each step is constrained by the low-rank structure of AB⊤AB . 2. Identity initialization. With B initialized to zero, BA⊤=0BA =0 and Y′=Y =Y at the start of training. Postmix begins as a no-op and gradually learns useful cross-group mixing patterns, avoiding disruption to the early training dynamics. 3. Low-rank constraint. The bottleneck rank r≪Hvr H_v limits the expressiveness of the cross-head mixing, acting as a structural regularizer. This prevents Postmix from overfitting, consistent with the ablation finding that larger ranks (r=32,128r=32,128) degrade evaluation performance despite lower training loss.