November 4, 2026 中 #Attention & MLA #KV Cache #LLM Training

VHA：训练时扩展虚拟头，推理时为什么还能变回普通 GQA？

笔记日期： 2026-11-04 · 作者： Zhongzhu Zhou
论文： Virtual Head Attention
论文作者： Xibo Ding、Guoxia Wang、Jinle Zeng、Jiabin Yang、Sirui Tao、Chao Yang、Dianhai Yu
arXiv： 2610.17959 · 版本与日期： v1，2026-10-13

1. 先把问题问清楚：既然最后能合并，训练时为什么要拆开？

读 VHA 时，我最感兴趣的不是“虚拟头”这个名字，而是一个看起来有些矛盾的设计：训练时加入两组额外的矩阵运算，推理前却又把它们乘回原来的投影矩阵。既然部署时可以用普通的 GQA 表达同一个函数，为什么训练时不直接优化那个最终矩阵？论文的核心就在这里：前向函数相同，不代表优化过程相同。把参数拆成相互共享的因子，会改变梯度怎样流动、不同头怎样共同学习，以及模型在有限训练预算下找到什么解。

先看工程上的出发点。减少 KV 头数可以直接缩小缓存，而且不用改变普通 dense attention 的接口。但很多 query 头开始共用很少的 key/value 表示，质量可能下降。如果顺便把 query 头也砍掉一半，投影计算少了，注意力分布的数量却也少了。VHA 想同时保住这两件事：第一步 query 投影较窄，真正参与 attention 的 query 头数仍然与原来一致。[1，第 3 节]

具体做法分成两步。Premix 先把每个物理 query 头分别变换给两个 KV 组使用，于是原来的一半物理头扩展成完整数量的虚拟头。Postmix 再通过一个低秩残差分支，把 attention 输出沿着头的维度混合。两个模块都保持线性、参数不随 token 动态变化，因此训练完成后可以分别吸收到 query 投影和输出投影中，得到只有两个 KV 头的标准 GQA。

这里的“可精确合并”有严格的对象：它保持的是已经按 VHA 训练出来的模型在实数运算下的函数。论文没有给出一个办法，把任意现成的八 KV 头模型直接无损变成两 KV 头模型；也不是在已有缓存上做不改变任何结果的存储压缩。减小 KV 表示、重新组织训练参数和合并已训练矩阵，是三个不同的问题。

论文在 1.7B 实验中，把八项任务均分从 GQA 的 54.09 提高到 54.50，同时把每层每 token 的 BF16 KV 缓存从 4 KiB 降到 1 KiB。训练计算量的估算从 10.61 降到 9.78 GFLOPs/token。另一个长上下文实验中，检索任务收益更大。不过这三组数字分别属于质量测量、存储公式和算术估计，不能直接拼成一个端到端加速倍数。全文没有提供相应的实测解码吞吐或延迟表。

这篇笔记依据 21 页完整论文及附录 A–H。下面先把张量维度、矩阵合并和成本推导写清楚，再看不同实验到底验证了什么。文中的重绘数据来自原文表格；缓存例子、算术比例和性能上界模型则明确标为我的推导，不当作作者做过的新实验。

图 1：根据论文方法重画的训练与导出流程。导出后额外混合运算消失，但 query 投影会从物理头数量展开到虚拟头数量；图中两条路径的矩阵形状并不相同。

2. 前置知识：物理头、虚拟头和缓存头，分别数的是什么？

2.1 先统一三个容易混淆的数量

设输入为 $X\in\mathbb R^{S\times h}$ ，其中 $S$ 是序列长度， $h$ 是模型隐藏维度，每个头的宽度为 $d$ 。我用 $G$ 表示 KV 组数，用 $m$ 表示物理 query 头数，用 $H_v=Gm$ 表示最后参与注意力计算的虚拟 query 头数。注意，这里的 $G$ 是“有多少个 KV 组”，不是“一个 KV 组里有多少个 query 头”。两种定义在注意力论文里都常见，不先说明就很容易算错。

VHA 固定 $G=2$ ，并让 $m=H_q/2$ ，其中 $H_q$ 是原来的 query 头数。因此 $H_v=H_q$ 。以 1.7B 配置为例，第一步得到八个物理 query 头，扩展之后有十六个虚拟头，分别使用两个 KV 头。论文里“query 头减半”说的是第一步投影，并不意味着最终只计算八张注意力分布。

普通 GQA 中，同组 query 共用 $K_g,V_g$ ，但各自的 query 不同，所以仍然可以得到不同的注意力权重。暂时省略位置编码，第 $(i,g)$ 个头的结果是

Y_{i,g}=\operatorname{softmax}\!\left( \frac{Q_{i,g}K_g^\top}{\sqrt d}+C \right)V_g, \qquad Y_{i,g}\in\mathbb R^{S\times d}. \tag{1}

$C$ 是因果 mask，未来位置对应负无穷。共享 KV 不是要求所有头关注同一批位置，而是限制它们用什么 key 表示做匹配、用什么 value 表示取回信息。减少 KV 头数之后，query 仍可不同，但它们面对的共享表示空间变窄了。VHA 的问题意识，是怎样让这种更强的共享仍然容易训练。

2.2 缓存为什么会少四倍？

BF16 下，每个数占两个字节。一层里，每个历史 token 要存 $Gd$ 个 key 数值和 $Gd$ 个 value 数值。因此，若 batch 为 $B$ 、全部 $L$ 层都是完整注意力、缓存长度为 $S$ ，则

M_{\mathrm{KV}} =2\times2\times BLSGd =4BLSGd\quad\text{bytes}. \tag{2}

当 $d=128$ 时，八个 KV 头对应每层每 token 4,096 字节，两个头对应 1,024 字节。原文把它们写成 4 KB 和 1 KB，这里按二进制单位写作 4 KiB 和 1 KiB。两者的比例就是四，不需要任何压缩率估计。

如果取一个说明性的例子：batch 为一、28 个完整注意力层、131,072 个缓存 token，那么总 KV 分别为 14 GiB 和 3.5 GiB。这只是代入公式的结果，不是作者报告的显存容量测试；权重、临时激活、分配器碎片和算子工作区都还没算进去。理解这一点，才能知道“缓存少四倍”到底覆盖显存的哪一部分。

标准自回归解码一般不缓存历史 query，只缓存 key/value。因此，把当前 token 的 query 扩展成更多虚拟头，不会直接把历史 KV 缓存再撑大。但 query 投影、临时张量和 attention 计算仍要付出成本。不能把“不增加历史 query 缓存”简化成“多出来的头完全免费”。

2.3 这仍然是 dense attention

扩展之后仍然有 $H_v$ 个 query 头，每个头都对全部因果可见位置做注意力。 $QK^\top$ 与概率乘 $V$ 的密集计算没有因为物理 query 减半而自动减半，长序列下的二次复杂度也没有消失。省下来的主要是某些投影计算与 KV 存储，而不是把所有注意力乘法都换成线性复杂度。

因此，VHA 与 top- $k$ 选 token、永久驱逐旧 token、低比特量化缓存是不同层面的技术。基础方法仍保留每个位置的两个 KV 头，没有在读取历史时筛掉一大部分位置。后续可以考虑组合这些技术，但组合后的误差与性能需要单独验证，不能借用本篇的结果自动作保证。

2.4 RoPE 的位置决定能不能直接合并

论文结构图把 Premix 放在 RoPE 之前。对位置 $t$ 的行向量，用 $R_t$ 表示旋转矩阵，那么 query 实际是 $(x_tW_iP_g)R_t$ 。这里可以把两个常量矩阵先乘成 $W_iP_g$ ，再按原顺序做位置旋转。

不能随意把 $P_g$ 穿过 RoPE，因为通常 $P_gR_t\ne R_tP_g$ 。如果把位置相关矩阵与普通投影交换，得到的就未必是同一个注意力分数。后面写伪代码时，我会把 Premix、RoPE 和 attention 的顺序明确列出来。合并之所以成立，是保留了这个顺序，而不是假设任何线性矩阵都能在网络中自由移动。

3. Premix：让一套 query 基底服务两个 KV 组

3.1 扩展发生在特征维度

把物理 query 投影拆成 $m$ 个矩阵 $W_i\in\mathbb R^{h\times d}$ ；给每个 KV 组配置一个 $P_g\in\mathbb R^{d\times d}$ 。于是

Q_i=XW_i,\qquad \widetilde Q_{i,g}=Q_iP_g=X(W_iP_g), \qquad i=1,\ldots,m,\quad g=1,\ldots,G. \tag{3}

这里有两个共享关系。一个物理 query 头会服务所有 KV 组；一个组内的 $P_g$ 又会被所有物理 query 头共用。因此它不是把一串头简单复制并分别训练，也不是沿着头维度做平均。 $P_g$ 变换的是每个头内部的 $d$ 个特征通道，虚拟头由“物理头编号 $i$ 、KV 组编号 $g$ ”共同标识。

把两步投影看成一个有效矩阵，就是

W^{\mathrm{eff}}_{i,g}=W_iP_g. \tag{4}

每个有效矩阵本身仍可能有秩 $d$ ，所以不宜把 Premix 简单叫成“每个 query 都做低秩压缩”。真正的约束是多个有效矩阵共用因子。显式的低秩残差出现在后面的 Postmix。分清这两个位置，才能理解 VHA 为什么既有投影共享，又有另一套低秩头间混合。

从直觉上说，同一个物理 query 提供一个共同基底，两个组变换让它用不同的坐标组合去询问各自的 key/value。能不能用更少的基础投影得到足够丰富的注意力行为，取决于任务是否适合这种共享结构。矩阵分解本身不会自动创造质量收益，它只是规定了训练时可以怎样组织这些有效头。

图 2：用两个物理 query、两个 KV 组说明虚拟头编号。每个物理头分别乘两个组变换，得到四个虚拟头。箭头表示复用和特征变换，不表示把物理头平均。

3.2 为什么从单位阵附近开始？

论文用 $P_g=I+\alpha\epsilon_g$ 初始化 Premix，默认 $\alpha=0.1$ ， $\epsilon_g$ 为高斯扰动。开始时，同一个物理 query 在不同组中的版本较为接近，但又有机会逐渐分化。不同组的 K/V 投影本来就不同，所以 query 接近也不等于两个组一开始必须输出完全相同的结果。

这个设计在自由度与共享之间做了选择。完全独立的虚拟 query 投影更灵活，却放弃了第一步投影节省；完全一样的变换又可能让共享过强。单位阵加扰动希望保住原始 query 的含义，同时允许组特定的结构生长出来。附录实验支持默认初始化在质量上更有利，但没有证明其它初始化都不能稳定训练。

这里也要把“靠近单位阵”的直觉与严格的矩阵范数界区分开。元素级扰动较小，不等于对任意输入都能保证算子扰动很小。本文不依赖这种额外假设推导合并公式；合并只需要矩阵固定且线性。这样就不会把一个初始化经验误当作整个方法成立的数学前提。

3.3 梯度是怎样把不同头绑在一起的？

令 $D_{i,g}$ 表示损失对虚拟 query $\widetilde Q_{i,g}$ 的梯度。对式 3 求导，得到

\frac{\partial\mathcal L}{\partial W_i} =\sum_g X^\top D_{i,g}P_g^\top, \qquad \frac{\partial\mathcal L}{\partial P_g} =\sum_i(XW_i)^\top D_{i,g}. \tag{5}

第一个式子说明，同一个物理 query 矩阵接收多个 KV 组的梯度；第二个式子说明，一个组变换接收多个物理头的梯度。各个虚拟头并不像独立 dense 投影那样各学各的。即使某一步的前向可以写成同一个有效矩阵，下一步参数更新的方向也会受共享因子影响。

这种耦合可能是有益的正则化：多个头被迫利用共同基底，少一些互不相关的自由度。但它也可能造成梯度干扰，尤其当不同组希望修改同一物理头的方向相互冲突时。实验要回答的是，在所选模型、数据和优化设置中，这种约束究竟帮忙还是添乱；矩阵乘法的结合律不能替实验回答这个问题。

4. Postmix：先压到少量混合通道，再加回各个头

4.1 低秩的是残差，不是整个输出变换

对某一个 token，把所有 attention 头的输出按行排成 $Y\in\mathbb R^{H_v\times d}$ 。再取 $A,B\in\mathbb R^{H_v\times r}$ ，按论文正文的头维度约定写成

Y'=Y+B(A^\top Y)=MY, \qquad M=I_{H_v}+BA^\top. \tag{6}

先由 $A^\top$ 把 $H_v$ 个头汇入 $r$ 个中间混合通道，再由 $B$ 分配回各头，最后加上原来的输出。秩受限的是 $M-I$ ，也就是额外的残差分支。因为还有单位阵，整个 $M$ 通常仍是满秩矩阵。看到“低秩 Postmix”就认为所有头的信息必须被压成 $r$ 维，忽略的正是这条恒等残差。

默认设置为 $r=m/2$ 。在 1.7B 配置中， $m=8,H_v=16,r=4$ ，两个因子总共只有 $2H_vr=128$ 个参数。它们沿头维度工作，对所有位置和所有头内通道共享。这个结构不能为每个 token 任意选择一套新的混合权重，却能以很小的参数量改变不同头的输出组合。

初始化时， $A$ 是小随机数， $B$ 为零，所以一开始 $M=I$ 。这有一个细节：开始时 $B$ 通常已经能获得梯度，而 $A$ 的这条分支梯度会因为乘到零矩阵 $B$ 而为零。等 $B$ 开始变化， $A$ 才随之学习。如果两个因子都初始化成零，简单双线性残差的两侧梯度都会卡住。因此，这种不对称初始化有具体的优化原因。

4.2 逐步推导：为什么能吸收到输出投影里？

定义 $\operatorname{rowflat}(Y)$ 为按头的行顺序拼接成一个长行向量。头维度乘 $M$ 之后，有

\operatorname{rowflat}(MY) =\operatorname{rowflat}(Y)(M^\top\otimes I_d). \tag{7}

这里的转置不能省略。一个更容易手算的推导，是把原输出矩阵拆成 $H_v$ 个块 $W_{O,j}\in\mathbb R^{d\times h}$ 。混合后的第 $j$ 个头为 $\sum_iM_{ji}Y_i$ ，总输出于是为

\begin{aligned} \sum_j\left(\sum_iM_{ji}Y_i\right)W_{O,j} &=\sum_iY_i\left(\sum_jM_{ji}W_{O,j}\right),\\ W'_{O,i}&=W_{O,i}+\sum_j(BA^\top)_{ji}W_{O,j}. \end{aligned} \tag{8}

也就是说，可以把每个旧输出块重新线性组合成一个新块。沿行向量约定写成整体矩阵，就是

W'_O=\bigl((I+AB^\top)\otimes I_d\bigr)W_O. \tag{9}

式 6 中是 $BA^\top$ ，式 9 中出现 $AB^\top$ ，原因正是式 7 的转置。实际写公式时，先约定 head 的排列方式，再谈吸收规则，比只背一句“乘进输出投影”可靠得多。交换两个因子的命名当然也能得到等价表述，但一套推导里必须保持一致。

导出时不必真的构造巨大的 Kronecker 矩阵。先算 $C_k=\sum_jB_{jk}W_{O,j}$ ，再算 $W'_{O,i}=W_{O,i}+\sum_kA_{ik}C_k$ 即可。这既保留了分块解释，也说明吸收过程只是乘法与加法，不要求 $M$ 可逆。即使训练后 $M$ 恰好变成奇异矩阵，仍能照样乘进输出投影。

4.3 为什么不顺手加一个激活函数？

一旦在中间加上依赖输入的非线性，通常就不能用一个固定输出矩阵替代整个分支。VHA 保持线性，是为了最终能回到普通 GQA 结构。注意，保持线性的是额外的 Premix/Postmix；attention 的 softmax 仍然是非线性的，整个层绝不是线性模型。

论文做了 SiLU 消融。在 1.7B 设置中，把 SiLU 加到 Premix、Postmix、两者同时加入，平均分分别是 53.83、53.79、53.67，而纯线性版本是 54.50。因此在这些具体变体里，非线性既破坏简单吸收，也没有带来质量回报。这支持作者的选择，但不能扩展成“任何非线性头间交互都不好”的一般结论。

5. 两套编号算法：训练时怎么走，导出时怎么合并？

5.1 算法 1：显式保留虚拟头的前向过程

下面按论文数学方法写成伪代码，省略 batch 维度。head 排列采用“先 KV 组，再物理 query 编号”；query、attention 输出和输出投影必须使用同一顺序。代码块是算法说明，不代表某个仓库的内部实现。

01 输入：X[S,h]，物理投影 W[i]，KV 投影，A、B
02 设 G=2，物理头数为 m，虚拟头数 Hv=G*m
03 K=X*WK，V=X*WV，分别拆成 G 个宽度为 d 的头
04 对每个物理头 i=1,...,m：
05     Q[i]=X*W[i]
06 对每个 KV 组 g=1,...,G：
07     Krot[g]=RoPE(K[g], positions)
08     对每个物理头 i=1,...,m：
09         j=(g-1)*m+i
10         Qvirtual[j]=Q[i]*P[g]
11         Qrot[j]=RoPE(Qvirtual[j], positions)
12         Scores[j]=Qrot[j]*transpose(Krot[g])/sqrt(d)
13         Prob[j]=softmax(Scores[j]+causal_mask)
14         Y[j]=Prob[j]*V[g]
15 对每个 token t：
16     U[t]=transpose(A)*Y[t]             # r × d
17     Ymixed[t]=Y[t]+B*U[t]              # Hv × d
18     Output[t]=rowflat(Ymixed[t])*WO
19 返回 Output，训练时正常反向传播

Premix 位于 attention 之前，Postmix 位于概率乘 value 之后。两者都不改变因果 mask。对当前 token 来说，各头输出本来就只依赖可见历史，再沿头维度线性组合，不会引入未来信息。因果性来自原来的 attention mask 和逐位置混合，不需要发明另一套特殊 mask。

这份伪代码为了可读而列出了中间量，并不意味着训练时必须把所有中间张量完整写回显存。融合、分块和重计算会影响峰值显存与运行时间。仅凭算法里出现一个变量，不能推断底层一定多分配了一整块张量；同样，仅凭理论 FLOPs 少了，也不能推断实际少了相同比例的访存。

5.2 算法 2：把训练好的参数导出为普通 GQA

导出时参数固定。接下来只合并那些位置无关的线性映射，RoPE、因果 mask 和其余网络层保持原来的语义与相对顺序。

01 输入：训练好的 W[i]、P[g]、WK、WV、WO、A、B
02 固定头编号 j=(g-1)*m+i
03 对每个组 g 和物理头 i：
04     WQ_export[j]=W[i]*P[g]             # h × d
05 将 WO 拆成 Hv 个块 WO[j]，每块为 d × h
06 对每个中间通道 k=1,...,r：
07     C[k]=sum_j B[j,k]*WO[j]
08 对每个虚拟头 i=1,...,Hv：
09     WO_export[i]=WO[i]+sum_k A[i,k]*C[k]
10 按记录的头顺序拼接 WQ_export
11 按相同头顺序拼接 WO_export
12 保留 WK、WV、RoPE、因果 mask 及其余网络层
13 将标准 GQA 配置为 Hv 个 query 头、G 个 KV 头
14 删除显式 Premix 与 Postmix 运算
15 导出固定的推理参数

这个算法保证的是实数算术下的函数相等。浮点矩阵乘法重新结合后，舍入顺序可能改变；若再量化，合并后权重的数值范围与量化尺度也可能变化。所以“可精确吸收”不应被解释成在任何 BF16 实现中 logits 都逐比特一致，更不能自动保证随机采样生成的整段文本完全一样。这是数学结论与数值实现之间的边界，不是本文声称观察到了作者实验中的故障。

另一个关键变化是投影宽度。导出后的 $W_Q$ 有 $H_v$ 个头，而不是训练第一步的 $m$ 个物理头。展开后的形状使它能接入普通 GQA 推理路径。若部署时选择保留因子，可以继续享受较窄物理投影的算术优势，但也保留了额外操作。它们是两种执行方案，成本必须分别计算。

6. 成本拆开算：参数量、FLOPs、缓存和延迟不是一个比例

6.1 同样叫“省一半”，参数和计算其实不一样

原本 dense query 投影有 $Gmhd$ 个参数。分解后，物理投影有 $mhd$ 个参数，组变换有 $Gd^2$ 个参数。两者之比为

\rho_{\mathrm{param}} =\frac{mhd+Gd^2}{Gmhd} =\frac1G+\frac{d}{mh}. \tag{10}

若一次乘加计两次 FLOPs，物理投影需要 $2Smhd$ ，组变换作用到全部物理头上需要 $2SGmd^2$ ，因此 query 前向计算比例为

\rho_{\mathrm{Q,FLOPs}} =\frac{2Smhd+2SGmd^2}{2SGmhd} =\frac1G+\frac dh. \tag{11}

参数公式里有 $m$ ，计算公式对应项却没有，原因是一个 $P_g$ 虽然只存一次，却要应用到 $m$ 个物理头。参数共享减少存储，不会让每次应用都免费。

代入 $h=2048,d=128,m=8,G=2$ ：分解 query 有 2,129,920 个参数，原 dense query 有 4,194,304 个，比例约 50.78%；每 token 前向 query FLOPs 则是 4,718,592 对 8,388,608，比例为 56.25%。这意味着 query 路径节省约 43.75% 的前向算术，而不是整个 Transformer 快了 43.75%。KV 投影、attention 矩阵乘法、输出投影、FFN 与反向传播都还在外面。

图 3：按式 10–11 计算的示例，只统计 query 路径。分解后的参数和 FLOPs 比例不同；全部折叠后，标准 query 投影又恢复为十六头的宽度。这里没有模型实测时间。

6.2 放到整个模型里，训练收益就小得多

论文表 5 把其它 Transformer 运算算进去后，相对所选 GQA 基线报告的训练 FLOPs 降幅为 8.1%、7.9%、7.6%、6.5%，分别对应 0.6B、1.7B、4B、8B-A1B。1.7B、长度 4,096 的一行是 10.61 对 9.78 GFLOPs/token。直接用这两个已经四舍五入的数计算，约为 7.8%；讨论表中 7.9% 时应保留原文精度背景，不把最后一位当作独立的新结果。

不同基线的成本也不是齐刷刷地比 VHA 高。MHA 在若干配置里的训练 FLOPs 更低，但缓存更大。因此作者强调的价值在于缓存、质量和计算量三者的综合权衡，而不是每一项资源都对所有方法绝对最优。MFA 是另一个有用参照：1.7B 时同样是 1 KiB 缓存，但平均分 54.06，训练估计为 15.90 GFLOPs/token。

最值得保留的一条边界是：折叠成普通展开的 query GEMM 后，式 11 的因子化 query 节省不再原样存在。相比八 KV 头基线，两 KV 头仍能节省 KV 投影；但训练/未折叠 prefill 图与完全折叠的推理 prefill 图，不应共用同一套 query 成本。论文的可部署性论证是清楚的，读者却仍需要对具体执行图分别计时。

6.3 缓存少四倍，为什么不等于总延迟少四倍？

假设基线时间中有比例 $f$ 与读取 KV 缓存直接成正比，其余部分暂时不变。若 KV 部分减少 $c$ 倍，一个简化的 Amdahl 模型给出

\mathrm{speedup}_{\mathrm{ideal}} =\frac{1}{(1-f)+f/c}. \tag{12}

当 $c=4$ 时，如果 KV 读取占原时间的 80%，理想整体加速是 2.5 倍；如果只占 50%，就是 1.6 倍。这些是解释公式的假设例子，不是 VHA 实测结果。权重读取、FFN、MoE 通信、kernel 启动和实际内存访问效率，都可能限制收益。

缓存更小还有另一种价值：在同一显存预算中容纳更多并发请求。这可能提升吞吐，却不一定同比例降低单请求延迟。说“理论上有两到四倍的内存受限解码收益”时，应说明依赖的瓶颈假设；不能把它直接改写成“论文测得整体最快四倍”。本文没有这样的测量证据。

6.4 混合滑窗模型不能照搬完整注意力的总缓存比例

长上下文实验的完整注意力层把 KV 从四头减到两头，但滑窗层仍保留八头；完整层与滑窗层比例为 1:3。设完整上下文长度为 $S$ ，滑窗为 $W$ ，忽略相同的层数、batch、头宽和字节因子，则

\frac{M_{\mathrm{VHA}}}{M_{\mathrm{GQA}}} =\frac{\tfrac14(2S)+\tfrac34(8W)} {\tfrac14(4S)+\tfrac34(8W)} =\frac{0.5S+6W}{S+6W}. \tag{13}

只有当完整历史占主导时，这个比例才趋近二分之一。举例取 $S=W$ ，比例约为 0.929，因为大量未变化的滑窗缓存稀释了收益。真实总缓存还要看窗口长度和分配策略。完整注意力部分少一半，不能不带条件地宣传成整个混合模型的缓存少一半。

图 4：本文推导的说明性曲线，并非作者性能实验。左图是式 12 的时间占比模型；右图是式 13 的混合缓存比例。两者都说明局部缩减率与整体收益之间还隔着一个分母。

7. 实验拆开读：哪些结果属于同一组比较？

7.1 小模型实验说明了怎样的缓存—质量权衡

主实验采用 Qwen3 系列结构，规模为 0.6B、1.7B、4B、8B-A1B，训练数据为 50B FineWeb-Edu token，上下文长度 4,096，使用 BF16。表 1 报告八项任务的平均分。下面的增量是平均分之差，也就是百分点，不是相对百分比。

规模	GQA 均分	VHA 均分	VHA − GQA	GQA / VHA 缓存 KiB
0.6B	50.42	51.04	+0.62	4 / 1
1.7B	54.09	54.50	+0.41	4 / 1
4B	56.90	57.71	+0.81	4 / 1
8B-A1B	54.86	54.91	+0.05	2 / 1

四个规模的平均分方向一致，但幅度差别很大。8B-A1B 只高 0.05 个百分点，不能与后面长上下文的两位数收益用同样的确定语气讨论。各单项任务也不是全赢。例如 1.7B 上，VHA 的 BoolQ 为 60.64，低于 GQA 的 62.11；Winogrande 为 54.70，也低于 55.25。平均分改善说明整体权衡更好，不等于所有能力同时增强。

同一张表还给出 MHA、MLA、MFA。1.7B 时三者平均分分别是 53.24、52.49、54.06。把质量与缓存画在一起，VHA 在小缓存一侧的位置确实很有吸引力。但这些是为目标参数规模配置的完整架构，不同方法的内部头数并非完全相同。因此图中比较的是不同设计方案，而不是在其余每个内部维度都固定时，只替换一个算子的纯消融。

图 5：论文表 1 的 1.7B 数据重绘。横轴是每层每 token 的 BF16 KV 缓存，纵轴为八任务平均分。原表没有提供这些点的置信区间，因此图中也不添加误差条。

7.2 只改头数，与同时加入两种混合，结果不同

表 2 很适合回答“是不是把头数配置换一下就行”。16 query、8 KV 的普通 GQA 得到 54.09；保留 16 query、把 KV 改成 2，分数降到 53.09；再把 query 也改成 8，得到 52.38。VHA 使用 8 个物理 query、2 个 KV，但扩展到 16 个虚拟 query，最终为 54.50。

16-query/2-KV 基线与 VHA 的比较，能说明针对两 KV 头推理架构，训练参数化是否有帮助。8-query/2-KV 基线则提醒我们，只缩窄第一步投影并不是完整方法。不能因为它与 VHA 的“物理头数”相同，就说两者最终 attention 容量完全一样：VHA 恢复了十六个注意力分布，这本来就是它要付出的计算和获得的表达能力。

表 4 又拆开 Premix 和 Postmix。1.7B 上，只有 Premix 为 53.31，只有 Postmix 为 53.95，两者都有为 54.50，原 GQA 为 54.09。4B 上对应为 56.62、57.22、57.71、56.90。完整设计在两组中都最好，但 4B 的 Postmix 单独使用已经超过 GQA。因此不能写成“两部分单独用都没有效果”，更准确的说法是它们组合之后效果更稳定、更强。

图 6：表 2、表 4 的头数与组件消融。GQA 标签写的是 query/KV 数量；VHA 则是八个物理 query、十六个虚拟 query。虚线为 1.7B 的八 KV 头 GQA 基线。

7.3 30B 结果来自另一套训练环境

较大的实验使用 ERNIE 风格 30B-A3B MoE，内部语料 500B token，长度 8,192，FP8 训练。GQA 配置是 48 query、4 KV；VHA 是 24 物理 query、48 虚拟 query、2 KV。比较表中的缓存仍按 BF16 计算，不能因为训练用了 FP8，就自行把缓存数字再除以二。

表 3 把十二项 benchmark 分成四类，再对四个类别均分做平均。因此 47.76、49.66、50.36 这三个 GQA、MLA、VHA 总分，是“类别宏平均”，不是十二个任务一律等权的平均。VHA 比 GQA 高 2.60 点，比 MLA 高 0.70 点；但 coding 类别只有 53.23，低于 MLA 的 54.37。

详细结果更能说明这种权衡。HumanEval+ 上，VHA 为 47.56，低于 GQA 的 50.00 和 MLA 的 52.44；MATH 上却是 48.84，高于 38.18 和 43.42。只读总分会把能力之间的取舍抹平。数据、结构、精度和 token 预算又都不同，因此也不应把这组大模型分数直接接到前面小模型曲线上，声称得到了统一的 scaling law。

图 7：左图重绘表 1 的平均分增量，右图重绘表 5 的训练 FLOPs 降幅。右边是算术估算，不是训练时间或推理延迟。

图 8：30B-A3B 的四个类别均分，来自表 3。VHA 在三个类别领先，但 coding 不如 MLA；原文总分对四类等权。

7.4 长上下文不是拿同一个 30B 模型简单拉长

长上下文实验又是一套独立设置：10B-A1B，MiMo-2.5-Flash 风格的混合注意力，完整层和滑窗层为 1:3。只在完整注意力层使用两 KV 对四 KV 的改变，滑窗仍是八 KV。若把它描述成“30B 模型延长到 128K 的结果”，就把两组不同实验混在一起了。

训练分三段：8K 阶段使用 500B token，32K 阶段再训练 100B，128K 阶段再训练 100B，同时按阶段调整 RoPE base。最终 128K checkpoint 在 8K、16K、32K、64K、128K 五个评估长度的 RULER 总分，VHA 为 90.1、88.5、85.3、81.3、76.8；GQA 为 79.1、78.2、72.6、69.5、65.7。五个长度平均相差约 11.4 点。[1，表 6]

其中 variable tracking 特别突出：128K 时 VHA 为 83.0，GQA 只有 3.0。RULER 这里有十三项任务；如果总体按任务等权，仅这一项的 80 点差距就贡献约 $80/13=6.15$ 个总体分数点。这不是否定总体收益，而是帮助定位收益的来源：后续值得研究的是哪一类跨位置依赖被保留下来了，而不是把总体差距想象成十三项任务都均匀提高了同样多。

图 9：表 6 最终 checkpoint 的 RULER 数据。右图的 variable tracking 只是十三项任务之一，却贡献了相当一部分总分优势。曲线连接原始测量点，不是外推规律。

其它长上下文结果没有这么整齐。最终 checkpoint 上，LongBench-v2 是 24.2 对 18.9，BABILong 是 38.2 对 31.1，但 InfiniteBench 是 31.5 对 31.8，VHA 略低。HELMET 的 RAG、长问答和总结类别更有利于 VHA，in-context learning 却不总是如此。前一个 32K checkpoint 甚至在 LongBench-v2 上是 22.3 对 28.0，在 InfiniteBench 上是 27.3 对 29.9。

这说明训练阶段很重要。最终长上下文结果确实可观，但收益并非在每个中间阶段、每种任务上都已经出现。保留这些反例，能避免把“最终选择的一组 checkpoint 很好”写成“结构本身在所有阶段都必然更强”。

7.5 鲁棒性附录补了什么，又没有补什么？

附录 B 在另一套高效训练设置中重复了 0.6B 实验。四次 VHA 的平均分为 50.40、标准差 0.07；三次 GQA 的平均分为 50.00、标准差 0.34。其中一次 GQA 达到 50.39，已经接近 VHA 均值。这为小规模的排序增加了支持，但也表明单次结果有波动；由于训练设置不同，不能把这些分数直接和主表 1 混起来算一个更大的样本量。

初始化实验比较了默认单位阵扰动、更大扰动和完全随机的 Premix。报告的前 2,000 步稳定性检查中都没有出现 NaN/Inf，但默认设置的下游质量较好。因此可以说默认初始化更合适，不能说其它初始化一定数值不稳定。早期没有溢出，也不等于最终泛化一定好；稳定性与质量是不同层面的观察。

8. 局限：目前哪些结论还不能从论文里拿到？

第一，理论资源收益与系统实测之间还有缺口。论文对 FLOPs 和缓存做了比较细的核算，但没有给出不同设备、batch、上下文长度和导出策略下，端到端 prefill/decode 的完整延迟矩阵。兼容普通 GQA kernel 能降低接入成本，却不能单独决定加速幅度。极少的 KV 头还可能影响算子的并行工作划分，单纯数读取字节并不足以预测所有硬件行为。

第二，训练覆盖面仍有限。最大模型实验是 500B token，长上下文分阶段累计约 700B。它们已经是有分量的实验，但尚不能代表更长的万亿级训练预算、所有模型家族，或经过大量强化学习后的行为。论文也把更大规模训练，以及与量化、稀疏注意力的组合留给后续。不能因为这些组合在直觉上兼容，就把其质量和速度写成已被验证。

第三，架构比较不只改变一个变量。近似匹配参数规模、token 数和缓存大小，不等于同时匹配训练 FLOPs、优化器状态、函数表达能力和每个基线最适合的超参数。头数消融和模块消融缓解了这个问题，但没有消除全部替代解释。某种参数化也许在当前学习率下特别合适，而直接训练最终 GQA 矩阵的基线需要另一套调参预算。

第四，不同结果的随机性信息不均衡。小模型有附录重复，大模型与长上下文没有同样完整的多 seed 覆盖。0.05 点与 11.4 点的差异，不应使用同一种置信语气。评测如何平均也很重要：单个大幅获益任务可以推动总分，类别宏平均又与任务宏平均不同。比较前要先确认分母。

最后，“精确合并”不能无限扩展。固定线性映射可以组合，token 依赖的混合、非线性分支或错误的 RoPE 顺序通常不能直接套用同一证明。浮点和量化又有独立的数值问题。这篇论文提供的是一个可回到常规部署结构的训练架构，而不是所有压缩前后 checkpoint 都等价的保证。

9. 批判性分析：更有说服力的解释应该落在优化过程上

9.1 输出投影本来就在组合不同头

普通 attention 的 dense 输出投影已经能把不同头的信息组合起来。Postmix 不是第一次给这一层创造“跨头通信”，因为它最后完全可以吸收到 $W_O$ 中。对于不受约束的 dense 输出矩阵，它没有凭空扩大最终线性映射的函数集合。真正值得解释的是：额外因子是否给优化提供了更好的路径、是否施加了合适的正则化，以及它怎样配合共享 query 起作用。

可以用一个简化推导看到差异。设有效矩阵 $F=WP$ ，损失对它的梯度为 $D=\partial\mathcal L/\partial F$ 。用普通 SGD 同时更新两个因子，步长为 $\eta$ ，得到

W_+=W-\eta DP^\top,\qquad P_+=P-\eta W^\top D. \tag{14}

把新因子相乘，展开后是

\begin{aligned} F_+={}&F-\eta\left(DP^\top P+WW^\top D\right)\\ &+\eta^2DP^\top W^\top D. \end{aligned} \tag{15}

它显然不同于直接对 $F$ 做 $F-\eta D$ 。因子的尺度和方向改变了有效矩阵的更新；VHA 的跨头共享还会引入更多耦合。这个例子只解释为什么“前向一样”并不意味着“训练一样”，不是作者 AdamW 更新的精确公式，更没有证明这种变化总是有利。

对我来说，这比“额外混合增加表达力”更接近论文真正有价值的部分。一个部署时看起来普通的矩阵，可能通过不同的参数空间走出完全不同的训练轨迹。如何定量区分预条件、正则化和共享先验的贡献，是后续比继续堆一个平均分更有解释力的问题。

9.2 观察到矩阵不同，不等于证明功能不同

论文展示 Premix 逐渐偏离单位阵、不同组矩阵逐渐分化。这些可视化描述了学习到的参数，值得看；但因子分解存在换基自由度。对任意共同的可逆矩阵 $T$ ，都有

W_i\longmapsto W_iT,\qquad P_g\longmapsto T^{-1}P_g \quad\Longrightarrow\quad W_iP_g\ \text{不变}. \tag{16}

有效 query 投影完全一样时，原始因子的数值距离仍可能改变。初始化与 weight decay 会偏好某些分解，但它们不自动把因子之间的距离变成一个功能多样性的直接度量。更强的分析应该看有效 query–key 算子、相同输入上的注意力分布，以及有针对性的因子干预。这里是在提出后续实验方向，没有声称已经做了这些实验。

Postmix 的非对角元素也一样：它们证明残差分支确实在混合头坐标，却不能单独证明获得了原 dense 输出投影无法表达的新能力。更准确的问题是，低秩残差是否让训练更容易到达一个好的有效矩阵，而不是它是否首次允许两个头的信号相遇。

9.3 bottleneck 变宽，矩阵秩也不可能无限增长

论文在 1.7B 上把 $r$ 设为 4、32、128，平均分为 54.50、54.24、54.15。较大设置的训练 loss 可以更低，评测却更差，这与正则化解释相容。但这里头数是十六，残差矩阵只有 $16\times16$ ，因此

\operatorname{rank}(BA^\top)\leq\min(r,16). \tag{17}

$r=32$ 与 $r=128$ 不可能分别得到秩 32 和 128 的头间矩阵；两者都已经可以表示满秩的十六维残差。继续加宽改变的是过参数化程度和优化几何，不是矩阵最大秩还在按同样速度增长。把整组实验都称为“增大秩”，会混淆两种效应。

如果后续在 $r=1,2,4,8,16$ 附近做多 seed 比较，再与更宽的冗余因子化比较，就更容易分开“限制矩阵秩”的影响和“改变参数化宽度”的影响。当前结果支持小 bottleneck 是一个不错的选择，但其机制解释仍有进一步收紧的空间。

9.4 可合并架构最适合谁？

如果能够改预训练结构，又希望部署侧继续使用常规 GQA，VHA 很有吸引力：训练阶段试验共享与重参数化，服务阶段不一定要维护新算子。相反，只有一个任意现成 checkpoint 的部署团队，不能通过补上恒等模块再合并，就凭空得到论文里的质量—缓存优势。训练必须先学会利用两 KV 头与共享因子。

即使拿到了训练好的 VHA 模型，最优执行图也可能随工作负载变化。全部折叠方便接入与算子调度；保留因子则可能保住较窄 query 投影的算术优势。更完整的系统证据应同时报告两条路径的 prefill 延迟、decode 延迟、batch 容量与数值一致性，而不是把其中一条路径的优点全部算到另一条上。

我认为这篇论文最值得留下的思想，就是把训练参数化与推理参数化分开设计。它的经验收益可以由不同的学习轨迹解释，并不需要宣称矩阵合并凭空赠送了额外非线性能力。把这一点说清楚，既能保留方法的价值，也能让下一步该验证什么更加明确。

10. 结论

VHA 用较少的物理 query 投影，通过组特定的线性变换恢复完整的注意力头数，再以低秩残差混合头输出。训练完成后，两部分都能合并进 query 与输出投影，部署为两个 KV 头的普通 GQA。推理结构看起来熟悉，但它经历过的训练参数空间与直接训练最终矩阵不同。

现有结果支持更有利的缓存—质量权衡、适度减少的训练 FLOPs 估计，以及某些长上下文检索任务上的明显进步。它尚未证明端到端实测加速两到四倍，也没有证明任意已有模型能无损转换，或者每个任务、每个训练阶段都会提升。后续最有价值的证据，是把这套漂亮的合并代数与具体工作负载的系统测量连接起来，并更细地拆解共享和重参数化究竟怎样帮助优化。

参考资料

Xibo Ding、Guoxia Wang、Jinle Zeng、Jiabin Yang、Sirui Tao、Chao Yang、Dianhai Yu：Virtual Head Attention，arXiv:2610.17959v1，2026-10-13。论文记录 · 本次精读 PDF · 全文。本文全部 VHA 实验数字均来自这一版本及其附录。
作者提供的 VHA 官方资源入口。