门控网络和专家网络是高层组成部分,它们可以被组合成一个精确的数学模型。这个模型详细说明了单个token通过稀疏MoE层的完整前向计算过程,概述了具体的计算步骤。门控网络:从输入到路由权重该过程始于门控网络,也称作路由器。它的作用是确定哪些专家应处理当前的输入token。输入是一个token嵌入,表示为一个向量 $x \in \mathbb{R}^d$,其中$d$是模型的隐藏维度。门控网络本身是一个简单的线性层,由一个权重矩阵 $W_g \in \mathbb{R}^{d \times N}$定义,其中$N$是专家总数。该层将输入token投射到一个$N$维空间,为每个专家生成一个logit值。$$ h(x) = x \cdot W_g $$得到的向量$h(x)$包含$N$个原始得分。为了将这些得分转换为概率分布,我们应用softmax函数:$$ g(x) = \text{softmax}(h(x)) $$输出$g(x)$是一个稠密的$N$维向量,其中每个元素$g(x)_i$表示路由器将token分配给专家$i$的置信度。$g(x)$中所有元素的和为1。通过Top-K门控实现稀疏性一个稠密的$g(x)$向量意味着每个专家都会对输出做出贡献,这与MoE计算效率的目标相悖。为了强制稀疏性,我们采用TopK操作。不使用所有专家,我们选择得分最高的专家中的一小部分固定数量$k$。对于给定的token,我们识别出$g(x)$中Top $k$个值的索引,并将所有其他门控值设为零。这就创建了一个稀疏门控向量$G(x)$。$k$的选择是一个重要超参数。在Switch Transformers中,$k=1$,意味着每个token被路由到一个专家。更常见的选择是$k=2$,这为学习更复杂的函数提供了途径,并增加了一定程度的冗余。这一操作有效地修剪了每个token的计算图。如果$k=2$且我们有$N=64$个专家,我们只需对其中2个进行前向计算,忽略其余62个。专家网络$N$个专家中的每一个通常是一个独立的自前馈网络(FFN)。尽管它们都共享相同的架构,但它们不共享权重。每个专家$E_i$都有一套自己的参数。一个标准的双层FFN专家可以表示为:$$ E_i(x) = \text{ReLU}(x \cdot W_{1,i}) \cdot W_{2,i} $$这里,$W_{1,i}$和$W_{2,i}$分别是专家$i$的第一和第二线性层的权重矩阵。正是这种独立的专家权重集合导致模型总参数量大幅增加。完整的前向计算我们现在可以结合这些步骤来定义MoE层的最终输出$y(x)$。输出是来自所选专家的输出的加权和,使用TopK操作得到的稀疏门控权重。$$ y(x) = \sum_{i=1}^{N} G(x)_i \cdot E_i(x) $$因为$G(x)$是稀疏的,只有$k$个非零值,这个求和是计算高效的。我们只需计算路由器选择的$k$个专家的$E_i(x)$。单个token的完整数据流可以如下所示:digraph G { rankdir=TB; splines=ortho; node [shape=box, style="rounded,filled", fontname="sans-serif", margin="0.2,0.1"]; edge [fontname="sans-serif", fontsize=10]; subgraph cluster_input { label="输入Token"; style=invis; x [label="输入x", fillcolor="#a5d8ff"]; } subgraph cluster_gating { label="门控网络(路由器)"; bgcolor="#ffec99"; style=rounded; gating_linear [label="线性层 (Wg)", fillcolor="#ffe066"]; softmax [label="Softmax", fillcolor="#ffe066"]; topk [label="Top-K 选择\n(k=2)", fillcolor="#ffe066"]; } subgraph cluster_experts { label="专家网络"; bgcolor="#e9ecef"; style=rounded; e1 [label="专家 1 (FFN)", fillcolor="#dee2e6", style="dashed"]; e2 [label="专家 2 (FFN)", fillcolor="#b2f2bb"]; e_dots [label="...", shape=none]; eN [label="专家 N (FFN)", fillcolor="#b2f2bb"]; } subgraph cluster_output { label="最终输出"; style=invis; sum_op [label="加权和 Σ", shape=circle, fillcolor="#ffc9c9"]; y [label="输出 y(x)", fillcolor="#a5d8ff"]; } // Connections x -> gating_linear [label="d维向量"]; gating_linear -> softmax [label="N维Logits"]; softmax -> topk [label="N维权重"]; topk -> sum_op [label="G(x) (稀疏权重)", color="#f03e3e", fontcolor="#f03e3e"]; x -> e1 [style=dashed]; x -> e2; x -> e_dots [style=invis]; x -> eN; e1 -> sum_op [label="E₁(x)", style=dashed]; e2 -> sum_op [label="E₂(x)"]; eN -> sum_op [label="Eɴ(x)"]; sum_op -> y; // Invisible edges for alignment e_dots -> sum_op [style=invis]; // Annotations {rank=same; e1; e2; e_dots; eN;} note [label="在此示例中,专家2和专家N\n由Top-K门控选中。\n专家1未激活(虚线所示)。", shape=note, fillcolor="#fff9db", style=filled, align=left]; topk -> note [style=invis]; }单个token通过MoE层的数据流。输入x被送往门控网络以生成稀疏权重,并同时送往所选专家进行处理。关于可微分性和权重重新归一化的说明TopK函数是不可微分的,这给反向传播带来了问题。实际中,这通过直通估计器处理。在前向计算时,我们应用离散的TopK选择。在反向传播时,我们使梯度通过顶部$k$个门,就好像选择是一个简单的乘法。稠密门控输出$g(x)$用于计算门控权重$W_g$的梯度。此外,从初始softmax输出$g(x)$中选择Top $k$个值后,它们的和不再保证为1。为了形成一个合适的凸组合,这些$k$个值通常会被重新归一化。这通常通过仅对$h(x)$中选定的Top $k$个logit应用第二次softmax来完成。这确保了最终求和中使用的权重准确反映其相对重要性并和为1。这种表述提供了一个参数量庞大的模型,但每个token的计算成本是恒定的,由$k$而非$N$决定。然而,这种精巧的结构带来了一个重大挑战:如果门控网络学习到将大多数token路由到少数几个专家,其他专家将不会收到训练信号。这导致了专家退化问题,我们接下来通过引入负载均衡损失来解决。