解读一下MOE架构,部分内容图片参考自油管。
首先来简单了解一下什么是MoE(Mixture of Experts,专家混合)
MoE(Mixture of Experts)是一种深度学习架构,其核心思想是通过**多个专家网络(Experts)和一个门控网络(Gate Network)**来动态地选择和组合不同的专家,从而提高计算效率和模型能力。
MoE 的核心组成
-
专家网络(Experts)
- MoE 由多个**子神经网络(Experts)**组成,每个专家都是一个独立的子模型(通常是 MLP 或 Transformer 层)。
- 每个专家专注于学习不同的数据模式或特定任务的一部分。
-
门控网络(Gate Network)
- 这是一个路由机制,用于决定输入样本应该由哪些专家来处理。
- 通常是一个softmax 层,为每个专家分配一个权重(即该样本有多大概率被某个专家处理)。
- 例如,输入一个数据后,门控网络可能只激活 2~3 个专家,而不是所有专家都参与计算。
-
加权组合
- 选定的专家将对输入进行计算,然后按门控网络的权重加权求和,得到最终输出。
MoE 的主要优势
-
计算效率高
- 由于每个输入样本只会激活一部分专家,MoE 可以减少计算量,使得训练和推理更高效,尤其在大模型中非常重要。
- 例如,GPT-4 采用了一种 MoE 变体,使其可以拥有超大参数量,但每次计算只用到其中一小部分参数。
-
提升模型容量
- 传统神经网络的计算复杂度与参数量线性相关,而 MoE 允许使用超大规模模型,但每次计算仅使用一部分参数,因此不会导致推理开销过大。
-
任务分工更明确
- 不同的专家可以专注于不同类型的数据模式(如某些专家擅长处理长文本,另一些擅长短文本)。
MoE 可能的挑战
-
负载不均衡(Load Imbalance)
- 如果门控网络总是倾向于选择某几个专家,可能导致某些专家过载,而其他专家几乎不被使用,降低了并行计算效率。
-
训练不稳定
- MoE 由于路由机制的存在,训练过程中可能会出现梯度不稳定,甚至导致某些专家完全无法被训练。
-
通信开销
- 在分布式训练中,MoE 可能会带来额外的通信成本,影响整体效率。
MOE基本原理:
如图示,MOE的基本结构包括一个门控单元(也叫做路由),多个专家网络构成。门控单元的作用是根据输入数据的特征,动态地选择适合的专家网络进行计算。一般来说,门控单元采用softmax 机制,为每个专家分配一个权重,并根据这些权重决定激活哪些专家。这样,MoE 在保持大模型容量的同时,减少了计算开销,使得每个输入数据仅经过一部分专家的处理。
MoE 的计算流程
MoE 的计算过程在主流的decoder-only架构LLM中,可以分为以下几个步骤(假设在decoder block中的FFN中使用MOE):
- 输入数据处理(经典transformers流程)
- 输入数据首先通过嵌入层(Embedding Layer),转换成合适的表示形式后进入self-attention计算以及layer norm之后进入FFN层。
-
门控单元(Gate Network)计算
- 门控网络根据输入数据计算 softmax 权重,决定哪些专家网络应该被激活,以及激活的权重是多少。
-
专家网络计算
- 选中的专家网络对输入进行计算,每个专家独立地执行前向传播(Forward Pass)。
-
加权求和输出
- 门控单元的 softmax 权重用于加权求和所有选中的专家的输出,得到最终的合成输出。
数学上,MoE 的计算可以表示如下:
y
=
∑
i
=
1
N
G
(
x
)
i
E
i
(
x
)
y = \sum_{i=1}^{N} G(x)_i E_i(x)
y=i=1∑NG(x)iEi(x)
其中:
- G ( x ) G(x) G(x) 是门控网络计算出的 softmax 权重(决定专家的选择)。
- E i ( x ) E_i(x) Ei(x) 是第 i i i 个专家的输出。
-
N
N
N 是专家的总数。
在传统的 Transformer 结构中,每个前馈层(FFN)都会对所有 Token 进行计算,而MoE 只选择 部分专家 进行计算,使得推理过程中只计算一小部分参数,从而减少计算量,推理更高效。
在使用 MoE 进行推理时,模型会根据每个专家学习到的具体知识,在内部形成一条专家推理路径。这一路径由**门控网络(Gating Network)**动态决定,确保输入数据能够被路由到最适合的专家进行处理。
MoE 推理示例
假设我们使用一个 MoE 语言模型来回答问题:
“地球上最大的动物是什么?”
在 MoE 结构中,可能的专家推理路径如下:
- 门控网络分析输入,发现它涉及生物学领域。
- 分配专家:
- 专家 3(专精于生物知识) → 负责动物相关信息
- 专家 7(专精于常识问答) → 负责常识性回答
- 专家进行计算,输出可能的答案:
- 专家 3 可能输出:“蓝鲸”
- 专家 7 可能输出:“地球上最大的哺乳动物是蓝鲸”
- 综合专家结果,最终生成回答:
“地球上最大的动物是蓝鲸。”
MoE 的门控机制原理
在 Mixture of Experts (MoE) 结构中,门控机制(Gating Mechanism)负责决定哪些专家(Experts)将被激活,以及它们的计算权重。门控网络(Gating Network)通过学习输入数据的分布,智能地分配计算资源,使得不同输入流经不同的专家网络。
门控机制的计算流程
假设有一个 MoE 层,包含 N N N 个专家(Experts),一个输入 x x x 经过 MoE 层的计算流程如下:
1. 计算专家得分
门控网络通常是一个简单的全连接层,它的作用是将输入 x x x 映射到一个长度为 N N N 的分数向量 g ( x ) g(x) g(x):
g ( x ) = W g ⋅ x + b g g(x) = W_g \cdot x + b_g g(x)=Wg⋅x+bg
其中:
- W g W_g Wg 和 b g b_g bg 是门控网络的可学习参数。
- g ( x ) g(x) g(x) 是一个 N N N 维向量,其中每个值代表某个专家的得分。
2. 计算专家分配权重
为了获得专家的选择概率,通常对 g ( x ) g(x) g(x) 进行Softmax 归一化:
p i = exp ( g i ) ∑ j = 1 N exp ( g j ) p_i = \frac{\exp(g_i)}{\sum_{j=1}^{N} \exp(g_j)} pi=∑j=1Nexp(gj)exp(gi)
其中 p i p_i pi 代表第 i i i 个专家的激活概率。
3. 选择 Top-k 个专家
由于 MoE 是稀疏激活的,我们通常不会让所有专家都参与计算,而是选取Top-k 个专家(通常 k = 1 k=1 k=1 或 k = 2 k=2 k=2)。选择方法:
S = Top-k ( p ) S = \text{Top-k}(p) S=Top-k(p)
- 这里 S S S 是被选中的专家索引。
- 仅计算选中的专家,并忽略其余专家。
4. 计算专家输出
被选中的专家会使用输入 x x x 进行计算,并输出结果:
y i = E i ( x ) , ∀ i ∈ S y_i = E_i(x), \quad \forall i \in S yi=Ei(x),∀i∈S
其中:
- E i ( x ) E_i(x) Ei(x) 是第 i i i 个专家的计算输出。
5. 加权融合专家输出
MoE 需要结合多个专家的输出,我们根据 Softmax 权重 p p p 对选中的专家进行加权求和:
y MoE = ∑ i ∈ S p i y i y_{\text{MoE}} = \sum_{i \in S} p_i y_i yMoE=i∈S∑piyi
这样,最终输出 y MoE y_{\text{MoE}} yMoE 作为 MoE 层的结果,并传递给后续 Transformer 层。
门控机制的核心特点
- 动态选择专家:
- 每个 Token 经过 MoE 层时,都会由门控网络动态决定激活哪些专家。
- 稀疏计算:
- 由于只计算 Top-k 个专家,MoE 在计算效率上远远优于传统的稠密模型。
- 可学习性:
- 门控网络在训练过程中会不断优化,使得不同专家能够学习不同的任务,提高模型的表达能力。
MoE 的负载均衡挑战
虽然 Mixture of Experts (MoE) 结构能够有效地提升计算效率,但在实际应用中会遇到一个重要问题——负载不均衡(Load Imbalance)。负载均衡问题会导致计算资源浪费、模型训练不稳定,甚至影响最终模型的性能。
为什么会出现负载不均衡?
MoE 的负载不均衡主要来源于门控网络(Gating Network)对专家的选择不均衡,具体表现如下:
-
专家选择分布偏差
- 在训练过程中,某些专家可能被频繁选择,而其他专家几乎不被使用,导致计算资源无法有效利用。
-
部分专家过载
- 由于计算仅发生在被选中的专家上,如果某些专家被选中的概率远高于其他专家,它们的计算负担会过大,从而成为瓶颈。
-
梯度更新不均衡
- 如果某些专家长期不被选中,它们的梯度更新会较少,导致训练效果下降,使得这些专家在模型中的作用变弱,形成恶性循环。
-
训练不稳定
- 如果专家的负载极端不均衡,模型的训练过程会变得不稳定,甚至影响最终的收敛效果。
- 如果专家的负载极端不均衡,模型的训练过程会变得不稳定,甚至影响最终的收敛效果。
负载均衡的优化策略
针对上述问题,研究人员提出了多种优化方法,以改善 MoE 的负载均衡情况。
KeepTopK 机制:解决 MoE 负载均衡问题
在 Mixture of Experts (MoE) 结构中,负载均衡问题是一个关键挑战。KeepTopK 机制是一种常见的专家选择方法,它能够在保持计算效率的同时,避免过载问题。
什么是 KeepTopK?
KeepTopK 指的是在 MoE 层中,每个 Token 只会选择得分最高的 K 个专家,并将计算负载均摊到这些专家上。
KeepTopK 机制:基于加噪声的专家选择
在 Mixture of Experts (MoE) 结构中,KeepTopK 是一种优化专家选择的方法,能够在减少计算量的同时,缓解负载均衡问题。它的核心思想是:对专家得分进行加噪声处理后,选择得分最高的 K 个专家,以此增强专家选择的随机性,避免个别专家过载。
KeepTopK 计算流程
-
计算专家得分
对于输入 Token,MoE 层的门控网络(Gating Network)计算所有 N N N 个专家的得分:
g = W g ⋅ x g = W_g \cdot x g=Wg⋅x
其中:- W g W_g Wg 是门控网络的权重矩阵。
- x x x 是输入 Token 的向量表示。
- g g g 是一个长度为 N N N 的向量,表示该 Token 对应的专家得分(未归一化)。
-
加噪声扰动
在选择专家之前,向得分 g g g 添加一定的随机高斯噪声 ϵ \epsilon ϵ,增强选择的随机性:
g ′ = g + ϵ , ϵ ∼ N ( 0 , σ 2 ) g' = g + \epsilon, \quad \epsilon \sim \mathcal{N}(0, \sigma^2) g′=g+ϵ,ϵ∼N(0,σ2)
其中:- ϵ \epsilon ϵ 服从均值为 0 0 0,方差为 σ 2 \sigma^2 σ2 的正态分布。
- 该噪声可以防止门控网络过度偏向某些专家,使专家选择更加均衡。
-
选择得分最高的 K 个专家
- 通过 TopK 选择策略,选取得分最高的 K 个专家:
S = TopK ( g ′ , K ) S = \text{TopK}(g', K) S=TopK(g′,K) - 这里的
S
S
S 是一个包含
K
K
K 个索引的集合,对应被选中的专家编号。假设
K
=
2
K=2
K=2,如图:
- 通过 TopK 选择策略,选取得分最高的 K 个专家:
-
计算专家加权输出
- 仅激活
S
S
S 中的专家,并根据门控得分
g
g
g 进行加权求和:
y = ∑ i ∈ S g i ⋅ Expert i ( x ) y = \sum_{i \in S} g_i \cdot \text{Expert}_i(x) y=i∈S∑gi⋅Experti(x) - 这样可以保证计算量受控,同时避免完全随机选择专家带来的信息损失。
- 仅激活
S
S
S 中的专家,并根据门控得分
g
g
g 进行加权求和:
加噪声的作用
增强随机性:加噪声可以防止门控网络过度集中在少数专家上,提高模型的泛化能力。
缓解负载不均衡:由于不同 Token 选择专家时会受到噪声影响,专家的使用频率更均衡,降低了过载风险。
提升训练稳定性:避免专家梯度更新不均衡,使得所有专家都能充分训练,提高整体模型性能。
负载均衡优化
虽然 KeepTopK 通过加噪声的方式减少了过载问题,但仍然可能出现负载不均衡。
从 KeepTopK 过渡到辅助损失(Auxiliary Loss)
虽然 KeepTopK 机制通过在专家得分上加噪声,使专家选择更加均衡,但它仍然无法完全避免负载不均衡问题。某些专家仍可能被过度选择,而其他专家则可能很少被使用。为了进一步优化专家分配,我们可以引入 辅助损失(Auxiliary Loss),鼓励模型均衡使用所有专家。
为什么需要辅助损失?
即使采用了 KeepTopK 的加噪选择策略,仍然可能出现:
- 某些专家过载:部分专家频繁被选中,计算开销增大,并可能导致过拟合。
- 部分专家闲置:某些专家很少被使用,使得它们的参数训练不足,影响整体性能。
- 梯度更新不均衡:如果大部分 Token 仅流经少数专家,则这些专家的梯度更新过快,而其他专家的梯度更新不足,影响模型收敛。
为了解决这些问题,我们引入辅助损失,使得 MoE 层更加均衡地使用所有专家。
** 重要性分数(Importance Score)**
重要性分数(Importance Score) 用于衡量每个专家在训练过程中被选中的频率。
假设门控网络对每个 Token 计算了
N
N
N 个专家的得分,并通过 Softmax 归一化后得到专家选择概率:
g
=
softmax
(
W
g
⋅
x
)
g = \text{softmax}(W_g \cdot x)
g=softmax(Wg⋅x)
其中:
- W g W_g Wg 是门控网络的权重矩阵。
- x x x 是输入 Token 的表示。
- g ∈ R N g \in \mathbb{R}^{N} g∈RN 是专家得分向量。
每个专家
i
i
i 的重要性分数 定义如下:
I
i
=
∑
x
∈
batch
p
i
(
x
)
I_i = \sum_{x \in \text{batch}} p_i(x)
Ii=x∈batch∑pi(x)
其中:
- I i I_i Ii 表示专家 i i i 在当前 batch 内的重要性分数。
- p i ( x ) p_i(x) pi(x) 是门控网络分配给专家 i i i 的概率。
理想情况下,所有专家的 I i I_i Ii 应该接近均匀分布,以确保 MoE 层充分利用所有专家,并且不会导致计算负载不均衡。
4. 变异系数(Coefficient of Variation, CV)
变异系数(CV) 是衡量数据分布均衡性的重要指标,定义如下:
CV
=
σ
I
μ
I
\text{CV} = \frac{\sigma_I}{\mu_I}
CV=μIσI
其中:
- σ I \sigma_I σI 是所有专家的重要性分数的标准差。
- μ I \mu_I μI 是所有专家的重要性分数的均值。
变异系数用于衡量专家使用频率的离散程度,CV 越大,表示专家使用的不均衡程度越高。因此,我们可以定义一个基于 CV 的损失:
L
cv
=
CV
2
=
(
σ
I
μ
I
)
2
L_{\text{cv}} = \text{CV}^2 = \left(\frac{\sigma_I}{\mu_I}\right)^2
Lcv=CV2=(μIσI)2
该损失项的目标是减少专家之间的重要性分数的方差,使其分布更加均匀。
专家容量
在 MoE 结构中,除了使用辅助损失平衡专家负载外,我们还需要 专家容量(Expert Capacity) 的机制来限制每个专家能处理的最大 Token 数量。
专家容量的作用
- 避免个别专家过载:如果不加限制,部分专家可能会接收过多 Token,导致计算资源不均衡。
- 提高计算效率:设定最大容量可以减少计算开销,避免单个专家成为瓶颈。
- 防止过拟合:过载的专家可能学习过于细粒度的信息,而闲置专家未能得到充分训练,影响泛化能力。
专家容量的定义
在 MoE 机制中,通常会设定一个 专家容量上限
C
C
C,即每个专家最多能处理的 Token 数量:
C
=
k
×
B
N
C = \frac{k \times B}{N}
C=Nk×B
其中:
- k k k:超参数,表示每个 Token 选择的专家数量(通常是 1 或 2)。
- B B B:当前 batch 的总 Token 数。
- N N N:MoE 层中的专家数。
如果某个专家收到的 Token 超过
C
C
C,那么 超出部分的 Token 需要被丢弃或者重新分配,这会引入新的问题,例如 Token 丢失或延迟计算。
通常,在实际应用中,会结合 负载均衡损失(Load Balancing Loss) 和 专家容量限制,确保计算效率和负载均衡的最优权衡。
如何结合专家容量和负载均衡?
在 MoE 训练过程中,专家容量和负载均衡是相互影响的:
- 如果容量 C C C 过小,可能导致大量 Token 被丢弃,影响信息传输。
- 如果容量 C C C 过大,则容易导致某些专家过载,影响计算效率。
- 负载均衡损失可以帮助均匀分配 Token,使得所有专家的负载尽可能接近容量限制 C C C,提高计算稳定性。