专家混合(MoE)层是一种稀疏的神经网络架构,与密集架构形成对比。基本MoE层所涉及的计算被正式阐述。这种数学描述准确阐明了输入数据如何流经系统、路由决策如何做出以及专家输出如何组合。理解此过程对于后续考察更复杂的架构变体和训练优化很重要。考虑一个集成在更大网络(例如Transformer块)中的MoE层。该层为序列中的每个标记接收一个输入表示。令 $x \in \mathbb{R}^d$ 表示单个标记的输入嵌入,其中 $d$ 是模型的隐藏维度。MoE层由 $N$ 个专家网络(表示为 $E_1, E_2, \dots, E_N$)和一个门控网络 $G$ 组成。门控网络的工作方式门控网络的主要功能是决定哪个或哪些专家应处理输入标记 $x$。一个常见实现是使用简单的线性变换,之后是softmax函数。线性变换: 输入标记 $x$ 使用一个可学习的权重矩阵 $W_g \in \mathbb{R}^{d \times N}$ 进行投影。这会产生logit值 $h(x) \in \mathbb{R}^N$,每个专家对应一个: $$h(x) = x W_g$$Softmax激活: Logit值随后通过一个softmax函数,生成归一化分数或概率,表明对每个专家的倾向性: $$G(x) = \text{softmax}(h(x))$$ 结果向量 $G(x) = [G(x)_1, G(x)_2, \dots, G(x)_N]$ 包含非负值且总和为1,其中 $G(x)_i$ 表示门控网络将标记 $x$ 分配给专家 $E_i$ 的分数。digraph G { rankdir=LR; node [shape=box, style="filled", fontname="Arial", fontsize=10, margin=0.1]; edge [fontname="Arial", fontsize=9]; // 定义输入标记子图 subgraph cluster_input { label = "输入标记"; bgcolor="#e9ecef"; x [label="x (d维)", shape=ellipse, style=filled, fillcolor="#a5d8ff"]; } // 定义门控网络子图 subgraph cluster_gating { label = "门控网络 G"; bgcolor="#e9ecef"; Wg [label="线性变换 (Wg)\n(d x N)", style=filled, fillcolor="#bac8ff"]; Softmax [label="Softmax", style=filled, fillcolor="#bac8ff"]; x -> Wg [label=" "]; Wg -> Softmax [label="Logit值 h(x)"]; } // 定义门控分数子图 subgraph cluster_scores { label = "门控分数"; bgcolor="#e9ecef"; G_scores [label="G(x) = [G(x)₁, ..., G(x)ɴ]\n(N维)", shape=ellipse, style=filled, fillcolor="#74c0fc"]; Softmax -> G_scores [label=" "]; } }数据流经门控网络,将输入标记表示 $x$ 转换为专家分配分数 $G(x)$。专家网络每个专家 $E_i$ 通常是一个神经网络模块,常见的是一个标准前馈网络(FFN),类似于Transformer块中的FFN。重要的是,每个专家 $E_i$ 都拥有自己独特的参数集。对于一个输入标记 $x$,第 $i$ 个专家的输出表示为 $E_i(x)$。尽管结构上相似,但参数差异使得每个专家能够专门处理不同类型的输入模式或执行不同的子任务。稀疏路由与输出计算虽然门控网络为每个专家生成一个分数 $G(x)_i$,但稀疏MoE模型通过仅为每个标记激活这些专家中的一部分来运用条件计算。这通常通过Top-k路由实现,即只选择门控分数最高的 $k$ 个专家(通常 $k=1$ 或 $k=2$)。专家选择: 根据 $G(x)$ 中的分数,确定Top-k个专家的索引 $I$: $$I = \text{TopKIndices}(G(x), k)$$ 例如,如果 $k=2$ 且专家 $j$ 和 $l$ 具有最高分数,则 $I = {j, l}$。门控分数重新归一化(可选但常见): 选定专家($i \in I$)的门控分数通常被保留并可能重新归一化。为此处简化,我们将用于加权专家 $i$ 输出的分数表示为 $g_i = G(x)_i$。一些实现可能只对Top-k分数重新应用softmax。专家计算: 输入标记 $x$ 仅由选定的专家处理: $$O_i(x) = E_i(x) \quad \text{for } i \in I$$加权输出组合: MoE层对标记 $x$ 的最终输出 $y$ 计算为选定专家输出的加权和,使用它们对应的门控分数作为权重: $$y = \sum_{i \in I} g_i O_i(x) = \sum_{i \in I} G(x)_i E_i(x)$$这种表述确保对于任何给定的输入标记,仅有专家参数的一小部分($k/N$)被使用,相比于激活所有 $N$ 个专家或使用单个大型专家,这在前向传播过程中显著节省了计算资源。digraph G { rankdir=TB; node [shape=box, style="filled", fontname="Arial", fontsize=10, margin=0.1]; edge [fontname="Arial", fontsize=9]; x [label="输入标记 x", shape=ellipse, style=filled, fillcolor="#a5d8ff"]; subgraph cluster_routing { label = "路由逻辑"; bgcolor="#e9ecef"; Gating [label="门控网络 G(x)\n分数 [G(x)₁, ..., G(x)ɴ]", style=filled, fillcolor="#74c0fc"]; TopK [label="Top-k 选择 (k=2)", style=filled, fillcolor="#ffc9c9"]; Gating -> TopK; } subgraph cluster_experts { label = "专家网络 (共N个)"; bgcolor="#e9ecef"; node [style=filled, fillcolor="#96f2d7"]; E1 [label="E₁"]; E2 [label="E₂"]; Edots [label="...", shape=plaintext]; EN [label="Eɴ"]; } Combiner [label="加权和\nΣ G(x)ᵢ Eᵢ(x)", style=filled, fillcolor="#ffd8a8"]; y [label="输出 y", shape=ellipse, style=filled, fillcolor="#ffec99"]; x -> Gating [style=dashed]; // 假设 TopK 选择了 E1 和 EN 作为示例 TopK -> E1 [label="G(x)₁", color="#fa5252", fontcolor="#fa5252"]; TopK -> EN [label="G(x)ɴ", color="#fa5252", fontcolor="#fa5252"]; // 输入仅流向选定的专家 x -> E1 [label=" ", style=solid, color="#1c7ed6"]; x -> EN [label=" ", style=solid, color="#1c7ed6"]; // 选定专家的输出流向组合器 E1 -> Combiner [label="E₁(x)"]; EN -> Combiner [label="Eɴ(x)"]; Combiner -> y; // 指示未选定的专家(可选的视觉提示) { rank=same; E1; E2; Edots; EN; } TopK -> E2 [style=dotted, color="#adb5bd", arrowhead=none]; } 单个标记 $x$ 在采用 $k=2$ 路由的MoE层中的前向传播。门控网络生成分数,Top-k选择专家(此处显示E1和EN),输入仅由选定的专家处理,它们的输出使用门控分数进行组合。批量处理的考量实际中,MoE层对标记批次进行操作。对于一个包含 $T$ 个标记的批次,输入通常表示为一个矩阵 $X \in \mathbb{R}^{T \times d}$。门控网络同时计算所有标记的分数 $G(X) \in \mathbb{R}^{T \times N}$。重要的一点是,Top-k选择对批次中的每个标记(行)独立执行。这意味着同一批次中的不同标记可以被路由到不同的专家组。尽管高效,但这种每个标记的独立路由带来了挑战,特别是在负载均衡方面:确保专家在整个批次中获得大致相等的计算量。如果门控网络持续将大多数标记路由到一小部分专家,其他专家则利用不足,削弱了专业化的优势,并可能导致训练不稳定。我们将在第3章讨论这些训练动态和相关的优化方法。这个数学框架提供了必要的构建模块,用来理解基本MoE层如何执行条件计算。门控网络、稀疏路由机制和专家网络之间的关联构成了MoE方法的主要组成部分。