图解 MoE 动态路由与负载均衡

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """单个专家网络（FFN 结构）"""
    def __init__(self, d_model, d_expert):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_expert)
        self.fc2 = nn.Linear(d_expert, d_model)
    
    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))

class GatingNetwork(nn.Module):
    """门控网络：计算路由概率"""
    def __init__(self, d_model, num_experts):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts)
    
    def forward(self, x):
        # 输出形状：[batch, num_experts]
        return F.softmax(self.gate(x), dim=-1)

class MoELayer(nn.Module):
    """基础 MoE 层"""
    def __init__(self, d_model, num_experts, d_expert, top_k=1):
        super().__init__()
        self.experts = nn.ModuleList([
            Expert(d_model, d_expert) for _ in range(num_experts)
        ])
        self.gate = GatingNetwork(d_model, num_experts)
        self.top_k = top_k
    
    def forward(self, x):
        # x: [batch, d_model]
        gate_probs = self.gate(x)  # [batch, num_experts]
        
        # Top-K 选择
        topk_probs, topk_indices = torch.topk(gate_probs, self.top_k, dim=-1)
        
        # 将选中的 token 分发给对应专家
        outputs = []
        for i in range(x.shape[0]):
            expert_outputs = sum(
                topk_probs[i, k] * self.experts[topk_indices[i, k]](x[i:i+1])
                for k in range(self.top_k)
            )
            outputs.append(expert_outputs)
        
        return torch.cat(outputs, dim=0)

class NoisyTopKGating(nn.Module):
    """带噪声注入的 Top-K 门控网络"""
    def __init__(self, d_model, num_experts, top_k=2, noise_epsilon=1e-2):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts)
        self.top_k = top_k
        self.noise_epsilon = noise_epsilon
    
    def forward(self, x, training=False):
        # 计算门控 logits
        logits = self.gate(x)  # [batch, num_experts]
        
        # 训练时添加高斯噪声
        if training and self.noise_epsilon > 0:
            noise = torch.randn_like(logits) * self.noise_epsilon
            logits = logits + noise
        
        # Top-K 选择
        topk_logits, topk_indices = torch.topk(logits, self.top_k, dim=-1)
        
        # 创建掩码：只保留 Top-K 的 logits
        mask = torch.zeros_like(logits).scatter_(1, topk_indices, 1.0)
        masked_logits = logits * mask
        
        # 对选中的专家进行 Softmax
        gate_probs = F.softmax(masked_logits, dim=-1)
        
        return gate_probs, topk_indices

def compute_load_balance_loss(gate_probs, num_experts):
    """
    计算负载均衡损失
    
    Args:
        gate_probs: [batch, num_experts] - 门控概率
        num_experts: int - 专家数量
    
    Returns:
        loss: float - 负载均衡损失值
    """
    # P_j: 每个专家的平均路由概率
    # 对 batch 维度取平均，得到每个专家被选择的平均概率
    P = gate_probs.mean(dim=0)  # [num_experts]
    
    # T_j: 每个专家实际处理的 token 比例
    # 使用硬分配：每个 token 只算它选中的专家
    # 对于 Top-K，这里用 gate_probs 近似（因为 Top-K 后概率集中在选中的专家）
    T = (gate_probs > 0).float().mean(dim=0)  # [num_experts]
    
    # 负载均衡损失：N * sum(P_j * T_j)
    # 当分布均匀时，P_j = T_j = 1/N，loss 最小
    loss = num_experts * (P * T).sum()
    
    return loss

class MoELayerWithBalance(nn.Module):
    """带负载均衡的完整 MoE 层"""
    def __init__(self, d_model, num_experts, d_expert, top_k=2, aux_alpha=0.01):
        super().__init__()
        self.experts = nn.ModuleList([
            Expert(d_model, d_expert) for _ in range(num_experts)
        ])
        self.gate = NoisyTopKGating(d_model, num_experts, top_k)
        self.num_experts = num_experts
        self.top_k = top_k
        self.aux_alpha = aux_alpha  # 平衡损失权重
    
    def forward(self, x, return_aux_loss=False):
        gate_probs, topk_indices = self.gate(x, training=self.training)
        
        # 专家计算
        outputs = []
        for i in range(x.shape[0]):
            expert_output = sum(
                gate_probs[i, k] * self.experts[topk_indices[i, k]](x[i:i+1])
                for k in range(self.top_k)
            )
            outputs.append(expert_output)
        output = torch.cat(outputs, dim=0)
        
        # 计算辅助损失
        aux_loss = compute_load_balance_loss(gate_probs, self.num_experts)
        
        if return_aux_loss:
            return output, aux_loss
        return output

# 构建使用 MoE 的分类器
class MoEClassifier(nn.Module):
    def __init__(self, input_dim, d_model, num_experts, d_expert, num_classes, top_k=2):
        super().__init__()
        self.embed = nn.Linear(input_dim, d_model)
        self.moe = MoELayerWithBalance(d_model, num_experts, d_expert, top_k)
        self.classifier = nn.Linear(d_model, num_classes)
    
    def forward(self, x):
        x = self.embed(x)
        x, aux_loss = self.moe(x, return_aux_loss=True)
        return self.classifier(x), aux_loss

# 训练循环
def train_moe(model, dataloader, optimizer, num_epochs, device):
    aux_alpha = 0.01
    
    for epoch in range(num_epochs):
        total_loss = 0
        expert_stats = {i: 0 for i in range(model.moe.num_experts)}
        
        for batch_x, batch_y in dataloader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)
            
            optimizer.zero_grad()
            
            # 前向传播
            logits, aux_loss = model(batch_x)
            
            # 总损失 = 分类损失 + alpha * 负载均衡损失
            task_loss = F.cross_entropy(logits, batch_y)
            total_loss_batch = task_loss + aux_alpha * aux_loss
            
            # 反向传播
            total_loss_batch.backward()
            optimizer.step()
            
            total_loss += total_loss_batch.item()
            
            # 统计专家利用率
            _, topk_indices = model.moe.gate(model.embed(batch_x))
            for idx in topk_indices.flatten().tolist():
                expert_stats[idx] += 1
        
        # 打印训练日志
        print(f"Epoch {epoch+1}: Loss={total_loss/len(dataloader):.4f}")
        print(f"专家利用率：{[expert_stats[i]/sum(expert_stats.values()) for i in range(len(expert_stats))]}")

# 使用示例
if __name__ == "__main__":
    model = MoEClassifier(
        input_dim=128,
        d_model=256,
        num_experts=8,
        d_expert=512,
        num_classes=10,
        top_k=2
    ).to("cuda")
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    
    # 创建假数据
    X = torch.randn(1000, 128)
    y = torch.randint(0, 10, (1000,))
    dataset = torch.utils.data.TensorDataset(X, y)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=32)
    
    train_moe(model, dataloader, optimizer, num_epochs=10, device="cuda")