[论文解读] Scaling Laws for Fine-Grained Mixture of Experts

Jakub Krajewski, Jan Ludziejewski+10 more|arXiv (Cornell University)|Feb 12, 2024

Advanced Statistical Process MonitoringDecision Sciences被引用 3

一句话总结

本文引入粒度（granularity）作为优化混合专家（MoE）模型的新超参数，实现对专家规模的细粒度控制。通过推导包含模型规模、训练样本数和粒度的缩放定律，作者证明MoE模型在效率方面始终优于稠密Transformer模型，即使在大规模计算预算下，计算量节省也超过40倍——这与以往认为MoE效率随规模增大而减弱的观点相悖。

ABSTRACT

Mixture of Experts (MoE) models have emerged as a primary solution for reducing the computational cost of Large Language Models. In this work, we analyze their scaling properties, incorporating an expanded range of variables. Specifically, we introduce a new hyperparameter, granularity, whose adjustment enables precise control over the size of the experts. Building on this, we establish scaling laws for fine-grained MoE, taking into account the number of training tokens, model size, and granularity. Leveraging these laws, we derive the optimal training configuration for a given computational budget. Our findings not only show that MoE models consistently outperform dense Transformers but also highlight that the efficiency gap between dense and MoE models widens as we scale up the model size and training budget. Furthermore, we demonstrate that the common practice of setting the size of experts in MoE to mirror the feed-forward layer is not optimal at almost any computational budget.

研究动机与目标

• 通过优化混合专家（MoE）架构，解决大规模语言模型（LLMs）训练的高计算成本问题。
• 挑战在固定训练时长和专家规模下，MoE效率随规模增大而减弱的假设。
• 在不同计算预算下，识别MoE模型的最优超参数（尤其是粒度）。
• 提供一个统一的框架，利用包含可变训练时长和专家粒度的缩放定律，实现MoE模型的计算最优训练。

提出的方法

• 引入粒度（G）作为控制MoE模型中专家规模的新超参数，实现对专家容量的细粒度调节。
• 推导新的缩放定律，建模模型规模、训练样本数、粒度和训练时长之间的关系，以FLOPs为主要成本度量指标。
• 采用改进的MoE架构，包含共享专家和分段专家，以支持在不同粒度下的高效训练。
• 在多个模型规模和计算预算下开展大量实验，以验证缩放定律并识别最优配置。
• 使用实际运行时间和困惑度作为评估训练效率的指标，表明更高的粒度可显著缩短达到相同性能所需的训练时间。
• 对扩展率（E）进行消融研究，证明结果在不同E值（包括E=16和E=64）下均成立。

实验结果

研究问题

• RQ1 调整MoE模型中的粒度如何影响其在不同计算预算下的训练效率和性能？
• RQ2 随着模型规模和训练预算的增加，MoE模型是否仍能保持甚至提升其相对于稠密Transformer模型的效率优势？
• RQ3 将专家规模设置为与前馈层规模相等（即G=1）是否是MoE模型的最优实践？
• RQ4 可变训练时长与缩放定律如何相互作用，以确定计算最优的MoE配置？
• RQ5 细粒度MoE模型是否能在极端计算预算下（例如10^25 FLOPs）仍实现优于稠密Transformer模型的效率？

主要发现

• 在10^25 FLOP预算下，采用最优粒度的MoE模型在计算效率方面始终优于稠密Transformer模型，计算节省超过40倍。
• 在几乎所有计算预算下，将专家规模设置为与前馈层规模相等（即G=1）的常规做法均非最优。
• 随着模型规模和训练预算的增加，MoE与稠密模型之间的效率差距进一步扩大，这与以往认为收益递减的观点相矛盾。
• 最优训练时长随模型规模和粒度而变化，将可变训练时间纳入缩放定律对于实现计算最优配置至关重要。
• 细粒度MoE模型能更快达到更低的困惑度，与标准MoE和稠密基线相比，在实际运行时间上展现出显著优势。
• 包含粒度、模型规模、训练样本数和可变训练时长的缩放定律，能够准确预测任意给定FLOP预算下的最优超参数。