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

上个月,我们发布了迄今为止最大的开源模型 K3。
作为 K2 的继任者,K3 并不是一次从零开始的重新设计,而是沿着我们过去一系列工作自然演化而来,并融合了我们对效果、效率、稳定性的一些最新理解和改进。可以说,它是一项持续的、“集大成”的研究成果,而非孤注一掷的豪赌。
这篇文章,我们来聊聊 K3 在架构上的一些设计思路。
写在前面
简单来说,在架构方面 K3 = KDA + MLA + Stable LatentMoE + AttnRes,训练优化器依然是 Moonlight版 Muon,但 Attention 部分的权重改成了 Per-Head 形式来优化。
在这些东西之中,AttnRes 我们之前已经在《Attention Residuals 回忆录》做过详细介绍,而 KDA 我们也已经分享过详细的技术报告《KimiLinear: 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 Mixtureof 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 的输入)这两个位置,都加上一个 RMSNorm。但进一步消融发现,后一个位置的 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 设计,它至少要满足如下条件:
效果不能差于 MLA(保证效果); 训练和 Prefill 成本至少不能超过 MLA(计算效率); KV Cache 比 MLA 小(超长文本); 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 的每一处改动都不算激进和花哨,背后基本都有明确的动机和实验支撑。效果、效率与稳定性之间的协调,依然是架构设计的主旋律。
参考链接

扫描二维码添加小助手微信
关于我们
