Unified Sparse Mixture of Experts

📄 arXiv: 2503.22996v2 📥 PDF

作者: Giang Do, Hung Le, Truyen Tran

分类: cs.CL

发布日期: 2025-03-29 (更新: 2025-10-27)

备注: 26 pages


💡 一句话要点

提出统一稀疏专家混合模型(USMoE),解决传统SMoE路由的局限性。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 稀疏专家混合模型 线性规划 路由算法 大型语言模型 计算机视觉

📋 核心要点

  1. 传统SMoE方法依赖固定k值,导致无法有效路由重要专家或token,并易产生表征崩溃。
  2. USMoE框架通过线性规划视角,统一整合专家和token维度信息,并采用统一评分函数。
  3. 实验表明,USMoE在多种任务中性能提升高达10%,或降低推理成本高达14%。

📝 摘要(中文)

稀疏专家混合模型(SMoEs)通过扩展模型容量,同时保持恒定的计算开销。早期的设计通常依赖于固定的k值,其中k代表每个token选择的专家数量或每个专家分配的token数量。然而,这些方法存在三个关键限制:可能无法路由到重要的专家或token,可能分配不相关的专家或token,并且经常遭受专家之间的表征崩溃。本文通过线性规划的视角重新审视SMoEs,并提出了一个统一稀疏专家混合模型(USMoE)框架,以解决这些限制。具体来说,我们的方法引入了一种统一机制,整合了来自专家和token维度的信息,以及一个统一的评分函数,线性组合了专家和token之间的相似性得分。我们提供了理论证明和经验证据,证明了USMoE在克服传统路由方法局限性方面的有效性。通过在大型语言模型和视觉任务的干净和损坏设置下,在无训练和训练场景下的全面评估,USMoE实现了比标准方法高达10%的性能提升,或降低高达14%的推理成本,同时保持了具有竞争力的准确性。

🔬 方法详解

问题定义:现有的稀疏专家混合模型(SMoE)在路由决策上存在局限性。它们通常依赖于固定的k值,即每个token选择的专家数量或每个专家分配的token数量是固定的。这种固定分配策略可能导致重要专家或token未被选中,或者将不相关的专家或token分配到一起,从而影响模型的性能。此外,专家之间的表征容易发生崩溃,降低模型的表达能力。

核心思路:USMoE的核心思路是将SMoE的路由问题视为一个线性规划问题,并设计一个统一的框架来优化token到专家的分配。通过整合来自专家和token两个维度的信息,并使用统一的评分函数来衡量它们之间的相似性,USMoE能够更有效地进行路由决策,从而克服传统SMoE方法的局限性。这种设计允许模型根据token和专家的实际情况动态地调整分配策略,避免了固定k值带来的问题。

技术框架:USMoE框架主要包含以下几个关键模块:1) 特征提取模块:用于提取token和专家的特征表示。2) 相似度计算模块:计算token和专家之间的相似度得分。3) 统一路由模块:基于线性规划优化目标,将token分配给专家。该模块是USMoE的核心,它整合了token和专家维度的信息,并使用统一的评分函数进行路由决策。4) 专家网络:每个专家都是一个独立的神经网络,用于处理分配给它的token。5) 聚合模块:将各个专家的输出进行聚合,得到最终的预测结果。

关键创新:USMoE最重要的技术创新在于其统一的路由机制。与传统的SMoE方法不同,USMoE不再依赖于固定的k值,而是通过线性规划来优化token到专家的分配。这种方法能够更有效地利用模型容量,并避免了传统方法中存在的路由问题。此外,USMoE还引入了一个统一的评分函数,该函数线性组合了token和专家之间的相似度得分,从而更准确地衡量它们之间的相关性。

关键设计:USMoE的关键设计包括:1) 统一评分函数:该函数线性组合了token和专家之间的相似度得分,可以根据具体任务进行调整。2) 线性规划优化目标:该目标旨在最大化token和专家之间的相似度,同时满足一定的约束条件,例如每个专家的容量限制。3) 专家网络的结构:每个专家都可以采用不同的网络结构,例如Transformer或MLP,具体取决于任务的需求。4) 损失函数:除了传统的交叉熵损失外,还可以引入正则化项来防止专家之间的表征崩溃。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,USMoE在大型语言模型和视觉任务中均取得了显著的性能提升。在某些任务中,USMoE的性能比标准方法提高了高达10%,同时在保持竞争力的准确性的前提下,推理成本降低了高达14%。这些结果验证了USMoE在克服传统SMoE方法局限性方面的有效性。

🎯 应用场景

USMoE具有广泛的应用前景,可应用于自然语言处理、计算机视觉等领域。例如,在大型语言模型中,USMoE可以提高模型的表达能力和泛化能力,从而改善文本生成、机器翻译等任务的性能。在计算机视觉领域,USMoE可以用于图像分类、目标检测等任务,提高模型的准确性和效率。此外,USMoE还可以应用于推荐系统、金融风控等领域,具有重要的实际价值和未来影响。

📄 摘要(原文)

Sparse Mixture of Experts (SMoEs) models scale the capacity of models while maintaining constant computational overhead. Early designs typically relied on a fixed value of $k$, where $k$ represents either the number of experts selected per token or the number of tokens assigned per expert. However, these approaches encounter three key limitations: they may fail to route to important experts or tokens, may assign irrelevant ones, and often suffer from representation collapse among experts. This paper reexamines SMoEs through the lens of \textit{Linear Programming}, and proposes a Unified Sparse Mixture of Experts (USMoE) framework that addresses these limitations. Specifically, our approach introduces a unified mechanism that integrates information from both the expert and token dimensions, and a unified scoring function that linearly combines similarity scores between experts and tokens. We provide both theoretical justification and empirical evidence demonstrating USMoE's effectiveness in overcoming the limitations of traditional routing methods. Through comprehensive evaluations on both clean and corrupted settings for large language models and vision tasks, under both training-free and training scenarios, USMoE achieves up to a 10\% performance improvement over standard approaches or reduces inference costs by up to 14\%, while maintaining competitive accuracy.