Training Matryoshka Mixture-of-Experts for Elastic Inference-Time Expert Utilization
作者: Yaoxiang Wang, Qingguo Hu, Yucheng Ding, Ruizhe Wang, Yeyun Gong, Jian Jiao, Yelong Shen, Peng Cheng, Jinsong Su
分类: cs.CL
发布日期: 2025-09-30
💡 一句话要点
提出Matryoshka MoE，实现MoE模型在推理时专家利用的弹性调整。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 混合专家模型 弹性推理 模型训练 专家排序 计算效率
📋 核心要点
- 现有MoE模型采用Top-K路由，推理时专家数量改变会导致性能急剧下降，缺乏弹性。
- M-MoE通过训练时改变激活专家数量，使模型学习粗到细的专家排序，提升弹性推理能力。
- 实验表明，M-MoE模型在不同专家数量下性能接近专门训练的模型，且训练成本更低。
📝 摘要（中文）
混合专家模型(MoE)已成为有效扩展大型语言模型的一种有前景的范例，它无需成比例地增加计算成本。然而，Top-K路由器的标准训练策略阻碍了MoE模型充分发挥其弹性推理的潜力。当在推理时改变激活专家的数量时，这些模型会表现出急剧的性能下降。本文介绍Matryoshka MoE (M-MoE)，一个训练框架，将粗到细的结构直接灌输到专家集成中。通过在训练期间系统地改变激活专家的数量，M-MoE迫使模型学习有意义的排序：排名靠前的专家协同提供必要的基础能力，而后续的专家则逐步添加更精细的细节。我们在多个粒度上探索了这一原则，确定了逐层随机化策略是最有效的。实验表明，单个M-MoE模型实现了卓越的弹性，其在各种专家数量下的性能与整个专家模型套件的性能非常匹配，但仅需总训练成本的一小部分。这种灵活性不仅解锁了弹性推理，而且还能够通过为不同的模型层分配不同的计算预算来优化性能。这项工作为大规模MoE模型更实用和更具适应性的部署铺平了道路。
🔬 方法详解
问题定义：现有MoE模型依赖Top-K路由，在训练时固定激活的专家数量。当推理时需要调整激活的专家数量以适应不同的计算资源或延迟要求时，模型的性能会显著下降。这种缺乏弹性的问题限制了MoE模型在实际部署中的应用。
核心思路：M-MoE的核心思路是在训练过程中引入专家数量的随机变化，从而迫使模型学习一种粗到细的专家排序。排名靠前的专家负责提供模型的基础能力，而后续的专家则逐步添加更精细的细节。通过这种方式，模型能够适应不同数量的激活专家，并在不同的计算预算下保持较好的性能。
技术框架：M-MoE的整体框架与标准的MoE模型类似，主要区别在于训练阶段。在每个训练步骤中，M-MoE会随机选择一个激活专家数量，然后使用Top-K路由选择相应的专家进行计算。通过这种方式，模型能够学习在不同专家数量下的最优策略。在推理阶段，可以根据实际需求选择激活的专家数量，而无需重新训练模型。
关键创新：M-MoE的关键创新在于其训练策略，即在训练过程中引入专家数量的随机变化。这种策略迫使模型学习一种粗到细的专家排序，从而提高了模型的弹性。与传统的MoE模型相比，M-MoE能够更好地适应不同的计算资源和延迟要求。
关键设计：M-MoE的关键设计包括：1) 逐层随机化策略，即在不同的模型层使用不同的专家数量；2) 合适的专家数量采样策略，以确保模型能够充分学习不同专家数量下的最优策略；3) 损失函数的设计，以鼓励模型学习粗到细的专家排序。
🖼️ 关键图片
📊 实验亮点
实验结果表明，单个M-MoE模型在不同专家数量下的性能与专门训练的模型套件的性能非常接近，但仅需总训练成本的一小部分。例如，在某个实验中，M-MoE模型在激活不同数量的专家时，性能仅下降了几个百分点，而专门训练的模型套件则需要数倍的训练成本。
🎯 应用场景
M-MoE可应用于各种需要弹性推理的大型语言模型部署场景，例如云端服务、边缘设备等。它能够根据实际的计算资源和延迟要求，动态调整激活的专家数量，从而在性能和效率之间取得平衡。此外，M-MoE还可以用于优化模型层的计算预算分配，提高整体性能。
📄 摘要（原文）
Mixture-of-Experts (MoE) has emerged as a promising paradigm for efficiently scaling large language models without a proportional increase in computational cost. However, the standard training strategy of Top-K router prevents MoE models from realizing their full potential for elastic inference. When the number of activated experts is altered at inference time, these models exhibit precipitous performance degradation. In this work, we introduce Matryoshka MoE (M-MoE), a training framework that instills a coarse-to-fine structure directly into the expert ensemble. By systematically varying the number of activated experts during training, M-MoE compels the model to learn a meaningful ranking: top-ranked experts collaborate to provide essential, coarse-grained capabilities, while subsequent experts add progressively finer-grained detail. We explore this principle at multiple granularities, identifying a layer-wise randomization strategy as the most effective. Our experiments demonstrate that a single M-MoE model achieves remarkable elasticity, with its performance at various expert counts closely matching that of an entire suite of specialist models, but at only a fraction of the total training cost. This flexibility not only unlocks elastic inference but also enables optimizing performance by allocating different computational budgets to different model layers. Our work paves the way for more practical and adaptable deployments of large-scale MoE models.