MoE（Mixture of Experts）混合专家模型笔记

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

# 构造最小示例：4 个 token，d_model=8，共 4 个专家，Top-2
torch.manual_seed(0)
N_tokens, d_model, num_experts, top_k = 4, 8, 4, 2

x = torch.randn(N_tokens, d_model)                        # 输入 token 向量
W_router = torch.randn(d_model, num_experts)               # Router 权重矩阵
experts = [nn.Linear(d_model, d_model, bias=False) for _ in range(num_experts)]

# 第一步：Router 计算每个 token 对各专家的得分
scores = F.softmax(x @ W_router, dim=-1)                   # [4, 4]

# 第二步：选 Top-2 专家
topk_scores, topk_indices = scores.topk(k=top_k, dim=-1)  # 各形状 [4, 2]

# 第三步：对 Top-K 权重重新归一化，使每个 token 的权重之和为 1
topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)

# 第四步：对每个 token，加权求和被选中专家的输出
output = torch.zeros(N_tokens, d_model)
for k_pos in range(top_k):
    expert_ids = topk_indices[:, k_pos]        # 每个 token 在第 k_pos 位选的专家编号，形状 [4]
    weights = topk_scores[:, k_pos]            # 对应权重，形状 [4]
    for token_i in range(N_tokens):
        expert_id = expert_ids[token_i].item()
        output[token_i] += weights[token_i] * experts[expert_id](x[token_i])

print(output.shape)   # torch.Size([4, 8])
print(topk_indices)   # 每个 token 选中的专家编号，如 tensor([[2,1],[0,3],[1,2],[3,0]])

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


class ExpertFFN(nn.Module):
    """单个专家：标准的两层 FFN"""

    def __init__(self, d_model, d_ffn):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ffn, bias=False)
        self.w2 = nn.Linear(d_ffn, d_model, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)))


class MoELayer(nn.Module):
    """MoE 层：N 个专家 + Router，每个 token 激活 Top-K 个专家"""

    def __init__(self, d_model, d_ffn, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # N 个独立的专家 FFN
        self.experts = nn.ModuleList([
            ExpertFFN(d_model, d_ffn) for _ in range(num_experts)
        ])

        # Router：线性层，输出每个专家的得分
        self.router = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # [batch*seq_len, d_model]

        # 路由：计算每个 token 对应各专家的得分
        router_logits = self.router(x_flat)                          # [N_tokens, N_experts]
        router_scores = F.softmax(router_logits, dim=-1)             # 归一化为概率

        # 选 Top-K 个专家
        topk_scores, topk_indices = router_scores.topk(self.top_k, dim=-1)
        # 对 Top-K 的权重重新归一化（使权重之和为 1）
        topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)

        # 对每个 token，加权求和被选中专家的输出
        output = torch.zeros_like(x_flat)
        for expert_idx in range(self.num_experts):
            # 找出哪些 token 选了这个专家
            token_mask = (topk_indices == expert_idx).any(dim=-1)
            if not token_mask.any():
                continue

            expert_input = x_flat[token_mask]
            expert_output = self.experts[expert_idx](expert_input)

            # 取该专家对应的权重
            expert_weight = topk_scores[token_mask][
                (topk_indices[token_mask] == expert_idx)
            ].unsqueeze(-1)

            output[token_mask] += expert_weight * expert_output

        return output.view(batch_size, seq_len, d_model)


# 验证
d_model, d_ffn = 512, 2048
moe = MoELayer(d_model=d_model, d_ffn=d_ffn, num_experts=8, top_k=2)

x = torch.randn(2, 16, d_model)   # batch=2, seq_len=16
out = moe(x)
print(out.shape)  # torch.Size([2, 16, 512])