DeepSeek实现低成本训练，原来是靠它！

DeepSeek推出的最新推理模型，以500万美元的训练成本，比肩数亿美元成本的OpenAI o1，离不开各种优化策略，除了之前提到的“知识蒸馏”以外，还包括今天的主角MoE。

在机器学习和深度学习领域，模型的设计和优化一直是研究的核心。近年来，一种名为Mixture of Experts (MoE) 的模型架构逐渐引起了广泛关注。MoE模型通过结合多个“专家”模型的优势，能够在处理复杂任务时表现出色。本文将详细介绍MoE模型的基本概念、工作原理、优势以及应用场景。

1. MoE模型的基本概念

Mixture of Experts (MoE)，中文译为“专家混合模型”，是一种由多个子模型（称为“专家”）组成的集成模型。每个专家模型通常专注于处理输入数据的某一部分或某一特定特征。MoE模型的核心思想是通过一个门控机制（Gating Network）来决定每个输入数据应该由哪个或哪些专家模型来处理。

MoE是前馈神经网络的一种替代方案。在前馈神经网络中，对于某个问题，需要全部的权重参与计算，而通过MoE，可以根据问题的不同，只让部分权重参与计算，从而实现高效的推理。

DeepSeek致力于MoE的研究，自2024年1月的DeepSeekMoE模型开始，到2025年初的DeepSeek R1模型，都有MoE的身影。

下图是DeepSeekMoE中的MoE。

 DeepSeekMoE论文：2401.06066https://arxiv.org/pdf/2401.06066

下图是DeepSeekV3中的MoE。

 我们在DeepSeek的官方代码中，也可以看到MoE的代码实现。

MoE代码实现：

DeepSeek-V3/inference/model.py at main · deepseek-ai/DeepSeek-V3 · GitHubhttps://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py#L633

DeepSeekV3论文：

DeepSeek-V3/DeepSeek_V3.pdf at main · deepseek-ai/DeepSeek-V3 · GitHubhttps://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf

2. MoE模型的工作原理

MoE模型的工作流程可以分为以下几个步骤：

1. 输入数据：模型接收输入数据，并将其传递给门控网络和各个专家模型。

2. 门控网络：门控网络根据输入数据生成一个权重向量，该向量决定了每个专家模型对最终输出的贡献程度。门控网络通常是一个简单的神经网络，如全连接层或softmax层。门控网络的原理与LSTM相似，都是让模型自动学习如何对问题做分类

3. 专家模型：每个专家模型独立处理输入数据，并生成自己的输出。专家模型可以是任何类型的模型，如神经网络、决策树等。若想减少计算量，可以根据门控网络的输出，只取前K个专家的计算结果。

4. 加权求和：最终输出是各个专家模型输出的加权和，权重由门控网络决定。

数学上，MoE模型的输出可以表示为：

其中：

• x是输入数据，

• fi(x)是第u个专家模型的输出，

• gi(x)是门控网络为第i个专家模型生成的权重，

• n是专家模型的数量

 3.MoE的PyTorch实现

前文提到DeepSeek V3中已经给出了MoE的代码实现，但里面涉及到许多变量，难以阅读，我对官方的代码进行了总结，提取出里面核心的思想，简化后的模型代码实现如下：

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


class Expert(nn.Module):
    """
    定义一个专家模块，用于在 MoE 中进行特征变换。
    每个专家模块包含两层全连接网络，中间使用 ReLU 激活函数。
    """
    def __init__(self, dim, inter_dim):
        super().__init__()
        self.w1 = nn.Linear(dim, inter_dim)  # 第一层全连接层，输入维度为 dim，中间维度为 inter_dim
        self.w2 = nn.Linear(inter_dim, dim)  # 第二层全连接层，输出维度为 dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播函数，实现专家模块的功能。
        :param x: 输入张量，形状为 [batch_size, dim]
        :return: 输出张量，形状为 [batch_size, dim]
        """
        return self.w2(F.relu(self.w1(x)))  # 先通过第一层全连接层，再通过 ReLU 激活函数，最后通过第二层全连接层


class MoE(nn.Module):
    """
    定义一个混合专家（MoE）模块。
    MoE 模块包含多个专家模块和一个门控网络，用于根据输入动态选择专家。
    """
    def __init__(self, dim, inter_dim, n_expert, top_k):
        super().__init__()
        self.top_k = top_k  # 每次激活的专家数量
        self.n_expert = n_expert  # 总专家数量

        self.gate = nn.Linear(dim, n_expert)  # 门控网络，输入维度为 dim，输出维度为 n_expert
        self.experts = nn.ModuleList([Expert(dim, inter_dim) for _ in range(n_expert)])  # 创建 n_expert 个专家模块

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播函数，实现 MoE 模块的功能。
        :param x: 输入张量，形状为 [batch_size, dim]
        :return: 输出张量，形状为 [batch_size, dim]
        """
        # 门控网络计算每个专家的权重
        gate_scores = self.gate(x)  # [batch_size, num_experts]，计算每个样本对每个专家的得分
        gate_weights = F.softmax(gate_scores, dim=-1)  # 归一化为概率分布，形状为 [batch_size, num_experts]

        # 选择 top-k 专家
        top_k_weights, top_k_indices = torch.topk(gate_weights, self.top_k, dim=-1)  # [batch_size, top_k]，选择 top-k 专家的权重和索引
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)  # 重新归一化，确保 top-k 权重和为 1

        # 初始化输出
        output = torch.zeros_like(x)  # [batch_size, output_dim]，初始化输出张量

        # 只激活 top-k 专家
        for i in range(self.top_k):
            expert_idx = top_k_indices[:, i]  # 当前批次的第 i 个专家索引，形状为 [batch_size]
            expert_mask = F.one_hot(expert_idx, num_classes=self.n_expert).float()  # [batch_size, num_experts]，创建 one-hot 掩码
            expert_output = torch.stack([expert(x) for expert in self.experts], dim=1)  # [batch_size, num_experts, output_dim]，计算所有专家的输出
            expert_output = torch.sum(expert_output * expert_mask.unsqueeze(-1), dim=1)  # [batch_size, output_dim]，根据掩码选择对应专家的输出
            output += top_k_weights[:, i].unsqueeze(-1) * expert_output  # 加权求和，将当前专家的输出加到总输出中

        return output

# 示例用法
if __name__ == "__main__":
    # 定义模型参数
    input_dim = 32  # 输入维度
    num_experts = 4  # 专家数量
    top_k = 2  # 每次激活 2 个专家
    batch_size = 8  # 批量大小

    # 创建 MoE 模型
    moe = MoE(input_dim, 16, num_experts, top_k=top_k)

    # 随机生成输入数据
    x = torch.randn(batch_size, input_dim)

    # 前向传播
    output = moe(x)
    print("Output shape:", output.shape)  # 应为 [batch_size, input_dim]