稀疏MoE密集反向传播：Default MoE

稀疏激活的专家混合模型（MoE）在大规模训练中面临独特挑战：路由器仅从未激活的专家处接收稀疏梯度更新。本文介绍Default MoE方法，通过为路由器提供密集梯度来解决这一问题，同时保持稀疏计算的优势。¹

---

1. 问题背景：稀疏MoE的训练困境

1.1 MoE架构回顾

稀疏MoE架构通过路由器选择性地激活Top-K专家：

$$y=\sum _{i=1}^K{G}_i\cdot E_i(x),G=\text{TopK}(\text{Router}(x))$$

其中 $E_i$ 为第 $i$ 个专家，$K$ 为激活专家数量（通常 $K\ll N$，$N$ 为总专家数）。

1.2 训练稳定性问题

稀疏梯度问题：路由器仅从未激活的专家处接收梯度信号！

具体例子：对于1个激活专家 + 7个未激活专家的情况，路由器仅获得 $12.5\%$ 的梯度信息。

后果：

• 路由决策收敛缓慢
• 专家利用率不均衡（某些专家几乎不被激活）
• 训练不稳定

1.3 现有解决方案

方法	策略	问题
Dense MoE	训练时激活所有专家	计算量激增
Switch Transformer	简化路由器	梯度仍不完整
Load Balancing	辅助损失	增加训练复杂度

1.4 Default MoE的核心思想

核心思想：为未激活专家提供”默认输出”，使路由器能够接收来自所有专家的梯度，同时保持推理时的稀疏性。

---

2. Default MoE方法详解

2.1 默认输出定义

定义：对于每个专家 $E_i$，维护一个指数移动平均（EMA）的默认输出：

$$D_i^{(t)}=\alpha \cdot O_i^{(t)}+(1-\alpha )\cdot D_i^{(t-1)}$$

其中：

• $O_i^{(t)}$ 是专家 $E_i$ 在时刻 $t$ 的实际输出（当被激活时）
• $D_i^{(t)}$ 是默认输出
• $\alpha$ 是EMA系数（通常 $0.01\sim 0.1$）

2.2 路由器梯度计算

修改后的路由器前向传播：

class DefaultMoELayer(nn.Module):
    def __init__(self, d_model, n_experts, top_k, alpha=0.1):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        self.alpha = alpha
        
        self.router = nn.Linear(d_model, n_experts)
        self.experts = nn.ModuleList([Expert(d_model) for _ in range(n_experts)])
        
        # 默认输出缓冲区
        self.default_outputs = [None] * n_experts
    
    def forward(self, x):
        # 路由器计算
        router_logits = self.router(x)
        top_k_logits, top_k_indices = torch.topk(router_logits, self.top_k, dim=-1)
        
        # 计算路由器输出
        router_output = torch.zeros_like(router_logits)
        
        # 对激活的专家
        for i, expert_idx in enumerate(top_k_indices[0]):
            expert_output = self.experts[expert_idx](x)
            
            # 更新EMA默认输出
            if self.default_outputs[expert_idx] is None:
                self.default_outputs[expert_idx] = expert_output.detach()
            else:
                self.default_outputs[expert_idx] = (
                    self.alpha * expert_output.detach() + 
                    (1 - self.alpha) * self.default_outputs[expert_idx]
                )
            
            router_output[0, expert_idx] = top_k_logits[0, i]
        
        # 对未激活的专家，使用默认输出
        for i in range(self.n_experts):
            if i not in top_k_indices[0]:
                router_output[0, i] = self.router(x)[0, i] + self._default_signal(i)
        
        return router_output, top_k_indices
    
    def _default_signal(self, expert_idx):
        """计算未激活专家的默认信号"""
        if self.default_outputs[expert_idx] is None:
            return 0.0
        
        # 计算当前输入与默认输出的相似度
        # 作为额外的路由器梯度信号
        return torch.norm(self.default_outputs[expert_idx]).item() * 0.1

2.3 密集梯度更新原理

关键洞察：通过默认输出，路由器现在可以计算关于所有专家的梯度：

$$\frac{\partial \mathcal{L}}{\partial W_{router}}=\sum _{i=1}^N\frac{\partial \mathcal{L}}{\partial E_i}\cdot \frac{\partial E_i}{\partial W_{router}}$$

其中 $\frac{\partial E_i}{\partial W_{router}}$ 通过默认输出链式传播。

2.4 与标准Top-K路由的对比

特性	标准Top-K	Default MoE
路由器梯度来源	仅激活的K个专家	所有N个专家
梯度稀疏性	稀疏	密集
推理计算量	O(K·E)	O(K·E)
训练计算量	O(K·E)	O(K·E + α·N·E)

---

3. 理论分析

3.1 梯度方差减少

定理：设 ${\sigma }_{sparse}^2$ 为标准稀疏路由的梯度方差，${\sigma }_{default}^2$ 为Default MoE的梯度方差。则：

$$\frac{{\sigma }_{default}^2}{{\sigma }_{sparse}^2}\le \frac{1}{N}+\left(1-\frac{1}{N}\right)\cdot \frac{\text{Var}(D)}{\text{Var}(O)}$$

其中 $\text{Var}(D)$ 和 $\text{Var}(O)$ 分别为默认输出和实际输出的方差。

直觉：当默认输出方差接近实际输出方差时，梯度方差显著减少。

3.2 收敛速度分析

定理：在温和条件下，Default MoE的收敛速率满足：

$$\mathbb{E}[\Vert {\theta }_k-{\theta }^{\ast }\Vert ]\le {\left(1-\eta \cdot {\lambda }_{eff}\right)}^k\Vert {\theta }_0-{\theta }^{\ast }\Vert$$

其中有效学习率 ${\lambda }_{eff}$ 满足：

$${\lambda }_{eff}\ge {\lambda }_{standard}\cdot \left(1+\frac{N-K}{K}\cdot \rho \right)$$

$\rho$ 是默认输出与实际输出相关性的度量。

3.3 默认输出的质量保证

引理：EMA默认输出 $D_i$ 满足：

$$\mathbb{E}[\Vert D_i-{\bar{E}}_i\Vert ]\le \frac{\alpha }{2-\alpha }\cdot {\sigma }_{noise}$$

其中 ${\bar{E}}_i$ 是专家 $i$ 的真实平均输出，${\sigma }_{noise}$ 是输出噪声。

推论：通过适当选择 $\alpha$，可以控制默认输出与真实平均输出的偏差。

---

4. 实验验证

4.1 训练稳定性

在1B参数MoE模型上的实验：

方法	训练损失方差	梯度范数标准差	专家利用率标准差
标准Top-K	0.42	8.3	0.31
Default MoE	0.18	4.1	0.12

发现：Default MoE显著降低了训练的不稳定性。

4.2 下游任务性能

在多种下游任务上的评估：

任务	标准Top-K	Default MoE	改进
语言建模 (PPL)	18.2	16.7	+8.2%
问答 (Accuracy)	72.1%	74.8%	+3.7%
推理 (Accuracy)	45.3%	48.1%	+6.2%

4.3 消融实验

EMA系数 $\alpha$ 的影响：

$\alpha$	训练稳定性	最终性能
0.001	差	中
0.01	良	优
0.1	优	良
0.5	不稳定	差

推荐：$\alpha =0.01\sim 0.05$

---

5. 与其他MoE改进方法的比较

5.1 SparseMixer

SparseMixer使用中点法（ODE求解器）估计稀疏路由的梯度：

• 共同点：都解决稀疏梯度问题
• 差异：SparseMixer使用解析近似，Default MoE使用EMA实际输出

# SparseMixer的核心思想
def midpoint_approx(expert_outputs, router_logits, top_k):
    """使用中点估计"""
    # ... SparseMixer实现
    pass
 
# Default MoE的核心思想
def ema_approx(expert_outputs, default_buffers):
    """使用EMA近似"""
    # ... Default MoE实现
    pass

5.2 混合使用

两种方法可以互补：

class HybridMoE(nn.Module):
    """结合SparseMixer和Default MoE"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.default_moe = DefaultMoELayer(*args, **kwargs)
        self.sparse_mixer = SparseMixerRouter(*args, **kwargs)
    
    def forward(self, x):
        # 使用SparseMixer计算精确梯度估计
        mixer_signal = self.sparse_mixer(x)
        
        # 使用Default MoE提供稳定的梯度流
        router_output, indices = self.default_moe(x)
        
        # 融合两种信号
        final_output = router_output + 0.1 * mixer_signal
        return final_output, indices

---

6. 实践指南

6.1 实现步骤

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class DefaultMoEImplementation:
    """Default MoE的完整实现"""
    
    def __init__(self, d_model, n_experts, top_k, alpha=0.01):
        self.n_experts = n_experts
        self.top_k = top_k
        self.alpha = alpha
        
        # 专家
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, 4 * d_model),
                nn.GELU(),
                nn.Linear(4 * d_model, d_model)
            ) for _ in range(n_experts)
        ])
        
        # 路由器
        self.router = nn.Linear(d_model, n_experts)
        
        # 默认输出缓冲区（CPU上维护，节省GPU内存）
        self.register_buffer('default_outputs', torch.zeros(n_experts, d_model))
        self.register_buffer('default_counts', torch.zeros(n_experts))
    
    def forward(self, x):
        batch_size = x.shape[0]
        
        # 路由器计算
        router_logits = self.router(x)
        
        # 获取Top-K
        top_k_logits, top_k_indices = torch.topk(router_logits, self.top_k, dim=-1)
        
        # 准备输出
        output = torch.zeros_like(x).repeat_interleave(self.top_k, dim=0)
        
        # 处理每个激活的专家
        for i in range(self.top_k):
            expert_idx = top_k_indices[:, i]
            expert_output = self.experts[expert_idx](x)
            
            # 更新默认输出
            self._update_default_outputs(expert_output, expert_idx)
            
            output[batch_size * i:batch_size * (i + 1)] = expert_output
        
        # 计算路由器损失（可选）
        router_loss = self._compute_router_loss(router_logits, top_k_indices)
        
        return output, top_k_indices, router_loss
    
    def _update_default_outputs(self, expert_output, expert_idx):
        """更新EMA默认输出"""
        # 在GPU上计算
        with torch.no_grad():
            mask = (self.default_counts[expert_idx] > 0).float()
            self.default_outputs[expert_idx] = (
                mask * (self.alpha * expert_output.mean(dim=0) + 
                       (1 - self.alpha) * self.default_outputs[expert_idx]) +
                (1 - mask) * expert_output.mean(dim=0)
            )
            self.default_counts[expert_idx] += 1
    
    def _compute_router_loss(self, router_logits, top_k_indices):
        """计算路由器辅助损失"""
        # 确保负载均衡
        gates = F.softmax(router_logits, dim=-1)
        
        # 专家利用率
        expert_counts = torch.zeros(self.n_experts, device=router_logits.device)
        for idx in top_k_indices:
            for i in idx:
                expert_counts[i] += 1
        
        # 辅助损失
        aux_loss = self.n_experts * torch.var(expert_counts / expert_counts.sum())
        
        return aux_loss

6.2 训练配置建议

# 推荐的超参数配置
config = {
    'd_model': 4096,
    'n_experts': 32,
    'top_k': 8,  # 或 2
    'alpha': 0.01,  # EMA系数
    'router_z_loss': 0.001,  # 路由器数值稳定性损失
    'aux_loss_weight': 0.01,  # 辅助损失权重
}

6.3 调试技巧

def debug_default_moe(model, dataloader):
    """调试Default MoE的训练"""
    
    # 1. 检查默认输出方差
    print("Default outputs variance:")
    for i, do in enumerate(model.default_outputs):
        print(f"  Expert {i}: {do.var():.4f}")
    
    # 2. 检查专家利用率
    total_counts = model.default_counts.float()
    utilization = total_counts / total_counts.sum()
    print(f"\nExpert utilization std: {utilization.std():.4f}")
    print(f"Min/Max utilization: {utilization.min():.4f} / {utilization.max():.4f}")
    
    # 3. 检查梯度范数
    router_grad_norm = model.router.weight.grad.norm()
    print(f"\nRouter gradient norm: {router_grad_norm:.4f}")

---

7. 总结与展望

7.1 主要贡献

1. 问题识别：明确稀疏梯度是MoE训练不稳定的主要原因
2. 解决方案：提出Default MoE，通过EMA默认输出提供密集梯度
3. 理论分析：建立梯度方差减少和收敛加速的理论保证
4. 实验验证：在多种任务上验证方法的有效性

7.2 局限性

• 需要维护默认输出缓冲区，增加内存开销
• EMA系数需要调优
• 不适用于极端稀疏的设置（如top-1路由）

7.3 未来方向

• 自适应$\alpha$调度
• 多层默认输出池化
• 与其他优化技术的深度整合

---

参考资料

Footnotes

1. Dense Backpropagation Improves Training for Sparse Mixture-of-Experts. arXiv:2504.12463. ↩