基于自适应专家权重的MoE动态路由负载均衡方案：提升混合专家模型效率与性能的新策略

本综述提出一种基于专家权重的动态路由负载均衡方案，针对混合专家（MoE）模型中的负载失衡问题，摒弃传统辅助损失函数方法，通过直接调整专家权重实现高效路由。该方法在自然语言理解（NLU）任务中显著降低负载标准差（如MNLI从12.25降至1.18），平均准确率提升0.7%，且无需超参数调优，避免梯度扰动，为大规模模型计算资源优化提供新视角。

　　

1 引言

Transformer模型在自然语言处理、计算机视觉和多模态领域取得显著成功，但随着模型规模扩大，计算需求急剧增加。混合专家（Mixture of Experts, MoE）模型通过稀疏结构扩展模型容量而不显著增加计算开销，成为解决密集模型高计算需求的有效方法。MoE基于Transformer架构，用专家层替换原始前馈网络（FFN），由门控网络和多个专家组成，通过稀疏门控实现计算效率提升。然而，MoE系统面临专家负载失衡的固有挑战，由训练数据分布不均和专家初始参数偏好差异引起。负载失衡导致资源浪费、梯度更新不均和优化陷入局部极小值，同时专家容量限制可能引发令牌溢出和数据丢失。当前主流方法通过添加负载均衡损失函数来缓解该问题，但存在超参数敏感、数据分布偏差和延迟反馈等局限性。此外，固定激活专家数量的top-k路由无法适应不同任务复杂度，可能降低模型处理困难任务的置信度。

2 相关工作

2.1 混合专家模型架构

MoE是一种稀疏门控深度学习模型，由多个专家模型和门控网络构成。在基于Transformer的MoE模型中，专家网络由若干相同的前馈网络（FFN）组成，MoE结构通常置于自注意力子层之后，由门控网络选择Transformer块内的前馈网络。由于模型扩展时FFN计算需求增加，MoE结构被用于减少激活参数量，如540B参数PaLM模型中90%参数位于FFN层。Shazeer等人引入稀疏门控策略，通过仅计算top-k专家输出的加权和来降低计算开销。

2.2 MoE模型中的top-k路由

门控网络是MoE模型的核心和实现稀疏化的方法，负责将令牌与专家匹配。top-k路由是最广泛使用的路由算法，但对不同难度任务选择固定数量专家激活。激活专家数量k作为超参数直接影响模型在不同任务上的性能，需要大量实验确定最优值。随着模型规模增大，这一限制导致计算资源浪费。Huang等人提出基于阈值的动态路由方法，可根据令牌复杂度动态调整激活专家数量。该方法使用Switch Transformer的损失函数，但动态路由中更频繁激活的专家更容易出现负载失衡，且损失函数引入额外超参数和梯度扰动。

路由过程的数学表示显示，门控权重线性组合E个专家的输出，其中P表示考虑负载均衡的路由概率，ε为负载损失，专家权重W可与其他网络参数一起训练。MoE层通过将输入X映射到两个张量并通过计算获得最终专家层输出，其核心是路由器将X映射到调度张量D_X和组合张量C_X。在没有正则化约束时，最终路由得分矩阵每行的最大值可能集中在某些索引，导致大多数令牌被路由到少数固定专家，造成专家利用不足和令牌丢失。当前工作通过添加辅助损失L_Balance来缓解，该损失函数量化每个训练批次内专家间激活频率差异并使用方差指标，随后对门控网络输出施加惩罚项。

2.3 动态路由

传统top-k路由限制激活专家数量，在某些情况下导致不必要浪费。实验表明MoE模型性能随top-k中k值不同而显著差异，需要大量计算资源验证不同下游任务的最优k值。此外，top-k门控方法假设每个令牌必须激活相同数量专家，不符合实际任务需求。固定专家数量在处理某些困难任务时可能产生较低置信度，影响输出结果。Guo等人提出DYNMOE算法，在训练和测试期间自动确定激活专家数量，通过将门控机制建模为多标签分类问题，独立计算每个专家的门控分数，激活所有超过阈值的专家。但该方法不加限制最大激活专家数量，可能导致令牌激活所有专家或仅少数特定专家，引发严重负载失衡或专家相似度高问题，同时添加删除专家带来额外计算成本和内存需求。

3 方法设计

针对top-k路由的局限性，采用基于阈值的动态路由策略，通过判断当前激活专家数量是否达到阈值来决定是否激活更多专家以提高令牌处理可靠性。该方法首先对路由概率值P排序获得排序索引I，找到累积超过阈值p的最小专家集合S，其中p为超参数。传统负载均衡通过添加辅助损失函数实现，但损失函数引入需要调优的超参数，不当选择可能导致模型过度关注负载均衡而削弱主任务性能。

本文从MoE门控网络路由算法原理出发，通过专家权重调整缓解负载失衡问题。具体地，在每批数据路由后，根据当前训练轮次的负载情况设置惩罚项，减少高负载专家的权重，增加低负载专家的权重。该方法避免引入额外辅助损失函数，消除超参数对模型处理不同下游任务的影响。模型通过路由概率激活专家，其中专家权重在路由过程中起关键作用，不同权重决定动态路由中激活的专家子集。动态路由设计思想是根据任务难度激活相应数量专家，动态专家激活提高模型效率但也带来新问题：在模型底层和训练早期，动态路由会激活更多专家以确保输出可靠性，更多专家意味着更高激活率和更容易失衡。

为更好实现负载均衡，本文直接根据负载情况调整专家权重。该设计的优势在于专家权重在路由中起决定性作用，每个令牌需要与专家权重计算以构建路由矩阵。每批数据路由后统计所有专家的负载情况，通过算法更新权重。在语言任务中，每批数据通常来自相同句子或段落且存在前后关联，过于频繁调整会丢失上下文联系，违反专家模型自然分工机制并额外消耗计算资源。

对于输入集X，首先执行动态路由并为动态路由设置较大阈值。在模型训练初期，随机分配的令牌被分配给专家，专家权重会随着模型共同训练逐渐显示专业化，因此需要较大阈值让更多专家参与训练。这有助于模型获取更多浅层表征以供后续深度专家激活。获得每个令牌的路由概率后按降序排序，将累积概率超过阈值的专家子集作为本轮路由专家。随后统计分配给每个专家的令牌数量并确定平均值，计算负载惩罚和令牌级别每个专家的负载权重，然后使用公式更新专家权重。

算法描述该方法的流程：通过优化动态路由中的负载均衡方案，实现基于输入的动态专家选择同时平衡专家间任务负载，利用模型优势提高效率。由于ε表示公式中的负载损失量，Fedus等人采用的负载均衡方案使ε = L_Balance，本文方案不使用损失函数，因此在公式中使ε = 0。在当前批数据路由结束时获得所有激活专家的集合，e_i = ā_i - a_i表示第i个专家处理令牌数量与平衡状态相比的多少（a_i为每个专家分配的令牌数量，ā_i为每个专家分配令牌的平均数量），表示令牌级别第i个专家的负载权重偏差。通过将令牌级别偏差映射到专家权重级别，并通过公式更新现有专家权重来实现下一轮路由的负载均衡。

目标是通过动态路由在下一轮路由计算中让专家包含上一轮的负载约束，在下一轮专家激活概率P = softmax(w_ix)中，经公式更新后专家权重包含上一轮负载条件。权重更新直接影响专家对输入令牌的激活概率，其中w_i/∑_i=1^Ew_i表示当前专家权重在总权重中的比例。当t_i高于负载平均值时取负值，减少当前专家权重并抑制下一轮激活，反之促进专家间负载均衡。

算法设计中，专家权重在每批路由后更新，因为在语言任务中令牌存在一定相关性。实时响应负载平衡会消耗显著计算资源并影响语言模型的任务准确性。例如，句子中的代词通常需要更多专家协同处理，参与专家也应包括处理被指代代词的专家。这种负载失衡是由于特殊任务激活更多专家，从而增加模型处理某些困难任务的置信度。

4 实验设计

基于上述设计，采用Mixtral-8×7B作为模型核心架构，这是基于Transformer的混合专家模型，每层包含8个前馈网络（即专家），每个专家参数大小为7B，隐藏层数和注意力头数均为32。实验在自然语言理解数据集上进行，主要关注模型的负载均衡、专家平均激活次数和每个子任务的准确性。在多个基准测试中，其性能达到或超过Llama 2-70B，尤其在数学和多语言理解任务中展现出卓越能力。

4.1 数据集

自然语言理解（NLU）是自然语言处理（NLP）的子领域，专注于使计算机理解人类语言的含义。本文在公开数据集上验证模型有效性，使用来自纽约大学、华盛顿大学等机构创建的多任务自然语言理解数据集GLUE（通用语言理解评估）。GLUE包含九个自然语言理解任务，均使用英语，用于评估模型在各种现有NLU任务中的性能。GLUE的九个任务涉及自然语言推理、文本蕴含、情感分析、语义相似度等多个任务，可分为单句分类任务、相似性任务和推理任务三类。

单句分类任务包括：情感分类（SST-2）、判断是否合乎语法（CoLA），均可抽象为判断句子是否合乎语法的二元任务；相似性任务包括：判断两个句子是否表达相同含义（MRPC、QQP）、判断两个句子的相互关联性（STS-B）；推理任务包括：判断句子是否语义矛盾和隐含（MNLI、RTE）、句子中是否存在某个问题的答案（QNLI）、代词指代哪个对象（WNLI）。

4.2 模型结构和参数设置

Mixtral-8×7B基于Transformer架构，处理上下文长度最高达32k令牌，前馈块被MoE层替换。Transformer层数设置为32，前馈网络（FFN）的嵌入大小为4096，查询（Q）向量数为32，键（K）和值（V）向量数为8。每层包含8个专家，每个专家有7B参数。每次激活top-2专家，FFN层使用SwiGLU激活函数。阈值p为超参数，初始值为0.4。所有其他可学习参数在预训练前随机初始化，标准差为0.006。在GLUE数据集中，对于训练数据量较大的任务（MNLI、SST-2、QQP、QNLI），epoch数设置为6；对于数据有限的任务（CoLA、STS-B、MRPC、RTE、WNLI），epoch数增加至10。预热步数值设置为16。此外，由于WNLI训练数据仅约600个样本，该任务的批大小单独设置为16。

4.3 基线

在实验设计中，设置一个密集模型和两个使用top-1和top-2路由的混合专家（MoE）模型。密集模型遵循标准Transformer架构，每个Transformer层由多头注意力层和标准前馈网络组成。基线MoE模型通过加入辅助损失函数实现负载均衡，具体地，top-1路由采用Switch Transformer模型，top-2路由采用Gshard模型。为公平比较，所有三种基线方法配置与本文提出方法相同，并使用学习率{1e-3, 2e-3, 5e-3}进行训练，仅使用最佳实验结果进行比较。对于基线MoE模型，使用Fedus等人提出的损失函数，涉及超参数α=0.001。

4.4 评估指标

对于GLUE数据集中的三类下游任务，使用准确率（ACC）作为评估指标。准确率计算方法如公式所示。负载均衡使用每个专家接收令牌比例P_i的标准差σ作为评估指标，标准差越小表示负载越均衡。负载均衡的评估指标如公式所示。

5 结果

实验基于GLUE数据集提供的训练和测试数据评估模型有效性，使用每个任务的测试准确率作为指标。表1展示了四个模型（每个模型有8个专家）在GLUE数据集三种不同类型下游任务上使用不同路由策略和辅助函数的测试准确率。

实验结果表明，在绝大多数任务中，MoE模型优于传统密集模型。这是因为MoE模型在相同参数设置下可增加模型容量，且每步激活的参数数量显著少于密集模型，从而降低推理成本。与其他模型相比，本文提出的方法平均准确率比密集模型高6.1%，比当前广泛使用的top-2路由高0.7%，表明所提出模型适用于一般自然语言理解任务并能取得良好性能。然而，在RTE和MRPC数据集上，准确率略低于top-2路由，表明所提出模型在数据有限场景下无法充分发挥动态路由优势。对于数据充足的数据集（如QQP和MNLI），模型可提高准确率超过1%。MRPC和QQP都存在正负样本分布不平衡问题，在QQP中负样本占63%，正样本占37%。通过使用本文提出的负载均衡方法，模型可避免在正负样本分布不平衡场景下因仅选择部分专家而导致的令牌丢失，因此在分类任务中能取得比top-2更好的结果，证明该方法对大数据集中不平衡样本分布的鲁棒性。

模型在小数据集上表现较差，主要因为RTE和MRPC的训练集分别仅包含2500和3700个样本，其中负样本占比高达65%。为进一步研究小数据集对模型性能的影响，进行以下对比实验：首先逐步减少训练集大小，使用训练集数据的100%、75%和50%训练模型。在batch_size=16设置下，计算第一层、第八层和第十六层专家负载的平均方差，以ACC为评估指标。实验结果如表2所示。

随着训练数据减少，模型准确率也逐渐恶化，最高下降3.27%。从路由角度看，MoE的核心特征是不同专家专注于不同输入子空间，但在小样本数据集中，路由器在早期阶段将有限样本随机分配给多个专家，导致每个专家的有效样本量较小。由于样本量不足，专家无法训练稳定模式，导致后续阶段令牌无法准确路由到相关专家。随后调整训练期间的批大小，设置为8、16和32，并包含数据充足的大数据集作为对比（QNLI）。实验结果如表3所示。在epoch不变条件下，小数据集在batch_size=16时表现最佳。这是因为较小的批大小增加了批次内所有样本属于同一类别的概率，导致某些专家的激活频率增加。动态路由随后减少这些专家的权重，影响后续路由轮次的结果。较大的batch_size设置减少了模型迭代次数，导致泛化性能差。通过上述对比实验验证了模型在小数据集上表现不佳的主要原因是数据量有限和数据分布不均，但通过调整批大小可以缓解模型准确率差的问题，从而增强模型对有限数据的适应性。

Jiang等人发现，在MoE模型的较高层级，连续分配现象显著高于随机分配。这意味着在没有负载均衡控制的情况下，模型更可能仅激活少数专家，严重影响模型的训练和推理，并导致空闲专家资源浪费。通过可视化Mixtral-8×7B模型中同一句子在第一MoE层和最后MoE层的专家分配，验证了本文提出的方法在平衡高层专家负载方面的有效性，如图3所示。

图3a,b分别显示了MoE第一层和最后一层中八个专家对同一句子的分配情况，发现最后一层在令牌分配中更容易出现专家的连续分配，统计每个专家的激活频率，与第一层相比存在严重负载失衡，其中专家4的激活率从1.45%增长到34.78%，而专家2和专家4仅占5.8%。图3c表示采用本文负载均衡方案后最后一层专家的负载，专家负载的标准差值为1.76，而图3b中为9.33。此外，分析其他层（从第1层到第7层）的专家负载方差，发现除前两层外，本文提出的方法在所有层上都优于Fedus等人提出的损失函数方法。这是因为两种方法在第一层都使用随机分配，导致性能相似并实现均匀分配。因此，本文设计的方法对第一层专家权重调整影响较小。但随着MoE层数增加，损失函数无法有效处理连续分配现象，导致负载失衡随着层数增加而恶化。

为验证阈值对模型中动态路由的影响，根据数据量大小从GLUE中选择QQP、SST-2和RTE数据集。它们的训练集大小分别约为360k、67k和3k，阈值大小分别设置为0.2、0.4和0.6。通过统计分析不同阈值下的专家激活数量和实验结果，验证阈值的合理性。图4显示了三个数据集在不同阈值下使用动态路由的MoE层中专家激活数量，虚线表示top-2路由的激活数量。随着数据集大小减小，同一数据集内不同阈值下的专家激活数量也减少，且减少速率变小。这是因为随着数据集大小减小，快速形成高度专业化专家变得越来越困难。动态路由策略必须在满足激活阈值要求的同时平衡负载均衡。因此在较大数据集中，提出的动态路由策略在减少专家激活数量方面非常有效。在模型训练早期或浅层，专家尚未达到专业化水平，路由概率分布通常相对均匀。当阈值设置为p=0.2时，阈值容易满足，即使所有专家具有相同激活概率，也仅需要两个专家就能达到阈值。因此随着训练进行或在模型更深层，低阈值导致路由退化为传统top-1路由，且这种现象随着数据集减小而更加明显。随着阈值增加，动态路由根据路由概率激活相应数量专家以促进专业化。在上述三个数据集中，更深层的激活专家数量显著少于浅层。这是因为训练后，具有较高专业化水平的专家无需激活额外专家进行计算就能满足阈值约束。表4显示了不同阈值下动态路由在不同大小数据集上的实验结果。在三个数据集中，p=0.4的设置产生最佳结果。因此在该任务中也验证了Huang等人将p值设置为0.4的有效性。

通常，模型在下游任务上的性能取决于参数数量。例如，实验中top-2路由表现出显著优于top-1路由的性能。Zhou等人证明模型的更深表示容易出现过拟合现象，因此良好的浅层表示比复杂的深层表示更有价值。这两种现象表明参数数量和浅层表示都能有效提高模型性能。本文提出的方案通过多个专家在浅层实现全面特征提取，并通过影响专家权重和动态路由在深层激活较少专家，从而在减少模型参数数量的同时实现专家间负载均衡。

WNLI是一个二元分类任务数据集，包含634个训练样本和146个测试样本。训练集正负样本分布平衡，而65%的测试样本为负样本。动态激活专家数量一致的主要原因是对于更具挑战性的推理任务，模型会激活更多专家参与。有限的训练数据、正负样本分布不均和任务难度都导致模型在该任务的平均激活专家数量上没有优势。在模型顶层，也激活约两个专家，且由于数据分布和数量不足，该方案在平均激活数量上退化为top-2路由。通过在九个数据集上的实验，本文提出的方案不仅在激活较少专家（即较少参数）的同时实现负载均衡，而且通过激活模型较低层的更多专家获得更好的浅层表示。随着层数增加，每个令牌激活的专家数量逐渐减少，有效防止模型过于复杂并减少不必要的计算。除少数数据集外，提出的方案取得比top-2路由更好的结果。

6 消融实验

消融实验根据数据量大小从GLUE中选择五个数据集，验证本文提出的方法对专家负载均衡的效果。记录训练过程中八个专家的负载情况，证明提出的负载均衡方案在相同数据集上优于使用损失函数的方法。实验结果如图5所示。图5a显示了使用损失函数实现负载均衡时不同数据集上八个专家的激活率，可见每个数据集都存在较严重的负载均衡问题。灰色虚线表示理想的专家激活率。在一些数据量较大的数据集（如QQP和MNLI）上负载失衡更加严重。这是因为负载均衡损失通过对门控网络的选择概率分布进行正则化实现，通常使用softmax或其他归一化函数，其梯度在极值附近可能变得非常小或非常大。这种梯度问题可能导致某些专家的选择概率快速增加，加剧失衡。其次，采用动态路由后，每次激活的专家数量不固定。在训练开始时，模型激活更多专家以获取浅层表示，如图4所示。这减少了旨在改善专家间负载均衡的惩罚项的权重。因此损失函数无法有效改善专家负载均衡，导致某些专家的激活率在整个训练过程中远高于理想值。图5b显示了使用基于专家权重的动态路由方案的动态路由策略，在不同数据集上取得良好的负载均衡效果。通过直接修改专家权重，避免了损失函数带来的梯度影响，且不随训练批次增加而增加负载损失。在较大数据集上也表现良好，验证了提出的方法能有效缓解专家间负载失衡。

对上述五个数据集，计算两种负载损失下专家激活的标准差作为验证模型有效性的标准，如表5所示。所有五个数据集的标准差都降低，其中MNLI数据集的负载失衡改善最显著。相应地，在表1中可见，本文提出的负载均衡算法在MNLI数据集上的实验准确率提高了2.7%，这是所有任务中最大的准确率提升。这也证实了负载失衡影响模型的准确率和训练效果。此外，我们的方法采用动态路由；实验结果表明该方法使模型能够激活更少专家（参数），在保持专家间负载均衡的同时减少模型的计算需求，进一步优化了混合专家（MoE）模型。

在MoE模型中，通过门控网络G(x)实现稀疏性而不显著增加计算成本。虽然负载失衡直接影响模型处理大数据集和复杂任务时的最优性能，但在某些情况下负载失衡可能具有意义。例如，如果某批输入数据大多属于一种类型，混合专家模型最初设计意图是不同专家处理不同种类的输入，而不是将相同种类数据平均分配给不同专家，强制均衡可能违反模型的自然分工机制。当前负载均衡机制未考虑训练数据本身的不平衡，仅通过专家权重设置达到平均负载均衡状态，并要求在模型训练中所有类型训练数据均匀分布以避免数据本身分布带来的负载失衡。如果数据自然分布不平衡，追求绝对均衡可能影响模型效果。未来工作将专注于更现实场景，解决数据本身不平衡时如何避免上述问题。此外，实验表明模型在标签分布不均匀的数据集上测试时性能不理想。需要进一步研究解决小样本大小、样本分布不均和其他相关挑战引起的负载均衡问题。

7 总结

本文提出一种基于专家权重的动态路由算法用于负载均衡，根据先前路由中专家间的负载情况调整专家权重以影响下一次路由分配。该方法优势在于避免损失函数中超参数的训练设置和梯度变化对专家负载的影响。实验验证表明，基于专家权重的动态路由算法不仅提高任务准确性，而且与top-k路由相比减少模型参数数量，同时取得良好的负载均衡效果。尤其在数据充足的数据集中，该路由算法能更好发挥其优势。