Turn Waste into Worth: Rectifying Top-$k$ Router of MoE

📄 arXiv: 2402.12399v2 📥 PDF

作者: Zhiyuan Zeng, Qipeng Guo, Zhaoye Fei, Zhangyue Yin, Yunhua Zhou, Linyang Li, Tianxiang Sun, Hang Yan, Dahua Lin, Xipeng Qiu

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-02-17 (更新: 2024-02-21)

💡 一句话要点

提出Rectify-Router以解决MoE模型中的冗余计算问题

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

关键词: 稀疏混合专家 路由机制 计算效率 深度学习 模型优化

📋 核心要点

  1. 现有的top-$k$路由机制在稀疏MoE模型中存在冗余计算和内存成本高的问题，导致部分专家过载而丢弃令牌。
  2. 本文提出了Rectify-Router，通过Intra-GPU Rectification和Fill-in Rectification解决丢弃令牌和填充问题，提升模型性能。
  3. 实验结果显示，Intra-GPU Rectification和Fill-in Rectification分别有效处理丢弃和填充问题，组合使用时准确率提升4.7%。

📝 摘要（中文）

稀疏专家混合模型（MoE）因其计算效率而受到广泛关注。然而，现有的top-$k$路由机制存在冗余计算和内存成本高的问题，导致部分专家过载而丢弃超出令牌，另一些专家则因未使用而填充零，影响模型性能。为了解决这些问题，本文提出了Rectify-Router，包括Intra-GPU Rectification和Fill-in Rectification。Intra-GPU Rectification有效处理丢弃的令牌，将其高效路由到所在GPU内的专家，避免跨GPU通信；Fill-in Rectification通过用高路由分数的令牌替换填充令牌来解决填充问题。实验结果表明，这两种方法有效改善了模型性能，组合使用时准确率比传统top-1路由提升了4.7%。

🔬 方法详解

问题定义：本文旨在解决稀疏MoE模型中top-$k$路由机制导致的冗余计算和内存浪费问题。现有方法在处理令牌时，部分专家过载而丢弃令牌，另一些专家则因未使用而填充零，影响模型性能。

核心思路：论文提出的Rectify-Router通过两种主要的修正机制来优化路由过程。Intra-GPU Rectification专注于处理丢弃的令牌，确保令牌在其所在GPU内高效路由；Fill-in Rectification则通过用高路由分数的令牌替换填充令牌，减少无效计算。

技术框架：整体架构包括两个主要模块：Intra-GPU Rectification和Fill-in Rectification。前者处理丢弃的令牌，后者解决填充问题。两者结合使用时，能够显著提升模型的准确性和效率。

关键创新：最重要的技术创新在于提出了Intra-GPU Rectification和Fill-in Rectification两种修正机制，分别针对丢弃和填充问题进行优化，显著提升了模型性能。与传统top-1路由相比，本文方法在处理效率和准确性上均有显著改进。

关键设计：在设计中，参数设置和损失函数经过精心调整，以确保高效的令牌路由和最小化冗余计算。网络结构方面，采用了适应性路由策略，以便在不同的计算场景下灵活应对。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Intra-GPU Rectification和Fill-in Rectification有效处理了丢弃和填充问题，组合使用时模型的准确率比传统top-1路由提升了4.7%。这一显著提升展示了新方法在优化稀疏MoE模型中的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、图像识别和其他需要高效计算的深度学习任务。通过优化MoE模型的路由机制，能够在保持高性能的同时降低计算资源的消耗，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Sparse Mixture of Experts (MoE) models are popular for training large language models due to their computational efficiency. However, the commonly used top-$k$ routing mechanism suffers from redundancy computation and memory costs due to the unbalanced routing. Some experts are overflow, where the exceeding tokens are dropped. While some experts are vacant, which are padded with zeros, negatively impacting model performance. To address the dropped tokens and padding, we propose the Rectify-Router, comprising the Intra-GPU Rectification and the Fill-in Rectification. The Intra-GPU Rectification handles dropped tokens, efficiently routing them to experts within the GPU where they are located to avoid inter-GPU communication. The Fill-in Rectification addresses padding by replacing padding tokens with the tokens that have high routing scores. Our experimental results demonstrate that the Intra-GPU Rectification and the Fill-in Rectification effectively handle dropped tokens and padding, respectively. Furthermore, the combination of them achieves superior performance, surpassing the accuracy of the vanilla top-1 router by 4.7%.