arXiv:2410.10456cs.CL2024-10被引 23

动态调整专家数量,让大模型更省力还更准。

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

  • 根据上下文自动决定每词激活的专家数,不再固定
  • 比传统方法减少25%计算量,推理快20%以上
  • 适合所有主流MoE大模型,训练仅需8小时

在大语言模型时代,混合专家(MoE)架构为扩展模型参数的同时控制计算成本提供了可行路径。传统MoE模型采用静态Top-K路由,对每个令牌均激活固定且相等数量的专家,忽略其语境重要性。本文提出一种新型自适应动态选择机制(Ada-K),可根据每个令牌的上下文需求动态调整激活专家数量,实现效率与性能的更好平衡。该策略引入可学习、轻量级分配模块,负责定制化资源分配,且设计为全插件式,可广泛适配主流MoE模型。通过近端策略优化(PPO)算法实现非可微决策过程的端到端训练。在四个主流基线模型上的实验证明,相比传统Top-K,Ada-K显著提升性能:减少超过25%的浮点运算量(FLOPs),推理速度提升超20%,且保持或提高多任务基准表现。训练高效,即使对于参数量超1400亿的Mixtral-8x22B模型,训练时间也仅需8小时。详细分析显示,复杂任务、中间层及内容词更倾向激活更多专家,为未来自适应MoE系统设计提供洞见。训练代码与模型权重将公开共享。

原文摘要 · Abstract (English)

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.

MoE大模型路由优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。