Gated MLA 计算公式：输出门与 FFN 的区别

一、MLA 先压缩 KV cache

Multi-head Latent Attention（MLA）的核心任务是减少长序列推理时需要保存的 K/V。它为每个历史 token 缓存一个低维 latent，再从 latent 重建各注意力头使用的内容 key 和 value。

对位置 s 的隐藏状态 $\mathbf{x}_s$，可以用下面的简化形式表示：

$$\mathbf{c}_s=W_c\mathbf{x}_s,\qquad \mathbf{k}^{(h)}_s=W_k^{(h)}\mathbf{c}_s,\qquad \mathbf{v}^{(h)}_s=W_v^{(h)}\mathbf{c}_s.$$

其中 $\mathbf{c}_s$ 是缓存的低维 KV latent，$h$ 表示注意力头。实际高效推理可以把部分上投影吸收到查询与输出投影中，不必显式重建完整 K/V；上式用于说明数据关系。

二、Gated MLA 的完整计算公式

第 $h$ 个头先执行因果 softmax attention：

$$a^{(h)}_{t,s}=\operatorname{Softmax}_{s\le t}\!\left(\frac{(\mathbf{q}^{(h)}_t)^\top\mathbf{k}^{(h)}_s}{\sqrt{d_h}}\right),\qquad \mathbf{o}^{(h)}_t=\sum_{s\le t}a^{(h)}_{t,s}\mathbf{v}^{(h)}_s.$$

连接所有头，得到尚未门控的 MLA 输出：

$$\widetilde{\mathbf{o}}_t=\operatorname{Concat}_{h=1}^{H}\!\left(\mathbf{o}^{(h)}_t\right).$$

Kimi K3 再从当前输入计算一个逐通道门，并在最终输出投影之前相乘：

$$\boxed{\mathbf{y}_t=W_o\!\left[\operatorname{Sigmoid}(W_g\mathbf{x}_t)\odot\widetilde{\mathbf{o}}_t\right]}.$$

门控发生在输出端。

它不修改 attention score，也不筛选某个 KV latent；它调节的是 MLA 已经检索到的各通道信息以多大幅度进入输出投影和残差流。

三、FFN 已经使用 SiLU，为什么这里还要加门？

因为它们作用在不同子层，回答的是两个不同问题。以常见的 SwiGLU FFN 为例：

$$\operatorname{FFN}(\mathbf{x})=W_{\mathrm{down}}\!\left[\operatorname{SiLU}(W_{\mathrm{gate}}\mathbf{x})\odot W_{\mathrm{up}}\mathbf{x}\right].$$

FFN 的 SiLU 是特征变换中的非线性：它帮助当前 token 在通道维度生成和重组特征。Gated MLA 的 Sigmoid 则是注意力路径的输出门：它决定从其他 token 检索来的信息应当放行多少。

比较项	Gated MLA 的 Sigmoid	FFN 的 SiLU / SiTU
所在位置	attention 输出投影之前	FFN 内部
控制对象	跨 token 检索结果	当前 token 的通道特征
主要作用	控制信息写回	提供非线性变换
数值范围	0 到 1，适合作为通过比例	可为负且无固定上界，更像特征值

因此，后面的 FFN 即使含有 SiLU，也无法在 attention 写入残差之前替它做选择。两条路径的参数也不同：公式里虽然都可能出现“gate”这个词，但语义不能混用。

四、把两条子层放回完整数据流

$$\mathbf{x}_t\xrightarrow{\mathrm{MLA}}\widetilde{\mathbf{o}}_t\xrightarrow{\operatorname{Sigmoid}(W_g\mathbf{x}_t)\odot}\mathbf{y}_t\xrightarrow{\mathrm{Residual}}\mathbf{x}'_t\xrightarrow{\mathrm{FFN/SiTU}}\mathbf{z}_t.$$

这条顺序揭示了两者不能替代的原因：attention gate 管理“从序列别处读回来的信息”，FFN 管理“已经进入当前 token 状态的特征如何继续计算”。

Kimi K3 的 MLA 层还采用 NoPE，即 query 和 key 不显式加入位置编码。位置敏感和近因信息主要由层间穿插的 KDA 提供，Gated MLA 则保留不受局部窗口限制的全局内容交互。

五、这个门能说明什么，不能说明什么

能说明：门值是输入相关、逐 token、逐通道的；同一层可以对不同 token 采用不同写回强度。
不能说明：门值大不等于 token 更重要，也不等于某个历史位置获得了更高注意力权重。
额外成本：模型增加一个全秩投影 $W_g$ 和逐元素乘法；它换来更细的输出控制，但并非零成本。
训练风险：Sigmoid 若长期靠近 0 或 1，梯度会变小，因此初始化和输入尺度仍然重要。

为什么不用 SiLU 直接充当注意力门？

SiLU 满足 $\operatorname{SiLU}(z)=z\operatorname{Sigmoid}(z)$，输出可以为负且没有固定上界，更适合表达经过非线性变换的特征。Sigmoid 的输出限定在 0 到 1，更直接地表示某个通道的通过比例。

一、MLA 先压缩 KV cache

二、Gated MLA 的完整计算公式

三、FFN 已经使用 SiLU，为什么这里还要加门？

四、把两条子层放回完整数据流

五、这个门能说明什么，不能说明什么

参考资料

继续阅读