MOE 详解:结合 Transformer 架构的完整分析

MOE 详解:结合 Transformer 架构的完整分析

本文基于 Google 的论文《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》,结合 Transformer 架构,详细解释 MOE 的工作原理、数学公式和训练过程。


目录

  1. MOE 概述
  2. MOE 在 Transformer 中的位置
  3. MOE 层的结构
  4. 前向传播详解
  5. 反向传播详解
  6. 训练中的负载均衡
  7. 完整数值示例
  8. MOE 的优势与挑战

一、MOE 概述

1.1 什么是 MOE?

MOE(Mixture of Experts) 是一种条件计算技术,核心思想是:

  • 多个专家网络(Experts)替换单一的前馈网络
  • 使用门控网络(Gating Network)为每个输入动态选择少量专家
  • 只有被选中的专家参与计算,实现稀疏激活

1.2 核心优势

特性 传统 MLP MOE
参数总量 受限于计算量 可以非常大
每次推理激活的参数 100% 仅 top-k 专家
模型容量 有限 大幅增加
计算成本 与参数量成正比 相对恒定

示例:DeepSeek-V3

  • 总参数:6710 亿
  • 每次推理激活参数:370 亿(仅 5.5%)
  • 专家数量:256 个
  • 每个 token 激活专家:9 个(1 个共享 + 8 个选中)

1.3 发展历程

1991: Jacobs 等人提出 Mixture of Experts 概念
2017: Google 提出 Sparsely-Gated MOE(本文重点)
2022: Switch Transformer、DeepSpeed-MoE
2024: Mixtral、DeepSeek-MoE、DeepSeek-V3

二、MOE 在 Transformer 中的位置

2.1 回顾标准 Transformer Block

Input X
    │
    ├──────────────────────────────┐
    ↓                              │ (Residual)
Layer Norm → Multi-Head Attention  │
    │                              │
    ↓                              │
    + ←────────────────────────────┘
    │
    ├──────────────────────────────┐
    ↓                              │ (Residual)
Layer Norm → MLP (FFN)             │  ← 【MOE 替换这里】
    │                              │
    ↓                              │
    + ←────────────────────────────┘
    │
    ↓
Output

2.2 MOE Transformer Block

MOE 替换 MLP 层

Input X
    │
    ├──────────────────────────────┐
    ↓                              │ (Residual)
Layer Norm → Multi-Head Attention  │
    │                              │
    ↓                              │
    + ←────────────────────────────┘
    │
    ├──────────────────────────────┐
    ↓                              │ (Residual)
Layer Norm → MOE Layer             │  ← 【MOE 层】
    │        ┌─────────────┐       │
    │        │ Gating Net  │       │
    │        │    ↓        │       │
    │        │ [E1][E2]... │       │
    │        │ [En]        │       │
    │        └─────────────┘       │
    ↓                              │
    + ←────────────────────────────┘
    │
    ↓
Output

2.3 参数对比

组件 标准 Transformer (GPT-2) MOE Transformer
Embedding 50257 × 768 50257 × 768
Attention 4 × 768² = 2.4M 4 × 768² = 2.4M
MLP/MOE 768 × 3072 × 2 = 4.7M n × (768 × 3072 × 2) + 门控
激活参数 全部 top-k 专家

三、MOE 层的结构

3.1 整体架构

                        Input x (768维)
                             │
              ┌──────────────┴──────────────┐
              │                             │
              ↓                             ↓
     ┌─────────────────┐          ┌─────────────────────────┐
     │  Gating Network │          │      Experts            │
     │    G(x)         │          │  E1, E2, ..., En        │
     │                 │          │  (n个前馈网络)            │
     └────────┬────────┘          └─────────────────────────┘
              │                             │
              │ 门控权重                     │ 专家输出
              │ [g1, g2, ..., gn]          │ [E1(x), E2(x), ..., En(x)]
              │                             │
              └──────────────┬──────────────┘
                             │
                             ↓
                      加权求和 (稀疏)
                   y = Σ G(x)_i × E_i(x)
                             │
                             ↓
                      Output y (768维)

3.2 核心组件

3.2.1 专家网络(Experts)

每个专家是一个独立的前馈网络,结构与标准 MLP 相同:

\[E_i(x) = W_2^{(i)} \cdot \text{GELU}(W_1^{(i)} \cdot x + b_1^{(i)}) + b_2^{(i)} \]

参数:

  • \(W_1^{(i)} \in \mathbb{R}^{d \times d_{ff}}\):第一层权重(768 × 3072)
  • \(W_2^{(i)} \in \mathbb{R}^{d_{ff} \times d}\):第二层权重(3072 × 768)
  • 每个专家约 4.7M 参数

3.2.2 门控网络(Gating Network)

门控网络决定每个输入应该由哪些专家处理:

简单 Softmax 门控

\[G_\sigma(x) = \text{Softmax}(x \cdot W_g) \]

其中 \(W_g \in \mathbb{R}^{d \times n}\)\(n\) 是专家数量。

Noisy Top-K 门控(推荐):

\[G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k)) \]

\[H(x)_i = (x \cdot W_g)_i + \text{StandardNormal}() \cdot \text{Softplus}((x \cdot W_{noise})_i) \]

\[\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in top-k of } v \\ -\infty & \text{otherwise} \end{cases} \]

3.3 MOE 输出公式

\[y = \sum_{i=1}^{n} G(x)_i \cdot E_i(x) \]

稀疏性:由于 \(G(x)\) 是稀疏的(只有 top-k 个非零),我们只需计算被选中的 k 个专家。


四、前向传播详解

4.1 输入设置

假设:

  • 输入 \(x \in \mathbb{R}^{768}\):一个 token 的表示
  • 专家数量 \(n = 8\)
  • top-k = 2(每个 token 选择 2 个专家)

4.2 Step 1: 计算门控 logits

\[H(x) = x \cdot W_g + \text{noise} \]

x 形状: (768,)
W_g 形状: (768, 8)

H(x) = x × W_g
形状: (8,)  ← 每个专家一个分数

示例:
H(x) = [1.2, 0.5, -0.3, 2.1, 0.8, -0.5, 1.5, 0.2]
        ↑                   ↑
      专家0               专家3

4.3 Step 2: Noisy Top-K 选择

添加噪声(训练时):
noise = StandardNormal() × Softplus(x × W_noise)

H(x) + noise = [1.25, 0.48, -0.28, 2.15, 0.82, -0.52, 1.48, 0.18]

Top-2 选择:
- 专家3: 2.15 ← 最高
- 专家0: 1.25 ← 第二高

KeepTopK(H(x), 2) = [1.25, -∞, -∞, 2.15, -∞, -∞, -∞, -∞]

4.4 Step 3: Softmax 得到门控权重

\[G(x) = \text{Softmax}(\text{KeepTopK}(H(x), 2)) \]

Softmax([1.25, -∞, -∞, 2.15, -∞, -∞, -∞, -∞])

= [e^1.25, 0, 0, e^2.15, 0, 0, 0, 0] / (e^1.25 + e^2.15)
= [3.49, 0, 0, 8.58, 0, 0, 0, 0] / 12.07
= [0.289, 0, 0, 0.711, 0, 0, 0, 0]

门控权重:
G(x) = [0.289, 0, 0, 0.711, 0, 0, 0, 0]
         ↑              ↑
       专家0          专家3

4.5 Step 4: 计算被选中专家的输出

只计算 G(x)_i ≠ 0 的专家:

专家0:
E_0(x) = W_2^(0) × GELU(W_1^(0) × x + b_1^(0)) + b_2^(0)
形状: (768,)

专家3:
E_3(x) = W_2^(3) × GELU(W_1^(3) × x + b_1^(3)) + b_2^(3)
形状: (768,)

其他专家 (1,2,4,5,6,7): 不计算! ← 节省计算

4.6 Step 5: 加权求和

\[y = G(x)_0 \cdot E_0(x) + G(x)_3 \cdot E_3(x) \]

\[y = 0.289 \cdot E_0(x) + 0.711 \cdot E_3(x) \]

y 形状: (768,)

这就是 MOE 层的输出!

4.7 批处理情况

对于一个 batch 的 tokens(形状 (batch_size, seq_len, 768)):

输入: X 形状 (B, S, 768)  ← B个样本,S个token

对于每个 token x_{b,s}:
1. 计算门控权重 G(x_{b,s})
2. 选择 top-k 专家
3. 计算选中专家的输出
4. 加权求和

输出: Y 形状 (B, S, 768)

注意: 不同 token 可能选择不同的专家!

4.8 前向传播流程图

输入 x (batch_size × seq_len × 768)
        │
        ├─────────────────────────────────────────────────┐
        │                                                 │
        ↓                                                 ↓
┌───────────────────────┐                    ┌─────────────────────────┐
│    Gating Network     │                    │       8 个专家           │
│                       │                    │                         │
│  H(x) = x × W_g       │                    │  E_0: MLP (768→3072→768)│
│  + noise              │                    │  E_1: MLP (768→3072→768)│
│                       │                    │  E_2: MLP (768→3072→768)│
│  TopK + Softmax       │                    │  ...                    │
│                       │                    │  E_7: MLP (768→3072→768)│
└───────────┬───────────┘                    └───────────┬─────────────┘
            │                                            │
            │ G(x) = [0.289, 0, 0, 0.711, 0, 0, 0, 0]    │
            │                                            │
            │ 选中专家: {0, 3}                            │
            │                                            │
            └──────────────────┬─────────────────────────┘
                               │
                               ↓
                    只计算 E_0(x) 和 E_3(x)
                               │
                               ↓
                    y = 0.289 × E_0(x) + 0.711 × E_3(x)
                               │
                               ↓
                    输出 y (batch_size × seq_len × 768)

五、反向传播详解

5.1 总体梯度流向

                                    ∂L/∂y
                                      │
              ┌───────────────────────┴───────────────────────┐
              │                                               │
              ↓                                               ↓
     ∂L/∂G(x) (门控梯度)                           ∂L/∂E_i(x) (专家梯度)
              │                                               │
              ↓                                               ↓
     ∂L/∂W_g, ∂L/∂W_noise                         ∂L/∂W_1^(i), ∂L/∂W_2^(i)

5.2 MOE 输出公式回顾

\[y = \sum_{i=1}^{n} G(x)_i \cdot E_i(x) \]

对于稀疏情况(只有 top-k 个专家被激活):

\[y = \sum_{i \in \text{TopK}} G(x)_i \cdot E_i(x) \]

5.3 Step 1: 对专家输出的梯度

\[\frac{\partial \mathcal{L}}{\partial E_i(x)} = G(x)_i \cdot \frac{\partial \mathcal{L}}{\partial y} \]

关键:只有被选中的专家才有非零梯度!

假设 ∂L/∂y 形状: (768,)

对于专家0 (G(x)_0 = 0.289):
∂L/∂E_0(x) = 0.289 × ∂L/∂y
形状: (768,)

对于专家3 (G(x)_3 = 0.711):
∂L/∂E_3(x) = 0.711 × ∂L/∂y
形状: (768,)

对于未选中的专家 (G(x)_i = 0):
∂L/∂E_i(x) = 0  ← 没有梯度更新!

5.4 Step 2: 对门控权重的梯度

\[\frac{\partial \mathcal{L}}{\partial G(x)_i} = E_i(x)^T \cdot \frac{\partial \mathcal{L}}{\partial y} \]

对于专家0:
∂L/∂G(x)_0 = E_0(x)^T × ∂L/∂y
形状: 标量

对于专家3:
∂L/∂G(x)_3 = E_3(x)^T × ∂L/∂y
形状: 标量

5.5 Step 3: 通过 Softmax 反向传播

Softmax 的反向传播:

\[\frac{\partial \mathcal{L}}{\partial H(x)_i} = G(x)_i \cdot \left( \frac{\partial \mathcal{L}}{\partial G(x)_i} - \sum_{j} \frac{\partial \mathcal{L}}{\partial G(x)_j} \cdot G(x)_j \right) \]

注意:只有 top-k 个专家参与 Softmax,所以求和只在这些专家上进行。

5.6 Step 4: 对门控网络权重的梯度

\[\frac{\partial \mathcal{L}}{\partial W_g} = x^T \cdot \frac{\partial \mathcal{L}}{\partial H(x)} \]

x 形状: (768,)
∂L/∂H(x) 形状: (8,)  ← 只有 top-k 个非零

∂L/∂W_g = x^T × ∂L/∂H(x)
形状: (768, 8)

5.7 Step 5: 对每个专家网络的反向传播

对于被选中的专家(如专家0),其内部的 MLP 反向传播:

专家0 的结构:
h_1 = GELU(x × W_1^(0) + b_1^(0))
E_0(x) = h_1 × W_2^(0) + b_2^(0)

反向传播:
∂L/∂W_2^(0) = h_1^T × ∂L/∂E_0(x)
∂L/∂h_1 = ∂L/∂E_0(x) × (W_2^(0))^T
∂L/∂(pre_gelu) = ∂L/∂h_1 ⊙ GELU'(pre_gelu)
∂L/∂W_1^(0) = x^T × ∂L/∂(pre_gelu)

5.8 Step 6: 对输入 x 的梯度

梯度从两条路径汇合:

\[\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial x}\bigg|_{\text{from gating}} + \frac{\partial \mathcal{L}}{\partial x}\bigg|_{\text{from experts}} \]

从门控网络

\[\frac{\partial \mathcal{L}}{\partial x}\bigg|_{\text{gating}} = \frac{\partial \mathcal{L}}{\partial H(x)} \cdot W_g^T \]

从专家网络

\[\frac{\partial \mathcal{L}}{\partial x}\bigg|_{\text{experts}} = \sum_{i \in \text{TopK}} G(x)_i \cdot \frac{\partial E_i(x)}{\partial x} \]

5.9 完整反向传播流程图

                                    ∂L/∂y (768,)
                                        │
              ┌─────────────────────────┴─────────────────────────┐
              │                                                   │
              ↓                                                   ↓
┌─────────────────────────────────────┐    ┌─────────────────────────────────────┐
│         对门控权重的梯度               │    │          对专家输出的梯度              │
│                                     │    │                                     │
│  ∂L/∂G(x)_0 = E_0(x)^T × ∂L/∂y      │    │  ∂L/∂E_0(x) = 0.289 × ∂L/∂y         │
│  ∂L/∂G(x)_3 = E_3(x)^T × ∂L/∂y      │    │  ∂L/∂E_3(x) = 0.711 × ∂L/∂y         │
│                                     │    │                                     │
└──────────────┬──────────────────────┘    └──────────────┬──────────────────────┘
               │                                          │
               ↓                                          ↓
┌─────────────────────────────────────┐    ┌─────────────────────────────────────┐
│         通过 Softmax 反向            │    │        对每个专家 MLP 反向             │
│                                     │    │                                     │
│  ∂L/∂H(x) = Softmax_backward(...)   │    │  专家0: ∂L/∂W_1^(0), ∂L/∂W_2^(0)     │
│                                     │    │  专家3: ∂L/∂W_1^(3), ∂L/∂W_2^(3)     │
│  (只在 top-k 专家上计算)              │    │  (只更新被选中的专家!)                 │
│                                     │    │                                     │
└──────────────┬──────────────────────┘    └──────────────┬──────────────────────┘
               │                                          │
               ↓                                          ↓
┌─────────────────────────────────────┐    ┌─────────────────────────────────────┐
│      对门控网络权重的梯度               │   │       对输入 x 的梯度 (专家部分)        │
│                                     │    │                                     │
│  ∂L/∂W_g = x^T × ∂L/∂H(x)           │    │  ∂L/∂x|_exp = Σ G_i × ∂E_i/∂x        │
│                                     │    │                                     │
└──────────────┬──────────────────────┘    └──────────────┬──────────────────────┘
               │                                          │
               │                                          │
               └──────────────────┬───────────────────────┘
                                  │
                                  ↓
                    ∂L/∂x = ∂L/∂x|_gating + ∂L/∂x|_experts
                                  │
                                  ↓
                        传递给前一层 (Attention)

六、训练中的负载均衡

6.1 问题:专家不均衡

如果不加约束,门控网络会倾向于:

  • 总是选择某几个专家
  • 这些专家训练更多,变得更好
  • 门控网络更倾向于选择它们
  • 形成恶性循环
训练过程中的不均衡示例:

训练初期:
专家使用次数: [120, 115, 118, 122, 125, 110, 108, 112]
分布相对均匀

训练后期 (无负载均衡):
专家使用次数: [450, 380, 25, 15, 10, 8, 7, 5]
只有前两个专家被使用,其他专家几乎废弃!

6.2 解决方案:辅助损失函数

6.2.1 重要性损失(Importance Loss)

定义专家 \(i\) 对一个 batch \(X\) 的重要性:

\[\text{Importance}(X)_i = \sum_{x \in X} G(x)_i \]

重要性损失:

\[\mathcal{L}_{\text{importance}}(X) = w_{\text{importance}} \cdot \text{CV}(\text{Importance}(X))^2 \]

其中 CV 是变异系数(标准差/均值)。

直觉:如果所有专家的重要性相等,CV = 0,损失为 0。

6.2.2 负载损失(Load Loss)

重要性相等不意味着接收的样本数相等。定义负载:

\[\text{Load}(X)_i = \sum_{x \in X} P(x, i) \]

其中 \(P(x, i)\) 是在重新采样噪声时,专家 \(i\) 被选中的概率。

\[\mathcal{L}_{\text{load}}(X) = w_{\text{load}} \cdot \text{CV}(\text{Load}(X))^2 \]

6.3 总损失函数

\[\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{CE}} + \mathcal{L}_{\text{importance}} + \mathcal{L}_{\text{load}} \]

其中 \(\mathcal{L}_{\text{CE}}\) 是标准的交叉熵损失。

6.4 负载均衡的效果

损失配置 Test Perplexity CV(Importance) CV(Load) 最大负载/平均负载
无辅助损失 39.8 3.04 3.01 17.80
\(w_{imp}=0.2\) 35.6 0.06 0.17 1.47
\(w_{load}=0.2\) 35.7 0.22 0.04 1.15
\(w_{imp}=0.1, w_{load}=0.1\) 35.6 0.06 0.05 1.14

结论:辅助损失显著改善了模型质量(perplexity 从 39.8 降到 35.6)和负载均衡。


七、完整数值示例

7.1 设置

输入: x ∈ R^768 (一个 token 的表示)
专家数量: n = 8
每个专家: MLP (768 → 3072 → 768)
top-k: 2

7.2 前向传播

Step 1: 计算门控 logits

# x 形状: (768,)
# W_g 形状: (768, 8)
# W_noise 形状: (768, 8)

H_clean = x @ W_g  # (8,)
# H_clean = [1.2, 0.5, -0.3, 2.1, 0.8, -0.5, 1.5, 0.2]

noise = torch.randn(8) * F.softplus(x @ W_noise)
# noise = [0.05, -0.02, 0.02, 0.05, 0.02, -0.02, -0.02, -0.02]

H = H_clean + noise
# H = [1.25, 0.48, -0.28, 2.15, 0.82, -0.52, 1.48, 0.18]

Step 2: Top-K 选择

# 选择 top-2
top_values, top_indices = torch.topk(H, k=2)
# top_values = [2.15, 1.48]
# top_indices = [3, 6]  # 专家3和专家6

# 构建稀疏门控
H_masked = torch.full((8,), float('-inf'))
H_masked[top_indices] = top_values
# H_masked = [-inf, -inf, -inf, 2.15, -inf, -inf, 1.48, -inf]

Step 3: Softmax

G = F.softmax(H_masked, dim=-1)
# G = [0, 0, 0, 0.663, 0, 0, 0.337, 0]
#            专家3 ↑         专家6 ↑

Step 4: 计算选中专家的输出

# 只计算专家3和专家6
E3_output = expert_3(x)  # (768,)
E6_output = expert_6(x)  # (768,)

# 每个专家内部:
# h1 = GELU(x @ W1 + b1)  # (3072,)
# output = h1 @ W2 + b2    # (768,)

Step 5: 加权求和

y = G[3] * E3_output + G[6] * E6_output
# y = 0.663 * E3_output + 0.337 * E6_output
# y 形状: (768,)

7.3 反向传播

假设上游梯度:

dL_dy = ...  # (768,) 从后续层传来

Step 1: 对专家输出的梯度

dL_dE3 = G[3] * dL_dy  # = 0.663 * dL_dy
dL_dE6 = G[6] * dL_dy  # = 0.337 * dL_dy

# 其他专家梯度为 0

Step 2: 对门控权重的梯度

dL_dG3 = (E3_output * dL_dy).sum()  # 标量
dL_dG6 = (E6_output * dL_dy).sum()  # 标量

Step 3: 通过 Softmax 反向

# 只在 top-k 专家上计算
G_topk = torch.tensor([G[3], G[6]])  # [0.663, 0.337]
dL_dG_topk = torch.tensor([dL_dG3, dL_dG6])

# Softmax 反向
s = (dL_dG_topk * G_topk).sum()  # 加权和
dL_dH_topk = G_topk * (dL_dG_topk - s)

# 映射回原始索引
dL_dH = torch.zeros(8)
dL_dH[3] = dL_dH_topk[0]
dL_dH[6] = dL_dH_topk[1]

Step 4: 对门控网络权重的梯度

dL_dWg = x.unsqueeze(1) @ dL_dH.unsqueeze(0)
# 形状: (768, 8)

Step 5: 对每个专家的反向传播

# 专家3 (被选中)
dL_dW2_3 = h1_3.unsqueeze(1) @ dL_dE3.unsqueeze(0)  # (3072, 768)
dL_dh1_3 = dL_dE3 @ W2_3.T  # (3072,)
dL_dpre_gelu_3 = dL_dh1_3 * gelu_derivative(pre_gelu_3)  # (3072,)
dL_dW1_3 = x.unsqueeze(1) @ dL_dpre_gelu_3.unsqueeze(0)  # (768, 3072)

# 专家6 类似...

# 其他专家: 梯度为 0,不更新!

Step 6: 对输入 x 的梯度

# 从门控网络
dL_dx_gating = dL_dH @ W_g.T  # (768,)

# 从专家网络
dL_dx_exp3 = G[3] * (dL_dpre_gelu_3 @ W1_3.T)  # (768,)
dL_dx_exp6 = G[6] * (dL_dpre_gelu_6 @ W1_6.T)  # (768,)

# 总梯度
dL_dx = dL_dx_gating + dL_dx_exp3 + dL_dx_exp6

7.4 参数更新统计

参数更新情况:

门控网络:
  W_g: (768, 8) ← 全部更新
  W_noise: (768, 8) ← 全部更新

专家网络:
  专家0: 不更新 (未被选中)
  专家1: 不更新
  专家2: 不更新
  专家3: W1_3, W2_3, b1_3, b2_3 ← 更新 ✓
  专家4: 不更新
  专家5: 不更新
  专家6: W1_6, W2_6, b1_6, b2_6 ← 更新 ✓
  专家7: 不更新

本次迭代:
- 门控网络参数: 全部更新
- 专家参数: 只有 2/8 = 25% 更新

八、MOE 的优势与挑战

8.1 优势

优势 说明
模型容量大 可以有数千亿参数
计算效率高 每次只激活少量专家
专家专业化 不同专家学习不同的模式
训练并行 不同专家可以在不同 GPU 上

8.2 挑战

挑战 解决方案
负载不均衡 辅助损失函数
通信开销 混合数据并行和模型并行
训练不稳定 Noisy Top-K, 负载均衡
显存占用 专家参数分布在多个 GPU

8.3 实验结果

来自原论文的结果:

模型 参数量 ops/timestep Test Perplexity
LSTM-2048-512 151M 151M 43.7
MoE-4096-h 4.3B 8.9M 34.1

结论

  • MOE 模型有 28倍 的参数量
  • 但只有 6% 的计算量
  • 同时获得了 22% 的性能提升

8.4 专家专业化示例

来自机器翻译实验的观察:

专家 专业化领域
专家 381 研究、创新相关词汇
专家 752 领导力、核心作用相关
专家 2004 速度、快速相关副词
专家 381 激活的上下文:
"... with researchers , ..."
"... to innovation . ..."
"... technology innovations is ..."

专家 752 激活的上下文:
"... plays a core ..."
"... plays a critical ..."
"... assume a leadership ..."

专家 2004 激活的上下文:
"... with rapidly growing ..."
"... to swiftly ..."
"... the fastest ..."

九、与标准 Transformer 的对比总结

9.1 结构对比

标准 Transformer Block:                MOE Transformer Block:

Input                                  Input
  │                                      │
  ↓                                      ↓
Layer Norm                             Layer Norm
  │                                      │
  ↓                                      ↓
Multi-Head Attention                   Multi-Head Attention
  │                                      │
  ↓                                      ↓
Residual Add                           Residual Add
  │                                      │
  ↓                                      ↓
Layer Norm                             Layer Norm
  │                                      │
  ↓                                      ↓
┌─────────────────┐                   ┌─────────────────────────────┐
│      MLP        │                   │          MOE Layer          │
│ (768→3072→768) │                   │  ┌─────────┐  ┌──────────┐  │
│                 │                   │  │ Gating  │  │ Experts  │  │
│  一个前馈网络   │                   │  │ Network │  │ E1...En  │  │
│                 │                   │  └────┬────┘  └────┬─────┘  │
│  ~4.7M 参数    │                   │       └──────┬─────┘        │
│  100% 激活     │                   │       加权求和 (稀疏)        │
└─────────────────┘                   │                             │
  │                                   │  n × ~4.7M 参数             │
  ↓                                   │  top-k/n × 100% 激活        │
Residual Add                          └─────────────────────────────┘
  │                                      │
  ↓                                      ↓
Output                                 Residual Add
                                         │
                                         ↓
                                       Output

9.2 关键公式对比

组件 标准 MLP MOE
前向传播 \(y = W_2 \cdot \text{GELU}(W_1 \cdot x)\) \(y = \sum_{i \in \text{TopK}} G(x)_i \cdot E_i(x)\)
参数量 \(2 \cdot d \cdot d_{ff}\) \(n \cdot 2 \cdot d \cdot d_{ff} + d \cdot n\)
计算量 \(2 \cdot d \cdot d_{ff}\) \(k \cdot 2 \cdot d \cdot d_{ff} + d \cdot n\)
激活比例 100% \(k/n\)

9.3 反向传播对比

步骤 标准 MLP MOE
梯度传播 传给所有参数 只传给被选中的专家
参数更新 所有参数 门控网络全部 + 选中专家
辅助损失 负载均衡损失

十、参考资料


本文结合 Transformer 架构,详细分析了 MOE 的工作原理。MOE 通过稀疏激活实现了"大参数量、小计算量"的目标,是构建超大规模语言模型的关键技术。

posted @ 2025-12-11 10:01  b1uesk9  阅读(217)  评论(0)    收藏  举报