Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization

📄 arXiv: 2603.21862v1 📥 PDF

作者: Weilin Wan, Jingtao Han, Weizhong Zhang, Cheng Jin

分类: cs.LG

发布日期: 2026-03-23


💡 一句话要点

提出MoE架构优化框架,通过联合约束和降维搜索,实现任意计算预算下的最优架构配置。

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

关键词: 混合专家模型 MoE 缩放定律 架构优化 计算效率

📋 核心要点

  1. 现有MoE模型缩放研究受限于实验预算,难以在庞大的架构设计空间中找到最优配置,且忽略了全局交互。
  2. 论文提出一种整体MoE架构优化框架,通过联合约束FLOPs、激活参数和总参数,并采用降维搜索策略。
  3. 实验表明,该框架能够根据计算预算生成完整的、最优的MoE架构,且近优配置范围随规模扩大。

📝 摘要(中文)

大型语言模型的缩放定律指导宏观资源分配,但将其转化为精确的混合专家(MoE)架构配置仍然是一个开放问题,因为设计空间巨大。现有的MoE缩放研究受限于实验预算,要么使用额外的MoE变量来扩充缩放公式,面临不可靠拟合的风险,要么固定所有非MoE因素,忽略全局交互。我们提出了一个可重用的MoE架构整体优化框架,弥合了这一差距。我们首先表明,仅凭每个token的FLOPs不足以作为MoE模型的公平指标,因为不同层类型的计算密度差异会增加参数量,而不会带来成比例的计算成本,并建立了每个token的FLOPs、激活参数和总参数的联合约束三元组。然后,通过代数约束和隐藏维度的秩保持特性,我们将16维的架构搜索空间简化为两个连续的低维阶段。经过对跨越六个数量级的计算量的数百个MoE模型的验证,我们的框架产生了鲁棒的缩放定律,可以将任何计算预算映射到完整的、最优的MoE架构。一个关键发现是,接近最优的配置范围随着规模的扩大而扩大,为从业者提供了定量的灵活性,以平衡缩放定律的建议与基础设施的约束。

🔬 方法详解

问题定义:现有MoE模型的缩放定律研究,难以在巨大的架构搜索空间中找到最优配置。现有方法要么引入过多MoE变量导致拟合不可靠,要么固定非MoE因素忽略全局交互。此外,仅使用FLOPs/token作为公平性指标是不够的,因为不同层计算密度不同,会导致参数量虚增。

核心思路:论文的核心思路是通过联合约束FLOPs/token、激活参数和总参数,来更准确地衡量MoE模型的计算成本和参数效率。同时,利用代数约束和隐藏维度的秩保持特性,将高维的架构搜索空间降维,从而降低搜索难度。

技术框架:该框架包含两个主要阶段。第一阶段,通过代数约束和秩保持特性,将16维的架构搜索空间简化为低维空间。第二阶段,在简化的搜索空间中,利用缩放定律来寻找最优的MoE架构配置。该框架能够根据给定的计算预算,输出完整的、最优的MoE架构。

关键创新:最重要的创新点在于提出了联合约束三元组(FLOPs/token、激活参数、总参数),解决了仅使用FLOPs/token作为公平性指标的不足。此外,通过降维搜索,显著降低了MoE架构优化的计算复杂度。

关键设计:关键设计包括:1) 联合约束三元组的构建方式,需要仔细选择合适的权重来平衡三个指标;2) 降维搜索的具体方法,包括代数约束和秩保持特性的应用;3) 缩放定律的具体形式,需要通过实验数据进行拟合和验证。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,该框架能够有效地将计算预算映射到最优的MoE架构。通过对数百个MoE模型进行验证,证明了该框架的鲁棒性。一个关键发现是,随着模型规模的扩大,近优配置的范围也随之扩大,为实际应用提供了更大的灵活性。

🎯 应用场景

该研究成果可应用于大规模语言模型的MoE架构设计,帮助研究人员和工程师在给定的计算预算下,快速找到最优的MoE模型配置。这有助于提升模型的性能和效率,降低训练成本,并加速大型语言模型的部署和应用。

📄 摘要(原文)

Scaling laws for Large Language Models govern macroscopic resource allocation, yet translating them into precise Mixture-of-Experts (MoE) architectural configurations remains an open problem due to the combinatorially vast design space. Existing MoE scaling studies are constrained by experimental budgets to either augment scaling formulas with extra MoE variables, risking unreliable fits, or fix all non-MoE factors, ignoring global interactions. We propose a reusable framework for holistic MoE architectural optimization that bridges this gap. We first show that FLOPs per token alone is an inadequate fairness metric for MoE models because differing computational densities across layer types can inflate parameters without proportional compute cost, and establish a joint constraint triad of FLOPs per token, active parameters, and total parameters. We then reduce the 16-dimensional architectural search space to two sequential low-dimensional phases through algebraic constraints and a rank-preserving property of the hidden dimension. Validated across hundreds of MoE models spanning six orders of magnitude in compute, our framework yields robust scaling laws that map any compute budget to a complete, optimal MoE architecture. A key finding is that the near-optimal configuration band widens with scale, giving practitioners quantitative flexibility to balance scaling law recommendations against infrastructure constraints.