Distributionally Robust Mixture-of-Experts Training¶

会议: NeurIPS 2026
arXiv: 2610.07207
代码: https://drmoet.github.io/
领域: 模型压缩
关键词: 混合专家模型, 分布鲁棒优化, 负载均衡, 路由鲁棒性, 语言模型预训练

一句话总结¶

提出分布鲁棒混合专家训练（DRMoET），将各层专家视为内生鲁棒性分组并采用激活加权信用与熵正则镜像上升优化高损失路由结果，在不改动前向路由开销的前提下显著缓解偶发误路由带来的性能退化。

研究背景与动机¶

混合专家（Mixture-of-Experts, MoE）架构通过稀疏门控为每个 token 动态激活少量专家子网络，实现了在保持单个 token 计算量近乎不变的前提下极大扩展模型参数容量。这一稀疏激活机制已成为现代大语言模型扩展的核心技术。然而，稀疏路由也引入了难以察觉的系统可靠性隐患：模型的最终输出极度依赖于所选专家的实际胜任能力。现实场景中，门控网络不可能完美无缺，再加上负载均衡惩罚或测试期的数据分布漂移，大量 token 难免被分流至并非全局最优的次级候选专家。在标准训练范式下，这些排位靠后或激活频率较低的次选专家获得的有效任务梯度更新相对贫乏，能力成长明显落后于头部专家，导致一旦遭遇次优路由就会带来不成比例的严重预测损失。

主流缓解方案主要依赖辅助负载均衡损失（如 Switch Transformer 与 GShard 中的方差或 KL 惩罚）或者免辅助损失的门控偏置调节（如 Loss-Free Balancing）。但这些常规手段的核心矛盾在于：它们只优化流量分配的平均度，而非专家被选中时的实际胜任力。机械地强制均衡流量无法保证次级专家在冷不丁被选中时具备足够的处理能力；不仅如此，专家特化本就是 MoE 模型的重要优势，盲目抹平各专家间的流量差异甚至会损害门控本身的专业化偏好。

本文的切入角度是跳出单纯调节流量分配的思路，将鲁棒优化的思想引入 MoE 内部。与其依赖先验的外部标签定义群组，不如直接把模型各层由门控内生诱导出的路由专家视作鲁棒性组，自适应增强高损失路由路径的梯度信号。核心 idea：将各层专家建模为内生分布鲁棒优化（DRO）分组，基于 EMA 平滑的激活加权损失动态更新专家对偶权重，通过对偶加权任务损失回传重点强化高损失次优路由路径，在保留天然特化分工的同时消除专家的能力短板。

方法详解¶

整体框架¶

DRMoET 整体遵循标准前向计算流程：给定输入 token，门控网络计算各专家的路由权重并选出 top-$k$ 专家进行前向传播与激活合并。在反向传播与状态维护阶段，框架不额外引入任何辅助损失项，也不修改门控的离散判定规则，而是通过“激活加权归因 $\to$ 跨 batch 指数移动平均（EMA）平滑 $\to$ 熵正则化 Softmax 对偶更新 $\to$ 任务损失重加权回传”这一极轻量闭环工作。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Token 序列"] --> B["前向计算与 Top-k 稀疏路由<br/>门控概率 p 与专家激活 h"]
    B --> C["激活加权信用归因<br/>截断梯度: c = sg(p · ||h||)"]
    C --> D["专家损失聚合与 EMA 追踪<br/>滑动平滑各专家历史损失"]
    D --> E["熵正则对偶权重更新<br/>Softmax 镜像上升自适应加权"]
    E --> F["对偶加权任务损失回传<br/>重点强化高损失次优专家"]

关键设计¶

1. 激活加权信用归因：精准度量专家在最终损失中的实际贡献

直接将 token 标量损失均摊给所有被激活的专家忽略了专家在前向合并中的贡献差异。如果仅使用门控概率 $p_{l,i}$，又会忽略各专家输出特征向量的实际模长变化。DRMoET 设计了结合门控概率与未合并激活特征 $L_2$ 范数的联合归因信用，并显式引入截断梯度（Stop-Gradient）算子： $$\tilde{c}_{l,i}(x_j^{(b)}) \triangleq \text{sg}\left(p_{l,i}(x_j^{(b)}) \cdot \|h_{l,i}(x_j^{(b)})\|_2\right)$$ 其中 $h_{l,i}$ 为第 $l$ 层第 $i$ 个专家的前向输出向量。使用 $\text{sg}(\cdot)$ 确保了归因系数纯粹作为后验权重，而不会把激活模长或门控概率本身变成一个被反向传播强行拉大的优化伪目标。第 $l$ 层专家 $i$ 在当前批次中承担的归一化损失标量定义为： $$\ell_{l,i}^{(t)} \triangleq \frac{1}{N_t} \sum_{(b,j) \in \mathcal{B}} \tilde{c}_{l,i}(x_j^{(b)}) \cdot \ell_j^{(b)}(\theta) \cdot \mathbf{1}[i \in \mathcal{T}_{j,l}^{(b)}]$$ 通过全 batch 的 token 总数 $N_t$ 归一化，使得归因量在不同的 micro-batch 划分下具有尺度一致性。

2. 专家损失的 EMA 平滑追踪：抑制微批次间内生组损失的高方差

与具有固定先验标签的传统 Group DRO 不同，MoE 的路由选择是高度动态且稀疏的。单个 micro-batch 中某些专家的分配 token 极少甚至为零，若直接基于单批次损失更新极易导致对偶变量剧烈震荡。为此，DRMoET 为每层所有专家维护跨步持久化的 EMA 状态： $$\hat{\ell}_{l,i}^{(t)} = \beta \hat{\ell}_{l,i}^{(t-1)} + (1 - \beta) \, \text{sg}(\ell_{l,i}^{(t)})$$ 衰减系数 $\beta \in (0, 1)$ 通常取值较接近 1（实验默认取 0.999）。EMA 有效吸收了 token 到达的随机性与采样方差，为下一阶段的对偶上升提供了平稳且能反映真实能力缺陷的历史统计指标。

3. 熵正则对偶更新与任务损失重加权：平滑极小化最差专家风险

为避免硬取最坏专家（Hard Max）带来的非光滑性与训练崩塌，DRMoET 在概率单纯形 $\Delta_E$ 上定义平滑的鞍点问题，并推导出了显式的概率空间 Softmax 镜像上升更新规则： $$\mu_{l,\cdot}^{(t)} = \text{softmax}\left(\mu_{l,\cdot}^{(t-1)} + \eta_t \hat{\ell}_{l,\cdot}^{(t)}\right)$$ 其中 $\eta_t = \eta_0 / \sqrt{E}$ 为控制对偶学习率的缩放步长。该规则在理论上等价于带有负熵与欧氏正则项 $\Omega(\mu_l) = -\sum_i \mu_{l,i}\log\mu_{l,i} + \frac{1}{2}\|\mu_l\|_2^2$ 的镜像上升，具备 $\widetilde{\mathcal{O}}(T^{-1/2})$ 的平稳点收敛率保证。最终构筑的鲁棒训练目标为： $$\mathcal{L}_{\text{DRO}}^{(t)} \triangleq \sum_{l=1}^L \sum_{i=1}^E \text{sg}(\mu_{l,i}^{(t)}) \, \ell_{l,i}^{(t)}$$ 通过这一重加权，表现不佳、损失持续偏高的专家在反向传播时将获得更高权重的任务梯度，被迫补齐短板，而原本就表现极佳的专家则不会受到多余干扰。

损失函数 / 训练策略¶

DRMoET 在训练时不破坏底层模型的标准配方。总损失函数在任务损失重加权项之外，完全保留原有 baseline 的辅助负载均衡项与路由器 $z$-loss： $$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{DRO}} + \lambda_{\text{aux}} \mathcal{L}_{\text{aux}} + \lambda_z \mathcal{L}_z$$ 在实现中，$\mu$ 的更新完全在 autograd 图外部完成，持久化状态仅需为每层保存 $2E$ 个标量（EMA 与对偶权重），计算额外开销仅为 $\mathcal{O}(L N_t k d)$，相对于前向 FFN 的 $\mathcal{O}(L N_t k d d_{\text{ff}})$ 几乎可忽略不计。

实验关键数据¶

主实验¶

在 FLAME-MoE 架构与 DataComp-LM (DCLM) 数据集上分别进行了 290M 激活/746M 总参数（33.5B tokens）与 1.7B 激活/10.3B 总参数（67B tokens）的从头预训练，下游 0-shot 评测包含 ARC-Challenge、ARC-Easy、HellaSwag、PIQA、WinoGrande、SciQ 及 ReCoRD（F1）。

模型规模与方法	ARC-C	ARC-E	HellaSwag	PIQA	WinoGrande	SciQ	ReCoRD (F1)	平均得分
290M-746M (33.5B tokens)
FLAME-MoE (基线)	0.2329	0.5295	0.3447	0.6697	0.5012	0.8020	0.7023	0.5403
Aux-free ($u=10^{-3}$)	0.2261	0.4878	0.3283	0.6670	0.4949	0.7100	0.6445	0.5084
Aux-free ($u=10^{-2}$)	0.2184	0.4895	0.3301	0.6627	0.4949	0.6990	0.6440	0.5055
DRMoET ($\eta=0.1$)	0.2338	0.5522	0.3472	0.6844	0.5209	0.7950	0.7115	0.5493
DRMoET ($\eta=0.01$)	0.2270	0.5400	0.3470	0.6839	0.4941	0.8180	0.7007	0.5444
DRMoET ($\eta=0.001$, 最佳)	0.2295	0.5682	0.3452	0.6888	0.5051	0.8190	0.7105	0.5523
DRMoET ($\eta=0.0001$)	0.2193	0.5556	0.3461	0.6970	0.5162	0.8060	0.7037	0.5491
1.7B-10.3B (67B tokens)
FLAME-MoE (基线)	0.3481	0.6970	0.4942	0.7563	0.5904	0.9050	0.8467	0.6625
Aux-free ($u=10^{-2}$)	0.3148	0.6843	0.4667	0.7459	0.5722	0.8930	0.8248	0.6431
DRMoET ($\eta=0.01$, 最佳)	0.3805	0.7210	0.4973	0.7666	0.6235	0.9020	0.8457	0.6767
DRMoET ($\eta=0.001$)	0.3652	0.7024	0.4996	0.7628	0.6148	0.9110	0.8452	0.6716

消融与鲁棒性验证实验¶

1. 机制验证：专家损失分布收敛性与强制次优路由探针

评测维度	评测指标	FLAME-MoE 基线	DRMoET (本文)	相对变化率
专家间损失均匀性	最差专家损失 (Worst loss)	2.6930	2.6379	-2.05%
	最优专家损失 (Best loss)	2.1149	2.0722	-2.02%
	平均专家损失 (Mean loss)	2.3712	2.3699	-0.05%
	极差 (Range)	0.5781	0.5656	-2.16%
	标准差 (Std.)	0.1516	0.1386	-8.58%
	变异系数 (CV)	0.0639	0.0585	-8.45%
强制次优路由探针	正常 Top-k 路由验证损失	3.187	3.182	-0.16%
(替换为排位靠中的中间专家)	强制中位专家误路由损失	7.249 (退化 2.27×)	7.071 (退化 2.22×)	超额损失降低 4.3%

2. 关键组件消融（信用归因、稀疏度与 EMA 衰减，5 项基准平均）

消融配置维度	具体参数设置	ARC-E	HellaSwag	PIQA	SciQ	ReCoRD (F1)	5项平均得分
归因信用机制	FLAME-MoE 基线	0.5295	0.3447	0.6697	0.8020	0.7023	0.6096
	激活加权归因 ($\eta=10^{-2}$)	0.5400	0.3470	0.6839	0.8180	0.7007	0.6179
	激活加权归因 ($\eta=10^{-3}$)	0.5682	0.3452	0.6888	0.8190	0.7105	0.6263
	仅纯门控概率归因 ($\eta=10^{-2}$)	0.5455	0.3384	0.6839	0.7920	0.6938	0.6107
	仅纯门控概率归因 ($\eta=10^{-3}$)	0.5421	0.3435	0.6861	0.7910	0.7065	0.6138
Top-k 路由密度	Top-12 路由基线	0.5459	0.3506	0.6991	0.7950	0.7033	0.6188
	Top-12 + DRMoET ($\eta=10^{-3}$)	0.5640	0.3523	0.6910	0.8130	0.7056	0.6252
	Top-12 + DRMoET ($\eta=10^{-2}$)	0.5619	0.3515	0.6991	0.8110	0.6998	0.6247
EMA 衰减系数 $\beta$	$\beta = 0.9$	0.5568	0.3477	0.6915	0.7970	0.7048	0.6196
($\eta=10^{-3}$ 固定)	$\beta = 0.99$	0.5408	0.3455	0.6937	0.7940	0.7022	0.6152
	$\beta = 0.999$ (默认推荐)	0.5682	0.3452	0.6888	0.8190	0.7105	0.6263
	$\beta = 1.0$ (不更新历史)	0.5467	0.3465	0.6844	0.7850	0.7068	0.6139

关键发现¶

最弱专家短板大幅补齐：DRMoET 将最差专家的平均损失压低 2.05%，专家间损失的标准差与变异系数下降超过 8%，同时整体平均损失基本保持不变（仅 -0.05%），说明模型并没有通过牺牲头部专家的上限来换取平庸，而是扎实提升了边缘专家的能力底线。
次优误路由惩罚显著降低：在强制使用排序靠中的中间专家（Mid-k misrouting）的严苛探针测试下，DRMoET 的超额损失较基线降低了 4.3%，有效证明了当实际部署或长尾输入导致路由器选择次优路径时，模型具备更强的兜底鲁棒性。
专业化分工不降反升：通过对 6 个功能领域的教师强制推理分析，DRMoET 的领域-专家互信息 $I(E; G)$ 提升了 12.9%，首选领域的胜任优势 $\Delta Q$ 大幅提升 30.2%，证明追求路由鲁棒性不仅没有抹杀专家的特化分工，反而促成了更高质量的专业化解耦。

亮点与洞察¶

内生化 DRO 组映射：创新性地将各层专家直接映射为 DRO 的对抗组，省去了传统 DRO 对外部领域或人口统计先验标签的依赖，使得鲁棒优化能无缝内嵌于现代神经网络的稀疏路由中。
激活模长截断信用分配：引入未合并特征激活范数与门控概率乘积并施加 stop-gradient，巧妙契合了 MoE 线性合并的物理本质，既避免了虚假的模长梯度膨胀，又比单纯依赖门控概率的加权提升了 1.25 个百分点的下游平均分。
零额外前向反向通信开销：完全复用前向已经计算好的激活张量与损失标量，对偶更新为极轻量的标量 Softmax，实测硬件吞吐（67.19 vs 64.17 TFLOP/s/GPU）甚至因消除了部分掉队节点（straggler）而略有上升。

局限与展望¶

模型族与超大专家数的泛化待验：论文主要在 FLAME-MoE 架构（32 与 64 专家）下完成验证，更深层次路由（如细粒度 256+ 专家 MoE 或多头 MoE）下的鲁棒性边界仍需更广阔的实证支撑。
极端长尾分布构造不足：当前评测仍基于通用预训练语料与学术标准评测集，未来可在高度倾斜的多语言、小众代码等极端长尾场景下进一步探究其对罕见路由崩溃的防御上限。
动态演化机理分析有限：强制误路由与专家层级分析仅在最终收敛权重上进行，尚未完全展示专家鲁棒性与特化分工在预训练全周期内的演变过程。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将层级专家视为内生 DRO 分组，突破了传统仅依靠流量平衡优化 MoE 的范式。
实验充分度: ⭐⭐⭐⭐⭐ 从 290M 到 10.3B 规模两套从头预训练，机制探针、误路由破坏实验、特化指标与消融完备。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰，算法描述与系统实现考量扎实，实验图表紧密支撑核心主张。
价值: ⭐⭐⭐⭐⭐ 为大模型稀疏化扩展提供了开箱即用、几乎零成本的可靠性与鲁棒性增强范式。