一文读懂 DeepSeek使用的 MoE 架构到底是什么 原创 精华

import torch
import torch.nn as nn


# 一个简单的专家网络
class SimpleExpert(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleExpert, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size)
        )


    def forward(self, x):
        return self.network(x)


# 简化版的MOE模型
class SimpleMOE(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_experts=4):
        super(SimpleMOE, self).__init__()


        # 创建多个专家
        self.experts = nn.ModuleList([
            SimpleExpert(input_size, hidden_size, output_size) 
            for _ in range(num_experts)
        ])


        # 路由器：决定用哪个专家
        self.router = nn.Linear(input_size, num_experts)


    def forward(self, x):
        # 1. 路由器计算每个专家的权重
        router_outputs = torch.softmax(self.router(x), dim=-1)


        # 2. 获取权重最高的两位专家
        top_k_weights, top_k_indices = torch.topk(router_outputs, k=2, dim=-1)


        # 3. 归一化权重，使它们的和为1
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)


        # 4. 初始化最终输出
        final_output = torch.zeros_like(x)


        # 5. 对于每个样本
        for i in range(x.shape[0]):
            # 6. 计算并组合所选专家的输出
            for j in range(2):  # 两位专家
                expert_idx = top_k_indices[i, j]
                weight = top_k_weights[i, j]


                # 获取专家输出
                expert_output = self.experts[expert_idx](x[i:i+1])


                # 加权组合
                final_output[i:i+1] += weight * expert_output


        return final_output
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.

def simple_load_balancing(router_probs, num_experts):
    """简单的负载均衡计算"""
    # 计算每个专家实际使用频率
    expert_usage = router_probs.sum(dim=0) / router_probs.shape[0]


    # 理想情况下每个专家使用率相等
    ideal_usage = torch.ones_like(expert_usage) / num_experts


    # 计算实际与理想的差距作为损失
    balance_loss = ((expert_usage - ideal_usage) ** 2).sum()


    return balance_loss
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.

# 简化版的DeepSeek MOE架构
class SimpleDeepSeekMOELayer(nn.Module):
def __init__(self, hidden_size, num_experts=8, k=2):
        super(SimpleDeepSeekMOELayer, self).__init__()


# 自注意力层（保持不变）
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=8)
        self.norm1 = nn.LayerNorm(hidden_size)


# MOE替代标准FFN
        self.moe = SimpleMOE(
            input_size=hidden_size,
            hidden_size=hidden_size*4,
            output_size=hidden_size,
            num_experts=num_experts
        )
        self.norm2 = nn.LayerNorm(hidden_size)


def forward(self, x):
# 自注意力部分
        residual = x
        x = self.attention(x, x, x)[0]
        x = self.norm1(x + residual)


# MOE部分
        residual = x
        x = self.moe(x)
        x = self.norm2(x + residual)


return x
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.