引言与一些碎碎念
没想到我的文章有这么多人看( ,深受感动,于是快马加鞭的开始赶这一篇了,下一篇和下下篇的主题也差不多确定了,现在的我文思泉涌(bushi
说起来今天是八一建军节呢( ,祝大家节日快乐捏
昨天不知道为什么失眠了熬穿了,等我写完这篇就去睡觉……
这篇主要是讲解MoE相比于FFN,在模型训练时 FLOPs / 参数对比以及推理时的显存有哪些变化,这个是常错题(大概,以及目前常用的各类MoE变体以及他们独特的机制。昨天有佬友想让我讲解一下Kimi k3的MoE框架,这篇文章也会讲到。
最后还有一个题外话,也就是上一篇所讲到的 Temperature = 0 的情况下,为什么每次输出还会不一样?这与 MoE 的路由机制有很大的关系,也牵扯到我们这篇会谈的一个 MoE 变体:Soft MoE。
如果你还没看过上半部分,这里直达!
上半部分直达!
MoE FLOPs/参数量,推理显存相比FFN的变化
FLOPs与参数量
我们先来讲下基础FFN的FLOPs/参数量作为baseline
我们假设一个FFN一共有L层感知机,我们把每层看作一个普通的 ax + b 公式,这样每一层就拥有乘与加两个浮点操作
而每一层token要经历两次变化:D \rightarrow d_{\text{ff}} \rightarrow D,这里 d_{\text{ff}} 指的是每层的维度,D指的是token的维度,我们这里假设 d_{\text{ff}} = 4D。因此每层每个 token 的计算量大致是:2D d_{\mathrm{ff}}。分别为一个 D \times d_{\mathrm{ff}} 的矩阵乘法和一个 d_{\mathrm{ff}} \times D 的矩阵乘法。
再乘以 T 个 token、L 层,以及乘加操作按 2 个 FLOPs 计算:
\mathrm{FLOPs}_{\mathrm{FFN}}^{\mathrm{dense}}
\approx
2 \cdot T \cdot L \cdot D \cdot d_{\mathrm{ff}} \cdot 2
代入 d_{\mathrm{ff}}=4D:
\mathrm{FLOPs}_{\mathrm{FFN}}^{\mathrm{dense}}
\approx 16TLD^2
接下来我们来看MoE的Flops是如何变化的
我们知道MoE本质上是多个FFN的集合,我们假设MoE 把一个 FFN 换成 N 个专家:
\mathrm{Expert}_1,\mathrm{Expert}_2,\ldots,\mathrm{Expert}_N
那么我们就需要在FFN的Flops基础上乘以一个N。
但是,看了上半部分的佬友们都知道,MoE路由不可能让token经过所有专家,而是会进行打分选取Top k 的专家,也就是说,每次token经过FFN时,只会激活k个专家参与前馈传播。因此我们还需要在Flops基础上乘以k/N 来筛除没有参与计算的专家。
不难分析,N和k/N相乘后N被消掉了,只剩下了k这个参数。也就是说,使用MoE后,Flops的计算量变为原来的k倍,与专家数量无关。
至于参数量,一个普通 FFN 的参数量约为:
D\cdot 4D+4D\cdot D=8D^2
D·4D和4D·D分别代表 D \rightarrow d_{\text{ff}} \rightarrow D 的两步计算所需要的矩阵的大小,两者相加就是FFN的参数量。
而MoE中由于同时存在N个专家,故参数量也需要乘以N,由于不管专家是否激活,其所携参数是固定存在的,故参数量的变化与k无关。
总结一下就是一句话:MoE 的"参数容量N倍, 计算量k倍"
推理显存
在MoE推理显存计算场景,经常会出现一个陷阱问题:MoE 推理时显存是按"active 参数"还是"总参数"算?有人可能会觉得MoE推理时只使用了激活专家的参数,因此GPU上只需要存储当前激活专家的参数即可,因此是按active参数算。但这显然是错误的,因为GPU没有办法提前预知路由会选择哪几个专家,下一个 token 可能路由到任何 expert,因此所有专家的参数必须常驻 GPU。
MoE各变体及其机制
写这章的过程中发现其实一些基础变体在上半部分中已经提过了……包括Expert Choice 版本的变体,使用Auxiliary-Loss-Free Balance的变体之类的。这些变体的机制变化相对比较简单且上半部分已经详细讲过其机制了,我们就略过它们挑些重要的变体讲吧!(绝对不是我太困了
细想了下发现MoE变体不能简单的按DeepSeek MoE,Kimi StableLatentMoE这样子的按模型厂家分类,因为他们也是多种MoE变体的融合,因此我们下面会按机制来更细致的来分类变体(我的中文怎么写出来这么糟糕
共享专家 (Shared Experts)
在传统MoE中,我们希望每个专家都有自己专长的部分,充分发挥MoE高参数容量所带来的优势,但就像现实生活中物化生的基础都是数学(大概)一样,不同专家都要掌握相同的基础知识(如语法),导致参数浪费,造成知识的冗余。同时还导致,比如一个资深的高等数学专家,由于自身要回答用户问题,必须抽出一定时间学习如何更通俗的讲解知识点,导致原本高质的高数能力下降,单个专家被迫学习差异巨大的多种知识(如数学和情感分析),导致其难以真正专业化,导致知识混杂。
为了解决上述问题,DeepSeekMoE等架构提出了“共享专家 + 路由专家”的混合架构。
- 共享专家 (Shared Experts):始终处于激活状态,处理所有输入的Token。它们负责捕捉跨领域的通用知识,保证基础信息的流动。
- 路由专家 (Routed Experts):仅在需要时被稀疏激活。模型通过路由器(Router)为每个输入动态选择最合适的少数专家来处理,负责学习特定领域的专业知识。
对于每个输入的Token,模型会将其同时送入所有共享专家和由路由器选出的Top-K个路由专家。最终将所有专家的输出(通常是加权求和)与原始输入相加,作为该层的最终输出。
细粒度专家 (Fine-Grained Experts)
传统的MoE模型(如经典的Top-2 MoE)通常只使用少数几个(如8或16个)大型专家网络。这种设计会带来两个主要问题:
知识混杂:每个专家被迫学习过多、过杂的知识。例如,同一个专家可能需要同时处理“数学计算”和“情感分析”两种差异巨大的任务,导致其难以在任一领域做到真正精通。
知识冗余:不同的专家可能都需要掌握一些基础知识(如语法),导致这些知识在多个专家的参数中被重复存储,造成了参数浪费。
我们不难发现,这些问题其实是和共享专家要解决的是相同的,这两个问题的根本原因就是传统MoE专家分工不够明确。
细粒度专家MoE将一个大型的专家网络拆分为数量更多、但规模更小的专家网络。例如,一个典型的配置可能拥有64到256个细粒度专家。对于每个输入的token,模型不再是只激活1-2个专家,而是选择6到8个最相关的细粒度专家来共同处理。
现在常用的DeepSeek MoE(Kimi k3所使用的MoE架构也是基于此的),就是通过同时采用共享专家和细粒度专家的变体的方式,改善了传统MoE知识冗余与混杂的问题。
潜在MoE (Latent MoE)
相比于MoE效果,Latent MoE主要是针对MoE的计算效率进行了优化。其核心思路为:在将任务分配给专家处理之前,先将输入的token“压缩”到一个更小的维度中,让专家在这个低维空间里进行计算,最后再将结果“解压”回原始维度。
我们这边引用一下k3技术报告中的原话:
“增加专家池和激活专家的数量会扩展专家专业化的空间,但在传统MoE中,每个选中的专家都会接收完整的d维度token表示,因此通信和专家权重流量会随着路由倍数增长。Latent MoE通过将完整模型宽度与路由专家宽度分离,使这种扩展变得可行”
不难发现,Latent MoE的主要目的就是为了解决传统MoE专家网络仍然在模型的完整大维度上运行导致的通信与内存瓶颈问题。
K3 的 Stable Latent MoE(这章比较长)
我大致看了下k3的技术报告,k3所使用的Stable Latent MoE本质上就是DeepSeek MoE的共享专家机制,细粒度专家机制和Latent MoE的杂交版,具体表现为:共享专家保留完整维度用于常见转换,而专门的路由专家则在维度为ℓ的紧凑潜在空间中操作。这使得KimiK3能够将通道混合扩展到896个路由专家,每个token选择16个激活专家。
这样的融合看似很美好,但是过多的矩阵连乘和过多的专家数量会带来许多问题带来了两类失败,分别为内部激活值爆炸和负载均衡失效。
激活值爆炸主要原因如下:在加入 Shared MoE和 Latent MoE 后,K3的MoE大致流程如下所示
输入 x
↓
降维投影 W↓
↓
896 个专家中选 16 个
↓
每个专家内部做 GLU/FFN 计算
↓
16 个专家输出加权聚合
↓
升维投影 W↑
↓
回到主 hidden space
而且k3一开始使用的普通SwiGLU门控,还会额外添加两个分支相乘(我有点后悔上半部分没有展开讲门控了)
我这里简略的讲解下GLU门控的大致机制吧
普通 FFN 通常是:
\text{FFN}(x)=W_2\,\sigma(W_1x)
也就是:
x → 线性层 W1 → 激活函数 → 线性层 W2
但 GLU / SwiGLU / SiTU-GLU 会先把输入送进两个不同的线性分支gate和up。
gate branch 负责产生一个“门”。
例如:
g = W_gx
然后过一个非线性函数:
\tilde g = \text{Swish}(g)
它的作用类似:
这个维度的信息该放大、压小,还是关闭?
所以叫 gate branch。
up branch 负责产生主要内容向量:
u = W_ux
它可以理解成普通 FFN 里的中间表示。
所以:
gate branch 决定“开关/强弱”;
up branch 提供“内容”。
两个分支输出的向量会逐元素相乘,这种乘法提升表达能力,但也会把两个大激活相乘成更大的激活。这就导致k3的MoE流程变成了一个超级雷霆巨大且猎奇的四重矩阵连乘
W↓ → W_gate / W_up → 乘法门控 → W_expert_out → W↑
再加上K3自身高达2.8T的参数量,中间只要有一点误差就会被这个矩阵连乘持续放大,最终导致模型的效果变差。
而负载均衡失效的原因在于,k3高达10³量级的路由专家数量超出了现有无辅助损失(也就是我们上半部分所讲的Auxiliary-Loss-Free Balance)的偏差更新保持负载均衡的能力范畴。因为普通 bias update 通常是固定步长:专家过载:b -= η,专家不足:b += η。专家数量少的时候还能工作。但 K3 有接近 1000 个专家,问题变得非常复杂:
- 每个专家的 token 数波动更大
- top-k 是 16,不是 简单的top-1/top-2
- 专家之间竞争边界非常密集
- 固定步长太小会调整慢,固定步长太大会振荡
- 热门专家和冷门专家数量多,系统整体更难稳定。
为了解决上述两个问题,k3提出了三重机制,上投影前的RMSNorm和用于抑制激活爆炸的SigmoidTanh单元GLU(SiTU‑GLU),以及用于负载均衡的分位数平衡(QB),我们一一讲解。
RMSNorm before up-projection
RMSNorm全称是 Root Mean Square Layer Normalization,可以翻译成 均方根归一化。其目标是控制“向量整体尺度”,位置放置于升维投影矩阵之前,防止计算过程中由第一类失败导致的参数偏大传播出去。我们来随着其数学公式和具体数值例子来看看它具体是怎么发挥作用的。
假设某一层里,一个 token 的 hidden state 是: x=[x_1,x_2,\ldots,x_d] ,这里的d是指token的维度。
RMSNorm 计算这个向量的均方根:
\text{RMS}(x)
=
\sqrt{
\frac{1}{d}
\sum_{i=1}^{d}x_i^2
+
\epsilon
}
然后把整个向量除以这个 RMS:
\hat x_i=
\frac{x_i}{\text{RMS}(x)}
最后再乘一个可学习缩放参数:
y_i=g_i\hat x_i
完整写作:
\text{RMSNorm}(x)
=
g\odot
\frac{x}{
\sqrt{
\frac{1}{d}
\sum_{i=1}^{d}x_i^2+\epsilon
}
}
假设:x=[3,4]
那么:
\text{RMS}(x)=
\sqrt{\frac{3^2+4^2}{2}}
=
\sqrt{\frac{25}{2}}
\approx3.536
归一化后:
\hat x=
[3/3.536,\ 4/3.536]
\approx
[0.849,\ 1.131]
归一化后,这个向量的 RMS 大约为 1:
\sqrt{\frac{0.849^2+1.131^2}{2}}
\approx1
所以 RMSNorm 做的事情很简单:
原始向量可能很大或很小
↓
除以它自己的 RMS
↓
把尺度调到大约 1
↓
再乘上可学习 gain
其中gain参数的作用是防止 RMSNorm 只是把所有维度都归一到统一尺度,限制模型表达。每个维度都有自己的 gain,这允许模型学习:第 1 维应该放大一点、第 2 维应该压小一点、第 3 维保持不变……等等
这样可能还不够直观,我们来通过具体数据来展示其作用
我们假设有这样一个token,在MoE流程的初步阶段由于一些误差导致分数偏高,在经过一系列矩阵连乘后,误差累积过大,成为了向量u,其模相比于正常token高了100倍!此时u进入RMSNorm
RMSNorm 对 u 做:
\tilde u
=
\frac{u}{
\sqrt{
\frac{1}{d}\sum_i u_i^2+\epsilon
}
}
\odot g
其中:
\sqrt{
\frac{1}{d}\sum_i u_i^2
}
就是 u 的 RMS 尺度。
如果 u 的整体尺度变大 100 倍:
u' = 100u
那么:
\text{RMS}(u')=100\text{RMS}(u)
所以:
\frac{u'}{\text{RMS}(u')}
=
\frac{100u}{100\text{RMS}(u)}
=
\frac{u}{\text{RMS}(u)}
可以看到 RMSNorm将异常大的token拉回了正常水平,防止经过升维投影矩阵后在后续造成更大的灾难。但RMSNorm同样有局限。如果问题是某个单独维度极端大,例如:
u=[1,1,1,10000]
RMSNorm 会把整体缩放下来,但方向仍然被那个极端维度主导。虽然 norm 稳住了,但表示方向已经被 outlier 污染。所以 K3 还需要 SiTU-GLU在内部继续解决第一类灾难,SiTU-GLU 在专家内部限制大值,减少这种单维或局部 outlier 的产生。
SiTU-GLU
K3一开始使用的是SwiGLU作为门控,有以下问题
SwiGLU 常见形式如下
\text{SwiGLU}(x)
=
\text{Swish}(W_gx)\odot W_ux
其中:
\text{Swish}(a)=a\sigma(a)
当 a 很大时:
\sigma(a)\approx 1
所以:
\text{Swish}(a)\approx a
也就是说 gate branch 在大值区域不会被限制住。
假设两个分支都因为误差产生了较大的偏移,如
W_gx=20
W_ux=30
那么:
\text{Swish}(20)\times30\approx600
这个中间激活值就会大的很变态
K3 的 SiTU-GLU 就是为了解决这个。
它会对两个分支做 soft cap。
可以粗略理解为
\text{softcap}(x,\beta)=\beta\tanh(x/\beta)
当 x 很小时:
\text{softcap}(x,\beta)\approx x
当 x 很大时:
\text{softcap}(x,\beta)\approx \beta
所以如果 gate 分支限制在大约 4,up 分支限制在大约 25,那么乘积最大规模就被控制住,原来可能:20 × 30 = 600.限制后约:4 × 25 = 100
原技术报告给出的公式如下:
\text{SiTU-GLU}(\boldsymbol{x}) = \left[\beta_1 \tanh\left(\frac{\mathbf{W}_g \boldsymbol{x}}{\beta_1}\right) \odot \text{Sigmoid}(\mathbf{W}_g \boldsymbol{x})\right] \odot \left[\beta_2 \tanh\left(\frac{\mathbf{W}_u \boldsymbol{x}}{\beta_2}\right)\right],
原文也给出了β参数的具体值
“对于KimiK3,我们将软截断超参数设置为β1=4用于门控分支,设置为β2=25用于上分支。缩放后的tanh在原点附近近似线性,并在大数值处有界,这使得SiTU‑GLU能够在控制乘积中两个因子的同时,保留SwiGLU的局部响应。”
QuantileBalancing(QB)
K3 有接近 1000 个专家,每个 token 选 16 个专家。如果 router 自由选择,很容易让少数专家过热、很多专家吃不到 token。传统的 auxiliary-loss-free bias 更新在这个规模下容易震荡或调不平,所以 K3 用“分位数”直接估计每个专家 bias 应该调到哪里。
QB不像传统bias按照步长逐步更新bias,而是直接计算:如果想让专家e恰好收到目标数量 q个 token,那么它的 bias 应该是多少。那么我们该如何计算这个bias呢?k3将每个token的topk + 1的分数作为阈值,因为它代表“入选门槛”。对于 token t,某个专家 e 要进入 topk,它的 biased score 必须超过当前边界,这个边界近似就是当前t排名第 k+1 个专家的分数
对于专家 e,定义它相对边界的 margin:
r_{t,e}=s_{t,e}-\tau_t
如果:
r_{t,e}+b_e>0
也就是:
s_{t,e}+b_e>\tau_t
那么专家 e 就有机会进入 token t\ 的 top-k。这个$$ b_e $$ 就是我们后续要求的偏置项
所以对于一个专家 e,它面对所有 token 都有一个 margin 分布:
token 1: r_1,e = -3.2
token 2: r_2,e = -0.7
token 3: r_3,e = 0.1
token 4: r_4,e = 1.5
如果想让恰好 q 个 token 满足这个条件,就要找 margin 分布中的某个分位点。
假设有 m 个 token,我们希望专家接收 q 个 token,那么目标比例是:
\rho=\frac{q}{m}
也就是说,希望大约 $$ \rho $$ 比例的 token 选择这个专家。
要做到这一点,就应该让阈值落在 margin 分布的对应分位点上。
举个具体的数值例子
某专家 e 对 10 个 token 的 margin:
-5.0, -3.0, -2.1, -1.0, -0.3, 0.2, 0.7, 1.1, 2.0, 3.4
如果希望它被 3 个 token 选中,就需要让最大的 3 个 margin 过线:
1.1, 2.0, 3.4
边界大概在 0.7 和 1.1 附近
于是 bias 应该设置到让这个边界刚好变成 0
这就是分位数 balancing
这一章终于结束了,好累(,怎么还有一章啊(
MoE机制和Temperature为0时模型输出仍然不一致的关系和Soft MoE
MoE是如何影响输出确定性的
在上半部分的结语部分,我说到" Temperature = 0 的情况下,为什么每次输出还会不一样?这与 MoE 的路由机制有很大的关系,也牵扯到我们后续会谈的一个 MoE 变体:Soft MoE。"今天我们就来讲解这些内容
按理来说,temperature=0 对应贪婪解码,即每一步都选择概率最高的那个词,结果应该是确定的,但现实结果并不符合上述理论。除了和计算机本身硬件的浮点数计算不确定性和批处理(Batching)的影响以外(关于这部分内容,感兴趣的佬友可以搜索下Thinking Machine Lab的一篇叫《Defeating Nondeterminism in LLM Inference》的文章,里面详细讲解了这个现象出现的原因,这里只讲解MoE相关部分)
MoE模型由多个“专家”子模型组成,每个输入Token会被路由(Route) 到最擅长处理它的1个或几个专家那里
这个路由过程引入了不确定性:
- 专家分配不确定,如果两个 expert 的 router 分数非常接近,极小数值扰动可能改变 token 被送去哪个 expert。导致最终结果不一致。
- Token 丢弃(Token Dropping):每个专家处理Token的数量有上限。如果某批次中太多Token都想找同一个专家,超出的Token就会被丢弃。这些Token会因此失去被该专家处理的机会,导致最终结果发生变化。
- 批次构成的随机性:你无法控制自己的请求和谁一起组成批次。这个批次的构成是随机的,因此是否会发生Token丢弃、丢弃哪些Token,都变得不可预测,最终导致输出的不确定性。
Soft MoE是什么
简单来说与传统的MoE不同,Soft MoE使用软分配 (Soft Assignment):每个Token会以不同权重分配给所有专家。而不是直接全部交给某个/某些专家。每个专家处理所有输入Token的加权组合(加权平均值)。通过让所有专家都参与每个Token的处理,从根本上避免了“Token丢弃”和“路由竞争”问题,从而提升了模型输出的确定性。
写在最后
说实话我感觉这篇写的没有很好(,包括kimi k3那部分写的有点更加怪怪的( ,今天熬穿了导致脑子昏昏沉沉的,好像还吃坏了肚子导致每写一会就要去窜,磨磨蹭蹭从早上七点半写到现在才写完(,佬友们还是要爱惜身体(,再次祝大家节日快乐。
下次文章的话,我比较想写一下RAG相关的东西,或者是Kimi k3的完整技术报告解析吧( 。暂时拟一个吸睛的标题吧:你真的了解RAG吗?知道从chunking、embedding、粗排精排,查询优化的所有内容吗?知道知识图谱RAG正在被逐渐发掘潜力吗?你真的认为RAG已死吗?我们下次再见
我c我要睡觉了