Virtual Head Attention

📄 arXiv: 2610.17959v1 📥 PDF

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

分类: cs.CL, cs.AI

发布日期: 2026-10-13

备注: Published as a main conference paper at NeurIPS 2026

💡 一句话要点

提出虚拟头注意力机制以解决长上下文模型的缓存与计算问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文模型 虚拟头注意力 KV缓存优化 计算效率 自然语言处理 模型推理 深度学习

📋 核心要点

  1. 现有长上下文模型在KV缓存、计算开销和模型质量之间存在权衡，导致性能受限。
  2. 提出虚拟头注意力（VHA），通过轻量级线性操作提升模型性能，减少KV缓存和计算需求。
  3. 在多个实验中，VHA在各个规模上均超越标准GQA基线，并显著降低KV缓存和计算开销。

📝 摘要（中文）

在长上下文的大型语言模型中，注意力机制的设计面临KV缓存大小、训练和预填充FLOPs以及模型质量之间的“不可调和三角”问题。GQA模型虽然质量较高，但KV缓存开销大；而MLA和MFA等替代方法虽然能压缩缓存，但MLA存在质量损失，MFA则显著增加训练和预填充计算。本文提出的虚拟头注意力（VHA）通过两个轻量级线性操作同时改善这三方面的性能。VHA在推理时可将Premix和Postmix操作折叠进查询和输出投影中，减少模型复杂度，并在多个实验中超越标准GQA基线，同时显著降低KV缓存和计算开销。

🔬 方法详解

问题定义：本文旨在解决长上下文大型语言模型中，KV缓存大小、训练和预填充FLOPs与模型质量之间的矛盾，现有方法如GQA、MLA和MFA各有不足，导致性能受限。

核心思路：提出虚拟头注意力（VHA），通过两个轻量级线性操作（Q Premix和Linear Postmix）来恢复虚拟头的多样性，优化KV缓存和计算效率。

技术框架：VHA的整体架构包括两个主要模块：Q Premix对查询进行近似恒等变换以恢复虚拟头多样性，Linear Postmix通过低秩残差结构融合头间特征。推理时，Premix和Postmix可折叠进查询和输出投影中，简化模型结构。

关键创新：VHA的创新在于引入了无非线性操作的轻量级线性变换，使得模型在保持高质量的同时，显著降低KV缓存和计算开销，与现有方法相比具有本质区别。

关键设计：VHA使用两个KV头，组大小可达64，充分利用Tensor Core并行性，训练和预填充FLOPs降低7.9%。在不同规模的实验中，VHA均表现出色，且在工业模型中超越了MLA和GQA-4。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

在多个实验中，VHA在4K上下文50B标记的预训练中，始终超越标准GQA基线，提升幅度在0.05%到0.81%之间。同时，VHA在密集规模下将KV缓存减少了4倍，在MoE规模下减少了2倍，训练和预填充FLOPs降低了7.9%。在工业30B-A3B MoE模型的评估中，VHA分别超越MLA和GQA-4达0.70%和2.60%。

🎯 应用场景

该研究的虚拟头注意力机制具有广泛的应用潜力，特别是在需要处理长上下文的自然语言处理任务中，如文本生成、对话系统和信息检索等。通过优化KV缓存和计算效率，VHA能够在资源受限的环境中实现更高效的模型推理，推动大型语言模型在实际应用中的落地与普及。

📄 摘要（原文）

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.