Ada-K Routing: Boosting the Efficiency of MoE-based LLMs

📄 arXiv: 2410.10456v2 📥 PDF

作者: Tongtian Yue, Longteng Guo, Jie Cheng, Xuange Gao, Jing Liu

分类: cs.CL

发布日期: 2024-10-14 (更新: 2024-10-15)

备注: Coauthors do not reach a consensus on submitting the current version

💡 一句话要点

提出Ada-K路由，通过动态调整专家激活数量提升MoE-LLM效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 MoE 动态路由 自适应资源分配 大语言模型 计算效率 近端策略优化

📋 核心要点

  1. 传统MoE模型采用静态Top-K路由，忽略了不同token对专家资源的需求差异，导致效率低下。
  2. Ada-K路由通过可学习的分配器模块，动态调整每个token激活的专家数量，实现资源优化。
  3. 实验表明，Ada-K在降低FLOPs和加速推理的同时，还能提升模型性能，且训练高效。

📝 摘要（中文）

在大语言模型（LLMs）时代，混合专家（MoE）架构为在扩展模型参数的同时管理计算成本提供了一种有前景的方法。传统的基于MoE的LLM通常采用静态的Top-K路由，即为每个token激活固定且数量相等的专家，而忽略了它们在上下文中的重要性。本文提出了一种新的Ada-K路由策略，该策略动态调整每个token激活的专家数量，从而提高计算效率和模型性能之间的平衡。具体来说，我们的策略结合了可学习的轻量级分配器模块，这些模块决定为每个token定制的专家资源分配，以满足上下文需求。这些分配器被设计为完全可插拔的，使其广泛适用于所有主流的基于MoE的LLM。我们利用近端策略优化（PPO）算法来促进这种不可微决策框架的端到端学习过程。在四个流行的基线模型上的广泛评估表明，我们的Ada-K路由方法明显优于传统的Top-K路由。与Top-K相比，我们的方法在各种基准测试中实现了超过25%的FLOPs减少和超过20%的推理加速，同时仍然提高了性能。此外，Ada-K的训练非常高效。即使对于Mixtral-8x22B，一个具有超过140B参数的基于MoE的LLM，训练时间也限制在8小时内。详细分析表明，更困难的任务、中间层和内容词倾向于激活更多的专家，这为未来的自适应MoE系统设计提供了有价值的见解。训练代码和模型检查点都将公开。

🔬 方法详解

问题定义：现有MoE模型中的Top-K路由策略为每个token分配固定数量的专家，没有考虑到不同token的重要性差异。这种静态分配方式导致计算资源的浪费，无法充分利用专家模型的潜力，尤其是在处理复杂任务时。

核心思路：Ada-K路由的核心思想是根据token的上下文信息，动态地调整激活的专家数量。通过引入可学习的分配器模块，模型可以自主地决定哪些token需要更多的专家资源，哪些token可以使用较少的专家资源，从而实现计算效率和模型性能的平衡。

技术框架：Ada-K路由的整体框架包括以下几个关键模块：1) 基础MoE模型：采用现有的MoE架构，如Mixtral。2) 可学习的分配器模块：该模块接收token的上下文表示作为输入，输出每个token应该激活的专家数量。3) 路由机制：根据分配器模块的输出，动态地选择激活的专家。4) 训练机制：使用PPO算法对分配器模块进行端到端训练，优化资源分配策略。

关键创新：Ada-K路由的关键创新在于引入了可学习的分配器模块，实现了动态的专家资源分配。与传统的静态Top-K路由相比，Ada-K能够根据token的上下文信息，自适应地调整激活的专家数量，从而提高计算效率和模型性能。

关键设计：分配器模块的设计至关重要，论文中采用轻量级网络结构，以减少额外的计算开销。PPO算法用于训练分配器模块，目标是最大化模型性能，同时约束激活的专家数量。损失函数的设计需要平衡模型性能和计算成本，鼓励模型在保证性能的前提下，尽可能减少激活的专家数量。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，Ada-K路由在四个流行的基线模型上显著优于传统的Top-K路由。与Top-K相比，Ada-K实现了超过25%的FLOPs减少和超过20%的推理加速，同时在各种基准测试中仍然提高了性能。例如，在Mixtral-8x22B模型上，Ada-K的训练时间仅为8小时，展示了其高效的训练能力。

🎯 应用场景

Ada-K路由具有广泛的应用前景，可以应用于各种基于MoE的LLM，例如机器翻译、文本生成、对话系统等。通过动态调整专家资源分配，Ada-K可以显著提高模型的计算效率和推理速度，降低部署成本，使其更适用于资源受限的场景。此外，Ada-K还可以促进更大规模MoE模型的训练和部署，推动LLM的发展。

📄 摘要（原文）

In the era of Large Language Models (LLMs), Mixture-of-Experts (MoE) architectures offer a promising approach to managing computational costs while scaling up model parameters. Conventional MoE-based LLMs typically employ static Top-K routing, which activates a fixed and equal number of experts for each token regardless of their significance within the context. In this paper, we propose a novel Ada-K routing strategy that dynamically adjusts the number of activated experts for each token, thereby improving the balance between computational efficiency and model performance. Specifically, our strategy incorporates learnable and lightweight allocator modules that decide customized expert resource allocation tailored to the contextual needs for each token. These allocators are designed to be fully pluggable, making it broadly applicable across all mainstream MoE-based LLMs. We leverage the Proximal Policy Optimization (PPO) algorithm to facilitate an end-to-end learning process for this non-differentiable decision-making framework. Extensive evaluations on four popular baseline models demonstrate that our Ada-K routing method significantly outperforms conventional Top-K routing. Compared to Top-K, our method achieves over 25% reduction in FLOPs and more than 20% inference speedup while still improving performance across various benchmarks. Moreover, the training of Ada-K is highly efficient. Even for Mixtral-8x22B, a MoE-based LLM with more than 140B parameters, the training time is limited to 8 hours. Detailed analysis shows that harder tasks, middle layers, and content words tend to activate more experts, providing valuable insights for future adaptive MoE system designs. Both the training code and model checkpoints will be publicly available.