PaperlayerView exact revision on arXiv

Source-linked AI summary

Virtual Head Attention

Xibo Ding, Guoxia Wang, Jinle Zeng, Jiabin Yang, Sirui Tao, Chao Yang, Dianhai Yu

arXiv:2610.17959v1cs.CLcs.AI

TL;DR

Long-context attention must balance KV cache size, training and prefill FLOPs, and model quality, while existing compression approaches trade among these goals. VHA combines query-side virtual head expansion with low-rank output mixing in two linear operations. Across reported scales, it improves on GQA quality while reducing cache use and training/prefill FLOPs, and its operations merge into standard GQA-2 at inference.

  • Problem

    Attention for long-context language models must balance KV cache size, training and prefill FLOPs, and model quality, while GQA, MLA, and MFA each leave trade-offs among these objectives.

    Source: §1 · pp. 1–2

  • Method

    VHA halves physical query heads and uses Q Premix for virtual query expansion and low-rank Linear Postmix for inter-head feature fusion, with both operations foldable into standard GQA-2 projections at inference.

    Source: §1 · p. 2

  • Results

    Across four scales, VHA exceeds standard GQA by +0.05% to +0.81%, reduces KV cache by 4× versus GQA-8 at dense scales and 2× versus GQA-4 at the MoE scale, and lowers FLOPs by 7.9%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    VHA's linear operations merge into standard GQA-2 at inference, allowing reuse of existing GQA inference optimizations.

    Source: §1 · p. 2, §6 · p. 10

    The largest validation is at 30B-A3B MoE trained on 500B tokens; scaling to larger MoE models and substantially more data, broader long-context architectures and recipes, and reinforcement-learning training remain to be explored.

    Source: Limitations and Future Work · p. 14

Abstract

from arXiv · show

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.

1 Introduction

Long-context attention must balance KV cache size, training and prefill FLOPs, and model quality, but existing approaches trade one of these objectives against another. VHA uses two linear operations to improve this balance, and the reported evaluations show quality gains alongside lower cache use and compute.

1 Introduction

  • The attention design matters because KV cache grows linearly with sequence length and attention FLOPs scale quadratically during training and prefill.The resulting cache use constrains serving concurrency and generation latency.Source: §1 · p. 1
  • GQA-8 preserves strong quality but uses 4 KB of KV cache per layer per token, while GQA-2 uses 1 KB and loses 1.00% quality at 1.7B.MLA compresses cache to 1.12 KB but has the lowest quality across multiple scales; MFA uses 1 KB but raises training and prefill FLOPs by 49.8% at 1.7B, S=4096.Source: §1 · p. 2
  • VHA halves physical query heads and uses Q Premix to recover virtual query-head diversity and Linear Postmix to fuse inter-head features through a low-rank identity-residual transformation.The added linear operations can be folded into the query and output projections, leaving standard GQA-2 at inference.Source: §1 · p. 2
  • Across four scales, VHA beats standard GQA by +0.05% to +0.81%, uses 4× less KV cache than GQA-8 at dense scales and 2× less than GQA-4 at the MoE scale, and lowers FLOPs by 7.9%.The FLOP reduction is reported at 1.7B with S=4096; the four-scale pretraining used 4K context and 50B tokens.Source: §1 · p. 2
  • On a 30B-A3B MoE model trained on 500B tokens at 8K context, VHA beats MLA and GQA-4 by +0.70 / +2.60 % on a 12-benchmark evaluation suite.The evaluation is described as a base-model suite.Source: §1 · p. 2

2 Related Work

Prior attention designs reduce cache use through KV sharing, low-rank compression, or head communication, but often retain quality, compute, or inference-pipeline trade-offs. VHA instead virtualizes queries and mixes output features with linear operations that merge into standard GQA at deployment.

2.1 KV Cache Efficient Attention

  • VHA differs from cache-compression alternatives by merging its richer training-time parameterization exactly into standard GQA at deployment, without custom kernels or specialized inference procedures.Earlier alternatives include MLA, which requires custom kernels and a non-standard pipeline, and MFA, which increases training and prefill FLOPs.Source: §2.1 · p. 3

2.2 Inter-Head Communication and Virtual Heads

  • Unlike inter-head methods that mix attention weights, VHA virtualizes queries with Premix and fuses output-space head features with a low-rank Postmix.The operations are strictly linear and absorbable at inference.Source: §2.2 · p. 3

2.3 Reparameterization

  • VHA extends reparameterization from a single weight matrix to attention's head dimension, with Premix factorizing per-virtual-head query projections and Postmix mixing the head axis.The decoupled parameterization has distinct training dynamics while merging back into standard GQA-2 at inference with zero overhead.Source: §2.3 · p. 3

3 Method: Virtual Head Attention

VHA reduces the number of physical query heads and uses two linear operations to recover virtual-head diversity and mix head outputs. Both operations can be absorbed into standard GQA projections at inference, while the design targets lower cache and compute with improved quality.

3.2 Overall Design

  • With Hkv = 2, VHA halves physical query heads while retaining Hq virtual heads, yielding a cache 4× smaller than GQA-8 at dense scales and 2× smaller than GQA-4 at MoE scales.Q Premix restores diversity, and Linear Postmix fuses inter-head features.Source: §3.2 · p. 4

3.5 Efficiency Analysis

  • Because both modules are linear, they can be absorbed into the standard GQA query and output projections, leaving a GQA-2 model at inference.VHA’s decoding cost is identical to standard GQA-2; the passage estimates a 2–4× theoretical speedup in memory-bound decoding from the smaller cache.Source: §3.2 · p. 4, §3.4 · p. 5, §3.5 · p. 6

3.3 Q Premix: Virtual Query Head Expansion

  • Q Premix expands physical query heads into virtual heads using per-group transformations initialized near identity, so the virtual heads begin as slight perturbations of physical heads.The factorized Q path uses roughly 50% fewer projection parameters and FLOPs than a fully decoupled virtual-head Q projection at 1.7B.Source: §3.3 · pp. 4–5

3.4 Linear Postmix: Low-Rank Identity-Residual Mixing

  • Linear Postmix applies a low-rank identity-residual transformation along the head dimension independently for each token and head-dimension channel.Its matrices are initialized so the transformation starts as the identity; the passage reports a +1.19% improvement at 1.7B when Postmix is added on top of Premix.Source: §3.4 · p. 5

3.5 Efficiency Analysis

  • VHA is Pareto-optimal on KV cache and training/prefill FLOPs at every scale, with the highest average accuracy and smallest cache across the four scales.The cache reductions are 4× versus GQA-8 at dense scales and 2× versus GQA-4 at MoE scales.Source: §3.2 · p. 4, §3.5 · p. 5, Table 1 · p. 6

4 Experiments

Experiments evaluate VHA across dense and MoE scales, including industrial-scale and long-context settings, and test its quality, cache, and compute trade-offs. VHA maintains quality gains with reduced KV cache, while ablations and efficiency results examine the design choices behind those gains.

4.1 Experimental Setup

  • The experiments use 50B-token, 4K-context pretraining at four scales, 500B-token, 8K-context industrial pretraining, and matched 10B-A1B long-context mid-training.The long-context comparison shares backbone, data, optimization, and sliding-window layers; only full-attention layers differ, with two KV heads for VHA and four for GQA.Source: §4.1 · p. 6

4.2 Main Results

  • Across dense and MoE architectures, VHA achieves the highest average accuracy at every scale, exceeding GQA by +0.05% to +0.81%, with 4× less KV cache than GQA-8 at dense scales and 2× less than GQA-4 at the MoE scale.Its KV cache is 1.00 KB per layer per token, tied with MFA.Source: §4.2 · p. 7
  • At equal 1.00 KB cache, VHA beats MFA at every scale (+0.10% to +0.77%) and reduces FLOPs 7.9% at 1.7B, while MFA increases them 49.8%.MLA underperforms at every scale in the 50B-token training setting.Source: §4.2 · p. 7

4.3 Large-scale and long-context validation

  • The 30B-A3B VHA model achieves the highest 12-benchmark average using 50% of GQA-4’s cache and 89% of MLA’s.The evaluation covers four task categories; the model was pretrained on 500B tokens at 8K context.Source: §4.3 · p. 7, Table 3 · p. 7
  • At the final 128K checkpoint, VHA’s RULER length-average is 84.4 versus GQA’s 73.0 across 8K–128K, and it remains on par on InfiniteBench (31.5 vs. 31.8).It also outperforms GQA on LongBench-v2 (24.2 vs. 18.9) and BABILong (38.2 vs. 31.1) at 128K.Source: §4.3 · pp. 7–8

4.4 Ablations

  • At both tested scales, only the joint Premix–Postmix configuration consistently gives the largest gains at 25% of GQA-8’s cache.The component ablation tests the modules individually and together at 1.7B and 4B.Source: §4.4 · p. 8, Table 4 · p. 8
  • Larger Postmix ranks lower training loss but worsen downstream evaluation; SiLU degrades quality and prevents inference-time absorption.The rank sweep reports Avg 54.50/54.24/54.15 for r ∈ {4, 32, 128}; SiLU underperforms the linear formulation at 1.7B.Source: §4.4 · p. 8

4.5 Computational Efficiency

  • Across four scales at S=4096, VHA matches the 1.00 KB smallest cache and reduces training/prefill FLOPs versus GQA by 8.1%, 7.9%, 7.6%, and 6.5%.MFA matches the cache but raises FLOPs by 48–74%; MLA has a slightly larger cache and raises FLOPs at three scales.Source: §4.5 · pp. 8–9, Table 5 · p. 8

5 Analysis: How VHA Reparameterization Works

Premix learns differentiated query transforms across KV groups, while Postmix redistributes attention across heads. Their complementary structures together outperform either component alone.

Premix learns non-trivial, group-differentiated query transforms

  • Premix learns substantially different, but not independent, query transforms for the two KV groups, creating group-conditioned virtual-head expansion.The inter-group distance peaks at 6.13, compared with a peak 5.85 deviation from identity.Source: §5 · p. 9

Postmix learns structured low-rank cross-head communication

  • Postmix shows pronounced off-diagonal flow, primarily redistributing information within KV groups while maintaining sparse bridges across groups.The analysis rules out simple per-head rescaling.Source: §5 · pp. 9–10

Joint reparameterization closes the doubly-compressed gap

  • +0.41% over GQA-8 at 25% of its KV cache; Premix and Postmix are complementary, and neither alone suffices.Premix-only has a −0.78% gap and Postmix-only a −0.14% gap at 1.7B.Source: §5 · p. 10

6 Conclusion

VHA uses Q Premix and Linear Postmix to address the trade-off among KV cache, training and prefill FLOPs, and model quality. It surpasses GQA across evaluated scales and maintains an advantage through 128K context.

Conclusion

  • VHA mitigates the trade-off among KV cache size, training and prefill FLOPs, and model quality using Q Premix and Linear Postmix.Halving physical query heads and recovering expressiveness through virtual-head expansion and low-rank cross-group mixing reduces the cache to its architectural minimum.Source: §6 · p. 10
  • The linear operations absorb into standard GQA-2 at inference, allowing reuse of existing serving optimizations with zero additional overhead.The conclusion describes Premix and Postmix as functionally complementary; neither alone suffices.Source: §6 · p. 10
  • VHA surpasses GQA across five evaluated scales, from 0.6B to 30B-A3B MoE, and maintains its advantage over GQA through 128K context.The long-context result comes from a matched 10B-A1B evaluation.Source: §6 · p. 10

Limitations and Future Work

The paper identifies limits to its validation scale and scope, alongside directions for further study and possible combinations with other compression methods.

Limitations and Future Work

  • The largest validation is a 30B-A3B MoE trained on 500B tokens; larger MoE models and substantially more data, such as 20T tokens, remain to be explored.This bounds the demonstrated model and training-data scale.Source: Limitations and Future Work · p. 14
  • VHA is validated through 128K on a 10B-A1B MoE, while broader long-context architectures and training recipes remain unexplored.Source: Limitations and Future Work · p. 14
  • VHA's behavior under reinforcement-learning training remains an open direction; combining it with quantization or sparsification may yield further inference savings.Source: Limitations and Future Work · p. 14

A.1 Model and Stage-wise Training

In a controlled 10B-A1B MoE comparison, VHA uses fewer KV heads in full-attention layers and leads RULER Overall across evaluated lengths. The appendix also reports complementary evaluations and notes that robustness experiments use a different training implementation.

A.1 Model and Stage-wise Training

  • The matched models share backbone, data, optimization, and sliding-window layers; VHA uses two versus GQA's four KV heads in full-attention layers.Both use eight KV heads in sliding-window layers, so the cache reduction applies to full-attention layers without changing the sliding-window component.Source: §A.1 · p. 14
  • Long-context mid-training adds 100B tokens at 32K and another 100B at 128K after 500B-token 8K pretraining; final results use the 128K checkpoint.The 32K checkpoint provides intermediate results.Source: §A.1 · p. 14

A.2 RULER and HELMET

  • RULER Overall averages 13 tasks, while HELMET results are reported by category because its categories measure distinct capabilities.Evaluation spans 8K to 128K.Source: §A.2 · p. 14
  • 84.4 versus 73.0 for GQA: VHA's average RULER Overall leads across evaluated lengths; it wins seven of eight additional reported metrics, with InfiniteBench nearly tied.On HELMET, VHA wins RAG, LongQA, and Summarization, ties Reranking, and trails on ICL.Source: §A.2 · p. 14, §A.3 · p. 14

B Additional Robustness Experiments

  • Robustness experiments use a more efficient training implementation than the main 50B-token runs, so absolute scores need not match, though methods within each comparison share a recipe.Source: B Additional Robustness Experiments · p. 14

B.1 Repeated-run Robustness at 0.6B

Repeated runs at 0.6B show VHA with a higher mean and lower observed dispersion than GQA. At 1.7B, all tested Premix initializations train stably, while the default near-identity initialization yields the highest downstream average.

B.1 Repeated-run Robustness at 0.6B

  • VHA achieves a higher mean with lower observed dispersion than GQA across repeated 0.6B runs.The comparison uses four VHA runs and three GQA runs, with the unweighted mean over eight tasks and sample standard deviation across run-level averages.Source: §B.1 · p. 14

B.2 Sensitivity to Premix Initialization

  • At 1.7B, all three Premix initializations remain stable, and the default near-identity initialization gives the highest downstream average.The tested settings were α=0.1, α=1.0, and fully random initialization without the identity anchor.Source: §B.2 · pp. 14–15

C Per-benchmark Scores for the 30B-A3B MoE Evaluation

The 30B-A3B evaluation reports per-benchmark scores for the evaluated models, with category averages presented in the main text. The listed training configurations include an ERNIE-style MoE architecture for the 30B-A3B model.

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, with bold marking the best score per row.Category-averaged results are reproduced in Table 3 in the main text; configurations are GQA 48-4-4, MLA 28-28-28, and VHA 24-2-2.Source: C Per-benchmark Scores for the 30B-A3B… · p. 15

D.1 Model Architecture and Training Hyperparameters

  • The 30B-A3B model uses an ERNIE-style MoE architecture with auxiliary-free routing and z-loss (λ=4×10−4).The listed optimizer is AdamW [Loshchilov and Hutter, 2017], with β1=0.9, β2=0.95, weight decay 0.1, and gradient clipping at 1.0.Source: §D.1 · p. 15

D.2 Evaluation Benchmarks

The evaluation section defines zero-shot accuracy benchmarks for small-scale models and accuracy benchmarks with specified shot counts for large-scale models.

D.2 Evaluation Benchmarks

  • Small-scale evaluation reports zero-shot accuracy on eight tasks: ARC-Challenge, BoolQ, HellaSwag, OpenBookQA, PIQA, Winogrande, SciQ, and Social IQA.The passage provides the benchmark bibliography markers alongside the task names.Source: §D.2 · pp. 15–16
  • Large-scale evaluation reports accuracy on 12 tasks, with shot counts listed in Table 11.The tasks include MMLU, MMLU-Redux, MMLU-Pro, BBH, GPQA, GSM8K, MATH, HumanEval+, MBPP, MBPP+, MMMLU, and INCLUDE.Source: §D.2 · p. 16

E Training Loss Curves

Across four scales, VHA has the lowest convergence loss at 4B and remains within 0.01 nats of the best-converging method at the other reported scales. The appendices explain its factorized query and output transformations and their absorption into standard GQA-2 inference.

F VHA Algorithm

G Premix as Structured Factorization of Virtual-Head Query Projections

  • Q Premix is forward-equivalent to virtual query heads whose projections lie on a structured low-dimensional manifold, reducing Q-projection parameters and FLOPs.For H′q=8 and Hkv=2, the factorized projection uses about 2.13M parameters versus about 4.19M unconstrained, a roughly 50% reduction; coupled gradients also change optimization dynamics.Source: G Premix as Structured Factorization… · pp. 16, 18, 19

E Training Loss Curves

  • At 4B, VHA achieves the lowest convergence loss; at 0.6B, 1.7B, and 8B-A1B, its gap from the best-converging method is within 0.01 nats.The curves compare MHA, GQA, MLA, MFA, and VHA on FineWeb-Edu after warm-up, using 50B tokens and 4K context.Source: E Training Loss Curves · p. 16, Figure 4 · p. 19

Inference Absorption (one-time precomputation)

H Proof of Postmix Linear Redundancy

Related papers

Loading 2610.17959v1…