虚拟头注意力（VHA）：兼顾 KV cache、算力与模型质量的长上下文注意力机制

查看 PDF HTML（实验性）

摘要：长上下文大语言模型中注意力机制的设计面临 KV 缓存大小、训练与预填充 FLOPs 以及模型质量之间的“不可能三角”。GQA 实现了较强的质量，但会带来高昂的 KV 缓存开销；MLA 和 MFA 等替代方法可以压缩缓存，但 MLA 存在质量损失，而 MFA 会大幅增加训练和预填充计算量。我们提出虚拟头注意力（Virtual Head Attention，VHA），通过两个轻量级线性操作同时改善这三个顶点：Q Premix 对每个 KV 组内的查询应用近似恒等变换，从减半的物理查询头中恢复虚拟头的多样性；Linear Postmix 通过低秩 $I + AB^\top$ 残差结构融合头间特征。由于 VHA 不引入非线性操作，推理时可将 Premix 折叠进查询投影、将 Postmix 折叠进输出投影，使模型退化为标准 GQA-2 架构，从而直接复用所有现有的 GQA 推理优化。仅使用两个 KV 头时，大模型中的组大小可轻松达到 64，在解码期间充分利用 Tensor Core 并行性。在 4K 上下文、50B token 预训练的四个规模上，VHA 在每个规模上都持续超越标准 GQA 基线（$+$0.05\% 至 $+$0.81\%），同时在稠密规模下将 KV 缓存较 GQA-8 减少 $4\times$、在 MoE 规模下较 GQA-4 减少 $2\times$，并将训练和预填充 FLOPs 降低 7.9\%（1.7B，$S{=}4096$）。在 8K 上下文、500B token 训练的工业级 30B-A3B MoE 模型上，VHA 在包含 12 项基准的基座模型评测套件上分别以 $+$0.70\,/\,$+$2.60\,\% 超越 MLA 和 GQA-4；在另一项匹配的 10B-A1B MoE 长上下文中期训练后对比中，VHA 在 8K--128K 上取得平均 RULER 分数 84.4，而 GQA 为 73.0，同时在全注意力层中使用两个而非四个 KV 头。

Comments:	作为主会议论文发表于 NeurIPS 2026
Subjects:	计算与语言（cs.CL）；人工智能（cs.AI）
Cite as:	arXiv:2610.17959 [cs.CL]
	（或 arXiv:2610.17959v1 [cs.CL] 用于此版本）
	https://doi.org/10.48550/arXiv.2610.17959 arXiv 颁发的 DOI 经由 DataCite（待注册）

提交历史

发件人：Guoxia Wang [查看邮箱]
[v1] 2026年10月13日星期二 08:39:20 UTC（486 KB）