MoE,混合专家
MoE的基础原理
MOE是一种神经网络的架构模式,它通过**条件化计算(Conditional Computation)**来增加模型的参数容量,而不成比例地增加其计算成本(FLOPs)。其核心在于将单一的、密集的处理模块(如前馈网络FFN)替换为一个由两个核心组件构成的系统:
-
专家网络(Expert Networks):
- 一组(
E
个)并行的子网络。在Transformer架构中,这些专家通常是结构相同的独立FFN层。 - 每个专家网络都拥有独立的参数,负责学习数据分布中的特定模式或子任务。
- 一组(
-
门控网络(Gating Network)/路由器(Router):
- 一个可训练的轻量级网络,其功能是为每个输入的token动态地生成一组路由决策权重(通常是logits)。
- 它的输入是token的向量表示,输出是一个维度为
E
的向量,其中每个元素对应一个专家的分数。
工作流程(前向传播)
- 路由决策: 对于输入的token表征向量
x
,门控网络计算出所有E
个专家的对数得分(logits):logits = Gating(x)
。 - 稀疏激活: 应用一个Top-K选择算法到
logits
向量上,其中K
是一个远小于E
的整数(例如K=2,E=8)。该算法选择得分最高的K
个专家,并记录它们的索引和得分。此步骤是实现“稀疏性”的核心,因为它确保只有一小部分参数子集被激活。 - 条件化计算: 输入
x
仅被分发(dispatch)到这K
个被激活的专家。这些专家并行地处理x
,生成K
个输出向量output_i = Expert_i(x)
。 - 输出聚合: 最终的输出向量
y
是通过对这K
个专家的输出进行加权求和得到的。权重通常由这K
个专家得分经过Softmax
函数归一化后得到。y = Σᵢ (Softmax(top_k_logits)ᵢ * outputᵢ)
。
通过此机制,MoE将模型的总参数量(与E
成正比)与单次前向传播的计算量(与K
成正比)解耦。
稠密模型 vs. 稀疏MoE模型
通常比较的是稠密模型(Dense Model)与稀疏MoE模型(Sparse MoE Model)。
稠密模型 (Dense Model)
- 定义: 指的是标准的、非条件化的神经网络架构,如标准的BERT或GPT。
- 激活方式: 无条件激活。在一次前向传播中,对于每一个输入token,模型的所有参数都会被激活并参与计算。
- 关系: 其计算成本(FLOPs)与模型参数总量呈线性正相关。增加参数容量必然导致计算成本的同等增加。
稀疏MoE模型 (Sparse MoE Model)
- 定义: 指的是集成了上述MoE层,并采用稀疏激活(
K < E
)的神经网络架构。 - 激活方式: 条件化激活。只有一小部分(K个)专家子集的参数会根据门控网络的动态决策被激活。
- 关系: 成功地解耦了参数容量与计算成本。总参数量可以随专家数量
E
的增加而大幅增长,而每个token的计算成本仅与被激活的专家数量K
成正比。 - 训练挑战: 其稳定训练依赖于一个辅助损失函数(Auxiliary Loss),通常用于实现负载均衡(Load Balancing),以确保所有专家在训练过程中被均匀地利用,防止模型仅依赖少数专家而导致其他专家参数失效的“模式坍塌”问题。
“稠密MoE” (Dense MoE) 的概念
这是一个理论上存在但实践中无意义的场景,它指的是MoE层中设置**K = E
**。
- 在这种情况下,门控网络会激活所有专家来处理每一个token。
- 其计算成本将是单个专家计算成本的
E
倍,与模型的总参数量重新变为线性正相关。 - 这种做法完全违背了MoE设计的初衷——即通过稀疏激活来节省计算,因此在实际应用中不会被采用。它实质上退化成了一个计算极其昂贵的集成模型。
特性 | 稠密模型 | 稀疏MoE模型 |
---|---|---|
架构 | 单一、整体的处理模块 | 多个并行的专家模块 + 一个门控网络 |
参数激活 | 稠密 / 无条件 (Dense / Unconditional) | 稀疏 / 条件化 (Sparse / Conditional) |
总参数量 | N | 约 E × N |
每Token计算量 | 约 N | 约 K × N (其中 K << E ) |
核心挑战 | 硬件资源对模型规模的物理限制 | 训练稳定性(需处理负载均衡) |
负载均衡
基于损失函数的软负载均衡
出自论文 《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-Of-Experts Layer》
(8)P(x,i)=Pr((x⋅Wg)i+StandardNormal()⋅Softplus((x⋅Wnoise)i)>kth_excluding(H(x),k,i))(8) P(x, i) = \text{Pr}\left( (x \cdot W_g)_i + \text{StandardNormal}() \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i) > \text{kth\_excluding}(H(x), k, i) \right)(8)P(x,i)=Pr((x⋅Wg)i+StandardNormal()⋅Softplus((x⋅Wnoise)i)>kth_excluding(H(x),k,i))
(9)P(x,i)=Φ((x⋅Wg)i−kth_excluding(H(x),k,i)Softplus((x⋅Wnoise)i))(9) P(x, i) = \Phi\left( \frac{(x \cdot W_g)_i - \text{kth\_excluding}(H(x), k, i)}{\text{Softplus}((x \cdot W_{\text{noise}})_i)} \right)(9)P(x,i)=Φ(Softplus((x⋅Wnoise)i)(x⋅Wg)i−kth_excluding(H(x),k,i))
(10)Load(X)i=∑x∈XP(x,i)(10) \text{Load}(X)_i = \sum_{x \in X} P(x, i)(10)Load(X)i=x∈X∑P(x,i)
(11)Lload(X)=wload⋅CV(Load(X))2(11) L_{\text{load}}(X) = w_{\text{load}} \cdot \text{CV}(\text{Load}(X))^2(11)Lload(X)=wload⋅CV(Load(X))2
通用符号定义
在深入每个公式之前,我们先定义一些在多个公式中出现的通用符号:
x
: 代表一个输入 token 的向量表示 (vector representation)。X
: 代表一个批次 (batch) 的所有输入 tokens 的集合。i
: 代表我们当前正在分析的第 i 个专家 (expert) 的索引。k
: 代表在 Top-K 路由策略中,需要选择的专家数量。E
: 代表专家的总数量。
公式 (8): 路由选择的概率事件定义
P(x,i)=Pr((x⋅Wg)i+StandardNormal()⋅Softplus((x⋅Wnoise)i)>kth_excluding(H(x),k,i))P(x, i) = \text{Pr}\left( (x \cdot W_g)_i + \text{StandardNormal}() \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i) > \text{kth\_excluding}(H(x), k, i) \right)P(x,i)=Pr((x⋅Wg)i+StandardNormal()⋅Softplus((x⋅Wnoise)i)>kth_excluding(H(x),k,i))
这个公式定义了一个概率事件,其目的是计算专家 i
被选入 Top-K 行列的概率。
符号 | 含义和作用 |
---|---|
P(x,i)P(x, i)P(x,i) | 目标变量:输入x 被路由到专家i 的最终概率。 |
Pr(...)Pr(...)Pr(...) | 概率算子:表示计算括号内逻辑事件为真的概率。 |
WgW_gWg | 门控权重矩阵 (Gating Weights):门控网络的核心参数,用于将输入x 转换为所有专家的原始路由分数。下标g 代表"gating"。 |
(x⋅Wg)i(x \cdot W_g)_i(x⋅Wg)i | 专家i的原始分数 (Logit):输入x 与门控权重矩阵相乘,得到一个长度为E 的分数向量,此项代表该向量中属于专家i 的分量。 |
StandardNormal()StandardNormal()StandardNormal() | 标准正态分布随机数生成器:一个函数,返回一个从均值为0、标准差为1的正态分布中采样的随机数。这是引入随机性的来源。 |
WnoiseW_{noise}Wnoise | 噪声权重矩阵 (Noise Weights):一个独立的、可学习的权重矩阵,专门用于学习如何根据输入x 来调整噪声的幅度。 |
Softplus(...)Softplus(...)Softplus(...) | Softplus激活函数:数学形式为log(1+e^z) 。其作用是确保其输出值永远为正,因为噪声的标准差不能为负数。 |
Softplus((x⋅Wnoise)i)Softplus((x \cdot W_{\text{noise}})_i)Softplus((x⋅Wnoise)i) | 噪声的标准差 (σ):此项计算出一个与输入x 和专家i 相关的正数值,用作StandardNormal() 随机数的缩放因子,即噪声的幅度。 |
H(x)H(x)H(x) | 所有专家的分数向量:这是 x \cdot W_g 的简写,即 H(x) = x \cdot W_g 。 |
kth_excluding(v,k,i)kth\_excluding(v, k, i)kth_excluding(v,k,i) | 竞争门槛函数:一个自定义函数,它在向量v (这里是H(x) )中,排除掉第i 个元素后,找到剩下元素中第k 大的那个值。这个值就是专家i 要想被选中,所必须超越的**“录取分数线”**。 |
公式 (9): 可微的概率计算
P(x,i)=Φ((x⋅Wg)i−kth_excluding(H(x),k,i)Softplus((x⋅Wnoise)i))P(x, i) = \Phi\left( \frac{(x \cdot W_g)_i - \text{kth\_excluding}(H(x), k, i)}{\text{Softplus}((x \cdot W_{\text{noise}})_i)} \right)P(x,i)=Φ(Softplus((x⋅Wnoise)i)(x⋅Wg)i−kth_excluding(H(x),k,i))
这个公式是公式(8)的解析解(analytical solution),它将一个不可微的概率事件转化成了一个平滑、可微的函数。
符号 | 含义和作用 |
---|---|
Φ(...)Φ(...)Φ(...) | 标准正态分布的累积分布函数(CDF):一个函数,输入一个Z分数,输出一个0到1之间的概率值,代表随机变量小于等于该Z分数的概率。 |
分子 | 竞争优势/净胜分 (Competitive Margin):专家i 的原始分数与其“录取分数线”之间的差值。 |
分母 | 不确定性/噪声标准差 (Uncertainty/Noise Std Dev):与公式(8)中定义相同,代表了决策过程的随机性幅度。 |
整个分数 | Z分数 (Z-score):一个标准化的值,衡量了专家i 的“竞争优势”是其自身决策“不确定性”的多少倍。 |
公式 (10): 专家负载的定义
Load(X)i=∑x∈XP(x,i)\text{Load}(X)_i = \sum_{x \in X} P(x, i)Load(X)i=∑x∈XP(x,i)
这个公式定义了单个专家的“软”负载或期望负载。
符号 | 含义和作用 |
---|---|
Load(X)iLoad(X)_iLoad(X)i | 专家i的期望负载:在一个批次X 中,专家i 被期望分配到的token的总数量。 |
ΣΣΣ | 求和算子:对批次X 中的每一个输入x 进行遍历求和。 |
P(x,i)P(x, i)P(x,i) | 单token概率:由公式(9)计算得出的,单个输入x 被路由到专家i 的概率。 |
公式 (11): 负载均衡损失函数
Lload(X)=wload⋅CV(Load(X))2L_{\text{load}}(X) = w_{\text{load}} \cdot \text{CV}(\text{Load}(X))^2Lload(X)=wload⋅CV(Load(X))2
这个公式构建了最终用于优化的辅助损失函数。
符号 | 含义和作用 |
---|---|
Lload(X)L_{load}(X)Lload(X) | 负载均衡损失:一个标量值,衡量整个专家团队负载分配的不均衡程度。这个值越小,说明负载越均衡。 |
wloadw_{load}wload | 损失权重:一个手动设置的超参数(标量),用于平衡负载均衡损失与模型主任务损失之间的重要性。 |
Load(X)Load(X)Load(X) | 负载向量:一个包含了所有专家期望负载的向量,即 [Load(X)₁, Load(X)₂, ..., Load(X)ₑ] 。 |
CV(...)CV(...)CV(...) | 变异系数 (Coefficient of Variation):一个统计量,定义为 标准差 / 平均值 。它是一个归一化的、衡量数据离散程度的指标,非常适合用于衡量负载是否均衡。 |
Q&A
-
公式9的具体含义
其实公式9就是公式8的可微表达,这里涉及到概率论的基础知识:
如果一个随机变量 Y 服从均值为 μ、标准差为 σ 的正态分布,即 Y ~ N(μ, σ²),那么,P(Y > a) = Φ( (μ - a) / σ )
我们假设括号内为z,Φ(z)\Phi(z)Φ(z)表示从标准正态分布中随机取一个值,该值小于z的概率。当(x⋅Wg)i=kth_excluding(...)(x\cdot W_g)_i = kth\_excluding(...)(x⋅Wg)i=kth_excluding(...)时,该概率值为0.5,表示当预测专家i分数刚好够着topk专家分数的门槛时,选择i专家的概率为0.5。 -
公式9的分母含义
分母在公式8中表示加上的正态分布噪声的标准差,因此该标准差越大(分母越大)表示分配专家的随机性越大,分母越小表示基本按照专家的分数进行分配。 -
公式9中分母的作用
鼓励探索:在训练早期或遇到不确定的输入时,模型可以学习到生成一个较大的分母(高噪声),这会使得专家的分数排名更容易发生变化,从而让更多的专家有机会得到训练,避免了“马太效应”。
学习确定性:对于模型已经充分学习过的模式,模型可以学习到生成一个很小的分母(低噪声),使得路由决策非常稳定和高效。 -
为什么通过优化损失11,就可以实现负载均衡
正向传播过程中的x*Wg这个量,指的是由各个专家处理某个输入的分数(近似理解为概率)。而Lload损失函数保证了整体样本X上,各个专家的负载(处理的x个数期望)是大概相同的。 -
为了得到某样本分配给各个专家的概率,我们计算了(x⋅Wg)(x\cdot W_g)(x⋅Wg),这些概率值什么时候需要用公式9处理,能不能直接用Softmax
分为训练和预测两种情况:
训练时,(x⋅Wg)(x\cdot W_g)(x⋅Wg)用于两处,第一处就是用于计算LloadL_{load}Lload,此时需要用公式9进行计算负载,加噪声(分母)的作用见上个问题;第二处是用于计算主任务损失LtaskL_{task}Ltask,此时直接使用softmax并选取topk个专家来处理,最终的结果与标签计算损失
测试时,(x⋅Wg)(x\cdot W_g)(x⋅Wg)直接使用softmax并选取topk个专家来处理得到最终结果
基于专家容量的硬负载均衡
出自经典论文GShard
专家容量
为了确保负载是均衡的,我们强制要求由一个专家处理的token数量必须低于某个统一的阈值,我们将这个阈值定义为专家容量 (expert capacity)。假设一个训练批次(batch)中的token总数是N
,并且每个token最多被分派给两个专家 (at most two experts),那么专家的容量就被设置为O(N/E)
的量级。门控函数GATE(·)
会为每个专家e
维持一个运行中的计数器cₑ
,记录有多少token被分派给了该专家。当一个token所选择的两个专家都已经超出了它们的容量时,这个token就被认为是一个溢出token (overflowed token),此时其门控向量G_s,E
会退化成一个零向量。这类token的表征向量x_s
会通过残差连接直接传递到下一层。
讲解
这是一种强制的、基于计数的硬性负载均衡机制,它与我们之前讨论的、基于辅助损失函数的“软”均衡机制形成了鲜明对比。这是GShard为了解决大规模并行计算中的实际工程问题而设计的一个非常重要的机制。
1. 核心思想:设置“最大容量”上限
- 问题: 即使有辅助损失函数,在训练的某些阶段,门控网络仍然可能“任性地”将大量token发送给少数几个专家。
- 解决方案: 与其用损失函数“软性地劝导”,不如直接给每个专家设置一个**“最大接诊量”,我们称之为“专家容量 (Expert Capacity)”**。一旦专家的“接诊”数量达到了这个上限,它就关门谢客,不再接收任何新的token。
2. 如何确定容量大小?
N
: 一个训练批次中的token总数。E
: 专家的总数。- 每个token最多选择2个专家 (k=2):这意味着在一个批次中,总的“专家处理槽位”大约是
2 * N
个。 - 理想情况: 在完美的均衡状态下,每个专家应该分摊到的槽位是
(2 * N) / E
。 - 设置容量: 因此,将专家容量设置为
O(N/E)
的量级是合理的。O
表示“数量级”,通常会再乘以一个大于2的容量因子(capacity factor),比如1.25,来提供一些缓冲空间。例如,Capacity = 1.25 * (2 * N) / E
。
3. 运行机制:计数与溢出
这是一个动态的、按顺序的过程:
- 初始化: 在一个batch开始时,所有专家的计数器
cₑ
都清零。 - 逐token处理: 模型按顺序处理batch中的每个token。
- 路由决策: 假设现在处理token
T
,门控网络为其选择了专家A和专家B。 - 检查容量:
- 检查专家A的计数器
c_A
是否已满(c_A >= Capacity
)。 - 检查专家B的计数器
c_B
是否已满(c_B >= Capacity
)。
- 检查专家A的计数器
- 分发或溢出:
- 情况一 (未满): 如果专家A和B中至少有一个还有空位,那么token
T
就被正常地分发过去,并将其对应的计数器加一。 - 情况二 (全满): 如果专家A和B都满了,token
T
就无法被任何一个它想去的专家处理。此时,它被标记为**“溢出token (overflowed token)”**。
- 情况一 (未满): 如果专家A和B中至少有一个还有空位,那么token
4. 如何处理“溢出token”?
这是这个机制的关键一环。如果一个token被“丢弃”,它的信息就会丢失,这会损害模型的学习。GShard的处理方式非常优雅:
- 门控向量归零: 这个溢出token的门控向量
G
(即决定专家输出权重的向量)会被强制设为零向量。 - MoE层失效: 这意味着,MoE层的输出对于这个token来说,将是
Σ (0 * Expert_output) = 0
。MoE层没有对这个token做任何有意义的计算。 - 通过残差连接“跳过”: 在Transformer的架构中,每个子层(如MoE层)的输出都会与该层的输入进行一个残差连接(
output = input + Sublayer(input)
)。- 对于溢出token
x_s
,这个公式就变成了output = x_s + 0
。 - 最终效果: 这个溢出token未经任何处理,原封不动地通过残差连接“跳过”了整个MoE层,直接进入了下一层网络。
- 对于溢出token
分组分配
我们讨论的“专家容量”是一个全局的、串行的计数过程,这在大规模分布式训练中会成为新的瓶颈。“本地分组分派”就是为了解决这个瓶颈,让容量检查也可以并行化。这里更多在工程实践,因此不过多描述。
辅助损失
原文
至关重要的是,门控函数不应该总是选择同样的那少数几个专家,因为这会导致少数几个专家的容量溢出,而其余专家则利用不足。我们定义一个辅助损失项l_aux
来强制施加这个约束。它被添加到一个常数乘数k
的模型总损失函数中:L = l_nll + k * l_aux
。在算法的(13)行中,辅助损失项l_aux
的具体形式是基于以下考量而设计的:cₑ/S
这一项代表了路由到每个专家的输入所占的比例,而我们想要最小化cₑ/S
的均方值。但是,因为cₑ
是从top-2操作中派生出来的,它是不可微的,所以我们使用每个专家的平均门控值mₑ
作为一个可微的近似,并用mₑ(cₑ/S)
来替换(cₑ/S)²
,这个新形式现在就可以用梯度下降来优化了。
讲解
这段话详细阐述了GShard中辅助损失函数的设计动机和具体实现,完美地连接了我们之前讨论的**“硬性容量机制”和“软性损失函数”**。
1. 为什么需要辅助损失?—— 硬性容量机制的局限
- 我们之前讨论了“专家容量(Expert Capacity)”机制,它通过设置硬性上限来防止任何一个专家过载。
- 但是,这个机制有一个问题:它只管“堵”,不管“疏”。它能防止最坏的情况发生(即系统崩溃),但它不能主动地去引导门控网络做出更均衡的选择。
- 门控网络仍然可能会“愚蠢地”持续将token发往已经满了的专家,导致大量的token被“溢出”(overflow)并跳过MoE层。这同样会损害模型的学习效率。
- 因此,我们需要一个“软性”的机制来从源头上教导门-控网络要“雨露均沾”,这就是辅助损失的作用。
2. 损失函数的设计目标
cₑ
: 指的是在一个本地组(local group)中,被实际分派给专家e
的token数量(容量检查得到的计数)。S
: 指的是一个本地组中的token总数。cₑ / S
: 这一项代表了专家e
在这个组中实际承担的负载比例。- 理想状态: 在完美的均衡状态下,每个专家承担的负载比例应该是相等的。
- 设计目标: 我们想要最小化这些负载比例的**方差(Variance)**或类似的离散程度度量。最小化均方值(Mean Square)
Σ(cₑ/S)²
是一个简化但有效的代理目标。
3. 遇到的核心问题:不可微
- 问题所在:
cₑ
这个计数值,是通过一个离散的、硬性的Top-2选择操作得到的。这个操作是不可微分的(non-differentiable)。 - 后果: 你无法对
cₑ
求梯度,因此也无法对包含cₑ
的损失函数进行梯度下降。梯度无法从这个损失函数回传到门控网络的权重Wg
上。这个损失函数就成了一个摆设。
4. 解决方案:寻找一个“可微的替身”
既然无法直接使用cₑ
,我们就需要构造一个新的量,我们称之为mₑ
,它必须满足两个条件:
- 它必须是可微的,能够为梯度下降提供平滑的梯度。
- 它必须与
cₑ
高度相关,以至于优化mₑ
能够间接地优化cₑ
。
mₑ
的定义与计算
mₑ
被定义为**“专家e的平均门控值” (mean gates per expert)**。
- 计算门控值
gₑ(x)
: 对于每一个tokenx
,我们将门控网络计算出的原始分数logits
通过一个Softmax
函数。gₑ(x) = Softmax(x · Wg)ₑ
。Softmax
的输出是一个0到1之间的连续浮点数,代表了专家e
的“概率权重”。 - 计算平均值
mₑ
: 将一个批次(或一个本地组)中所有token为专家e
计算出的门控值gₑ(x)
相加,再除以token的总数S
。
mₑ = (1/S) * Σ_x gₑ(x)
为什么 mₑ
是可微的?
我们来审视从Wg
到mₑ
的完整计算路径:
Wg
-> x · Wg
(线性变换,可微) -> Softmax
(可微函数) -> gₑ(x)
-> Σ
(求和,可微) -> (1/S) * ...
(标量乘法,可微)。
整条路径上的每一步操作都是可微的。根据链式法则,mₑ
对于Wg
是完全可微的。当Wg
发生微小变化ΔWg
时,mₑ
也会相应地发生一个微小、平滑的变化Δmₑ
。
5. mₑ
与 cₑ
的强相关性
一个专家e
被门控网络赋予的平均门控值mₑ
越高,直观上意味着它被Top-K
操作选中的频率也应该越高,因此其真实的token计数值cₑ
也越大。反之亦然。这种强烈的正相关性使得mₑ
成为了cₑ/S
(负载比例)的一个非常优秀的代理(Proxy)。
6. 最终的损失函数形式
现在,我们有了一个不可微的理想度量cₑ/S
和一个可微的代理度量mₑ
。我们可以构建一个新的损失函数。
- 理想损失:
Σ (cₑ/S)²
,但不可微。 - GShard的解决方案:
l_aux = Σ_e mₑ * (cₑ/S)
。
在这个形式中:
mₑ
提供了可微的路径。梯度可以通过mₑ
顺畅地反向传播到Wg
。∂l_aux / ∂Wg
的计算是围绕∂mₑ / ∂Wg
进行的。cₑ/S
在这里被当作一个常数缩放因子。虽然它本身不可微,但在计算梯度时,可以把它看作一个“权重”,用来衡量当前实际不均衡的严重程度,从而对梯度进行缩放。
这样,整个辅助损失项l_aux
就变得可以通过梯度下降进行优化了,优化它会促使mₑ
的分布变得均衡,并因为mₑ
和cₑ
的强相关性,间接地也实现了对cₑ
的负载均衡。
随机路由
原文
随机路由
直观上来说,因为y_s
是被选中的专家们返回结果的加权平均值,所以如果第二个专家(2nd expert)的权重非常小,我们可以简单地忽略掉这个专家以节省总体的专家容量。因此,除了遵守专家容量的约束之外,门控函数GATE(·)
会以一个与其权重g₂
成正比的概率来决定是否将token分派给第二好的那个专家(2nd-best expert)。
讲解
这段话介绍了一种非常巧妙的、用于节省和优化专家容量使用的随机化策略。它是在Top-K
(这里特指K=2)路由决策基础上的一个进一步优化。
1. 问题背景:被“浪费”的容量
- 回顾Top-2路由: 对于每个token,模型会选择得分最高的两个专家(我们称之为专家1和专家2)来处理它。
- 权重差异: 在很多情况下,门控网络对专家1的“信心”可能远大于对专家2的。例如,经过Softmax计算后,专家1的权重
g₁
可能是0.99,而专家2的权重g₂
只有0.01。 - 计算与结果:
- 即使
g₂
如此之小,模型仍然会把这个token发送给专家2,消耗掉专家2的一个宝贵的“容量槽位”。 - 在最终计算输出
y_s = g₁ * output₁ + g₂ * output₂
时,专家2的输出output₂
因为乘以了一个极小的权重g₂
,对最终结果y_s
的贡献微乎其微。
- 即使
- 结论: 我们为了一个几乎可以忽略不计的贡献,却支付了实实在在的计算成本和容量成本。这是一种浪费。
2. 解决方案:引入“概率性忽略”
GShard的作者们提出的解决方案非常符合直觉:对于那些贡献微不足道的专家,我们不如有一定概率直接忽略掉它。
这个机制具体是这样实施的(仅针对Top-2中的第二名):
- 选出Top-2: 像往常一样,门控网络选出得分最高的专家1和得分第二高的专家2。
- 计算权重: 计算出它们的门控权重
g₁
和g₂
(其中g₁+g₂=1
)。 - 决策一(专家1): 第一个专家(最好的专家)总是被选择。token一定会被发送给专家1。
- 决策二(专家2) - 随机门:
- 对于第二个专家,我们引入一个随机决策。
- 模型会以**概率
p
**将token发送给专家2。 - 这个概率
p
与专家2自身的权重g₂
成正比。在实际操作中,通常就直接令p = g₂
(或者一个与g₂
相关的函数)。 - 例子:
- 如果
g₂=0.01
,那么模型只有1%的概率会真的把这个token发给专家2。在99%的情况下,模型会“假装”没选它,直接跳过。 - 如果
g₂=0.4
,那么模型有40%的概率会把token发给专家2。
- 如果
3. 优点与效果
这种“随机路由”策略带来了几个显著的好处:
- 节省专家容量: 这是最直接的好处。通过有选择地忽略掉那些贡献小的专家,为其他更需要这个专家的token节省出了宝贵的容量槽位。这能有效降低“专家容量溢出”的token比例,提升训练效率。
- 节约计算资源: 在token被忽略的情况下,专家2不需要进行任何计算,从而节省了计算资源。
- 动态的“Top-1”与“Top-2”切换:
- 当门控网络非常自信时(
g₁
接近1,g₂
接近0),这个机制的行为就近似于Top-1路由,只使用一个专家。 - 当门控网络不太确定,两个专家的权重比较接近时(比如
g₁=0.6, g₂=0.4
),这个机制的行为就更接近于标准的Top-2路由。 - 它实现了一种在Top-1和Top-2之间的平滑、动态的过渡。
- 当门控网络非常自信时(
总结
“专家容量”机制是一个非常实用的工程解决方案,它:
- 强制负载均衡: 通过设置硬性上限,从物理上杜绝了任何一个专家被过度使用,从而保证了计算负载在不同设备间的均衡。
- 简化计算图: 在静态图的计算框架(如早期的TensorFlow)中,预先设定每个专家的缓冲区大小(即容量),可以极大地简化内存分配和计算图的构建。
- 优雅处理溢出: 通过利用残差连接,它为那些“运气不好”的溢出token提供了一条“绿色通道”,让它们可以跳过拥堵的MoE层,避免了信息丢失,保证了训练的稳定性。