Research map / Architectures and efficiency / Attention and transformers

Virtual Head Attention

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Virtual Head Attention (VHA), which halves physical query heads with 2 KV heads and adds linear Q Premix and low-rank Linear Postmix during training, then absorbs both into a standard GQA-2 at inference.

Problem

Attention designs face an impossible triangle among KV cache size, training/prefill FLOPs, and model quality; GQA has high cache, MLA loses quality, and MFA increases compute.

Approach

Recover virtual head diversity from halved query heads via a near-identity per-group query transform (Q Premix) and fuse cross-head features with a low-rank identity-residual output mix (Linear Postmix); both are purely linear so they fold into WQ and WO at inference, yielding a standard GQA-2 model.

Key results

• At 1.7B, VHA 8-2-2 (1 KB KV cache) reaches 54.50% average accuracy vs. 54.09% for GQA-8 (4 KB) and reduces training/prefill FLOPs by 7.9% at S=4096.
• On the 30B-A3B MoE trained on 500B tokens, VHA 24-2-2 averages 50.36 vs. 47.76 for GQA-4 and 49.66 for MLA over 12 benchmarks.
• After long-context mid-training on a 10B-A1B MoE, VHA averages RULER 84.4 over 8K-128K vs. 73.0 for GQA.

Surprising findings

• Reducing GQA to 8-2-2 shape costs -1.71% but VHA with the same shape exceeds the GQA-8 baseline by +0.41%.
• Linear Postmix is theoretically redundant (absorbable into WO), yet ablations show it still improves quality; larger Postmix rank lowers training loss but degrades downstream accuracy.

Limitations

• Gains over GQA at small scales are modest (+0.05% to +0.81%) and within reported run-to-run variance at 0.6B.
• Speedups are theoretical (memory-bound decoding estimates) rather than measured end-to-end.
• Postmix is theoretically redundant under linear absorption, so its benefit relies on empirical optimization effects rather than expressivity.

Experimental setup

Datasets: FineWeb-Edu internal multi-domain pretraining corpus ProLong

Benchmarks: ARC-Challenge BoolQ HellaSwag OpenBookQA PIQA Winogrande SciQ SIQA MMLU MMLU-Redux MMLU-Pro BBH

Metrics: accuracy average accuracy RULER score bits per byte KV cache size training FLOPs

Baselines: MHA GQA MLA MFA

Builds on: Grouped-Query Attention Multi-head Latent Attention Multi-matrix Factorization Attention RepVGG LoRA

Closest papers by meaning

• DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads (ICLR 2025 · similarity 0.89)
• QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.89)
• Q Cache: Visual Attention Is Valuable in Less than Half of Decode Layers for Multimodal Large Language Model (AAAI 2026 · similarity 0.89)
• Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.89)
• A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.89)
• Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention Calibration (ICML 2024 · similarity 0.88)