MoE数学架构 Mathematics
MOE（Mixture of Experts）架构¶
Mixture of Experts（MOE）是一种通过组合多个专家网络（Experts）来提高模型表达能力和计算效率的架构。每个专家负责处理输入数据的不同部分，而门控网络（Gating Network）则决定如何组合这些专家的输出。MOE架构广泛应用于大规模模型，如语言模型和图像处理模型，以实现参数高效和计算可扩展性。
MOE 的关键组件¶
| 组件 | 描述 |
|---|---|
| 专家网络（Experts） | 多个独立的子网络，每个网络专门处理输入数据的不同方面。 |
| 门控网络（Gating Network） | 决定每个输入应由哪些专家处理，以及各专家的权重。 |
| 稀疏激活（Sparse Activation） | 仅选择少数几个专家参与每次前向传播，提高计算效率。 |
| 负载均衡机制（Load Balancing Mechanism） | 确保各专家得到大致相同的负载，避免部分专家过载。 |
MOE 的数据流动¶
| 阶段 | 输入 | 操作 | 输出 | 输出去向 |
|---|---|---|---|---|
| 输入层 | 原始输入数据 | 数据预处理与嵌入 | \(X \in \mathbb{R}^{N \times d_{model}}\) | MOE层 |
| 门控网络 | \(X\) | 1. 线性变换：\(G = XW_g + b_g\) 2. Softmax：\(G_{ij} = \frac{\exp(G_{ij})}{\sum_{k=1}^{E} \exp(G_{ik})}\) 3. 选择 top-k 专家 |
选择的专家索引及权重 | 专家层 |
| 专家层 | 输入数据 \(X\) 和门控权重 \(G\) | 1. 将输入分配给选定的专家 2. 每个专家独立处理分配到的数据 3. 汇总专家输出 |
\(Y \in \mathbb{R}^{N \times d_{model}}\) | 聚合层 |
| 聚合层 | 专家输出 \(Y\) | 1. 根据门控权重加权组合专家输出 2. 叠加与后续层 |
\(Z \in \mathbb{R}^{N \times d_{model}}\) | 后续模型层 |
MOE 层详细结构¶
| 层 | 输入 | 操作 | 输出 |
|---|---|---|---|
| 门控网络 | 输入嵌入 \(X \in \mathbb{R}^{N \times d_{model}}\) | 1. 线性变换：\(G = XW_g + b_g\), \(W_g \in \mathbb{R}^{d_{model} \times E}\) 2. Softmax激活：\(G_{ij} = \frac{\exp(G_{ij})}{\sum_{k=1}^{E} \exp(G_{ik})}, \forall j \in \{1, ..., E\}\) 3. 选取 top-k 专家 |
选择的专家索引与权重 \(G' \in \mathbb{R}^{N \times k}\) |
| 专家处理 | 输入数据 \(X\), 选择的专家 \(G'\) | 1. 将输入分配到选定的专家 2. 每个专家独立执行前向传播：\(Y_e = \text{Expert}_e(X_e)\) 其中 \(e \in \{1, ..., E\}\) |
各专家的输出 \(Y_e \in \mathbb{R}^{N_e \times d_{model}}\) |
| 输出聚合 | 各专家输出 \(Y_e\), 门控权重 \(G'\) | 1. 根据门控权重加权：\(Y = \sum_{e=1}^{k} G'_{e} \cdot Y_e\) 2. 汇总所有专家的输出 |
聚合输出 \(Z \in \mathbb{R}^{N \times d_{model}}\) |
MOE 数据流动示例¶
| 层 | 输入 | 输出 | 输出去向 |
|---|---|---|---|
| 输入嵌入 | 原始输入数据 | \(X_0 \in \mathbb{R}^{N \times d_{model}}\) | MOE层 |
| MOE层 | \(X_0\) | \(Z \in \mathbb{R}^{N \times d_{model}}\) | 后续模型层 |
| 后续层 | \(Z\) | ... | ... |
注：
- \(N\) 是批量大小。
- \(d_{model}\) 是模型的隐藏维度。
- \(E\) 是专家的数量。
- \(k\) 是每次前向传播中激活的专家数量（通常较小，如2或4）。
MOE 层中的参数更新¶
门控网络¶
| 参数 | 维度 | 描述 |
|---|---|---|
| \(W_g\) | \(\mathbb{R}^{d_{model} \times E}\) | 门控网络的权重矩阵 |
| \(b_g\) | \(\mathbb{R}^{E}\) | 门控网络的偏置向量 |
专家网络¶
| 组件 | 参数 | 维度 | 描述 |
|---|---|---|---|
| 专家1 | \(W^{(1)}\) | \(\mathbb{R}^{d_{model} \times d_{ff}}\) | 专家1的权重 |
| \(b^{(1)}\) | \(\mathbb{R}^{d_{ff}}\) | 专家1的偏置 | |
| 专家2 | \(W^{(2)}\) | \(\mathbb{R}^{d_{model} \times d_{ff}}\) | 专家2的权重 |
| \(b^{(2)}\) | \(\mathbb{R}^{d_{ff}}\) | 专家2的偏置 | |
| ... | ... | ... | ... |
| 专家E | \(W^{(E)}\) | \(\mathbb{R}^{d_{model} \times d_{ff}}\) | 专家E的权重 |
| \(b^{(E)}\) | \(\mathbb{R}^{d_{ff}}\) | 专家E的偏置 |
注意：
- 每个专家都拥有独立的参数集。
- 专家数量 \(E\) 和每个专家的隐藏维度 \(d_{ff}\) 可根据具体需求调整。
- 通过稀疏激活（只激活少数专家），可以在保持高模型容量的同时，提高计算效率。