Stable LatentMoE：896 选 16 的极端稀疏怎么才不训崩

MoE 的账很好算：把一层大 FFN 拆成很多个专家，每个 token 只激活其中几个，参数量涨上去、单 token 的算力却几乎不变。省算力的关键就是稀疏——激活比例越低越省。但稀疏和稳定是反着来的：越稀疏，路由越容易崩、负载越容易偏、激活越容易炸。多数模型不敢把稀疏度拉太狠，就是怕训不动。

Kimi K3 的 Stable LatentMoE 把这条线推到了 896 个路由专家、每 token 只激活 16 个——激活比例不到 2%。它不是靠运气训下来的，而是靠三个稳定器：归一化、SiTU-GLU、Quantile Balancing。名字里的 "Stable" 不是形容词，是这套东西的全部重点。

先说 "Latent"：专家为什么要压一刀

普通 MoE 的每个专家就是一个完整的 FFN，宽度和主干一样。专家一多，光是这些大矩阵的参数和显存就吃不消。

"Latent" 指的是先把输入下投影到一个较低的潜在维度，让专家在这个压缩空间里算，再投影回来。等于给每个专家的大线性层压了一刀——和多头潜在注意力（MLA）压 KV 是同一个思路，只不过这里压的是专家的 FFN。好处是能在同样的显存预算里塞下多得多的专家（896 个才养得起），坏处是压缩本身会放大数值不稳定，这又把担子交回给下面三个稳定器。

极端稀疏到底难在哪

896 选 16，问题会集中爆发在三个地方：

  1. 路由崩塌：训练早期，路由器会偏爱几个碰巧先学好的专家，把 token 都发过去。这些专家越练越强、越强越被选，剩下几百个专家几乎收不到梯度，等于白养——稀疏度越高，这个马太效应越猛。
  2. 负载不均：就算不完全崩，token 在专家间的分布也很难均匀。分布一偏，专家并行下就有卡空转、有卡排队，静态形状假设被打破，还会引入 host 同步开销。
  3. 激活不稳：潜在压缩 + 极稀疏路由叠加，专家输出的数值尺度容易在层间累积漂移，深了就炸。

Stable LatentMoE 的三个组件，正好一对一地摁住这三件事。

稳定器一：归一化

在路由和专家输出的关键位置插归一化，把尺度拉回可控范围。它治的是第 3 个问题——不让潜在空间里的激活在层间累积漂移。听着朴素，但在 2% 激活比例、2.8T 参数的规模下，少一层归一化就可能是训到一半发散和训到底的区别。

稳定器二：SiTU-GLU

专家内部的激活用 SiTU-GLU——Sigmoid-Tanh Unit 的门控线性单元变体。GLU 家族（SwiGLU、GeGLU 那些）的套路都是用一路信号去门控另一路：

GLU(x)=(σgate(xWg))⊙(xWv)\text{GLU}(x) = \big(\sigma_{\text{gate}}(x W_g)\big) \odot (x W_v)

SiTU 换的是门控函数 σgate\sigma_{\text{gate}}，用 sigmoid 与 tanh 组合出更有选择性的门：既能像 sigmoid 那样把该关的通道干净地压到接近 0，又借 tanh 保留有符号的幅度信息。在极稀疏 MoE 里，专家本来就少、每个都得干净利落，激活的选择性直接关系到专家分工清不清楚——这是治"专家学不出分工"的一手。

（SiTU-GLU 的精确定义以官方报告为准，这里给的是它在 GLU 框架里所处的位置和作用，不是逐字公式。）

稳定器三：Quantile Balancing

这是三个里最关键的，专治路由崩塌和负载不均。

传统负载均衡有两条老路，都别扭：一是额外加一个 auxiliary loss 逼专家用得均匀，但这个 loss 会和主任务打架、拉低质量；二是拿固定阈值卡路由分数，但 896 个专家、分数分布还随训练漂移，固定阈值根本跟不上。

Quantile Balancing 换了个思路：不设固定阈值，而是从路由分数的分布里取分位数来定专家的选取/容量。分布往哪偏，分位数就跟着动，均衡标准是自适应的，不需要手调阈值，也不必靠一个和主任务对着干的 aux loss。配合全均衡的专家并行——静态形状、去掉 host 同步——训练既稳又快。这一手把"越稀疏越容易崩"的马太效应从根上摁住了。

合起来：K3 的三件套

Stable LatentMoE 不是孤立的，它和前两篇讲的东西共同撑起 K3：

  • KDA：3 混合的注意力，扛长上下文、免掉大部分 KV cache；
  • Attention Residuals：让深网络的层间信息按需流动、还训得稳；
  • Stable LatentMoE：把 MoE 稀疏度推到 896 选 16，用三个稳定器保住训练。

据 K3 技术报告，这套组合相比 Kimi K2 带来约 2.5 倍的整体 scaling 效率提升，模型规模来到 2.8T 总参数。这些是报告给出的数字，不是我实测的，具体条件与口径以官方为准。

踩坑

  • 激活参数别只按 16/896 估：除了 16 个路由专家，通常还有常驻的 shared expert 一起参与计算，实际激活参数比"16/896 × 总量"要多。官方没给全所有形状，照最简单的比例去估会偏低。
  • 潜在压缩不是免费午餐：down/up 投影省了显存，但多两次矩阵乘、也多了数值不稳的风险——它和归一化是绑在一起用的，别只抄压缩、不抄稳定器。
  • Quantile Balancing 依赖分布统计：它要在 batch/步之间统计路由分数分布，分布式训练下这块的实现和同步策略是能不能真的"去 host 同步"的关键，不是接个 loss 那么简单。
  • 别照搬 896 选 16：这个稀疏度是配着这三个稳定器、在这个规模上才立得住的。小模型直接抄极端稀疏度，大概率训崩。

小结

Stable LatentMoE 做的是一件很难的事：把 MoE 的激活比例压到 2% 以下还能稳定训练。它的答案不玄——潜在压缩腾出养 896 个专家的空间，归一化摁住激活漂移，SiTU-GLU 让专家分工干净，Quantile Balancing 用自适应分位数替掉打架的 aux loss 和跟不上的固定阈值。三件套里 "Stable" 才是主角。它和 KDA、Attention Residuals 一起，构成了 Kimi K3 那套"把模型往大、往深、往稀疏推，同时不让它散架"的架构。

参考：Kimi K3 技术报告（arXiv.24653）。文中对 SiTU-GLU、Quantile Balancing 的描述是对其作用的说明性重建，精确定义、公式与实验数字以官方报告为准。