路由架构：线性、非线性、基于注意力

门控网络（或路由）是混合专家（MoE）层的控制中心。它的主要作用是检查每个输入标记，然后决定哪个或哪些专家应该处理它。这个路由的设计不只是一个实现细节，反而对模型实现专业化、训练稳定性以及整体计算情况有很深远的影响。路由的不同架构选择在表达能力、计算开销和优化便利性之间有各自的取舍。以下分析三种常见的路由架构：线性的、非线性的和基于注意力的。

线性路由

最简单的路由形式采用单一线性变换，然后是一个函数来获得专家概率或分配，通常是使用softmax来生成概率，或直接基于logits进行top-k选择。

给定输入标记表示$x \in \mathbb{R}^d$，其中$d$是模型维度，并且有$N$个可用专家，路由使用学到的权重矩阵$W_g \in \mathbb{R}^{d \times N}$来计算logits $h \in \mathbb{R}^N$：

$$h = x W_g$$

这些logits $h$直接决定选择过程。例如，在概率路由场景中（在现代稀疏MoE中不太常见，但可用于说明），可以通过softmax计算门控概率$p \in \mathbb{R}^N$：

$$p = \text{softmax}(h)$$

更常见的是，对于top-k路由（选择$k$个专家，通常$k=1$或$k=2$），路由直接选择$h$中对应$k$个最大值的专家。另外，在top-k选择之前可以向logits添加噪声，特别是在训练期间，以促使试探并改善负载均衡：

$$h_{noisy} = h + \epsilon \cdot \text{softplus}(x W_{noise})$$

其中$\epsilon$是从标准正态分布中采样的，$W_{noise}$是另一个学到的投影。最终的专家选择使用$h$或$h_{noisy}$。

线性路由架构的基本数据流。

优点：

• 计算效率高： 只需一次矩阵乘法，速度很快，与专家计算本身相比，开销很小。
• 简单： 易于实现和理解。参数数量相对较少（$d \times N$）。
• 稳定性好： 通常比更复杂的路由更容易训练，是一个可靠的基准。

缺点：

• 表达能力有限： 线性变换可能不足以捕捉复杂条件逻辑，这对于复杂的专家专业化来说是必需的。路由只能学习输入空间中的线性分离来进行路由决策。
• 可能出现退化： 如果没有适当的正则化或负载均衡机制（在第3章讨论），线性路由有时可能导致表示退化，或出现只有少数专家被持续选择的情况。

非线性路由

为了提升路由学习复杂路由函数的能力，可以引入非线性，通常通过将路由构建为小型多层感知机（MLP）来实现。

非线性路由不使用单一线性层，而是可能使用一个或多个带激活函数（如ReLU、GeLU或Swish）的隐藏层。例如，一个单隐藏层MLP路由：

$$h_{hidden} = \text{Activation}(x W_{g1} + b_{g1})$$ $$h = h_{hidden} W_{g2} + b_{g2}$$

其中，$W_{g1} \in \mathbb{R}^{d \times d_{hidden}}$、$b_{g1} \in \mathbb{R}^{d_{hidden}}$、$W_{g2} \in \mathbb{R}^{d_{hidden} \times N}$和$b_{g2} \in \mathbb{R}^{N}$是学到的参数。最终的logits $h$接着用于top-k选择，和之前一样，也可能添加噪声。

非线性（MLP）路由架构的基本数据流。

优点：

• 表达能力增强： 能模拟标记表示和专家匹配度之间更复杂、非线性的关系。可能实现更有效的专业化。
• 提升专业化潜力： 路由能学习更复杂的决策边界来路由标记。

缺点：

• 计算开销增加： 引入额外的矩阵乘法和非线性函数评估，增加了路由的计算量。
• 参数量增加： 需要比线性路由更多的参数，特别是当隐藏维度$d_{hidden}$很大时。
• 训练复杂性： 可能略微更难训练和稳定。路由本身可能遇到更深层网络常见的优化难题，尽管路由MLP通常保持较浅（1-2层）。

基于注意力机制的路由

一种新的、更先进的方法是在路由中加入注意力机制。这使得路由在进行路由决策时，能衡量输入表示的不同部分，甚至考虑语境信息。

有几种设计方法：

1. 自注意力预路由： 在将输入标记表示$x$送入（可能是线性的）路由层之前，对其应用自注意力层。这使得路由能够处理标记的语境增强表示。
2. 专家查询注意力： 使用专家特定的查询来关注输入标记的表示（键和值从$x$导出）。注意力分数可以直接影响路由logits。例如，为每个专家$e$学习可查询向量$q_e$： $$\alpha_e = \text{Attention}(q_e, K_x, V_x)$$ 其中，$K_x, V_x$是输入标记$x$的投影。得到的结果$\alpha_e$值（可能经过进一步处理后）构成路由logits $h$。

基于注意力机制的路由架构数据流。

优点：

• 表达能力最高： 注意力机制能捕捉复杂关联和语境细节，可能实现高度复杂和动态的路由策略。
• 语境感知： 能根据更宽泛的语境调整路由（如果设计得当，例如关注多个标记，尽管这会极大增加复杂性）。

缺点：

• 计算开销巨大： 注意力机制计算量大，增加了显著的开销，特别是与线性或简单MLP路由相比。这可能成为一个瓶颈。
• 实现和训练复杂性： 设计、实现和稳定基于注意力机制的路由要复杂得多。它们引入更多超参数和潜在的故障模式。
• 延迟： 额外的计算会显著增加推理延迟。

选择合适的路由架构

路由架构的选择涉及平衡多个因素：

特点	线性路由	非线性（MLP）路由	基于注意力机制的路由
表达能力	低	中	高
计算开销	低	中	高 / 很高
参数数量	低	中	高
实现难度	简单	中等	复杂
训练稳定性	通常较好	中等	可能较难

• 从简开始： 线性路由常提供强劲的基准，且计算成本最低。通常建议从这里开始，只有当性能停滞或出现特定路由需求时才增加复杂性。
• 考虑任务： 更复杂的任务，如果能从高度专业化的专家那里受益，则可能更有效地使用非线性甚至基于注意力机制的路由。但收益必须大于成本。
• 预算限制： 训练和推理期间可用的计算资源是重要因素。基于注意力机制的路由在资源受限的环境中可能不适用。
• 训练情况： 密切关注负载均衡和专家专业化（如第3章所述）。如果较简单的路由导致专业化差或不均衡，可能需要更具表达能力的路由，但通常，处理辅助损失或正则化是更直接的解决办法。

实践中，许多大型MoE模型仍倾向于选择在训练期间添加噪声的简单线性路由，因为它在性能和效率之间有良好的平衡。非线性路由在需要时提供中等的表达能力提升。基于注意力机制的路由代表了更偏研究方向，承诺更高的能力，但代价是显著的复杂性和计算量。最佳选择通常需要针对特定任务和数据集进行实证验证。