跳到主要内容
返回时间线
arXiv来源发表:

循环混合专家模型首次拥有联合缩放定律:稀疏度把循环的有效参数增益抬高并延后饱和

核心概要

该工作提出 Loop Scaling Laws,首次把循环(recurrence)与 MoE 稀疏度同模型规模、训练数据一起纳入同一个缩放定律,用一个有界且以稀疏度为条件的循环映射刻画循环带来的有效参数增益及其饱和点,在留出集上比线性与幂律映射预测损失更准,并据此在算力与显存约束下选择循环次数与专家数,在万亿 token 规模上让 0.3B 激活/1.3B 总参数的循环 MoE 在推理基准上追平约 2 倍大的非循环 MoE。

AI-generated editorial illustration: Scaling Laws for Looped Mixture of Experts

深度剖析

论文提出 Loop Scaling Laws,把循环次数与 MoE 稀疏度(由专家数变化)和模型规模、训练数据放在同一个函数形式中联合建模,并给出以稀疏度为条件的循环映射:每个循环带来的有效参数增益递减,总增益趋近一个有限上界,而稀疏度既抬高这个上界又让它需要更多循环才饱和。 此前的缩放定律只单独处理循环或稀疏:Parcae 用线性、无界的展开参数计数,Iso-Depth 用幂律、无界映射,MoE 定律则不含循环;并行的循环 MoE 研究在主要缩放分析中把循环固定在两次。该工作把两者写进同一形式,并说明在循环次数为 1 时退化为稠密循环定律、在循环次数为 1 时退化为标准 MoE 定律、在两者同时取 1 时退化为标准稠密定律。 作者在模型规模、训练 token、专家数与循环次数四个轴上做联合扫描,用同一套拟合流程比较线性、幂律、有界以及稀疏度条件映射,并报告留出集 RMSE:稠密情形下有界映射为 0.0092,优于幂律 0.0313 与线性 0.2566;MoE 情形下稀疏度条件有界映射在多个留出切片上取得最低 RMSE(如 0.0100、0.0047、0.0050、0.0043)。拟合使用超过 1000 个损失观测,并以 100 次自助重采样给出系数与联合最优配置的 90% 置信区间。

拟合出的定律被用作设计工具:在固定稀疏度下求算力最优循环次数,在固定循环下求显存最优专家数,并在给定训练算力与权重显存预算时求两者的联合最优配置。 这把缩放定律从预测损失推进到资源约束下的架构选择,并明确区分两条互补的交换关系:固定权重显存时循环用额外算力换质量,固定算力时稀疏用额外权重显存换质量。 论文给出三个选择准则(式 11、12、13),以拟合定律的 RMSE 作为可分辨的最小损失改进阈值,并在五档模型规模、若干专家数与循环次数构成的 300 个候选配置上搜索;报告的趋势包括算力最优循环次数随训练算力预算上升、更小激活模型需要更高循环、更大稀疏度把算力最优区间推向更高循环,以及显存越紧越偏向高循环加小专家数(如 3 GB 下选 1.0B 激活、循环 3、专家数 8;10 GB 下选 1.6B 激活、循环 2、专家数 16)。

在 14 个下游基准上,稀疏与循环给出互补增益:稀疏带来激活参数效率,循环带来总参数效率且更偏向推理任务,两者同时放大时性能前沿进一步前移。 这为“两条轴互补”提供了下游任务层面的经验证据,而不只是损失预测层面的拟合结果。 评测覆盖推理(BBH、GSM8K)、科学(ARC-C/E、OpenBookQA)、常识(HellaSwag、PIQA、SIQA、WinoGrande)、阅读(BoolQ、DROP)与知识(MMLU、Natural Questions、TriviaQA)五类共 14 个基准,比较 0.3B、0.6B、1.0B 激活规模在不同专家数与循环次数下的表现;报告称在匹配训练 FLOPs 时,0.3B 的 MoE 模型可超过更大的 1.0B 稠密模型,而循环带来的增益可抵消总参数翻倍,同时提升稀疏度与循环的模型持续优于稠密基线。

在万亿 token 规模的实际案例中,用定律推导出的循环次数训练 0.3B 激活/1.3B 总参数的循环 MoE,在匹配训练算力下于推理基准上追平 0.6B 激活/2.9B 总参数的非循环 MoE,并可通过在推理时改变循环次数实现按需的测试时扩展。 这把定律的预测能力落到真实训练规模,并展示循环带来的额外推理算力可以换取参数效率,同时提供一种推理期可调的算力开关。 两个模型分别在万亿 token 量级训练(循环 MoE 训练 token 数少于非循环 MoE),对应约相同训练 FLOPs;表 2 显示循环 MoE 在 BBH、GSM8K 上追平更大的非循环 MoE,但在 Overall 上落后,且把推理循环从 1 提高到 5 可显著提升 Overall 表现。

启示与展望

该定律面向在给定训练算力与权重显存预算下设计循环 MoE 的场景,适用于作者采用的模型阶梯、数据配方与中段循环策略;它给出的是在观测循环范围内的渐近性质,而非超出该范围的保证。对读者而言,这提供了一套可操作的配方:先拟合定律,再按式 11 选算力最优循环、按式 12 选显存最优专家数、按式 13 求联合最优,从而在资源受限部署中权衡参数、算力与显存。论文也指出可探索的方向,包括更丰富的循环策略、从边缘小模型到云端大模型的更广规模,以及通过优化循环状态或加入提前退出门控来降低循环模型的显存与延迟。

定律的渐近上界是映射的性质而非超出观测循环范围的保证,因此超出拟合区间的循环次数与专家数仍需谨慎;显存优化只覆盖权重显存,KV 缓存与其他运行时状态未纳入;下游评测中循环 MoE 在 Overall 上仍落后于更大的非循环 MoE,只在推理类任务上追平;此外,定律在作者自己的模型阶梯与数据配方上拟合,换用其他循环策略或数据分布时其系数是否稳定仍是开放问题。

来源