Router（路由）门控网络，它是MoE（Mixture of Experts，混合专家模型）架构中的核心组件，主要作用是为每个输入token（或序列）选择最合适的"专家"网络（通常是FFN层）来处理，以此在不显著增加计算量的前提下提升模型容量。
一、Router门控网络的核心定位
在标准Transformer中，所有token共享同一个FFN层；而在MoE中，FFN层被拆分为多个独立的"专家"（Expert），Router的作用就是：
- 接收每个token的特征向量作为输入；
- 计算该token对所有专家的"匹配分数"；
- 根据分数选择Top-K个专家（通常K=1或2）来处理这个token；
- （可选）对选中专家的输出进行加权融合。
简单来说，Router就是MoE模型的"调度员"，让不同的token由最擅长处理它的专家来加工，实现"分而治之"。
二、核心结构与工作流程
1. 数学原理（以最常用的Top-1/Top-2 Router为例）
假设模型有 n_experts 个专家，输入token的特征维度为 d_model：
-
第一步：计算评分
Router通过一个线性层将输入向量映射为
n_experts维的分数：
logits(x)=Wr⋅x+br \text{logits}(x) = W_r \cdot x + b_r logits(x)=Wr⋅x+br其中 Wr∈Rnexperts×dmodelW_r \in \mathbb{R}^{n_{\text{experts}} \times d_{\text{model}}}Wr∈Rnexperts×dmodel 是Router的权重矩阵。
-
第二步：归一化与选择
用Softmax对logits归一化得到每个专家的"选中概率"，再选择概率最高的Top-K个专家：
pi=Softmax(logits(x))i p_i = \text{Softmax}(\text{logits}(x))_i pi=Softmax(logits(x))i（工程上常使用Gumbel-Softmax或直接硬选择，避免梯度消失）
-
第三步：加权输出
仅让选中的专家处理token，再将专家输出按概率加权求和：
output(x)=∑i∈Top-Kpi⋅Experti(x) \text{output}(x) = \sum_{i \in \text{Top-K}} p_i \cdot \text{Expert}_i(x) output(x)=i∈Top-K∑pi⋅Experti(x)
2. 关键约束：负载均衡
为了避免少数专家被过度调用（导致计算瓶颈），Router会加入负载均衡损失 ，强制每个专家被选中的概率接近 1/n_experts（或按预设比例），保证计算资源均匀分配。
三、PyTorch实现简单的Router门控网络（结合MoE-FFN）
下面是可直接运行的代码，实现Top-2 Router + MoE-FFN，清晰展示Router的核心逻辑：
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Expert(nn.Module):
"""单个专家网络（本质是FFN层）"""
def __init__(self, d_model: int, d_ff: int = 2048, dropout: float = 0.1):
super().__init__()
self.w1 = nn.Linear(d_model, d_ff)
self.w2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
self.act = nn.GELU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
out = self.w1(x)
out = self.act(out)
out = self.dropout(out)
out = self.w2(out)
return out
class Top2Router(nn.Module):
"""Top-2 Router门控网络：为每个token选择2个最优专家"""
def __init__(self, d_model: int, n_experts: int):
super().__init__()
# Router的核心：线性层映射为n_experts维评分
self.router_linear = nn.Linear(d_model, n_experts)
self.n_experts = n_experts
def forward(self, x: torch.Tensor):
"""
Args:
x: 输入张量 [batch_size, seq_len, d_model]
Returns:
gate_scores: 选中专家的权重 [batch_size, seq_len, 2]
expert_indices: 选中的专家索引 [batch_size, seq_len, 2]
"""
# 1. 计算每个token对所有专家的评分
logits = self.router_linear(x) # [bs, seq_len, n_experts]
# 2. 选择Top-2专家的索引和分数
gate_scores, expert_indices = torch.topk(logits, k=2, dim=-1) # 均为 [bs, seq_len, 2]
# 3. Softmax归一化分数（保证权重和为1）
gate_scores = F.softmax(gate_scores, dim=-1)
return gate_scores, expert_indices
class MoEFFN(nn.Module):
"""带Router的MoE-FFN层（Router + 多专家）"""
def __init__(self, d_model: int, n_experts: int, d_ff: int = 2048, dropout: float = 0.1):
super().__init__()
self.n_experts = n_experts
self.d_model = d_model
# 1. 初始化Router
self.router = Top2Router(d_model, n_experts)
# 2. 初始化多个专家（用ModuleList管理）
self.experts = nn.ModuleList([
Expert(d_model, d_ff, dropout) for _ in range(n_experts)
])
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: 输入张量 [batch_size, seq_len, d_model]
Returns:
输出张量 [batch_size, seq_len, d_model]
"""
batch_size, seq_len, _ = x.shape
# 1. Router选择专家
gate_scores, expert_indices = self.router(x) # [bs, seq_len, 2], [bs, seq_len, 2]
# 2. 初始化输出张量
output = torch.zeros_like(x)
# 3. 遍历每个专家，处理被分配到该专家的token
for expert_idx in range(self.n_experts):
# 找到所有选中当前专家的token位置（Top-1或Top-2）
mask = (expert_indices == expert_idx) # [bs, seq_len, 2]
# 收集这些token的权重和位置
for k in range(2): # 遍历Top-1和Top-2
# 维度：[bs, seq_len]
k_mask = mask[..., k]
k_scores = gate_scores[..., k]
if not k_mask.any():
continue # 无token选中该专家，跳过
# 提取需要当前专家处理的token
selected_x = x[k_mask] # [n_selected, d_model]
# 专家处理
expert_out = self.experts[expert_idx](selected_x) # [n_selected, d_model]
# 加权并回填到输出
output[k_mask] += k_scores[k_mask].unsqueeze(-1) * expert_out
return output
# ------------------- 测试代码 -------------------
if __name__ == "__main__":
# 超参数
BATCH_SIZE = 2
SEQ_LEN = 10
D_MODEL = 512
N_EXPERTS = 4 # 4个专家
# 创建MoE-FFN（含Router）
moe_ffn = MoEFFN(d_model=D_MODEL, n_experts=N_EXPERTS)
# 测试输入
test_input = torch.randn(BATCH_SIZE, SEQ_LEN, D_MODEL)
# 前向传播
output = moe_ffn(test_input)
print("输入形状:", test_input.shape) # [2, 10, 512]
print("输出形状:", output.shape) # [2, 10, 512]
print("专家数量:", len(moe_ffn.experts)) # 4四、代码关键解释
- Expert类：就是普通的FFN层，每个专家独立参数，负责处理特定类型的token；
- Top2Router类 ：核心是
router_linear线性层，输出每个token对所有专家的评分，再通过topk选择最优的2个专家，并用Softmax归一化权重； - MoEFFN类：整合Router和所有专家，遍历每个专家处理被分配的token，最后加权融合输出，保证每个token仅由2个专家处理（大幅降低计算量）。
五、Router的常见变体
- 硬Router vs 软Router ：
- 硬Router：直接选择Top-K专家（硬分配，无梯度），工程上常用Gumbel-Softmax近似梯度；
- 软Router：对所有专家的输出加权（软分配），但计算量接近全专家，很少用。
- 负载均衡策略 ：
- 加入"专家容量限制"（每个专家最多处理固定数量的token）；
- 增加负载均衡损失，惩罚被过度选中的专家。
- 动态Router：根据输入序列长度、token类型动态调整K值（选中专家数）。
总结
- Router门控网络是MoE模型的核心，核心作用是为每个token选择最优的Top-K个专家（FFN层），实现"分治"提升模型容量；
- Router的核心逻辑是"线性评分→Top-K选择→加权融合"，需配合负载均衡保证计算效率；
- 工程上最常用Top-2 Router，兼顾模型效果和计算成本，是大模型（如GPT-3、PaLM）扩容的关键技术。