S2MoE: Robust Sparse Mixture of Experts via Stochastic Learning

📄 arXiv: 2503.23007v1 📥 PDF

作者: Giang Do, Hung Le, Truyen Tran

分类: cs.CL

发布日期: 2025-03-29

备注: 4 pages

💡 一句话要点

S2MoE：基于随机学习的鲁棒稀疏专家混合模型，提升模型性能并降低推理成本

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

关键词: 稀疏专家混合模型 随机学习 表示崩溃 不确定性学习 路由策略

📋 核心要点

  1. 现有SMoE方法受限于专家嵌入维度过小和Top-K路由导致专家学习相似特征，易发生表示崩溃。
  2. S2MoE通过不确定性学习，使模型能够从确定性和非确定性输入中学习，从而提升模型的鲁棒性。
  3. 实验结果表明，S2MoE在保持性能的同时，显著降低了推理计算成本，降幅达28%。

📝 摘要（中文）

稀疏专家混合模型(SMoE)通过将输入token路由到少量专家，从而实现大型语言模型的高效训练。然而，由于表示崩溃问题，训练SMoE仍然具有挑战性。最近的研究主要集中在改进路由器以缓解这个问题，但现有方法面临两个关键限制：(1)专家嵌入远小于模型的维度，导致表示崩溃；(2)将每个输入路由到Top-K专家可能导致它们学习过于相似的特征。本文提出了一种名为基于随机学习的鲁棒稀疏专家混合模型(S2MoE)的新方法，它是一种通过不确定性学习从确定性和非确定性输入中学习的专家混合模型。在各种任务上的大量实验表明，S2MoE实现了与其他路由方法相当的性能，同时降低了28%的计算推理成本。

🔬 方法详解

问题定义：论文旨在解决稀疏专家混合模型(SMoE)训练过程中出现的表示崩溃问题。现有方法主要集中在改进路由策略，但存在两个痛点：一是专家嵌入维度远小于模型维度，导致信息瓶颈；二是Top-K路由使得专家学习到的特征过于相似，缺乏多样性。

核心思路：S2MoE的核心思路是通过引入随机学习机制，使模型能够从不确定性输入中学习，从而增强模型的鲁棒性和泛化能力。这种方法旨在鼓励专家学习更多样化的特征，并减轻表示崩溃问题。

技术框架：S2MoE的整体架构仍然是基于SMoE的框架，但其关键在于路由机制的改进。它包含以下主要模块：输入层、路由层、专家层和输出层。路由层负责将输入token分配给不同的专家。与传统的Top-K路由不同，S2MoE引入了随机性，使得token的分配具有一定的概率性。专家层由多个独立的神经网络组成，每个专家负责处理特定类型的输入。输出层将各个专家的输出进行聚合，得到最终的预测结果。

关键创新：S2MoE最重要的技术创新点在于其随机学习机制。通过引入随机性，S2MoE能够从不确定性输入中学习，从而增强模型的鲁棒性和泛化能力。与现有方法的本质区别在于，S2MoE不仅仅依赖于确定性的路由规则，而是允许模型探索更多的可能性，从而学习到更多样化的特征。

关键设计：S2MoE的关键设计包括：(1)随机路由策略：使用基于概率分布的路由方法，而不是简单的Top-K选择，鼓励专家学习不同的特征。(2)不确定性损失函数：引入额外的损失函数，用于衡量模型预测的不确定性，并鼓励模型学习更加鲁棒的表示。(3)专家容量平衡：采用容量平衡策略，确保每个专家接收到足够数量的token，避免某些专家过度饱和，而另一些专家则利用不足。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

S2MoE在多个NLP任务上进行了实验，结果表明其性能与现有SMoE方法相当，同时显著降低了推理计算成本，降幅达28%。这表明S2MoE在保持模型性能的同时，提高了模型的效率，使其更适用于实际应用。

🎯 应用场景

S2MoE适用于需要大规模模型和高效推理的各种自然语言处理任务，例如机器翻译、文本生成、问答系统等。该方法能够降低计算成本，使得在资源受限的环境中部署大型语言模型成为可能。未来，S2MoE可以应用于其他领域，例如计算机视觉和语音识别。

📄 摘要（原文）

Sparse Mixture of Experts (SMoE) enables efficient training of large language models by routing input tokens to a select number of experts. However, training SMoE remains challenging due to the issue of representation collapse. Recent studies have focused on improving the router to mitigate this problem, but existing approaches face two key limitations: (1) expert embeddings are significantly smaller than the model's dimension, contributing to representation collapse, and (2) routing each input to the Top-K experts can cause them to learn overly similar features. In this work, we propose a novel approach called Robust Sparse Mixture of Experts via Stochastic Learning (S2MoE), which is a mixture of experts designed to learn from both deterministic and non-deterministic inputs via Learning under Uncertainty. Extensive experiments across various tasks demonstrate that S2MoE achieves performance comparable to other routing methods while reducing computational inference costs by 28%.