这项由亚马逊研究团队完成的工作发表于2026年第43届国际机器学习大会ICML 2026会议地点为韩国首尔收录于PMLR第306卷。论文编号为arXiv:2606.01666有兴趣深入了解的读者可通过该编号查询完整论文。当你打开手机上的AI助手向它提问一个问题时背后那个大脑——也就是大型语言模型——需要在零点几秒内激活数十亿个参数来回答你。这就像一家万人企业每次接一个小订单都要全体员工放下手头工作同时上阵。效率低、耗电多、成本高这是当前AI大模型最真实的困境。亚马逊的研究团队决定对这种工作方式动一次根本性的手术。一、大公司的效率难题为什么AI模型这么费电现代大型语言模型比如驱动各种AI助手的那些系统通常拥有几十亿甚至上百亿个参数。每次你输入一句话模型就需要把所有这些参数都转动一遍来生成回答。这种方式被研究者称为稠密激活——不管你问的是今天天气怎么样还是请帮我推导量子力学公式同样数量的计算资源都会被调用。研究者很早就意识到这种方式过于浪费。以最近发布的Qwen3-30B-A3B模型为例它拥有305亿个参数但每次回答一个问题时实际被用到的只有33亿个。换句话说九成的员工在摸鱼。这种架构被称为混合专家模型Mixture of Experts简称MoE它的核心思想是把大模型拆分成许多小的专家小组每次只让最合适的几组专家来处理当前的问题。然而从头训练一个MoE模型并不是件容易的事。这个过程不稳定、数据需求量极大还需要复杂的负载均衡机制来防止某几个专家被过度使用而其他专家无事可做。于是研究者们开始探索另一条路能不能把一个已经训练好的普通大模型直接改造成MoE模型这个过程被称为MoE化。问题的核心在于改造的时候怎么决定哪些神经元归哪个专家管这个分组决策至关重要分得好改造后的模型依然聪明分得不好模型就会像一个被随意拆散的团队不知道该怎么协作。亚马逊的研究团队在这里找到了一个新颖的切入点他们把这个分组问题类比成了一个经典的物流问题——最优运输问题。二、前人的方法靠感觉分组的隐患在亚马逊团队提出新方法之前已有几种主流的MoE化方案但它们都有一个共同的软肋。最简单粗暴的方式是随机分组就像把一个公司的员工随机打散成若干个小队。LLaMA-MoE就采用了这种方式随机把神经元塞进不同的专家里然后再花费大量算力重新训练来弥补随机分组带来的损失。这种方式耗时耗力而且分组本身没有任何道理可言。更聪明一点的方式是根据神经元的性格相似性来分组。有的方法看神经元处理输入时的权重是否相近权重聚类有的方法看哪些神经元倾向于同时被激活共激活聚类。CMoE这个方法就用了平衡K均值聚类算法把同时活跃的神经元归为一组。这些方法听起来很合理但亚马逊团队指出了一个被所有人忽视的根本问题这些方法都在优化中间过程而不是最终结果。用一个具体的例子来理解这个问题。大模型里的前馈网络层FFN处理信息的方式是两步第一步把输入信号变换成一个中间表示第二步把这个中间表示再变换成输出。现有方法在分组神经元时只看了第一步里神经元的行为却没有考虑第二步——也就是这些神经元最终如何影响输出结果。这就好比给餐厅的厨师分组时只看谁的刀法相似却不看谁做出来的菜最受顾客欢迎。为了验证这个问题有多严重研究团队做了一个精准的单层重建实验把模型的某一层单独拿出来用不同方法分好专家组后直接测量改造后的输出和原始输出之间的误差。结果触目惊心。在LLaMA-2模型上随机分组的误差是亚马逊新方法的36倍而最好的对比方法CMoE的误差也有2倍多。在LLaMA-3模型上这个差距更是高达41倍对比2倍。这个实验清楚地证明看谁和谁长得像完全不够关键要看谁和谁搭档能产出最好的结果。三、核心创意把分组问题变成快递派单问题亚马逊团队的新方法叫做DOT-MoE其中DOT代表可微分最优运输Differentiable Optimal Transport。最优运输问题是数学领域的一个经典概念可以用一个非常直观的场景来理解假设城市里有若干个仓库和若干个零售商每个仓库有一批货物每个零售商需要特定数量的货物。最优运输问题就是要找到一个调货方案使得总的运输成本最低同时每个仓库的货物全部送完每个零售商的需求全部满足。在DOT-MoE里仓库变成了神经元零售商变成了专家组运输变成了神经元归属某个专家组的决定运输成本变成了这种归属方式导致的输出误差。研究团队的目标就是找到一个分配方案让每个神经元都有归属每个专家组的神经元数量恰好相等同时整体的输出误差最小。然而这个问题有一个技术上的挑战每个神经元只能属于一个专家组这是一个非此即彼的离散决策而神经网络的训练需要能够顺滑地计算梯度可以理解为调整方向的指示信号。一旦决策是离散的这个信号就会消失训练无法进行。研究团队用了一个叫做熵正则化的数学技巧来绕过这个障碍。简单说他们在最优化目标里加入了一个允许模糊的成分让分配方案在训练过程中可以是60%属于专家A40%属于专家B这样的模糊概率而不是非此即彼的硬性归属。这个模糊分配通过一种叫做Sinkhorn-Knopp迭代的算法来计算——这个算法的工作方式类似于反复校准交替保证每行的概率之和为1每个神经元只被分配一次和每列的概率之和为固定值每个专家组的容量相等。为了数值稳定整个计算在对数域进行避免了数值过小导致的下溢问题。等到训练完成后再把这个模糊的分配方案硬化成明确的归属按照分配概率从高到低排序依次确认每个神经元的最终归属同时确保每个专家组不超载。这个硬化过程本身不能传递梯度于是研究团队又用了一个叫做直通估计器Straight-Through EstimatorSTE的技巧前向计算时用硬化后的结果反向传播梯度时假装用的是软分配让优化信号能够顺利传回去更新分配逻辑。这整个设计的精妙之处在于它同时解决了三件事分配是可以学习的不靠启发式规则分配是平衡的每个专家组神经元数量相等通过Sinkhorn严格保证优化目标是输出级别的误差不是中间表示的相似度。四、路由器的设计谁来决定每个问题找哪个专家光有专家分组还不够还需要一个调度员来决定每个输入问题应该找哪几个专家组来处理。这个调度员在MoE架构里叫做路由器。DOT-MoE的路由器设计很直接给定一段输入文字路由器把每个词token映射到一个概率分布表示这个词应该由哪些专家组来处理的可能性有多大。然后选出概率最高的前k个专家组只激活这k组专家来处理这个词。路由器的训练同样面临离散决策的梯度问题——选前k个这个操作也是不可微的。于是研究团队对路由器也应用了直通估计器前向计算用真正的离散选择结果某个词就是被分配给了这k个专家反向传播时假装用的是softmax输出的连续概率值来传递梯度。最关键的创新点在于神经元分组的学习和路由器的学习是同时进行的相互影响共同优化。之前的方法通常是先固定分组再训练路由器而DOT-MoE允许两者在训练过程中互相适应。分组方式影响路由器应该把什么样的输入路由给哪个专家路由器的偏好反过来又会影响什么样的分组更合理。这种协同优化让最终结果更加自洽。五、训练过程在不改变大模型本身的情况下完成改造DOT-MoE的训练分为两个阶段对齐阶段和微调阶段。对齐阶段是整个方法的核心。在这个阶段大模型原有的所有权重全部冻结一个参数都不动只训练两组新引入的参数神经元与专家组之间的亲和度矩阵用于Sinkhorn计算的原始分数以及路由器的权重。训练用了一个叫做Dolmino-mix的数据集在8块H100 GPU上跑了3500步整个过程不到3小时。训练时的目标函数结合了多个成分。第一个成分是让改造后的MoE模型的输出概率分布尽可能接近原始稠密大模型的输出——这是知识蒸馏的思路让稠密模型作为老师MoE模型作为学生。第二个成分是标准的语言建模损失确保模型依然会说话。第三个成分是路由器Z-loss惩罚过大的路由分数防止数值不稳定。第四个成分是负载均衡损失鼓励所有专家组被大致均匀地使用防止出现少数专家被过度使用、其余专家闲置的专家崩塌现象。对齐阶段结束后提取最终的神经元分配矩阵把原来的稠密FFN权重真正拆分成E个独立的专家FFN模块得到一个标准的MoE架构模型。接下来可以选择进行微调阶段用更多数据进一步恢复和提升模型性能。研究团队统一用12亿个token进行了微调以便与其他基线方法做公平比较。在训练过程中模拟稀疏计算的方式也很巧妙并不真的把权重拆开而是通过一个掩码操作把不属于当前激活专家的神经元的中间激活值直接置零再进行后续的矩阵乘法。这样既模拟了稀疏MoE的计算效果又能在原有的稠密权重上完成梯度计算。六、实验结果数字背后的真实差距研究团队在三个主流开源大模型上验证了DOT-MoE的效果LLaMA-2-7B、LLaMA-3-8B和Qwen2.5-7B。评测覆盖了六个常识推理基准ARC-Challenge科学题、Winogrande常识填空、HellaSwag情境理解、PIQA物理常识、SciQ科学问答和BoolQ是非判断。在最严格的零样本评测中不给任何参考示例直接答题DOT-MoE的表现相当突出。以LLaMA-2-7B为例把模型压缩到原来一半参数量后DOT-MoE的六项平均准确率达到61.5%而同样设置下CMoE只有44.5%DISP-LLM一种结构化剪枝方法达到57.4%。在LLaMA-3-8B上DOT-MoE零样本均值59.8%CMoE只有41.8%。在Qwen2.5-7B上DOT-MoE在较大参数预算下达到72.3%远超CMoE的55.5%和DISP-LLM的66.7%。与结构化剪枝方法的比较同样引人关注。结构化剪枝是另一类压缩大模型的主流方法它的做法是直接永久性地删除一部分参数从而减小模型体积。在LLaMA-2-7B上DOT-MoE在WikiText-2语料上的困惑度衡量语言模型迷惑程度的指标越低越好达到7.99而当前最优的结构化剪枝方法DISP-LLM是9.84半结构化剪枝方法SparseGPT是10.17Wanda是11.02。数字越小意味着模型对语言的理解越好DOT-MoE在这里处于领先位置。加入12亿token的微调后差距进一步扩大。在LLaMA-3-8B上用相同的12亿token微调后DOT-MoE均值67.8%CMoE是64.7%。当微调数据量扩大到70亿token时DOT-MoE进一步提升到71.0%超过了同样用70亿token微调的LLaMA-MoE-v266.8%。在Qwen2.5-7B上较大参数量版本的DOT-MoE微调后达到73.4%而原始稠密模型是80.6%差距被大幅缩小。七、消融实验拆解每个设计决策的贡献研究团队做了一系列精心设计的消融实验验证每个设计选择背后的道理。第一组实验研究专家粒度的影响。总专家数量分别设为16、37、74、148、256同时保持每次激活的参数比例不变始终用25%的FFN参数。结果显示专家数量从16增加到148时模型性能稳步提升但超过148之后收益趋于饱和。这与OpenAI和AllenAI等机构在训练MoE模型时观察到的规律一致专家粒度越细路由决策就越精准但过于碎片化的专家会带来边际效益递减。对比方法CMoE在专家数量增加到37时性能直接崩溃WikiText困惑度超过5000而DOT-MoE在这个规模下依然运行良好说明其路由机制对高专家数量有更强的适应性。第二组实验研究推理时的灵活性。研究团队训练了两个版本的Qwen2.5-7B模型分别在50%和75%的FFN稀疏度下进行对齐训练然后分别在30%、50%、75%、90%四个不同的推理时稀疏度下进行评测。结果很有规律在75%稀疏度下训练的模型在所有推理稀疏度设置下都比在50%下训练的模型表现更好。在极端的90%稀疏度推理时只激活10%的FFN参数75%稀疏度训练版本的平均准确率是39.12%而50%稀疏度训练版本只有36.38%。研究团队解释这种现象的原因是当训练时可用的激活神经元更少每个专家被迫学会更高效地编码信息最终形成更紧凑、更具区分度的专家表示。这对实际部署意味着一个实用的好消息在高稀疏度下训练一次就能灵活调整推理时的计算预算无需重新训练。第三组实验研究专家粒度对推理速度的影响。使用vLLM这个高性能推理框架在4种专家数量8、16、74、148和3种批次大小1024、2048、4096下测量吞吐量每秒处理的token数。结果显示专家数量从8增加到148吞吐量几乎没有变化。原因在于vLLM的融合MoE内核实现所有专家的权重被拼接存储通过token重排序整个计算等效于一次大矩阵乘法专家数量的变化不影响GEMM运算的规模因此不影响速度。这意味着使用更细粒度的专家分组不会带来性能代价。第四组实验研究不同初始化策略对后续微调的影响。在相同训练数据和相同超参数下分别用DOT-MoE、CMoE和LLaMA-MoE-v2初始化LLaMA-3-8B然后追踪训练过程中的训练损失、WikiText困惑度和HellaSwag准确率。DOT-MoE从训练起点就有明显优势初始训练损失更低WikiText困惑度更低。随着训练进行CMoE和LLaMA-MoE-v2出现了一个令人担忧的现象训练损失继续下降但WikiText困惑度和HellaSwag准确率不再进步甚至略有下滑这是过拟合的典型信号。DOT-MoE则在两个维度上都持续改善。这说明更好的初始专家分组不仅提供了更好的起点还带来了更好的泛化能力。八、延伸应用同样的思路也能用在注意力机制上除了前馈网络层研究团队还把同样的最优运输框架扩展到了Transformer的另一个核心组件——多头注意力机制。注意力机制的工作原理是把输入的每个词同时与多个注意力头进行交互每个头关注输入序列的不同方面。类似地可以把这些注意力头分成若干个专家组每次只激活其中几组来处理当前输入。分组的优化逻辑与FFN完全相同引入亲和度矩阵用Sinkhorn迭代找到平衡分配用直通估计器处理离散选择。在Qwen2.5-7B上注意力层的50%稀疏度实验中OT分组方式的平均准确率是64.1%而随机分组加训练路由器的基线只有46.2%差距达到17.9个百分点。值得注意的是注意力参数大约只占模型总参数的三分之一所以仅对注意力层做MoE化的整体压缩效果有限将注意力MoE和FFN MoE结合起来是一个自然的扩展方向。九、规模扩展性验证为了证明DOT-MoE不只在7-8B参数规模的模型上有效研究团队还在Qwen2.5-32B这个更大的模型上进行了测试同样目标是激活25%的参数。结果显示DOT-MoE的零样本平均准确率达到73.1%而CMoE只有38.8%差距扩大到34.3个百分点。这说明随着模型规模增加OT分组方式的优势不仅没有缩小反而有所扩大。在序列长度的鲁棒性方面研究团队测试了从2048到32768个token的不同上下文窗口长度。DOT-MoE在所有长度下相比CMoE保持约2个困惑度点的优势说明每个token独立做路由决策的设计天然适合长上下文场景。在训练开销方面研究团队也做了详细的性能分析。Sinkhorn迭代只占总训练时间的约2%。DOT-MoE特有的所有操作包括亲和矩阵计算、硬化分配、直通估计器合计带来约15%的额外开销主要来源是贪婪取整操作目前在CPU上运行存在CPU到GPU的数据传输开销。一旦这个操作改用GPU专用内核实现开销可以大幅降低。更重要的是这些开销只在对齐训练阶段产生推理时的模型已经是标准MoE架构没有任何额外计算。说到底DOT-MoE做的事情可以用一句话概括它找到了一种有数学依据的方式把一个全体员工同时上岗的大公司改造成了一个根据订单智能派遣专属小团队的高效企业而且改造过程不需要重建公司冻结原有权重只需要重新规划部门划分和调度系统。这项研究的意义不只是一个学术上的方法改进。归根结底让AI模型变得更高效意味着它可以部署在算力更有限的设备上意味着更低的服务器成本意味着更低的碳排放。每次你用手机上的AI功能背后都需要远程服务器的大量计算——如果这些计算可以减少一半影响是相当现实的。当然这项研究也有明确的局限性和未来方向。亲和度矩阵目前是随机初始化的如果能利用权重相关性或预计算的激活统计量来初始化可能加速Sinkhorn收敛并得到更好的初始分组。另外对于训练中长期闲置的专家可以考虑直接剪除进一步压缩模型的内存占用。在更大规模如千亿参数量的模型上进行验证也是自然的下一步。对这些问题感兴趣的读者可以通过arXiv:2606.01666找到完整论文深入了解所有技术细节和实验设置。QAQ1DOT-MoE和普通剪枝方法有什么区别A普通剪枝是永久性地删除模型参数就像直接裁员被删的知识就消失了。DOT-MoE不删任何参数而是把所有参数分成若干个专家小组每次只激活几组来处理当前问题全部知识都保留着只是按需调用。两者都能减少推理计算量但DOT-MoE在保留模型能力方面更有优势特别是在长尾知识方面。Q2DOT-MoE改造一个大模型需要多少计算资源A资源需求相对较低。对齐训练阶段冻结了原始模型的全部权重只训练不到2%的新参数亲和度矩阵和路由器权重在8块H100 GPU上只需不到3小时就能完成LLaMA-3-8B的改造。后续微调阶段用了12亿个token规模也远小于从头训练。整体来看比从头训练MoE模型节省了大量资源。Q3改造后的模型在实际推理时速度会更快吗A是的激活参数减少50%意味着每个token的计算量减半实际推理速度有实质性提升。实验表明使用vLLM的融合MoE内核时专家数量的增加不会带来额外的速度损失因为所有专家的计算被合并成高效的矩阵运算。模型部署后是标准的MoE架构与主流推理框架完全兼容。
亚马逊团队“最优快递员“：把一个臃肿的AI大脑变成高效专家小组
这项由亚马逊研究团队完成的工作发表于2026年第43届国际机器学习大会ICML 2026会议地点为韩国首尔收录于PMLR第306卷。论文编号为arXiv:2606.01666有兴趣深入了解的读者可通过该编号查询完整论文。当你打开手机上的AI助手向它提问一个问题时背后那个大脑——也就是大型语言模型——需要在零点几秒内激活数十亿个参数来回答你。这就像一家万人企业每次接一个小订单都要全体员工放下手头工作同时上阵。效率低、耗电多、成本高这是当前AI大模型最真实的困境。亚马逊的研究团队决定对这种工作方式动一次根本性的手术。一、大公司的效率难题为什么AI模型这么费电现代大型语言模型比如驱动各种AI助手的那些系统通常拥有几十亿甚至上百亿个参数。每次你输入一句话模型就需要把所有这些参数都转动一遍来生成回答。这种方式被研究者称为稠密激活——不管你问的是今天天气怎么样还是请帮我推导量子力学公式同样数量的计算资源都会被调用。研究者很早就意识到这种方式过于浪费。以最近发布的Qwen3-30B-A3B模型为例它拥有305亿个参数但每次回答一个问题时实际被用到的只有33亿个。换句话说九成的员工在摸鱼。这种架构被称为混合专家模型Mixture of Experts简称MoE它的核心思想是把大模型拆分成许多小的专家小组每次只让最合适的几组专家来处理当前的问题。然而从头训练一个MoE模型并不是件容易的事。这个过程不稳定、数据需求量极大还需要复杂的负载均衡机制来防止某几个专家被过度使用而其他专家无事可做。于是研究者们开始探索另一条路能不能把一个已经训练好的普通大模型直接改造成MoE模型这个过程被称为MoE化。问题的核心在于改造的时候怎么决定哪些神经元归哪个专家管这个分组决策至关重要分得好改造后的模型依然聪明分得不好模型就会像一个被随意拆散的团队不知道该怎么协作。亚马逊的研究团队在这里找到了一个新颖的切入点他们把这个分组问题类比成了一个经典的物流问题——最优运输问题。二、前人的方法靠感觉分组的隐患在亚马逊团队提出新方法之前已有几种主流的MoE化方案但它们都有一个共同的软肋。最简单粗暴的方式是随机分组就像把一个公司的员工随机打散成若干个小队。LLaMA-MoE就采用了这种方式随机把神经元塞进不同的专家里然后再花费大量算力重新训练来弥补随机分组带来的损失。这种方式耗时耗力而且分组本身没有任何道理可言。更聪明一点的方式是根据神经元的性格相似性来分组。有的方法看神经元处理输入时的权重是否相近权重聚类有的方法看哪些神经元倾向于同时被激活共激活聚类。CMoE这个方法就用了平衡K均值聚类算法把同时活跃的神经元归为一组。这些方法听起来很合理但亚马逊团队指出了一个被所有人忽视的根本问题这些方法都在优化中间过程而不是最终结果。用一个具体的例子来理解这个问题。大模型里的前馈网络层FFN处理信息的方式是两步第一步把输入信号变换成一个中间表示第二步把这个中间表示再变换成输出。现有方法在分组神经元时只看了第一步里神经元的行为却没有考虑第二步——也就是这些神经元最终如何影响输出结果。这就好比给餐厅的厨师分组时只看谁的刀法相似却不看谁做出来的菜最受顾客欢迎。为了验证这个问题有多严重研究团队做了一个精准的单层重建实验把模型的某一层单独拿出来用不同方法分好专家组后直接测量改造后的输出和原始输出之间的误差。结果触目惊心。在LLaMA-2模型上随机分组的误差是亚马逊新方法的36倍而最好的对比方法CMoE的误差也有2倍多。在LLaMA-3模型上这个差距更是高达41倍对比2倍。这个实验清楚地证明看谁和谁长得像完全不够关键要看谁和谁搭档能产出最好的结果。三、核心创意把分组问题变成快递派单问题亚马逊团队的新方法叫做DOT-MoE其中DOT代表可微分最优运输Differentiable Optimal Transport。最优运输问题是数学领域的一个经典概念可以用一个非常直观的场景来理解假设城市里有若干个仓库和若干个零售商每个仓库有一批货物每个零售商需要特定数量的货物。最优运输问题就是要找到一个调货方案使得总的运输成本最低同时每个仓库的货物全部送完每个零售商的需求全部满足。在DOT-MoE里仓库变成了神经元零售商变成了专家组运输变成了神经元归属某个专家组的决定运输成本变成了这种归属方式导致的输出误差。研究团队的目标就是找到一个分配方案让每个神经元都有归属每个专家组的神经元数量恰好相等同时整体的输出误差最小。然而这个问题有一个技术上的挑战每个神经元只能属于一个专家组这是一个非此即彼的离散决策而神经网络的训练需要能够顺滑地计算梯度可以理解为调整方向的指示信号。一旦决策是离散的这个信号就会消失训练无法进行。研究团队用了一个叫做熵正则化的数学技巧来绕过这个障碍。简单说他们在最优化目标里加入了一个允许模糊的成分让分配方案在训练过程中可以是60%属于专家A40%属于专家B这样的模糊概率而不是非此即彼的硬性归属。这个模糊分配通过一种叫做Sinkhorn-Knopp迭代的算法来计算——这个算法的工作方式类似于反复校准交替保证每行的概率之和为1每个神经元只被分配一次和每列的概率之和为固定值每个专家组的容量相等。为了数值稳定整个计算在对数域进行避免了数值过小导致的下溢问题。等到训练完成后再把这个模糊的分配方案硬化成明确的归属按照分配概率从高到低排序依次确认每个神经元的最终归属同时确保每个专家组不超载。这个硬化过程本身不能传递梯度于是研究团队又用了一个叫做直通估计器Straight-Through EstimatorSTE的技巧前向计算时用硬化后的结果反向传播梯度时假装用的是软分配让优化信号能够顺利传回去更新分配逻辑。这整个设计的精妙之处在于它同时解决了三件事分配是可以学习的不靠启发式规则分配是平衡的每个专家组神经元数量相等通过Sinkhorn严格保证优化目标是输出级别的误差不是中间表示的相似度。四、路由器的设计谁来决定每个问题找哪个专家光有专家分组还不够还需要一个调度员来决定每个输入问题应该找哪几个专家组来处理。这个调度员在MoE架构里叫做路由器。DOT-MoE的路由器设计很直接给定一段输入文字路由器把每个词token映射到一个概率分布表示这个词应该由哪些专家组来处理的可能性有多大。然后选出概率最高的前k个专家组只激活这k组专家来处理这个词。路由器的训练同样面临离散决策的梯度问题——选前k个这个操作也是不可微的。于是研究团队对路由器也应用了直通估计器前向计算用真正的离散选择结果某个词就是被分配给了这k个专家反向传播时假装用的是softmax输出的连续概率值来传递梯度。最关键的创新点在于神经元分组的学习和路由器的学习是同时进行的相互影响共同优化。之前的方法通常是先固定分组再训练路由器而DOT-MoE允许两者在训练过程中互相适应。分组方式影响路由器应该把什么样的输入路由给哪个专家路由器的偏好反过来又会影响什么样的分组更合理。这种协同优化让最终结果更加自洽。五、训练过程在不改变大模型本身的情况下完成改造DOT-MoE的训练分为两个阶段对齐阶段和微调阶段。对齐阶段是整个方法的核心。在这个阶段大模型原有的所有权重全部冻结一个参数都不动只训练两组新引入的参数神经元与专家组之间的亲和度矩阵用于Sinkhorn计算的原始分数以及路由器的权重。训练用了一个叫做Dolmino-mix的数据集在8块H100 GPU上跑了3500步整个过程不到3小时。训练时的目标函数结合了多个成分。第一个成分是让改造后的MoE模型的输出概率分布尽可能接近原始稠密大模型的输出——这是知识蒸馏的思路让稠密模型作为老师MoE模型作为学生。第二个成分是标准的语言建模损失确保模型依然会说话。第三个成分是路由器Z-loss惩罚过大的路由分数防止数值不稳定。第四个成分是负载均衡损失鼓励所有专家组被大致均匀地使用防止出现少数专家被过度使用、其余专家闲置的专家崩塌现象。对齐阶段结束后提取最终的神经元分配矩阵把原来的稠密FFN权重真正拆分成E个独立的专家FFN模块得到一个标准的MoE架构模型。接下来可以选择进行微调阶段用更多数据进一步恢复和提升模型性能。研究团队统一用12亿个token进行了微调以便与其他基线方法做公平比较。在训练过程中模拟稀疏计算的方式也很巧妙并不真的把权重拆开而是通过一个掩码操作把不属于当前激活专家的神经元的中间激活值直接置零再进行后续的矩阵乘法。这样既模拟了稀疏MoE的计算效果又能在原有的稠密权重上完成梯度计算。六、实验结果数字背后的真实差距研究团队在三个主流开源大模型上验证了DOT-MoE的效果LLaMA-2-7B、LLaMA-3-8B和Qwen2.5-7B。评测覆盖了六个常识推理基准ARC-Challenge科学题、Winogrande常识填空、HellaSwag情境理解、PIQA物理常识、SciQ科学问答和BoolQ是非判断。在最严格的零样本评测中不给任何参考示例直接答题DOT-MoE的表现相当突出。以LLaMA-2-7B为例把模型压缩到原来一半参数量后DOT-MoE的六项平均准确率达到61.5%而同样设置下CMoE只有44.5%DISP-LLM一种结构化剪枝方法达到57.4%。在LLaMA-3-8B上DOT-MoE零样本均值59.8%CMoE只有41.8%。在Qwen2.5-7B上DOT-MoE在较大参数预算下达到72.3%远超CMoE的55.5%和DISP-LLM的66.7%。与结构化剪枝方法的比较同样引人关注。结构化剪枝是另一类压缩大模型的主流方法它的做法是直接永久性地删除一部分参数从而减小模型体积。在LLaMA-2-7B上DOT-MoE在WikiText-2语料上的困惑度衡量语言模型迷惑程度的指标越低越好达到7.99而当前最优的结构化剪枝方法DISP-LLM是9.84半结构化剪枝方法SparseGPT是10.17Wanda是11.02。数字越小意味着模型对语言的理解越好DOT-MoE在这里处于领先位置。加入12亿token的微调后差距进一步扩大。在LLaMA-3-8B上用相同的12亿token微调后DOT-MoE均值67.8%CMoE是64.7%。当微调数据量扩大到70亿token时DOT-MoE进一步提升到71.0%超过了同样用70亿token微调的LLaMA-MoE-v266.8%。在Qwen2.5-7B上较大参数量版本的DOT-MoE微调后达到73.4%而原始稠密模型是80.6%差距被大幅缩小。七、消融实验拆解每个设计决策的贡献研究团队做了一系列精心设计的消融实验验证每个设计选择背后的道理。第一组实验研究专家粒度的影响。总专家数量分别设为16、37、74、148、256同时保持每次激活的参数比例不变始终用25%的FFN参数。结果显示专家数量从16增加到148时模型性能稳步提升但超过148之后收益趋于饱和。这与OpenAI和AllenAI等机构在训练MoE模型时观察到的规律一致专家粒度越细路由决策就越精准但过于碎片化的专家会带来边际效益递减。对比方法CMoE在专家数量增加到37时性能直接崩溃WikiText困惑度超过5000而DOT-MoE在这个规模下依然运行良好说明其路由机制对高专家数量有更强的适应性。第二组实验研究推理时的灵活性。研究团队训练了两个版本的Qwen2.5-7B模型分别在50%和75%的FFN稀疏度下进行对齐训练然后分别在30%、50%、75%、90%四个不同的推理时稀疏度下进行评测。结果很有规律在75%稀疏度下训练的模型在所有推理稀疏度设置下都比在50%下训练的模型表现更好。在极端的90%稀疏度推理时只激活10%的FFN参数75%稀疏度训练版本的平均准确率是39.12%而50%稀疏度训练版本只有36.38%。研究团队解释这种现象的原因是当训练时可用的激活神经元更少每个专家被迫学会更高效地编码信息最终形成更紧凑、更具区分度的专家表示。这对实际部署意味着一个实用的好消息在高稀疏度下训练一次就能灵活调整推理时的计算预算无需重新训练。第三组实验研究专家粒度对推理速度的影响。使用vLLM这个高性能推理框架在4种专家数量8、16、74、148和3种批次大小1024、2048、4096下测量吞吐量每秒处理的token数。结果显示专家数量从8增加到148吞吐量几乎没有变化。原因在于vLLM的融合MoE内核实现所有专家的权重被拼接存储通过token重排序整个计算等效于一次大矩阵乘法专家数量的变化不影响GEMM运算的规模因此不影响速度。这意味着使用更细粒度的专家分组不会带来性能代价。第四组实验研究不同初始化策略对后续微调的影响。在相同训练数据和相同超参数下分别用DOT-MoE、CMoE和LLaMA-MoE-v2初始化LLaMA-3-8B然后追踪训练过程中的训练损失、WikiText困惑度和HellaSwag准确率。DOT-MoE从训练起点就有明显优势初始训练损失更低WikiText困惑度更低。随着训练进行CMoE和LLaMA-MoE-v2出现了一个令人担忧的现象训练损失继续下降但WikiText困惑度和HellaSwag准确率不再进步甚至略有下滑这是过拟合的典型信号。DOT-MoE则在两个维度上都持续改善。这说明更好的初始专家分组不仅提供了更好的起点还带来了更好的泛化能力。八、延伸应用同样的思路也能用在注意力机制上除了前馈网络层研究团队还把同样的最优运输框架扩展到了Transformer的另一个核心组件——多头注意力机制。注意力机制的工作原理是把输入的每个词同时与多个注意力头进行交互每个头关注输入序列的不同方面。类似地可以把这些注意力头分成若干个专家组每次只激活其中几组来处理当前输入。分组的优化逻辑与FFN完全相同引入亲和度矩阵用Sinkhorn迭代找到平衡分配用直通估计器处理离散选择。在Qwen2.5-7B上注意力层的50%稀疏度实验中OT分组方式的平均准确率是64.1%而随机分组加训练路由器的基线只有46.2%差距达到17.9个百分点。值得注意的是注意力参数大约只占模型总参数的三分之一所以仅对注意力层做MoE化的整体压缩效果有限将注意力MoE和FFN MoE结合起来是一个自然的扩展方向。九、规模扩展性验证为了证明DOT-MoE不只在7-8B参数规模的模型上有效研究团队还在Qwen2.5-32B这个更大的模型上进行了测试同样目标是激活25%的参数。结果显示DOT-MoE的零样本平均准确率达到73.1%而CMoE只有38.8%差距扩大到34.3个百分点。这说明随着模型规模增加OT分组方式的优势不仅没有缩小反而有所扩大。在序列长度的鲁棒性方面研究团队测试了从2048到32768个token的不同上下文窗口长度。DOT-MoE在所有长度下相比CMoE保持约2个困惑度点的优势说明每个token独立做路由决策的设计天然适合长上下文场景。在训练开销方面研究团队也做了详细的性能分析。Sinkhorn迭代只占总训练时间的约2%。DOT-MoE特有的所有操作包括亲和矩阵计算、硬化分配、直通估计器合计带来约15%的额外开销主要来源是贪婪取整操作目前在CPU上运行存在CPU到GPU的数据传输开销。一旦这个操作改用GPU专用内核实现开销可以大幅降低。更重要的是这些开销只在对齐训练阶段产生推理时的模型已经是标准MoE架构没有任何额外计算。说到底DOT-MoE做的事情可以用一句话概括它找到了一种有数学依据的方式把一个全体员工同时上岗的大公司改造成了一个根据订单智能派遣专属小团队的高效企业而且改造过程不需要重建公司冻结原有权重只需要重新规划部门划分和调度系统。这项研究的意义不只是一个学术上的方法改进。归根结底让AI模型变得更高效意味着它可以部署在算力更有限的设备上意味着更低的服务器成本意味着更低的碳排放。每次你用手机上的AI功能背后都需要远程服务器的大量计算——如果这些计算可以减少一半影响是相当现实的。当然这项研究也有明确的局限性和未来方向。亲和度矩阵目前是随机初始化的如果能利用权重相关性或预计算的激活统计量来初始化可能加速Sinkhorn收敛并得到更好的初始分组。另外对于训练中长期闲置的专家可以考虑直接剪除进一步压缩模型的内存占用。在更大规模如千亿参数量的模型上进行验证也是自然的下一步。对这些问题感兴趣的读者可以通过arXiv:2606.01666找到完整论文深入了解所有技术细节和实验设置。QAQ1DOT-MoE和普通剪枝方法有什么区别A普通剪枝是永久性地删除模型参数就像直接裁员被删的知识就消失了。DOT-MoE不删任何参数而是把所有参数分成若干个专家小组每次只激活几组来处理当前问题全部知识都保留着只是按需调用。两者都能减少推理计算量但DOT-MoE在保留模型能力方面更有优势特别是在长尾知识方面。Q2DOT-MoE改造一个大模型需要多少计算资源A资源需求相对较低。对齐训练阶段冻结了原始模型的全部权重只训练不到2%的新参数亲和度矩阵和路由器权重在8块H100 GPU上只需不到3小时就能完成LLaMA-3-8B的改造。后续微调阶段用了12亿个token规模也远小于从头训练。整体来看比从头训练MoE模型节省了大量资源。Q3改造后的模型在实际推理时速度会更快吗A是的激活参数减少50%意味着每个token的计算量减半实际推理速度有实质性提升。实验表明使用vLLM的融合MoE内核时专家数量的增加不会带来额外的速度损失因为所有专家的计算被合并成高效的矩阵运算。模型部署后是标准的MoE架构与主流推理框架完全兼容。
