今日新知
32 个栏目276 个主题
大模型

苏剑林：简单谈谈 K3 的 MoE 和 Attention

大模型智能 公众号

上个月，我们发布了迄今为止最大的开源模型 K3。

作为 K2 的继任者，K3 并不是一次从零开始的重新设计，而是沿着我们过去一系列工作自然演化而来，并融合了我们对效果、效率、稳定性的一些最新理解和改进。可以说，它是一项持续的、“集大成”的研究成果，而非孤注一掷的豪赌。

这篇文章，我们来聊聊 K3 在架构上的一些设计思路。

写在前面

简单来说，在架构方面 K3 = KDA + MLA + Stable LatentMoE + AttnRes，训练优化器依然是 Moonlight 版 Muon，但 Attention 部分的权重改成了 Per-Head 形式来优化。

在这些东西之中，AttnRes 我们之前已经在《Attention Residuals 回忆录》做过详细介绍，而 KDA 我们也已经分享过详细的技术报告《Kimi Linear: An Expressive, Efficient Attention Architecture》。至于 Per-Head Muon，它其实没有效果上的优势（当然也没有劣势），改用它更多是出于正确性考虑（每个 Head 本就是相对独立的，不应耦合在一块）。

所以接下来，我们主要围绕着 MoE 和 MLA 部分来讨论：前者聊聊我们是如何“驯服”LatentMoE 的；后者聊聊我们在 Attention 上的取舍。

混合专家

K3 所用的 MoE，我们称之为“Stable LatentMoE”，顾名思义，它是“Stable”的“LatentMoE”，其中 LatentMoE 不是新的，它出自《LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts》，好处是可以在大致相同的训练和推理成本下，实现更好的效果，但它的加入也引发了稳定性问题，所以我们提出了一些改进。

SiTU

当前主流 MoE 架构，单个 Expert 通常都是采用 SwiGLU 的形式：

其中 SiLU(x) = xσ(x)（Sigmoid Linear Unit，亦称 Swish），σ 是 Sigmoid 函数。作为主要非线性来源，SwiGLU 经常出现的问题是：W₁ 的某一行 w 与某个输入 x 同向（Align），导致输出 w·x 非常大。更极端的是，这种现象在 W₂x 也同时发生，并且发生的位置还一样，于是中间部分出现了 O(‖x‖⁴) 级别的异常值。

对此，我们先将 SiLU 换成了 SiTU（Sigmoid Tanh Unit）：

这样先将门控部分控制在 (-β, β) 内，其中 β = 4。进一步压测发现，这样还不能完全杜绝膨胀，所以我们干脆把线性也加上了 softcap 运算，形成了如今的 SiTU-GLU：

其中 β₁ = 4，β₂ = 25。给 SwiGLU 引入 Clip 操作已经不新鲜，GPT-OSS、DSV4 就已经引入过 Hard Clip 操作，但我们发现，在同样的界限下，softcap 往往能起到更好的效果，所以我们选择了 softcap。

Norm

LatentMoE 与 MoE 的区别是：

即 LatentMoE 会先降维，再做 2n 选 2k 的 MoE，最后升维，这样训推成本大致相同，但效果略好一些。其中降/升维都是线性投影，原版 LatentMoE 在这里没有加额外操作，所以加上中间的 MoE 就出现了 4 个矩阵连乘的模式，极不稳定。

为了稳定训练，一个朴素的想法是在降维之后（即 d → d/2 的输出）、升维之前（即 d/2 → d 的输入）这两个位置，都加上一个 RMS Norm。但进一步消融发现，后一个位置的 RMS Norm 作用最为本质，所以按照“最小改动原则”，我们只保留了最后一个 RMS Norm，形成了当前 Stable LatentMoE 的形式。

事后，我们还做了更仔细的对比实验，发现多加的这个 RMS Norm，不止能稳定训练，还对效果有神奇作用。具体表现是，在大家都能正常收敛的情况下，加不加这个 RMS Norm 对 Valid Loss 影响不大，但对于某些 Benchmark 来说，不加这个 RMS Norm 会稳定变差。

究其原因，可能是这个 Norm 更好地平衡了 Routed Expert 和 Shared Expert 的比例（实测加了这个 Norm 后，就不用额外加 Scaling Factor 了），也可能是因为 Norm 是个非线性运算（虽然非线性极弱），它的加入无形增加了 LatentMoE 的等效深度。

QB

K3 的 MoE，原本的设想是 448 选 8，而引入 LatentMoE 后，就变成了 896 选 16，虽然稀疏度不变，但总 Expert 数的增加，依然会加剧负载不均衡问题。

跟上一代模型 K2 一样，K3 同样采用 Loss-Free 的负载均衡方案，但之前所用的 SignSGD 式更新规则，在 K3 这么大的总 Expert 数下表现已经不大稳定了，所以我们引入了 QB（Quantile Balancing），优点是数学上更合理，且无额外超参，其细节我们在《MoE 环游记：6、最优分配促均衡》已做过详细介绍。

QB 的核心运算是求全局分位数，但分位数是非线性运算，如果朴素地计算，通信非常大。之前笔者提议的做法是局部求分位数然后全局平均，但后来发现在规模进一步扩大时该做法还是不够准确，所以 K3 最终用了分 bin 近似（直方图估计）：将要求分位数的分数压缩到 0～1 后，分 bin 估计分数的分布，再从分布中读出分位数。

对于分 bin 数，我们发现 10000 个 bin 相比 1000 个 bin，对负载均衡并没有更好的增益，所以推荐使用 1000 个 bin 就够了。由于分布的可加性，我们可以以极低的通信量跨机器、跨梯度累积地聚合分布信息，从而获得全局的近似分位数。

注意机制

K3 的注意力是“KDA + MLA”混合，这里我们主要谈谈 MLA。当然 MLA 早已经不新鲜，但有些细节可以拿出来讨论一下。可能有些读者会嗤之以鼻：DSV4 都放弃 MLA 了，你们居然还用 MLA，是没活了吗？当然不是。K3 用 MLA，依然是经过慎重考虑的结果。

MLA

一年前，笔者写了《Transformer 升级之路：20、MLA 好在哪里?（上）》和《Transformer 升级之路：21、MLA 好在哪里?（下）》，从实验和理论来探究 MLA 的好处，当时笔者给出的判断是：“在相同训练成本和推理成本下，MLA 可能是效果最好的 Full Attention 变体。”

这个判断在今天还对吗？基本对，但有一些小变化。在固定训练成本和 KV Cache 大小时，MLA 依然是近乎最优的 Attention，但现在除 KV Cache 外，Decoding 还有一个新的变数——MTP，或者说推测解码，其思想是计算换速度。然而 MLA 在 Decoding 时表现为 head_dims=512+ 的 MQA，已经提前消耗了大部分算力，所以“MLA+MTP”很容易吃亏。

但是，Attention 的选择是需要综合多方面考虑，MTP 只是一部分，换成别的设计，可能对 MTP 是友好了，但其他方面不见得更好。

MLA 在训练阶段是 192+128（qk_dims 和 v_dims）的 MHA，如果往小了改，比如 128+128 的 GQA8，效果上很难打赢 MLA。即便能打平，GQA8 的 KV Cache 也是 MLA 三倍多，并不现实。注意，MTP 的加入，只是让 Decoding 速度不只取决于 KV Cache 大小，并不意味着 KV Cache 可以随意大，在长文本场景下，KV Cache 依然是越小越好。

如果往大了改，比如换成 256+256 的 MFA（本质上是个 MQA），那么效果确实能追回来了，KV Cache 也不比 MLA 多，但是训练成本却上去了，在 Scaling Law 上大概率是要吃亏的。而且这时候 Prefill 成本也增加了，这一块同样不可忽略——因为在当前主流的 Agent/Coding 场景下，每一轮的 Prefill 长度也不短。

所以，理想的、比 MLA 更好的 Attention 设计，它至少要满足如下条件：

  1. 效果不能差于 MLA（保证效果）；
  2. 训练和 Prefill 成本至少不能超过 MLA（计算效率）；
  3. KV Cache 比 MLA 小（超长文本）；
  4. Decoding 的计算量比 MLA 小（MTP 友好）。

至少在笔者现在看来，目前没有一个简单优雅的 Attention 设计能同时满足以上特点，所以我们必须有所取舍，而在跟 KDA 混合的背景之下，MLA 的部分问题得以缓解，所以我们依旧选择了 MLA。

DSV4

这里我们加个插曲，简单讨论一下 DSV4 的 Attention。DSV4 看上去放弃了 MLA，重新设计了截然不同的 Attention，但如果我们细品，就会发现其实它还有 MLA 的影子，并且符合上述两篇 MLA 博客和刚才说的四个方向。

在《Transformer 升级之路：20、MLA 好在哪里?（上）》中，我们实验发现无形增大的 head_dims 是 MLA 效果的关键，然后在《Transformer 升级之路：21、MLA 好在哪里?（下）》中，我们指出给定 KV Cache 大小下，效果最好的 Attention 是“一个 head_dims 等于 KV Cache 大小、K 和 V 共享的 MQA”。

于是，DSV4 直接将 Attention 换成 head_dims=512、K=V 的 MQA（位置编码是 QKVO-RoPE），以保证效果，而这正是 MLA 的 Decoding 形式。但这样一来，训练和 Prefill 的计算暴涨，且 DSV4 没有线性注意力，每一层都有 KV Cache，即便每 Token 只有 512 维也吃不消。为此，DSV4 引入了 Sparse 和 Compress：Sparse 节省计算量，Compress 进一步压缩 KV Cache，也节省计算量。

所以，与其说 DSV4 放弃了 MLA，不如说它按照上一节提到的几个方向，将 MLA 推到了另一个极致。不过，这种推广不是毫无代价，首先是 Infra 上的复杂性，其次是如此激进地 Sparse 和 Compress，其最优性感觉还有待仔细检验。但总而言之，从 MLA 到 DSV4，更像是一次传承和升级，而不是抛弃和重造。

当然，“Linear+Full”路线也有它的不足之处，所以跟 Sparse 路线相比，究竟谁能走得更远，目前还不得而知。

NoPE

K3 的 MLA 还有另外一个细节：它在保持标准 MLA 结构的前提下，将 RoPE 直接移除，变成了 NoPE。这个改动在 Kimi Linear 中就已经出现了，但 K3 发布后又引来了一些讨论。

首先，K3 是可以把 RoPE 加回去的，只不过加 RoPE 之后效果看上去没啥变化，所以按照最简洁原则，就干脆不加了。但要注意，如果是 K2 这种全 MLA 模型，RoPE 是关键的，去掉 RoPE 会明显变差，K3 能用 NoPE，是因为它是“KDA+MLA”的混合模型。

为什么“KDA+MLA”就可以 NoPE 了呢？我们在《Transformer 升级之路：6、旋转位置编码的完备性分析》推导过，任意正交矩阵的幂都可以用作构建广义的 RoPE，通常说的 RoPE 选择的是简单的旋转矩阵，而 PaTH 尝试了另一种选择——Householder 矩阵，并且表现也不错。

在《线性注意力简史：从模仿、创新到反哺》中我们推导过，在正交情形下，PaTH 可以等价写成：

这个等价形式表明，给 Q、K 加 DeltaNet，也能起到类似 RoPE 的位置编码作用，我们知道 KDA 是更一般的 DeltaNet，所以说 KDA+MLA 的混合模型，本身就自带类似 RoPE、PaTH 的位置编码效果了。我们也可以说，K3 并不是没有 RoPE 了，而是 KDA 隐含提供了一种广义的 RoPE。

此外，还有一个审美上的问题：既然 MLA 不加 RoPE 了，为什么还保留着另外拼接 64 维的做法？

这个有很多原因，比如更好适配现有 MLA 基建，不用重写一套代码。更重要的是，如果直接投影出 576 维的 Latent，然后投影出 192+128 维的 K、V，这样确实是优雅了，但计算会增加，且效果没啥收益，综合起来反而是吃亏的。最后，K3 已经引入了 KDA、AttnRes 等新变量，所以也不想往 MLA 引入太多的变数，毕竟饭要一口一口地吃。

文章小结

本文简单聊了聊 K3 在 MoE 和 Attention 这两块的设计与取舍。总的来说，K3 的每一处改动都不算激进和花哨，背后基本都有明确的动机和实验支撑。效果、效率与稳定性之间的协调，依然是架构设计的主旋律。

参考链接

  • 本文由苏神授权发布，原文链接：https://kexue.fm/archives/11848
  • Kimi K3：https://www.kimi.com/blog/kimi-k3
  • 《Attention Residuals 回忆录》：https://kexue.fm/archives/11664
  • 《Kimi Linear: An Expressive, Efficient Attention Architecture》：https://papers.cool/arxiv/2510.26692
  • 《LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts》：https://papers.cool/arxiv/2601.18089
  • Swish：https://papers.cool/arxiv/1710.05941
  • GPT-OSS：https://papers.cool/arxiv/2508.10925
  • Loss-Free：https://kexue.fm/archives/10757
  • 《MoE 环游记：6、最优分配促均衡》：https://kexue.fm/archives/11619
  • 《Transformer 升级之路：20、MLA 好在哪里?（上）》：https://kexue.fm/archives/10907
  • 《Transformer 升级之路：21、MLA 好在哪里?（下）》：https://kexue.fm/archives/11111
  • MFA：https://papers.cool/arxiv/2412.19255
  • DSV4：https://papers.cool/arxiv/2606.19348
  • QKVO-RoPE：https://kexue.fm/archives/10862
  • 《Transformer 升级之路：6、旋转位置编码的完备性分析》：https://kexue.fm/archives/9403
  • PaTH：https://papers.cool/arxiv/2505.16381
  • 《线性注意力简史：从模仿、创新到反哺》：https://kexue.fm/archives/11033
  • Kimi K3 硬核开源，技术报告都说了啥
  • Kimi K3 技术报告英中对照 PDF 版来了