AI混合专家：这次怎么落地的

MoE 混合专家 稀疏路由 深度学习 模型架构
作者: 醉月思 发布日期: 2023-06-05 分类: 技术实践

AI混合专家我没按教科书顺序做。

别急着给AI混合专家：这次怎么落地的下定义，先看这次卡在哪。

为什么折腾这个

当时在做一个文本分类任务，数据量不大但类别不少，传统 Dense 模型很难同时兼顾各个类别的特征表现。增大模型参数吧，训练成本上去还不一定有提升；不增大吧，又感觉模型能力不够。

这时候了解到了混合专家的思路，正好解决两个痛点：一是可以增加参数总量但保持单次计算量可控，二是不同专家可以 specialize 到不同模式的数据上。

听起来很对路，于是就动手试了一下。

基础架构怎么搭

最基础的路，每个专家都是一个完整的前馈网络，通过一个路由器来决定样本走哪几个专家。用一个简化的代码来说明一下结构：

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

class Expert(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

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

class Router(nn.Module):
    def __init__(self, input_dim, num_experts, top_k=1):
        super().__init__()
        self.gate = nn.Linear(input_dim, num_experts)
        self.top_k = top_k

    def forward(self, x):
        logits = self.gate(x)
        if self.training:
            # 训练时加上一些噪声让路由更鲁棒
            logits += torch.randn_like(logits) * 0.01

        # 取 top-k 专家
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(top_k_logits, dim=-1)

        return weights, top_k_indices

class MoELayer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_experts, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList([
            Expert(input_dim, hidden_dim, input_dim)
            for _ in range(num_experts)
        ])
        self.router = Router(input_dim, num_experts, top_k)
        self.top_k = top_k

    def forward(self, x):
        batch_size, seq_len, input_dim = x.shape
        x_flat = x.view(-1, input_dim)

        weights, expert_indices = self.router(x_flat)

        # 为每个专家准备对应的样本
        expert_outputs = []
        for i in range(len(self.experts)):
            mask = (expert_indices == i).float()
            if mask.sum() > 0:
                expert_input = x_flat * mask.unsqueeze(-1)
                expert_output = self.experts[i](expert_input)
                expert_outputs.append(expert_output)
            else:
                expert_outputs.append(torch.zeros_like(x_flat))

        # 根据路由权重聚合
        output = torch.zeros_like(x_flat)
        for i, expert_output in enumerate(expert_outputs):
            weight_mask = (expert_indices == i).float()
            weight_for_expert = (weights * weight_mask).sum(dim=-1, keepdim=True)
            output += expert_output * weight_for_expert

        return output.view(batch_size, seq_len, input_dim)

这个实现很简单，但跑起来很快就遇到问题了。

踩过的第一个坑：负载均衡

训练没多久发现大部分样本都堆在了一两个专家身上，其他专家几乎没机会参与。这就是著名的"负载不均衡"问题，最后的效果跟 Dense 模型差不多，甚至更差，因为大部分专家参数都没被有效利用。

解决方案是加一个负载均衡损失，让专家使用更均匀：

def load_balance_loss(weights, expert_indices, num_experts):
    # 计算每个专家的样本比例
    expert_counts = torch.zeros(num_experts, device=weights.device)
    for i in range(num_experts):
        expert_counts[i] = (expert_indices == i).float().sum()

    expert_probs = expert_counts / expert_counts.sum()

    # 理想情况是每个专家 1/num_experts 的样本量
    ideal_probs = torch.ones(num_experts, device=weights.device) / num_experts

    # 用 KL 散度衡量分布差异
    loss = F.kl_div(
        expert_probs.log(),
        ideal_probs,
        reduction='batchmean'
    )

    return loss * 0.1  # 调整权重

这个损失加进去后情况好了很多，但又出现新问题：有些专家为了分担负载，强行接收一些不太适合的样本，反而降低了质量。

这个平衡点真的很微妙，需要根据具体任务反复调。

第二个坑：通信开销

在多卡训练时，MoE 的通信开销比预想的大很多。每个 batch 要把数据分发到不同专家所在的 GPU 上，再把结果聚合回来。这个过程中如果专家分配不均匀，有些 GPU 就会忙得要死，有些却在等数据。

专家负载分布不均的示例，红色表示过载的专家，绿色虚线表示理想的均匀负载线

从图上能明显看到前 4 个专家承担了大部分样本，后 4 个专家几乎闲着。这种情况下，就算你有 8 个专家，实际起作用的可能只有 2-3 个。

后来试了几种方案，效果最好的是"容量因子"限制：给每个专家设置一个容量上限，超过的部分要么丢弃要么转发到其他专家：

class CapacityLimitedRouter(nn.Module):
    def __init__(self, input_dim, num_experts, top_k=2, capacity_factor=1.5):
        super().__init__()
        self.gate = nn.Linear(input_dim, num_experts)
        self.top_k = top_k
        self.capacity_factor = capacity_factor
        self.expert_count = num_experts

    def forward(self, x):
        batch_size = x.shape[0]
        logits = self.gate(x)

        # 计算每个专家的容量
        capacity = int(batch_size * self.capacity_factor / (self.top_k * self.expert_count))

        # Softmax 并取 top-k
        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(probs, self.top_k, dim=-1)

        # 为每个专家限制样本数量
        expert_usage = torch.zeros(self.expert_count, device=x.device)
        valid_mask = torch.ones_like(top_k_indices, dtype=torch.bool)

        for i in range(batch_size):
            for k in range(self.top_k):
                expert_idx = top_k_indices[i, k]
                if expert_usage[expert_idx] < capacity:
                    expert_usage[expert_idx] += 1
                else:
                    valid_mask[i, k] = False

        # 只保留有效的路由
        final_probs = top_k_probs * valid_mask.float()
        final_probs = final_probs / (final_probs.sum(dim=-1, keepdim=True) + 1e-10)

        return final_probs, top_k_indices

这个方案能保证每个专家不会过载，但代价是有些样本可能走不到最优专家。这也是个权衡问题。

第三个坑：专家塌陷

训练到一定阶段后，发现有些专家开始"偷懒"：它们不再学习区分性的特征，而是输出一些接近零或者接近均值的惰性结果。这样既不会增加太多 loss，又能节省计算资源。

这个问题比较棘手，最后是通过几个手段组合解决：

  1. 给每个专家加一个独立的输出层，让它们不能简单地通过输出均值来"逃避责任"
  2. 在损失函数中引入专家间差异的正则项，鼓励专家学习不同模式
  3. 定期重新初始化表现最差的专家
def expert_diversity_loss(expert_outputs, num_experts):
    """
    鼓励专家输出之间的差异性
    """
    pairwise_similarity = 0.0
    count = 0

    for i in range(num_experts):
        for j in range(i+1, num_experts):
            # 计算两个专家输出的余弦相似度
            output_i = expert_outputs[i].view(-1)
            output_j = expert_outputs[j].view(-1)

            similarity = F.cosine_similarity(
                output_i.unsqueeze(0),
                output_j.unsqueeze(0)
            )

            pairwise_similarity += similarity.abs()
            count += 1

    # 相似度越高，损失越大
    return pairwise_similarity / count * 0.05

加了这些约束后，专家的分化明显好了很多，但训练时间也明显增加了。

实际效果如何

折腾了这么久，最终效果还算可以。跟同参数量的 Dense 模型相比：

  • 在准确率上提升了约 2-3 个百分点
  • 训练时间增加了约 40%（主要来自路由和通信开销）
  • 推理速度在 batch 比较大时接近 Dense 模型，小 batch 时会慢一些

这些数字不是什么通用结论，只是在我们这个具体任务上的观察。如果你的任务特性不同，结果可能完全不一样。

Dense 模型与 MoE 模型在不同 batch size 下的吞吐量对比，MoE 在小 batch 时明显较慢，大 batch 时接近 Dense 模型

从图上可以看出，MoE 在小 batch 时开销相对明显，但随着 batch 增大，并行化的优势开始体现。

一些实践建议

如果你也想试一下混合专家，这些建议或许能帮你少踩点坑：

  1. 先从小规模开始验证：不要一上来就搞超大模型和超多专家，先用 4-8 个专家验证思路是否适合你的任务
  2. 密切关注负载均衡：这是最容易出问题的地方，训练过程中定期检查专家使用情况
  3. 容忍一定的不均匀：完全均匀的负载往往意味着损失了质量，找到适合自己的平衡点
  4. 注意硬件条件：MoE 对多卡训练比较友好，单卡场景下通信开销可能得不偿失
  5. 监控专家分化情况：定期检查不同专家是否真的学到了不同的模式，不要等到训练结束才发现问题

回过头来看

混合专家不是什么银弹，它只是在某些场景下的一个选择。它能在增加参数总量的同时控制计算量，但代价是更复杂的工程实现和更多的超参数调优。

对我来说，这次实践最大的收获不是提升了多少准确率，而是对模型架构设计有了更深的理解：任何架构选择都是一系列权衡的结果，没有通用的最优解。

有时候最简单的 Dense 模型反而是最实用的选择，复杂架构的收益未必能抵消其带来的维护成本。但在资源充足且确实需要区分不同模式时，混合专家确实提供了一个有意思的方向。

就像开头说的，当所有人都往一个方向挤的时候，或许该想想是不是还有别的门。但也要记得，开一扇门的成本有时候比挤过原来那扇门还要高。

版权声明: 本文首发于 指尖魔法屋-AI混合专家：这次怎么落地的（https://blog.thinkmoon.cn/post/245-moe-sparse-router-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章