AdaMoE：面向混合专家语言模型的基于空专家的词元自适应路由方法
AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models
Zihao Zeng, Yibo Miao, Hongcheng Gao, Hao Zhang, Zhijie Deng
arXiv (Cornell University) · 2024-06-19
摘要
混合专家模型因能够在不显著增加计算开销的前提下提升模型容量，已成为构建生产级大语言模型的标准范式。然而，现有的混合专家模型方法通常对所有词元采用固定的top-k路由策略，这种限制显然较为僵化——因为不同词元（如“ ”与“apple”）在特征提取时所需的专家数量往往存在差异。打破这一限制有助于更高效地利用有限的计算资源，从而释放模型在下游任务中的潜力。基于此，我们提出AdaMoE以实现混合专家模型的词元自适应路由，允许不同词元选择不同数量的专家。AdaMoE对传统的top-k路由混合专家模型仅做了极小的修改：它仅在专家集中加入固定数量的空专家，这些空专家不会消耗任何浮点运算量，同时相应增大k值。AdaMoE并不强制每个词元占用固定数量的空专家，而是通过负载均衡损失函数确保空专家的平均利用率，进而使每个词元所使用的空专家与真实专家数量呈自适应状态。AdaMoE在运作机制上与采用专家选择路由的混合专家模型极为相似，同时能够轻松支持自回归建模。该方法实现简单，可有效应用于预训练的混合专家大语言模型。大量实验表明，AdaMoE能在降低平均专家负载（浮点运算量）的同时取得更优的性能。例如，在ARC-C数据集上，将该方法用于Mixtral-8x7B的微调时，浮点运算量可减少14.5%，而准确率却能提升1.69%。
本摘要由 AI（Hy-MT2-Pro）翻译生成，仅供参考，请以原文为准。
AI 综述
基于全文生成- 主要观点
- 论文质疑现有MoE对全部token一律采用固定top-k路由的做法，认为“<EOS>”与“apple”等不同token进行特征抽象所需的专家数量并不相同，固定top-k限制了有限资源的利用与模型在下游任务上的潜力。为此作者提出AdaMoE：在专家集合中引入固定数量、不消耗任何FLOPs的空专家（null experts，取零映射或恒等映射），同时提高top-k中的k值，使不同token实际选用的真专家数量可以不同。其关键机制是不强制每个token占用固定数量的空专家，而是用负载均衡损失约束空专家的平均使用率——m个空专家共享同一个平均负载、彼此不作区分——从而得到自适应的真/空专家配比。作者主张该设计在效果上接近expert choice routing，但避免了后者依赖未来token、不适用于自回归因果语言建模的缺陷，且可plug-and-play地用于预训练普通LLM与MoE-LLM。核心经验论断是token自适应路由能在降低平均专家负载（FLOPs）的同时提升精度，例如在ARC-C上微调Mixtral-8x7B使FLOPs下降14.5%、准确率提升1.69%；支撑这一论断的前置证据是对SocialIQA逐层路由分布锐度（累计路由概率超过50%所需的top专家数）在token间显著差异的统计。
- 研究方法
- 论文首先做诊断性统计：将SocialIQA送入采用top-2路由的Mixtral-8x7B，逐MoE层记录全部token的路由分布，统计累计路由概率超过50%所需的最少top专家数，据此把token分为四类并给出各层所占比例。方法实现上，AdaMOE层包含n个真专家与m个空专家，路由器输出维度由n扩为n+m并取更大的k；按式(6)改写负载均衡损失，对m个空专家取平均负载，避免在等价空专家之间做区分性均衡；路由概率归一化只在被选中的真专家上进行，以保持与vanilla MoE输出尺度一致；训练中采用由大到小（α1=0.02、α2=0.0001）的负载均衡权重退火。对普通LLM的验证使用Llama2-7B搭配Mo-LoRA架构（n=4），基线为top-1/top-2路由，AdaMoE测试m5k2、m5k4、m7k4、m9k4四种配置，数据集为GLUE的RTE、COLA与常识推理的SQA、CQA、OQA，优化器AdamW、学习率3e-4、LoRA rank为8、每个数据集训练2个epoch、3个随机种子取平均。对MoE-LLM的验证以Mixtral-8x7B（每层8个FFN专家、top-2路由）为基座，采用4-bit量化与QLoRA微调，LoRA rank为8、学习率5e-5、每个数据集取1000条训练样本训练2个epoch；由于原router不便直接扩展，另定义由原gate参数派生初始化的gate2，与gate共同构成覆盖真/空专家的路由器。评测数据集包括WINO、HELLA、PIQA、SIQA、OQA、ARC-C，主报配置为m=8、k=3，准确率统一用OpenCompass评测，指标含Acc.、%FLOPs↓与Load（每层平均使用的专家数）。此外还设了遍历不同m与k的消融实验（摘录中未给出完整消融结果）。
- 学术贡献
- 在路由机制层面，论文提出用零FLOPs空专家加增大的k把固定top-k改造为token自适应路由，对vanilla MoE的改动仅限于专家集合、k值与负载均衡损失，属于低侵入式的路由算法设计。对MoE路由策略这一子领域，论文给出了质疑固定top-k合理性的逐层路由分布锐度统计证据，并提供了一条与expert choice routing功能相近但保持因果/自回归建模可用的替代路径。对Mo-LoRA与MoE-LLM微调子领域，论文给出可插拔的兼容方案：为普通LLM添加随机初始化的router与LoRA专家，为MoE-LLM扩展router输出维度并由原gate参数派生出新增维度的初始权重，使微调起始阶段真/空专家的负载仍保持均衡。论文还提供了一项可操作的资源控制手段：通过设定空专家数量m把真专家负载（FLOPs）预先调到目标水平，再由负载均衡损失保证其被实际使用，从而在算力预算与精度之间做显式权衡。实验材料覆盖Llama2-7B（语义理解与常识推理共5个数据集）与Mixtral-8x7B（6个数据集），报告了平均15.21%的FLOPs下降、Load由2.00降至1.66，为token级自适应路由提供了跨模型、跨任务的证据。
- 存在局限
- 负载均衡权重的退火策略只给出α1=0.02、α2=0.0001两个阶段取值，摘录中未说明调度长度、切换时点以及与其他调度方案的对照，削弱了对该超参数设计必要性的论证。MoE-LLM实验每个数据集仅用1000条样本训练2个epoch，作者也承认MoE-LLM微调易过拟合，这一训练量限制使表1结论对数据抽样敏感，且表1未报告多种子方差。结果内部存在不一致：PIQA与OQA上AdaMOE低于微调基线，论文仅以平均准确率占优作结，未分析这两类任务上自适应路由失效的原因。验证场景限于监督微调，普通LLM上的专家只能以LoRA形式注入，未检验全参数微调或预训练阶段下空专家机制的行为；同时全文只报告FLOPs与Load等理论计算指标，未给出实际推理延迟与显存占用测量，“全部选中空专家即绕过该层”的加速效果未被实测确认。对“语义token应获得更多专家”这一核心直觉，论文仅有图1的单个SIQA样例与路由分布锐度统计作为支持，未对自适应专家数与token的语言学类别（标点、子词片段、功能词、实义词）做系统性对应分析。材料中未明确提及空专家数量m的选取准则与最优取值规律（仅给出若干配置与主配置m=8、k=3），也未见对空专家取恒等映射与零映射两种选择做对比。
- 进一步研究方向
- 在固定FLOPs预算下对(m, k)组合做系统网格扫描，逐点绘制精度—平均负载前沿曲线，检验其单调性并据此提炼m的可复用选择规则。将MoE-LLM微调的训练样本由1000条扩至全量并使用多随机种子报告均值与置信区间，直接检验“小样本+2 epoch”是否造成了表1中PIQA、OQA上的退化。对负载均衡权重α做对照实验，比较单阶段固定α、线性/余弦衰减与两阶段切换，记录训练过程中真/空专家使用比例与精度的演化轨迹，以判定退火是否必要。导出各层路由日志，按token类型（标点、子词片段、功能词、实义词）统计被选真专家数的分布并做显著性检验，以验证或修正“语义token需要更多专家”的假设。在推理引擎中实测wall-clock延迟、吞吐与显存占用，并与可跳层方案（如MoD）在同条件、同硬件下比较，量化“整层绕过”的实际收益。将AdaMoE用于小规模MoE模型的自监督预训练而非仅微调，观察空专家使用模式是否随训练阶段漂移，并对比零映射与恒等映射两类空专家在微调中的效果差异。
本综述由 AI（deepseek-flash）阅读开放获取全文自动生成，仅供参考，不代表原文作者观点，请以原文为准。
参考文献
相关论文
话题：Topic Modeling