门控网络(或路由)是混合专家(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}$。digraph G { rankdir=LR; node [shape=box, style=rounded, fontname="sans-serif", color="#495057", fillcolor="#e9ecef", style=filled]; edge [color="#868e96"]; InputToken [label="输入标记 (x)"]; LinearLayer [label="线性层 (Wg)", shape=ellipse, color="#1c7ed6", fillcolor="#a5d8ff"]; Logits [label="Logits (h)", shape=ellipse]; Selection [label="Top-k 选择", shape=diamond, color="#1098ad", fillcolor="#99e9f2"]; Expert1 [label="专家 1"]; ExpertN [label="专家 N"]; Dots [label="..."]; InputToken -> LinearLayer; LinearLayer -> Logits; Logits -> Selection; Selection -> Expert1 [label=" 已选"]; Selection -> ExpertN [label=" 已选"]; Selection -> Dots [style=invis]; subgraph cluster_experts { label = "专家"; style = filled; color = "#dee2e6"; Expert1; ExpertN; Dots; } }线性路由架构的基本数据流。优点:计算效率高: 只需一次矩阵乘法,速度很快,与专家计算本身相比,开销很小。简单: 易于实现和理解。参数数量相对较少($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选择,和之前一样,也可能添加噪声。digraph G { rankdir=LR; node [shape=box, style=rounded, fontname="sans-serif", color="#495057", fillcolor="#e9ecef", style=filled]; edge [color="#868e96"]; InputToken [label="输入标记 (x)"]; Linear1 [label="线性层 1 (Wg1)", shape=ellipse, color="#1c7ed6", fillcolor="#a5d8ff"]; Activation [label="激活函数", shape=ellipse, color="#ae3ec9", fillcolor="#eebefa"]; Linear2 [label="线性层 2 (Wg2)", shape=ellipse, color="#1c7ed6", fillcolor="#a5d8ff"]; Logits [label="Logits (h)", shape=ellipse]; Selection [label="Top-k 选择", shape=diamond, color="#1098ad", fillcolor="#99e9f2"]; Expert1 [label="专家 1"]; ExpertN [label="专家 N"]; Dots [label="..."]; InputToken -> Linear1; Linear1 -> Activation; Activation -> Linear2; Linear2 -> Logits; Logits -> Selection; Selection -> Expert1 [label=" 已选"]; Selection -> ExpertN [label=" 已选"]; Selection -> Dots [style=invis]; subgraph cluster_experts { label = "专家"; style = filled; color = "#dee2e6"; Expert1; ExpertN; Dots; } }非线性(MLP)路由架构的基本数据流。优点:表达能力增强: 能模拟标记表示和专家匹配度之间更复杂、非线性的关系。可能实现更有效的专业化。提升专业化潜力: 路由能学习更复杂的决策边界来路由标记。缺点:计算开销增加: 引入额外的矩阵乘法和非线性函数评估,增加了路由的计算量。参数量增加: 需要比线性路由更多的参数,特别是当隐藏维度$d_{hidden}$很大时。训练复杂性: 可能略微更难训练和稳定。路由本身可能遇到更深层网络常见的优化难题,尽管路由MLP通常保持较浅(1-2层)。基于注意力机制的路由一种新的、更先进的方法是在路由中加入注意力机制。这使得路由在进行路由决策时,能衡量输入表示的不同部分,甚至考虑语境信息。有几种设计方法:自注意力预路由: 在将输入标记表示$x$送入(可能是线性的)路由层之前,对其应用自注意力层。这使得路由能够处理标记的语境增强表示。专家查询注意力: 使用专家特定的查询来关注输入标记的表示(键和值从$x$导出)。注意力分数可以直接影响路由logits。例如,为每个专家$e$学习可查询向量$q_e$: $$ \alpha_e = \text{Attention}(q_e, K_x, V_x) $$ 其中,$K_x, V_x$是输入标记$x$的投影。得到的结果$\alpha_e$值(可能经过进一步处理后)构成路由logits $h$。digraph G { rankdir=LR; node [shape=box, style=rounded, fontname="sans-serif", color="#495057", fillcolor="#e9ecef", style=filled]; edge [color="#868e96"]; InputToken [label="输入标记 (x)"]; AttentionMech [label="注意力机制\n(例如:自注意力或\n专家查询注意力)", shape=box, color="#f06595", fillcolor="#fcc2d7", width=2.5]; RouterHead [label="路由头部\n(例如:线性)", shape=ellipse, color="#1c7ed6", fillcolor="#a5d8ff"]; Logits [label="Logits (h)", shape=ellipse]; Selection [label="Top-k 选择", shape=diamond, color="#1098ad", fillcolor="#99e9f2"]; Expert1 [label="专家 1"]; ExpertN [label="专家 N"]; Dots [label="..."]; InputToken -> AttentionMech; AttentionMech -> RouterHead; RouterHead -> Logits; Logits -> Selection; Selection -> Expert1 [label=" 已选"]; Selection -> ExpertN [label=" 已选"]; Selection -> Dots [style=invis]; subgraph cluster_experts { label = "专家"; style = filled; color = "#dee2e6"; Expert1; ExpertN; Dots; } }基于注意力机制的路由架构数据流。优点:表达能力最高: 注意力机制能捕捉复杂关联和语境细节,可能实现高度复杂和动态的路由策略。语境感知: 能根据更宽泛的语境调整路由(如果设计得当,例如关注多个标记,尽管这会极大增加复杂性)。缺点:计算开销巨大: 注意力机制计算量大,增加了显著的开销,特别是与线性或简单MLP路由相比。这可能成为一个瓶颈。实现和训练复杂性: 设计、实现和稳定基于注意力机制的路由要复杂得多。它们引入更多超参数和潜在的故障模式。延迟: 额外的计算会显著增加推理延迟。选择合适的路由架构路由架构的选择涉及平衡多个因素:特点线性路由非线性(MLP)路由基于注意力机制的路由表达能力低中高计算开销低中高 / 很高参数数量低中高实现难度简单中等复杂训练稳定性通常较好中等可能较难从简开始: 线性路由常提供强劲的基准,且计算成本最低。通常建议从这里开始,只有当性能停滞或出现特定路由需求时才增加复杂性。考虑任务: 更复杂的任务,如果能从高度专业化的专家那里受益,则可能更有效地使用非线性甚至基于注意力机制的路由。但收益必须大于成本。预算限制: 训练和推理期间可用的计算资源是重要因素。基于注意力机制的路由在资源受限的环境中可能不适用。训练情况: 密切关注负载均衡和专家专业化(如第3章所述)。如果较简单的路由导致专业化差或不均衡,可能需要更具表达能力的路由,但通常,处理辅助损失或正则化是更直接的解决办法。实践中,许多大型MoE模型仍倾向于选择在训练期间添加噪声的简单线性路由,因为它在性能和效率之间有良好的平衡。非线性路由在需要时提供中等的表达能力提升。基于注意力机制的路由代表了更偏研究方向,承诺更高的能力,但代价是显著的复杂性和计算量。最佳选择通常需要针对特定任务和数据集进行实证验证。