AdaMoE:面向混合专家语言模型的带空专家令牌自适应路由方法
AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models
Zihao Zeng, Yibo Miao, Hongcheng Gao, Hao Zhang, Zhijie Deng
Findings of the Association for Computational Linguistics: EMNLP (Findings of EMNLP) · 2024-01-01
摘要
混合专家模型因其能够在不显著增加计算开销的前提下提升模型容量,已成为构建生产级大语言模型的标准方案。然而,现有混合专家模型方法通常对所有令牌采用固定的top-k路由策略,这种限制较为明显——因为不同令牌(如“<EOS>”与“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路由的约束,指出不同token(如“<EOS>”与“apple”)所需专家数量不同,固定k会限制有限资源利用与下游任务潜力。作者提出AdaMOE,在专家集中加入固定数量、零FLOPs的null experts并增大k,使每个token可选不同数量的真实专家。AdaMOE不强制每个token占用固定数量null experts,而是通过负载均衡损失约束null experts的平均使用率,从而形成自适应的null/true experts分配。该路由保持token-choice形式,区别于依赖未来token的expert-choice路由,可用于自回归建模,并可作为即插即用方法应用于预训练LLM与MoE-LLM的监督微调。实验主张其在降低平均专家负载(FLOPs)的同时取得更优性能,如Mixtral-8x7B在ARC-C上FLOPs降低14.5%、准确率提高1.69%。
- 研究方法
- 研究采用路由分布统计诊断、监督微调实验与消融实验相结合的方法路径。首先在SocialIQA上统计Mixtral-8x7B各MoE层中累计路由概率超过50%的top专家数量,以此质疑固定top-k的合理性。常规LLM实验以Llama2-7B为基座,在Mo-LoRA架构中加入LoRA experts与路由器,使用RTE、COLA、SQA、CQA、OQA五个数据集,设置n=4个真实专家,基线为top-1/top-2,AdaMOE比较m5k2、m5k4、m7k4、m9k4等配置,采用AdamW、学习率3e-4、LoRA rank 8、2个epoch、3个随机种子取平均。MoE-LLM实验以Mixtral-8x7B为基座,覆盖WINO、HELLA、PIQA、SIQA、OQA、ARC-C六个数据集,使用4-bit量化与QLoRA,LoRA目标模块为gate、gate2、w1、w2、w3,rank 8、学习率5e-5、每个数据集1000个训练样本、2个epoch,并将负载均衡损失权重从0.02退火至0.0001,主要报告m=8、k=3配置并以OpenCompass评测准确率。材料还报告了不同m与k的消融结果,但节选未展开完整表格。
- 学术贡献
- 贡献主要在方法创新与应用验证层面:提出以null experts实现token级自适应路由,对vanilla top-k MoE只需增加零FLOPs专家并扩大k,降低了改造门槛。其修改后的负载均衡损失对null experts取平均负载而不区分个体,避免对路由器施加不必要约束;路由概率仅在真实专家内归一化,以保持与vanilla MoE层输出尺度一致。该方法在保持自回归建模可行性的同时,获得与expert-choice路由相似的变专家数量效果,对MoE路由策略、MoE-LLM微调与Mo-LoRA参数高效微调等子领域构成直接的方法贡献。实验层面给出了平均专家负载、FLOPs与多基准准确率的联合证据:Mixtral-8x7B平均负载由2降至1.66、平均FLOPs降低15.21%、平均准确率由84.64升至85.35。材料未明确提及新的理论建构或新数据集发布。
- 存在局限
- 实验主要在微调阶段展开,作者明确因资源限制未进行预训练,因此AdaMOE在从零预训练或持续预训练中的路由行为、收敛性与计算收益尚未被验证。Mixtral-8x7B实验每个数据集仅用1000个训练样本、训练2个epoch,并采用4-bit QLoRA与rank 8的LoRA,样本规模与参数高效设置限制了结论向全量数据、全参数微调或其他规模MoE-LLM的外推。常规LLM实验使用Llama2-7B与Mo-LoRA,材料指出LoRA expert负载对总FLOPs影响很小,因而该部分主要报告准确率,AdaMOE在非MoE LLM上的FLOPs收益缺乏直接度量。Null expert默认采用零映射,材料虽提到恒等映射也满足零FLOPs,但未明确展示两种实现的效果对比;路由概率归一化选择仅在真实专家内进行,材料未给出与全部top-k归一化的对照结果。论证以基准准确率、FLOPs与负载指标为主,未给出自适应路由数量与性能收益之间关系的理论刻画,也未报告实际硬件时延或吞吐。
- 进一步研究方向
- 可在较小规模MoE-LLM上开展从零预训练或持续预训练对照实验,将AdaMOE与vanilla top-k、expert-choice路由在相同token预算下比较验证损失、平均真实专家负载、FLOPs与下游准确率,以检验token自适应是否在预训练阶段稳定出现。可设计受控消融,系统比较零映射与恒等映射两类null experts、以及仅在真实专家内归一化与在全部top-k内归一化对路由分布、输出尺度和准确率的影响。可将评测从多项选择准确率扩展到开放式生成、长上下文、多语言与指令跟随任务,并在固定硬件上测量wall-clock时延、吞吐与显存,以验证平均专家负载下降能否转化为实际推理加速。可沿层记录每个token选择的真实/null专家数量与路由熵,结合token类型(如EOS、标点、语义词)和错误样本分析自适应路由的决策依据,并检验全选null experts时绕过MoE层的效果。还可将AdaMOE与Mo-LoRA、QLoRA等PEFT设置组合,在不同rank、量化精度和训练样本规模下重复实验,检验其即插即用性与超参数敏感性。
本综述由 AI(deepseek-flash)阅读开放获取全文自动生成,仅供参考,不代表原文作者观点,请以原文为准。
参考文献
相关论文
话题:Topic Modeling