简单谈谈K3的MoE和Attention

视频生成开源模型架构优化注意力机制

2026-08-05 14:08

45 浏览

0 点赞

AI 速览

月之暗面团队深度解读开源大模型 K3 的 MoE 与 Attention 架构设计，适合关注大模型底层原理与技术演进的开发者。核心事实 - 架构演进：K3 作为 K2 继任者，沿袭既有工作自然演化，融合效果、效率与稳定性改进 - MoE 优化：采用 Stable LatentMoE，引入 SiTU 激活函数抑制膨胀，通过 RMS Norm 稳定训练，使用 QB 分位数均衡实现负载均衡 - Attention 设计：采用 KDA + MLA 混合结构，在 KV Cache、Decoding 成本与效果间权衡，保留 MLA 而非激进替换 - 位置编码：移除 RoPE 变为 NoPE，因 KDA 隐含提供广义 RoPE 效果，同时保持代码兼容性价值点为研究人员和工程师提供 K3 架构设计的详细技术拆解与实验依据，理解大模型组件取舍的逻辑。

作者简介：苏剑林，科学空间（kexue.fm）博主，长期写作深度学习与自然语言处理技术博客，以数学推导见长。他提出的 RoPE（旋转位置编码）被 LLaMA、GPT 等主流大模型广泛采用，其博客中的《Transformer 升级之路》《线性注意力简史》等系列与本文内容直接相关。本文以 Kimi（月之暗面）团队视角写作（文中称"我们发布了 K3"），K3 架构中的 KDA、AttnRes 等设计在其博客中有持续讨论。

概述

本文是 K3 的架构解读，作者以团队视角拆解模型在 MoE 与 Attention 两块的设计取舍。

MoE 部分围绕「Stable LatentMoE」展开：在 LatentMoE 的基础上，用 SiTU（Sigmoid Tanh Unit）与 softcap 运算抑制 SwiGLU 的激活膨胀；在降维与升维之间补一个 RMS Norm 稳定训练；负载均衡从 SignSGD 式更新换成无额外超参的 QB（分位数均衡），用分 bin 近似实现全局分位数。

Attention 部分讨论「KDA + MLA」混合结构：为什么在 DSV4 放弃 MLA 的背景下 K3 仍选择 MLA——效果、训练与 Prefill 成本、KV Cache、Decoding 计算量四个条件的权衡；DSV4 的 Attention 如何被解读为 MLA 的另一种极致；以及 K3 去掉 RoPE 变成 NoPE 的原理：KDA 隐含提供了一种广义的 RoPE。

文章的总体结论是，K3 的每一处改动都不激进，背后都有明确的动机与实验支撑，效果、效率与稳定性之间的协调是架构设计的主旋律。

原文如下：

上个月，我们发布了迄今为止最大的开源模型 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 的形式：

W_{3} (SiLU (W_{1} x) ⊙ W_{2} x) \boldsymbol{W}_3(\mathop{\text{SiLU}}(\boldsymbol{W}_1 \boldsymbol{x}) \odot \boldsymbol{W}_2 \boldsymbol{x})

其中

SiLU (x) = x σ (x) \mathop{\text{SiLU}}(x) = x\sigma(x)

（Sigmoid Linear Unit，亦称 Swish），

σ \sigma

是 Sigmoid 函数。作为主要非线性来源，SwiGLU 经常出现的问题是：

W_{1} \boldsymbol{W}_1

的某一行

w \boldsymbol{w}

与某个输入

x \boldsymbol{x}

同向（Align），导致输出

w \cdot x \boldsymbol{w}\cdot\boldsymbol{x}

非常大。更极端的是，这种现象在

W_{2} x \boldsymbol{W}_2 \boldsymbol{x}

也同时发生，并且发生的位置还一样，于是中间部分出现了

O (∥ x ∥^{4}) \mathcal{O}(\Vert\boldsymbol{x}\Vert^4)

级别的异常值。

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

SiTU (x; β) = \underset{softcap (x; β)}{\underset{⏟}{β \tanh (\frac{x}{β})}} \cdot σ (x) \mathop{\text{SiTU}}(x;\beta) = \underbrace{\beta \tanh\left(\frac{x}{\beta}\right)}_{\mathop{\text{softcap}}(x;\beta)}\cdot\sigma(x)

这样先将门控部分控制在

(- β, β) (-\beta, \beta)

内，其中

β = 4 \beta=4

。进一步压测发现，这样还不能完全杜绝膨胀，所以我们干脆把线性也加上了 softcap 运算，形成了如今的 SiTU-GLU：

W_{3} (SiTU (W_{1} x; β_{1}) ⊙ softcap (W_{2} x; β_{2})) \boldsymbol{W}_3\Big(\mathop{\text{SiTU}}(\boldsymbol{W}_1 \boldsymbol{x};\beta_1) \odot \mathop{\text{softcap}}(\boldsymbol{W}_2 \boldsymbol{x};\beta_2)\Big)

其中

β_{1} = 4, β_{2} = 25 \beta_1=4,\beta_2=25

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

Norm

LatentMoE 与 MoE 的区别是：

MoE: \underset{n 选 k}{\underset{⏟}{d \to D \to d}} \text{MoE:}\qquad\qquad\underbrace{d \to D \to d}_{n\text{选}k}

LatentMoE: d \to \underset{2 n 选 2 k}{\underset{⏟}{d / 2 \to D \to d / 2}} \to d \text{LatentMoE:}\qquad d\to\underbrace{d/2 \to D \to d/2}_{2n\text{选}2k}\to d

即 LatentMoE 会先降维，再做

2 n 2n

选

2 k 2k

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

为了稳定训练，一个朴素的想法是在降维之后（即

d \to d / 2 d\to d/2

的输出）、升维之前（即

d / 2 \to d d/2 \to 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 可以等价写成

SoftmaxAttention (\underset{\tilde{Q}}{\underset{⏟}{Q - DeltaNet (Q, W, W)}}, \underset{\tilde{K}}{\underset{⏟}{K - DeltaNet (K, W, W)}}, V) \mathop{\text{SoftmaxAttention}}\Big(\underbrace{\boldsymbol{Q}-\mathop{\text{DeltaNet}}(\boldsymbol{Q},\boldsymbol{W},\boldsymbol{W})}_{\tilde{\boldsymbol{Q}}},\underbrace{\boldsymbol{K}-\mathop{\text{DeltaNet}}(\boldsymbol{K},\boldsymbol{W},\boldsymbol{W})}_{\tilde{\boldsymbol{K}}},\boldsymbol{V}\Big)

这个等价形式表明，给

Q, K \boldsymbol{Q}, \boldsymbol{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

好文章，需要你的鼓励