MoE模型中的“专家”与“调控”代码实现

《DeepSeek原生应用与智能体开发实践（人工智能技术丛书）》(王晓华)【摘要 书评 试读】- 京东图书

MoE模型作为一种深度学习结构，其核心思想在于通过集成多个专门化的网络组件（即“专家”）来提升模型的表示能力和泛化性能。MoE的两个主要组成部分（专家和门控网络）相互协作，共同实现这一目标。

1. 专家

首先是专家部分。在MoE架构中，传统的前馈神经网络层被扩展为一组可选择的专家网络。这些专家通常也是由前馈神经网络构成，但每个专家都专注于处理特定类型的数据或特征。这种设计使得模型能够在不同的情境下调用最合适的专家，从而更精确地捕捉数据的复杂性和多样性。通过专家的专门化学习，MoE能够在训练过程中更有效地分配计算资源，提高模型的效率和性能。

我们实现了一个专家层，代码如下所示。

class Expert(torch.nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(n_embd, 4 * n_embd),
            torch.nn.ReLU(),
            torch.nn.Linear(4 * n_embd, n_embd),
            torch.nn.Dropout(0.1),
        )

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

2. 门控网络

门控网络无疑是MoE架构中的核心组件，其重要性在于它掌控着在推理和训练过程中专家的选择权。简而言之，门控网络充当着智能调度员的角色，根据输入数据的特性，精准地调配各个专家的参与程度。

在最基础的操作层面，门控网络的运作可以概括为以下步骤：

首先，我们接收输入数据（X），这个数据可能是一个特征向量，包含了待处理任务的关键信息。如图10-3所示。

接下来，我们将这个输入数据与路由器的权重矩阵（W）进行矩阵乘法运算。这一过程实质上是在对输入数据进行线性变换，目的是提取出对于后续专家选择至关重要的特征。如图10-4所示。

           图10-3  门控网络       

图10-4  门控网络的调控

通过这种计算方式，门控网络能够为每个专家生成一个相应的得分或者称为“门控值”。这些门控值反映了在当前输入情境下，各个专家对于任务处理的适合程度。基于这些门控值，我们可以进一步应用Softmax函数等策略，来确定每个专家在最终输出中的贡献权重，从而实现MoE架构中灵活且高效的专家组合与调用。

在最后阶段，我们将利用门控网络（即路由器）生成的得分，与每个专家的输出进行结合，以选出最合适的专家贡献。具体来说，我们将每个专家的输出与其对应的门控得分相乘，这一步骤实质上是在对每个专家的预测结果进行加权。加权后的专家输出随后被相加，形成MoE模型的最终输出，如图10-5所示。

现在，让我们将整个过程综合起来，以更全面地理解输入数据在门控路由器和专家之间的流动路径，如图10-6所示。

图10-5  门控网络与专家模型的计算         

       图10-6  门控路由与专家模型的选择

• 输入数据的接收：首先，模型接收输入数据，这些数据可能是文本、图像或其他类型的特征向量，包含待处理任务的关键信息。
• 路由器的处理：输入数据随后被传递给路由器（即门控网络）。在这里，数据与路由器的权重矩阵进行矩阵乘法运算，生成一组得分。这些得分反映了在当前输入情境下，各个专家对于任务处理的适合程度。
• 专家的激活与输出：基于路由器的得分，一部分专家会被激活，而其余专家则保持休眠状态。被激活的专家会独立地处理输入数据，并生成各自的预测结果或输出向量。
• 输出的加权与合并：每个被激活的专家的输出都会与其对应的门控得分相乘，进行加权处理。随后，这些加权后的输出被相加，形成MoE模型的最终预测结果。

通过以上方式，MoE架构能够在不同的输入情境下动态地选择合适的专家组合，以实现更高效、更精确的数据处理与预测。我们实现了对于门控机制与专家模型的稀疏MoE，其代码如下所示。

import torch  

# 定义一个专家模型，它是一个简单的全连接网络  
class Expert(torch.nn.Module):  
    def __init__(self, n_embd):  
        super().__init__()  
        self.net = torch.nn.Sequential(  
            torch.nn.Linear(n_embd, 4 * n_embd),	  # 线性层，将维度扩大4倍  
            torch.nn.ReLU(),					  # ReLU激活函数  
            torch.nn.Linear(4 * n_embd, n_embd),	  # 线性层，恢复原始维度  
            torch.nn.Dropout(0.1),				  # Dropout层，防止过拟合  
        )  

    def forward(self, x):  
        return self.net(x)  # 前向传播  

# 定义一个Top-K路由器，用于选择前K个最佳专家  
class TopkRouter(torch.nn.Module):  
    def __init__(self, n_embed, num_experts, top_k):  
        super(TopkRouter, self).__init__()  
        self.top_k = top_k  # 选择前K个专家  
        self.linear = torch.nn.Linear(n_embed, num_experts)  # 线性层，输出为专家数量  

    def forward(self, mh_output):  
        logits = self.linear(mh_output)  # 通过线性层得到每个专家的得分  
        # 选择得分最高的前K个专家及其索引  
        top_k_logits, indices = logits.topk(self.top_k, dim=-1)  
        # 创建一个与logits形状相同且全为-inf的张量  
        zeros = torch.full_like(logits, float('-inf'))  
        # 将前K个专家的得分填充到zeros中对应的位置  
        sparse_logits = zeros.scatter(-1, indices, top_k_logits)  
        # 对sparse_logits进行softmax操作，使得得分转换为概率分布  
        router_output = torch.nn.functional.softmax(sparse_logits, dim=-1)  
        return router_output, indices  # 返回专家的概率分布和索引  

# 定义一个稀疏的混合专家（MoE）模型  
class SparseMoE(torch.nn.Module):  
    def __init__(self, n_embed, num_experts, top_k):  
        super(SparseMoE, self).__init__()  
        self.router = TopkRouter(n_embed, num_experts, top_k)  # 路由器，用于选择专家  
        # 创建一个专家列表，每个专家都是一个Expert实例  
        self.experts = torch.nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])  
        self.top_k = top_k  # 选择前K个专家  

    def forward(self, x):  
        # 通过路由器得到专家的概率分布和索引  
        gating_output, indices = self.router(x)  
        final_output = torch.zeros_like(x)  # 初始化最终输出为与输入形状相同的全零张量  

        # 将输入和路由器的输出展平，以便后续处理  
        flat_x = x.view(-1, x.size(-1))  
        flat_gating_output = gating_output.view(-1, gating_output.size(-1))  

        # 遍历每个专家，根据其概率分布对输入进行处理  
        for i, expert in enumerate(self.experts):  
            # 找出当前专家是前K个专家的token  
            expert_mask = (indices == i).any(dim=-1)  
            flat_mask = expert_mask.view(-1)  # 展平操作  
            
            # 如果当前专家对至少一个token是前K个专家之一  
            if flat_mask.any():  
                # 选出这些token的输入  
                expert_input = flat_x[flat_mask]  
                # 将这些token输入给当前专家进行处理  
                expert_output = expert(expert_input)  
                # 获取当前专家对这些token的概率分布  
                gating_scores = flat_gating_output[flat_mask, i].unsqueeze(1)  
                # 根据概率分布对专家的输出进行加权  
                weighted_output = expert_output * gating_scores  
                # 将加权后的输出累加到最终输出中对应的位置  
                final_output[expert_mask] += weighted_output.squeeze(1)  

        return final_output  # 返回最终输出

上面代码实现了一个稀疏的混合专家（Sparse Mixture of Experts，Sparse MoE）模型，其中包含一个路由器和多个专家。路由器负责根据输入选择前K个最佳专家，而每个专家则是一个简单的全连接网络。最终输出是所有选定专家的加权输出之和。