← 全部论文2610.17959v1 / r1
版本与下载

Virtual Head Attention:中文研读笔记

论文:虚拟头注意力(Virtual Head Attention,VHA)
作者:Xibo Ding、Guoxia Wang、Jinle Zeng、Jiabin Yang、Sirui Tao、Chao Yang、Dianhai Yu(百度、北京大学)
版本:arXiv:2610.17959v1,提交于 2026-10-13;本次获取与研读日期:2026-10-19。
材料:arXiv · 本地原文 PDF · 英文汇报 PDF · 可编辑 TeX
以下页码均为从 1 开始的 PDF 页码;本文原文印刷页码与其一致。

论文总览

分组查询注意力(Grouped-Query Attention,GQA)通过共享键值(Key–Value,KV)头降低缓存成本。VHA 研究一个具体问题:能否用两个 KV heads 的缓存,训练出质量接近或超过常规 GQA 的模型,同时保留标准 GQA 的部署路径?

作者把原本的 query heads 数减半,先计算一组共享的“物理 query heads”,再为两个 KV 组分别施加特征变换,得到与原 GQA 数量相同的“虚拟 query heads”。注意力计算后,再用低秩残差矩阵混合各头输出。两个新增模块都是线性的,因此训练结束后可以分别并入 query 和 output 投影。VHA 的关键贡献由此落在训练参数化、参数共享和优化路径上;最终部署模型仍是标准的两 KV 头 GQA。

主要证据分为三组,不能混用实验条件:

  • 50B tokens、4K 上下文的四规模预训练:VHA 八任务平均准确率比 GQA 高 0.05–0.81 个百分点。稠密模型的 KV 缓存降为 GQA-8 的 1/4;8B-A1B MoE 的缓存降为 GQA-4 的 1/2。
  • 30B-A3B MoE、500B tokens、8K 上下文:四类任务宏平均为 50.36,GQA 为 47.76,MLA 为 49.66;VHA 在编码类别上低于 MLA。
  • 另一套 10B-A1B 混合注意力模型:经过 32K、128K 长上下文续训,最终模型在 8K–128K 的 RULER 长度平均为 84.4,GQA 为 73.0。这里只把全注意力层的 KV 头从四个减至两个,滑窗层不变。

效率方面,论文报告 1.7B、序列长 4096 下的训练 FLOPs 从 10.61 降至 9.78 G/token,约减少 7.9%。这个数字对应保留因子化结构的计算;合并为标准 GQA 后,query 投影会展开。全文没有给出端到端训练吞吐、prefill 延迟或解码速度的实测表。[Sec. 3–4,PDF pp. 4–8,Tables 1–5;Appendix A,pp. 14–15]

结构地图

原文章节论证作用PDF 定位
1 Introduction缓存、计算量、质量的三方取舍pp. 1–2
2 Related Work;2.1–2.3KV 压缩、头间通信、重参数化的关系pp. 2–3
3 Method: Virtual Head Attention;3.1–3.5GQA 定义、Premix、Postmix、合并与复杂度pp. 3–6;Fig. 1;Eqs. (1)–(5)
4 Experiments;4.1–4.5实验设置、主结果、大规模验证、消融、效率pp. 6–9;Tables 1–5
5 Analysis: How VHA Reparameterization Works已训练混合矩阵的结构pp. 9–10;Figs. 2–3
6 Conclusion贡献归纳p. 10
Limitations and Future Work更大规模、训练量与组合应用p. 14
A Long-Context Evaluation;A.1–A.3混合骨干、长训配方、多套长上下文评测pp. 14–15;Tables 6–7
B Additional Robustness Experiments;B.1–B.2重复训练、Premix 初始化pp. 14–17;Tables 8–10
C Per-benchmark Scores for the 30B-A3B MoE Evaluation12 项分数与 few-shot 设置pp. 15、17;Table 11
D Training Configuration and Benchmark Citations;D.1–D.2模型配置、优化器与基准pp. 15–16、18;Table 12
E Training Loss Curves四规模训练曲线pp. 16、19;Fig. 4
F VHA Algorithm训练前向和一次性合并pp. 16、20;Algorithm 1
G Premix as Structured Factorization of Virtual-Head Query Projections前向等价性与梯度耦合pp. 16–19;Eqs. (6)–(13)
H Proof of Postmix Linear Redundancy输出投影吸收的 Kronecker 推导pp. 19–21;Eqs. (14)–(18)

附录中的表格、算法存在跨页浮动,表格页码与章节起始页不总相同。

一:动机——压缩 KV 后,怎样重新组织 query 的自由度?

1.1 三种成本由不同部分决定

对长度 $S$、隐藏维 $h$、每头维度 $d$ 的输入,GQA 使用 $H_q$ 个 query heads 和 $H_{kv}$ 个 KV heads。每个 KV 组服务 $G=H_q/H_{kv}$ 个 query heads。

KV 缓存按 token 累积,BF16 下每层每 token 占用

$$ C_{\mathrm{KV}}=2\ \text{(K,V)}\times 2\ \text{bytes}\times H_{kv}d =4H_{kv}d\ \text{bytes}. $$

当 $d=128$ 时,8、4、2 个 KV heads 分别对应 4096、2048、1024 bytes。论文把它们记作 4、2、1 KB,实际上采用 1024 bytes 的换算。减少 KV heads 可以直接减少缓存,但不会自动减少 query heads 对所有历史 token 的注意力计算。

普通 GQA 的线性投影成本与 $h,H_q,H_{kv},d$ 有关,注意力矩阵乘法仍有 $O(S^2H_qd)$ 成本。VHA 的虚拟头数恢复到原 $H_q$,因此没有把注意力的二次项变成线性,也没有把这个二次项减半。[Sec. 3.1–3.3,PDF pp. 3–5]

1.2 作者针对的经验缺口

在 1.7B 主实验中,GQA 16-8-8 的平均分是 54.09。仅把 KV 头降为两个,GQA 16-2-2 降至 53.09;进一步把 query 头降为八个,GQA 8-2-2 降至 52.38。作者希望同时利用少 KV 的缓存优势和少物理 query 的投影成本优势,再恢复质量。

多头潜注意力(Multi-head Latent Attention,MLA)与多矩阵分解注意力(Multi-matrix Factorization Attention,MFA)是不同的对照:MLA 压缩 KV 潜表示;MFA 通过更大的头维度维持容量。这里 MFA 使用 $d'=256$,与 VHA 同为 1 KB 缓存,却在 1.7B 的训练 FLOPs 上比 GQA 多 49.8%。这些结果定义了本实验中的取舍;论文所谓“不可能三角”是设计动机,并非一个不可能性定理。[Sec. 1、4.2,PDF pp. 2、7–8,Tables 2、5]

从相关工作的角度,VHA 接近“把训练时的结构折叠进部署权重”的思路。它在 query 特征空间和注意力输出空间操作,不像 Talking-Heads 那样直接混合 softmax 前后的注意力权重。[Sec. 2.2–2.3,PDF p. 3]

二:方法——两个可合并的线性模块

2.1 符号、维度和数据流

符号含义1.7B 示例
$X\in\mathbb R^{S\times h}$层输入$h=2048$
$d$每个 query/key/value 头的特征维度128
$H_q$原 GQA 的 query 头数16
$H_q'=H_q/2$VHA 物理 query 头数8
$H_{kv}$KV 组数2
$H_v=H_{kv}H_q'$真正参与注意力计算的虚拟 query 头数16
$W_g^{\mathrm{premix}}\in\mathbb R^{d\times d}$第 $g$ 个 KV 组共享的特征变换两个 $128\times128$ 矩阵
$A,B\in\mathbb R^{H_v\times r}$Postmix 的两个因子$16\times4$
$W_O\in\mathbb R^{H_vd\times h}$拼接虚拟头后的输出投影$2048\times2048$

原图 1:VHA 总体结构与资源取舍

图 1 来自 PDF p. 4。左侧应沿 $X\to W_Q\to$ Premix $\to$ RoPE $\to$ Attention $\to$ Postmix $\to W_O$ 阅读;K、V 分支各只有两个头。右侧两个气泡图的横轴是缓存、纵轴是训练 FLOPs,气泡面积对应平均准确率。方法图中 Q 路径的扩展发生在 RoPE 之前。图 1 右下气泡图把 4B MLA 标为 56.25,与 Table 1 的 56.38 不一致;本文结果表采用 Table 1 的数值。

1.7B 的直观例子:先从 $X$ 计算八个物理头 $Q_1,\ldots,Q_8$;组 1 用自己的变换生成八个虚拟头,组 2 也生成八个。两组分别读取各自的 $K_g,V_g$,总共得到十六个注意力输出。论文表里的 VHA “8-2-2”记录的是物理 Q、K、V 数量;部署后的标准 GQA 则是“16-2-2”。[Sec. 3.2,PDF p. 4;Algorithm 1,p. 20]

2.2 查询前混合(Q Premix):在特征维度上产生组条件化 query

对每个物理头 $i$ 和 KV 组 $g$:

$$ Q_i=XW_Q^{(i)},\qquad \widetilde Q_{i,g}=Q_iW_g^{\mathrm{premix}}. $$

这是原文 Eq. (1),PDF p. 4。矩阵 $W_g^{\mathrm{premix}}$ 在同组的所有物理头间共享;它变换每个头的 $d$ 维特征,并不是一个沿头轴直接混合 query heads 的矩阵。同一个物理头经过不同组的变换,会用不同的查询几何读取不同的 KV。

初始化为

$$ W_g^{\mathrm{premix}}=I+\alpha\epsilon,\qquad \epsilon_{ab}\sim\mathcal N(0,\sigma^2). $$

Table 12 给出的默认值是 $\alpha=0.1,\sigma=1$。作者称之为 near-identity 初始化,其含义是带单位阵锚点的扰动。相关实验并没有显示随机初始化无法训练:三种初始化都稳定,只是这个默认值取得更好的下游平均分,详见后文。

暂略 RoPE 和因果 mask,单个虚拟头的注意力为

$$ Y_{i,g}= \operatorname{softmax}\left( \frac{\widetilde Q_{i,g}K_g^\top}{\sqrt d} \right)V_g. $$

根据矩阵乘法结合律,定义

$$ \widetilde W_Q^{(i,g)} =W_Q^{(i)}W_g^{\mathrm{premix}} \in\mathbb R^{h\times d}, $$

便有 $\widetilde Q_{i,g}=X\widetilde W_Q^{(i,g)}$。所以虚拟头可以提前合并成普通 query 投影。这种等价性不需要把 Premix 穿过 softmax;先合并两个相邻的线性变换即可。RoPE 保留在合并后 query 的同一位置。[Eqs. (2)–(3),PDF pp. 4–5;Appendix G,pp. 17–19]

这个结构限制了训练时可达到的 query 权重集合:同一个 $W_Q^{(i)}$ 被不同组共享,同一个 $W_g^{\mathrm{premix}}$ 被不同物理头共享。它不会提供比任意独立的十六头 query 投影更大的函数集合;它提供的是有组织的参数共享。

2.3 为什么前向等价,训练结果仍可能不同?

令 $D_{i,g}=\partial\mathcal L/\partial\widetilde Q_{i,g}$。原文 Eqs. (12)–(13) 给出

$$ \frac{\partial\mathcal L}{\partial W_Q^{(i)}} =\sum_g X^\top D_{i,g} \left(W_g^{\mathrm{premix}}\right)^\top, $$
$$ \frac{\partial\mathcal L}{\partial W_g^{\mathrm{premix}}} =\sum_i (XW_Q^{(i)})^\top D_{i,g}. $$

第一式把各 KV 组的学习信号汇入同一个物理头;第二式把各物理头的信号汇入同一个组变换。这样的梯度耦合,与分别训练每个 $\widetilde W_Q^{(i,g)}$ 不同。作者据此提出隐式正则化解释,消融结果支持这种训练参数化有用;公式本身说明的是梯度结构,不直接推出泛化误差一定更低。[Appendix G,PDF pp. 18–19]

2.4 线性后混合(Linear Postmix):在头轴上进行低秩残差混合

对单个 token,令 $Y\in\mathbb R^{H_v\times d}$,各行是一个头的输出。Postmix 为

$$ Y'=Y+B(A^\top Y)=(I_{H_v}+BA^\top)Y, \qquad A,B\in\mathbb R^{H_v\times r}. $$

先用 $A^\top$ 把 $H_v$ 个头压到 $r$ 个共享通道,再用 $B$ 映回头空间;每一个特征通道使用同样的头混合矩阵。原始 $Y$ 沿残差路径保留。默认 $r=H_q'/2$,1.7B 用 $r=4$,4B 和 8B-A1B 用 $r=8$。

初始化令 $A$ 为小随机值、$B=0$,于是 $Y'=Y$。需要区分的是:低秩限制作用在增量 $BA^\top$,$I+BA^\top$ 本身一般是满秩矩阵。[Sec. 3.4,PDF p. 5,Eq. (4);Table 12,p. 18]

2.5 为什么 Postmix 也能消失?

令 $\operatorname{flat}(Y)$ 表示按“头、头内特征”顺序拼接的行向量。对任意头混合矩阵 $M$:

$$ \operatorname{flat}(MY) =\operatorname{flat}(Y)(M^\top\otimes I_d). $$

代入 $M=I+BA^\top$:

$$ \operatorname{flat}(Y')W_O =\operatorname{flat}(Y) \underbrace{\left((I+AB^\top)\otimes I_d\right)W_O}_{W_O^{\mathrm{abs}}}. $$

因此部署时只需保存

$$ W_Q^{\mathrm{abs}} =\operatorname{concat}_{g,i} \left(W_Q^{(i)}W_g^{\mathrm{premix}}\right), \qquad W_O^{\mathrm{abs}} =\left((I+AB^\top)\otimes I_d\right)W_O. $$

K、V 投影保持原样;query 的拼接顺序与对应 KV 组一致。前向的 $BA^\top$ 与合并式的 $AB^\top$ 因展平方式而互为转置,并不矛盾。[Eq. (5),PDF p. 5;Algorithm 1、Eqs. (14)–(18),p. 20]

普通 $W_O$ 原本就能混合所有头的特征,因此线性 Postmix 没有扩充最终输出映射的函数集合。它改变的是训练时对这个映射的参数化。另一个数学细节是,吸收等式不要求 $M$ 可逆;即便 $M$ 奇异,上面的乘积仍然成立。附录 H 的可逆性讨论不是完成合并所必需的条件。

“零额外开销”应理解为:相对于同形状的标准 GQA-2,部署图中不再增加 Premix/Postmix 算子。它不表示部署投影仍只有训练时的物理 query 头数。

2.6 计算量:参数节省、训练节省与部署展开

对每 token、每层的 Q 路径,忽略 bias,并以一次乘加计两 FLOPs:

项目普通 $H_v$ 头 QVHA 因子化 Q
参数量$H_vhd$$H_q'hd+H_{kv}d^2$
前向 FLOPs$2hH_vd$$2hH_q'd+2H_{kv}H_q'd^2$

Postmix 另需约 $4H_vrd$ FLOPs。注意 $W_g^{\mathrm{premix}}$ 虽然共享参数,应用时仍要作用于每个物理头,所以参数量和 FLOPs 不能使用同一个省略头数的表达式。

以下为按原文维度展开的教学计算。在 1.7B,$h=2048,d=128,H_q'=8,H_v=16,H_{kv}=2$:

  • Q 参数从 4,194,304 降到 2,129,920,减少约 49.2%。
  • Q 前向从 8,388,608 降到 $4,194,304+524,288=4,718,592$ FLOPs,减少 43.75%。
  • $r=4$ 的 Postmix 为 32,768 FLOPs/token/layer。
  • 合并后的 Q 又成为 $2048\times2048$ 的投影,前向成本回到 8,388,608 FLOPs。

因此“Q 参数约减半”与“Q 计算精确减半”不相同。总训练成本还包含 K/V/O、FFN、注意力矩阵乘法和反向传播,不能把上述 Q 路径比例当成全模型比例。

同样要区分两种 prefill 路径:保留因子化结构时可以使用它的算术节省;完全吸收并调用标准 GQA 内核时,Q 已展开,不能原封不动沿用 Table 5 的 7.9% 训练 FLOPs 降幅。后者仍保留两 KV 头的缓存与 KV 投影优势。[Sec. 3.5,PDF pp. 5–6;Table 5,p. 8;Appendix G,p. 18]

三:实验——三条验证线分别回答什么?

3.1 四规模受控预训练

论文采用 Qwen3 风格骨干,在 FineWeb-Edu 上训练 50B tokens,上下文 4096,BF16。比较 MHA、GQA、MLA、MFA、VHA。模型不是简单地把所有方法的 query 头数设成相同:作者为了参数匹配调整了各方法头数,例如 1.7B 的 MHA 为 12 个 query heads,GQA 为 16 个,VHA 为 8 个物理、16 个虚拟 heads。这会影响计算量比较。

规模层数隐藏维GQA Q/KV 头VHA 物理 Q/虚拟 Q/KVPostmix $r$
0.6B28102416/88/16/24
1.7B28204816/88/16/24
4B36256032/816/32/28
8B-A1B MoE24204832/416/32/28

八项零样本基准为 ARC-Challenge、BoolQ、HellaSwag、OpenBookQA、PIQA、Winogrande、SciQ、Social IQA,涵盖科学问答、阅读理解和常识推理。Avg 是八项分数的等权平均,越高越好,不是把所有题目合在一起计算的总体正确率。

优化使用 AdamW,$\beta_1=0.9,\beta_2=0.95$,weight decay 0.1,梯度裁剪 1.0;小规模实验使用 cosine 学习率,warmup 2000 steps。峰值学习率在 0.6B 为 $2.5\times10^{-4}$,其余三个规模为 $2\times10^{-4}$。[Sec. 4.1,PDF p. 6;Table 12,p. 18]

3.2 工业规模验证

30B-A3B 使用 ERNIE 风格 MoE,46 层、隐藏维 1536,128 个路由专家加 1 个共享专家,每 token 激活 8 个路由专家。训练数据是内部多领域语料,500B tokens、8K 上下文、FP8 混合精度。GQA 用 48 个 query heads/4 个 KV heads;VHA 用 24 个物理、48 个虚拟 query heads/2 个 KV heads,$r=12$。

12 项基准先分成 General、Math & STEM、Coding、Multilingual 四类,各类内部平均,再对四类求宏平均。类大小依次为 4、3、3、2,所以最终 Avg 不是 12 项直接等权平均。缓存列仍按 BF16 计算,不能因为训练使用 FP8 就把缓存数字再减半。[Sec. 4.3,PDF p. 7;Appendices C–D,pp. 17–18]

3.3 长上下文验证是另一套模型

10B-A1B 使用 MiMo-2.5-Flash 风格骨干,全注意力与滑窗注意力的层数比例为 1:3。两模型共享数据、优化设置与滑窗层;全注意力层分别使用 GQA-4 和 VHA-2,滑窗层都用 8 个 KV heads。

阶段上下文长度本阶段 tokens累计 stepsRoPE base
预训练8K500B60K$10^4$
中期长训32K100B72K$6.4\times10^5$
最终长训128K100B96K$5\times10^6$

因此最终结果包含总计 700B tokens 的训练,而不是把一个只在 8K 训练的模型直接外推到 128K。RULER Overall 是 13 项任务的平均:8 项 Retrieval、2 项 Aggregation、2 项 QA、1 项 Variable Tracking;各子类别在 Overall 中的权重不同。HELMET 分类别报告,另用 LongBench-v2、InfiniteBench、BABILong、短上下文八任务及 ProLong bits per byte(BPB,越低越好)补充评估。[Appendix A,PDF pp. 14–15]

四:结果——质量收益来自哪些条件和任务?

4.1 四个规模都提高平均分,幅度并不相同

下表摘录原 Table 1,PDF p. 6。所有分数为八任务平均准确率(%);最后一列为百分点差。

规模MHAGQAMLAMFAVHAVHA−GQA
0.6B50.2250.4249.5250.6651.04+0.62
1.7B53.2454.0952.4954.0654.50+0.41
4B56.6856.9056.3856.9457.71+0.81
8B-A1B54.6054.8653.7954.8154.91+0.05

VHA 在四个规模上均取得最高平均分,并与 MFA 并列最小缓存。这里的核心趋势是:在这个训练预算与配置下,两 KV 头方案可以保持有竞争力的质量。8B-A1B 的 0.05 个百分点差距很小;该规模没有配套重复训练误差条,不能把它解读为有充分统计证据的大幅提升。

各任务并非全部受益。例如 1.7B 的 BoolQ 是 GQA 62.11、VHA 60.64;HellaSwag 是 GQA 51.98、VHA 51.71。VHA 在 OpenBookQA 和 Social IQA 等任务上补回了优势。相同道理,这组 50B-token 实验中的 MLA 排名不能直接推广为 MLA 在其他训练规模上的结论;后面的 30B 实验中 MLA 明显超过了 GQA。

4.2 从压缩损失到恢复:最有解释力的对照

原 Table 2(PDF p. 7)把压缩步骤拆开:

1.7B 配置物理 Q 头注意力 Q 头KV 缓存Avg相对 GQA-8
GQA 16-8-816164 KB54.09—
GQA 16-2-216161 KB53.09−1.00
GQA 8-2-2881 KB52.38−1.71
VHA 8-2-28161 KB54.50+0.41

VHA 比普通 GQA 8-2-2 高 2.12 个百分点,但两者并不是相同注意力头数:VHA 恢复了十六个实际参与注意力的虚拟头。更贴近最终部署形状的基线是 GQA 16-2-2,VHA 比它高 1.41 个百分点。这个比较支持“用结构化训练得到更好的两 KV 头 GQA”,也避免把所有收益都归为额外线性模块在同一计算图上的纯优化收益。

4.3 30B-A3B:平均领先,但编码有明确例外

原 Table 3,PDF p. 7:

方法KV KB/层/tokenGeneralMath & STEMCodingMultilingual四类宏平均
GQA2.0052.8643.0951.8543.2647.76
MLA1.1255.5743.4954.3745.2049.66
VHA1.0056.2546.4053.2345.5650.36

VHA 相对 GQA 为 +2.60 个百分点,相对 MLA 为 +0.70。VHA 的缓存是 GQA 的 1/2,约为表中 MLA 数值的 89%。Math & STEM 的优势主要体现在 MATH;不能把类别平均领先解释为该类别每项都领先。

原 Table 11(PDF p. 17)的完整任务明细如下,保留 shot 设置以说明评价协议:

类别/任务设置GQAMLAVHA
General / MMLU5-shot55.9759.3460.24
General / MMLU-Redux5-shot63.0565.2566.10
General / MMLU-Pro5-shot36.2041.1240.03
General / BBH3-shot CoT56.2356.5758.61
Math / GPQA10-shot CoT34.6231.6234.11
Math / GSM8K4-shot CoT56.4855.4256.25
Math / MATH4-shot CoT38.1843.4248.84
Coding / HumanEval+0-shot50.0052.4447.56
Coding / MBPP3-shot45.0048.6049.20
Coding / MBPP+3-shot60.5462.0762.93
Multilingual / MMMLU5-shot39.8242.6942.95
Multilingual / INCLUDE5-shot46.6947.7148.16

HumanEval+ 是最明显的负向项:VHA 低于 GQA 2.44 点、低于 MLA 4.88 点。它解释了为什么 VHA 的整体均值最高,但 Coding 类仍落后 MLA 1.14 点。

4.4 长上下文:变量追踪最突出,综合任务存在反转

RULER。 下表来自最终 128K checkpoint,原 Table 6,PDF p. 15。每行只改变评测长度,使用同一最终模型:

评测长度GQA OverallVHA Overall差值(点)GQA Variable TrackingVHA Variable Tracking
8K79.190.1+11.044.694.0
16K78.288.5+10.328.493.4
32K72.685.3+12.73.292.6
64K69.581.3+11.84.287.2
128K65.776.8+11.13.083.0
跨长度平均73.084.4+11.4——

这里的 84.4 是五种长度的平均,不是 128K 那一行的分数。VHA 也随长度增长而下降,只是显著慢于 GQA。

128K 处的子类别分数分别是:Retrieval 75.3→79.9,Aggregation 92.5→92.5,QA 31.5→45.5,Variable Tracking 3.0→83.0。按 13 项任务权重作教学分解,变量追踪这一项给 Overall 带来约 $80/13=6.15$ 点,约占 11.1 点总差距的一半以上;检索约贡献 $8(79.9-75.3)/13=2.83$ 点,QA 约贡献 $2(45.5-31.5)/13=2.15$ 点。因为原表分数已四舍五入,分项加总略有舍入差异。

中间 32K checkpoint 同样在 RULER 领先:8K/16K/32K 的 GQA 为 80.9/73.7/69.7,VHA 为 85.7/83.2/73.5。但这种领先没有同步出现在所有综合基准。

HELMET。 根据 Table 6 最终 checkpoint 的五个长度,下面给出本次按原表计算的等权长度平均,保留两位小数:

类别GQAVHA解释
RAG28.1437.48稳定且较大的优势
LongQA15.7016.10小幅均值优势,32K 处 VHA 落后
Summarization6.947.20小幅均值优势,若干长度落后
ICL72.6868.20五个长度均低于 GQA
Reranking0.840.84按表中精度,均值持平

因此作者的“3 胜、1 平、1 负”说的是跨长度类别平均。它不表示每个长度上都赢三类,也不意味着 Reranking 已达到较高绝对能力。

综合基准和语言建模。 原 Table 7,PDF p. 15:

指标32K GQA32K VHA128K GQA128K VHA
LongBench-v2 ↑28.022.318.924.2
InfiniteBench ↑29.927.331.831.5
BABILong ↑36.537.431.138.2
Short 八任务 ↑62.763.762.664.1
arXiv BPB ↓0.5680.5680.5820.564
Books BPB ↓0.8120.8120.8190.804
Code BPB ↓0.2940.2950.3110.296
Web BPB ↓0.7630.7650.7680.760

最终 checkpoint 的确是八项中七项更好;但 LongBench-v2 的相对排序从中期到最终发生反转,InfiniteBench 仍略低。应把结论理解为“这套长训配方的最终结果有明显优势”,而不是所有训练阶段都一致领先。

此外,“缓存减半”只针对全注意力层。若以 $C_{\mathrm{full}}$ 和 $C_{\mathrm{SWA}}$ 表示 GQA 模型两部分的缓存,修改后总缓存为 $0.5C_{\mathrm{full}}+C_{\mathrm{SWA}}$,总比例取决于上下文长度、滑窗大小和层配置,不固定为 1/2。

4.5 模块、秩与非线性的消融

原 Table 4,PDF p. 8:

配置1.7B Avg相对 GQA4B Avg相对 GQA
GQA 基线54.09—56.90—
Premix only53.31−0.7856.62−0.28
Postmix only53.95−0.1457.22+0.32
完整 VHA54.50+0.4157.71+0.81

在 1.7B 的 Premix 上加入 Postmix 提高 1.19 点;在 4B 提高 1.09 点。组合在两种规模下都是最好的,支持模块互补。不过,原文“单个模块都不足以恢复压缩损失”的概括过强:4B 的 Postmix-only 已经超过 GQA 0.32 点。更准确的说法是,两模块联合获得最高分,而不是单模块在所有规模都不能超过基线。

固定 Premix 后,1.7B 的 Postmix 宽度 $r=4,32,128$ 对应平均分 54.50、54.24、54.15;作者报告训练 loss 随 $r$ 增大而下降。这个结果支持较小瓶颈有利于该设置的下游表现。需要注意,此时 $H_v=16$,所以 $\operatorname{rank}(BA^\top)\leq16$:32 和 128 是因子宽度,实际混合矩阵的秩不可能超过 16。它们还改变参数化冗余与优化成本,不能把这组实验解释为混合矩阵实际秩从 32 增加到 128。

SiLU 消融分别为 Premix 加非线性 53.83、Postmix 加非线性 53.79、两者都加 53.67,均低于线性默认的 54.50。它同时给出实践理由和部署理由:这些非线性在此实验中没有改善质量,并且阻止把模块折叠成相邻的固定权重矩阵。

4.6 效率:报告的是 FLOPs 和缓存,并非实测时间

原 Table 5,PDF p. 8,序列长度统一为 4096。FLOPs 为所有 Transformer 层聚合的训练 G/token:

规模GQA FLOPsVHA FLOPs原表报告降幅GQA KVVHA KV
0.6B7.576.96−8.1%4 KB1 KB
1.7B10.619.78−7.9%4 KB1 KB
4B24.3822.52−7.6%4 KB1 KB
8B-A1B10.279.60−6.5%2 KB1 KB

降幅沿用原表,显示值经过舍入。1.7B 的 MFA 为 15.90 G/token,与 VHA 相比仍使用相同 1 KB 缓存。MHA 在若干规模下的 FLOPs 更低,例如 1.7B 为 9.91,原因之一是参数匹配配置中它的 query 头较少;其缓存却为 6 KB。VHA 的贡献是这个实验集合中的联合取舍,不是每个单独资源指标都绝对最优。

作者在 Sec. 3.5(PDF p. 6)推测,访存主导解码中 KV 缓存减少 2–4 倍可带来相应理论速度收益。实际解码还要读模型权重、执行计算并承担内核开销,因此缓存比例不能直接当成端到端速度比。文中的 Tensor Core 大组并行优势也是机制层面的论述,全文未提供硬件、batch 与上下文条件下的实测解码吞吐曲线。

4.7 权重分析:学到了非平凡结构,但主要混合仍在组内

原图 2:Premix 偏离单位阵与组间差异

原 Fig. 2,PDF p. 9,分析的是训练后的 1.7B 模型。两面板横轴均为层号;左侧纵轴为 $\|W_g^{\mathrm{premix}}-I\|_F$,蓝、红表示两个 KV 组,紫线为组平均;右侧纵轴为两个组矩阵的 Frobenius 距离。偏离单位阵的峰值为 5.85,组间距离峰值为 6.13。它们表明所学参数既没有停留在单位阵,也没有变成两个完全相同的组变换。两个峰值发生在不同层,不应直接当成同层的比值。

原图 3:Postmix 残差头间混合

原 Fig. 3,PDF p. 9,分析的是 4B 模型的五个代表层:3、6、17、28、31。每个 Sankey 图从左侧源头流向右侧重组后的头;颜色区分源头所属 KV 组,线宽和饱和度与 $|AB^\top|_{j,i}$ 有关。这里绘制的是残差混合的绝对权重,未呈现正负号,也不包含单位阵路径的完整作用。

可见大量非对角连接,说明它不只是逐头缩放。主导模式是组内重新分配,加上稀疏的跨组连接;不能把图解读成主要依靠密集跨组通信。这与 Premix 的组条件化查询形成一个合理的互补解释。

这些图展示了已训练参数的组织方式。由于 Premix 存在因子分解的基变换自由度,Postmix 又能被 $W_O$ 吸收,矩阵形状不是唯一的功能表征;把这些结构进一步视为收益的因果来源,还需要针对结构的干预消融。原文的组件消融支持模块有用,但没有单独隔离跨组连接的因果贡献。[Sec. 5,PDF pp. 9–10]

4.8 重复训练与初始化:补强了什么证据?

附录 B 使用另一套更高效的训练实现。因此这里的绝对分数不应与主 Table 1 的分数直接相减;同一组内部保持训练与评价设置一致。

0.6B 重复训练每次八任务均分均值 ± 样本标准差
VHA,$n=4$50.42、50.48、50.35、50.3450.40 ± 0.07
GQA,$n=3$49.83、49.79、50.3950.00 ± 0.34

这里的误差是多次训练结果均值的样本标准差,不是标准误,也不是置信区间。它为 0.6B 的平均优势提供了重复运行证据;不能直接用这个误差范围衡量 8B 或 30B 结果。评价为 zero-shot、BF16,可用时使用 normalized accuracy。[Table 8,PDF p. 16]

1.7B 初始化对照中,$\alpha=0.1$、$\alpha=1.0$、完全随机三种方案都稳定,未出现 NaN、Inf 或梯度爆炸;前 2000 步的 loss 和裁剪前梯度范数相近。最终平均分为 54.04、53.45、53.11。默认值比大扰动高 0.59 点,比随机初始化高 0.93 点。这个实验说明单位阵锚定的小扰动是较好的经验选择,而非随机初始化必然失稳。[Tables 9–10,PDF pp. 16–17]

4.9 训练 loss 并不处处第一

原图 4:四种规模的训练 loss

原 Fig. 4,PDF p. 19。四个面板分别对应 0.6B、1.7B、4B、8B-A1B;横轴是训练 steps,纵轴是训练 loss,蓝色实线是 VHA。图从 4K step 开始,使用 $\alpha=0.93$ 的指数移动平均平滑。

VHA 在 4B 收敛 loss 最低;其余三个规模保持接近,作者称距最优方法在 0.01 nats 以内。结合“更大 Postmix 宽度降低 loss 却损害下游均分”的消融,论文证据更接近于有益的参数化约束,而不是单纯以更低训练 loss 解释所有质量提升。[Appendix E,PDF pp. 16、19]

五:结论——值得带走的设计原则与实际取舍

VHA 最值得关注的原则是:把部署算子保持在成熟的 GQA 家族内,把新的结构自由度放在训练参数化中。 Premix 共享 query 基底,并对 KV 组学习不同特征变换;Postmix 通过低秩残差因子改变输出映射的训练方式。训练结束后,这些结构都可被吸收到标准投影中。

证据最扎实的部分是两类结果:一类是可逐式验证的线性吸收;另一类是在所测预算下,VHA 用更少 KV 缓存取得有竞争力乃至更高的平均质量。1.7B 的 GQA 16-2-2 与 VHA 的比较尤其有价值:两者部署形状一致,平均分相差 1.41 点,直接体现训练方式的影响。

实际使用时,需要同时决定是否保留因子化 prefill 路径、如何将权重导出为标准 GQA,以及目标任务是否偏向长距离变量追踪、检索或编码。本文的结果支持继续研究这些选择;端到端速度仍需结合部署硬件与工作负载测量。扩展到更大 MoE、更长训练预算及 RL 训练属于作者提出的后续方向,量化和稀疏化组合也尚未由本文实验验证。[Limitations and Future Work,PDF p. 14]

适合在组会进一步追问的两个问题是:

  1. 在相同训练 FLOPs 而不只是相同 tokens 下,VHA 与直接训练 GQA-2 的质量差距如何变化?
  2. 只保留组内 Postmix 或只保留跨组连接时,质量与长上下文变量追踪表现分别如何变化?

前者决定预算归一化后的收益,后者能把权重可视化提出的机制解释推进到更直接的因果检验。