Router门控网络简单介绍

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）扩容的关键技术。