1
# 混合专家（Mixture of Experts）的核心思想
2

3
# 密集模型：所有输入经过所有参数
4
output = dense_layer(input)  # 激活所有参数
5

6
# 稀疏 MoE：每个输入只激活部分专家
7
experts = [expert1, expert2, expert3, expert4]
8
gating_scores = gate(input)  # 门控决定激活哪些
9

10
# 稀疏激活：只激活 top-k 个专家
11
top_k_indices = torch.topk(gating_scores, k=2).indices
12
output = sum(experts[i](input) for i in top_k_indices)

1
class GLaMGating(nn.Module):
2
    """GLaM 的门控机制"""
3

4
    def __init__(self, d_model, n_experts):
5
        super().__init__()
6
        self.n_experts = n_experts
7
        self.gate = nn.Linear(d_model, n_experts)
8

9
    def forward(self, x):
10
        # 计算每个专家的得分
11
        gating_scores = self.gate(x)
12

13
        # Softmax 归一化
14
        gating_probs = F.softmax(gating_scores, dim=-1)
15

16
        # Top-k 稀疏激活：只激活得分最高的 2 个专家
17
        k = 2
18
        top_k_probs, top_k_indices = torch.topk(gating_probs, k)
19

20
        # 重新归一化
21
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
22

23
        return top_k_probs, top_k_indices

1
class Expert(nn.Module):
2
    """单个专家网络"""
3

4
    def __init__(self, d_model, d_ff):
5
        super().__init__()
6
        self.w1 = nn.Linear(d_model, d_ff)
7
        self.w2 = nn.Linear(d_ff, d_model)
8
        self.act = nn.GELU()
9

10
    def forward(self, x):
11
        return self.w2(self.act(self.w1(x)))

1
def load_balancing_loss(gating_probs, top_k_indices, n_experts):
2
    """专家负载均衡损失，防止只激活少数专家"""
3

4
    batch_size, seq_len = top_k_indices.shape
5

6
    # 计算每个专家被选中的频率
7
    expert_counts = torch.zeros(n_experts)
8
    for idx in top_k_indices.view(-1):
9
        expert_counts[idx] += 1
10

11
    expert_counts = expert_counts / (batch_size * seq_len)
12

13
    # 计算平均门控概率
14
    avg_gating = gating_probs.mean(dim=[0, 1])
15

16
    # 负载均衡损失：最小化频率与门控概率的差异
17
    loss = n_experts * (avg_gating * expert_counts).sum()
18

19
    return loss

1
class MoELayer(nn.Module):
2
    """GLaM 的 MoE 层"""
3

4
    def __init__(self, d_model, n_heads, n_experts):
5
        super().__init__()
6
        self.n_experts = n_experts
7
        self.self_attn = MultiHeadAttention(d_model, n_heads)
8
        self.moe = MixtureOfExperts(d_model, n_experts)
9
        self.norm1 = nn.LayerNorm(d_model)
10
        self.norm2 = nn.LayerNorm(d_model)
11

12
    def forward(self, x):
13
        # 标准 Transformer 块 + MoE
14
        x = x + self.self_attn(self.norm1(x))
15
        x = x + self.moe(self.norm2(x))
16
        return x

1
# 密集模型：每个 token 激活所有参数
2
class DenseModel:
3
    def forward(self, x):
4
        for layer in self.layers:
5
            x = layer(x)  # 所有参数参与
6
        return x
7

8
# 稀疏 MoE：每个 token 只激活部分参数
9
class SparseMoEModel:
10
    def forward(self, x):
11
        for layer in self.layers:
12
            if isinstance(layer, MoELayer):
13
                x = layer(x, top_k=2)  # 只激活 2 个专家
14
            else:
15
                x = layer(x)
16
        return x