AIxiv专栏是呆板之心宣布学术、技巧内容的栏目。从前数年，呆板之心AIxiv专栏接受报道了2000多篇内容，笼罩寰球各年夜高校与企业的顶级试验室，无效增进了学术交换与传布。假如你有优良的任务想要分享，欢送投稿或许接洽报道。投稿邮箱：[email protected]；[email protected]本周，在阿里云通义千问 Qwen 团队提交的一篇论文中，研讨职员发明了现在最热点的 MoE（混杂专家模子）练习中存在的一个广泛要害成绩，并提出一种全新的方式——经由过程轻量的通讯将部分平衡抓紧为全局平衡，使得 MoE 模子的机能跟专家特异性都失掉了明显的晋升。论文：《Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models》论文链接：https://arxiv.org/abs/2501.11873MoE 模子练习中的要害成绩混杂专家模子（MoEs）经由过程路由机制静态并稀少地激活模子参数，使得能高效地增年夜模子参数范围。基于 TopK 机制的稀少激活会在练习中会碰到专家激活不平衡的成绩：多数被频仍抉择的专家会被优化得更多，进一步使得这些专家被更频仍地抉择，终极招致只抉择多数专家，形成残余专家的冗余。因而，MoE 在练习中须要引入额定帮助的负载平衡丧失（load balance loss，LBL）来激励专家的抉择趋于平衡。现在主流 MoE 练习框架中实现的 LBL 的优化目的是部分（micro-batch）的负载平衡，这使得模子须要将一个micro-batch的输入都平均调配给差别的专家。但是，一个micro-batch的输入每每只来自一般范畴，部分负载平衡会让模子将每个范畴的输入都平均调配。这种平均调配会妨碍某些专家更多处置特定范畴的数据，也即妨碍专家呈现范畴档次的分化特点。咱们发明，将部分的负载平衡抓紧到全局的负载平衡，能明显加强专家的特同化并进步模子机能。配景混杂专家（Mixture-of-Experts，MoE）是一种高效的在练习时扩大模子参数范围的技巧。平日，一个MoE层由一个路由器（平日是一个线性层）跟一组专家构成（对Transformer的模子，每个专家是一个前馈神经收集）。给定一个输入，只有局部专家会被激活，而后它们的输出会依据路由器调配的权重停止聚合。详细来说：负载平衡丧失负载平衡丧失是练习 MoE 收集中的一种主要正则化技巧，其中心思维是激励全部专家的平衡激活。它能够经由过程以下公式盘算：此中，?是专家??的激活频率，?是调配给专家??的均匀路由分数。但是，年夜少数现有的MoE练习框架（比方Megatron-core）实现的是部分（micro-batch）档次的平衡，这象征着在每个 micro-batch 内盘算 LBL ，而后在全局（global-batch）档次长进行均匀，即：此中??为 micro-batch 数，?是在第???个 micro-batch 上盘算的负载平衡丧失，?为在第??个 micro-batch 上统计出的激活频率跟路由分数。咱们存眷的要害点是，假如一个 micro-batch 中的数据不敷多样化，这种实现方法可能会妨碍专家的特同化。比方，假设一个 micro-batch 中只包括代码数据，上述负载平衡丧失依然会推进路由器将这些代码输入平均调配给全部专家。而幻想状态下，处置代码数据的专家收集应当对代码数占有更高的激活频率。在练习基于 MoE 的年夜型言语模子时，这种情形更罕见：一个较小的 micro-batch （平日为 1）中的数据平日来自统一范畴。这在必定水平上说明了为什么以后年夜少数基于 MoE 的年夜言语模子中都不察看到显明的范畴档次的专家特同化。这一毛病促使咱们将以后部分平衡的方式想措施扩大到全局（global-batch）平衡。从部分平衡到全局平衡得得益于 LBL 盘算的格局，咱们能够经由过程通讯差别节点的???来将部分???转化为全局的??：1）在全部 micro-batch 之间同步专家抉择频率??；2）在每个GPU上盘算负载平衡丧失；3）在全部 micro-batch 之间聚合丧失。详细来说：此中??是对全局统计的激活频率跟门控分数，第一个等式为???的盘算方法，第二个等式为全局路由分数能够由部分路由分数均匀而来，第三个等式表现用全局激活频率参加部分盘算后再均匀聚合等价于全局平衡丧失。由于???只是一个专家数巨细的向量，即便是在全局通讯的情形下也不会带来显明的开支。别的因为 LBL 的盘算与模子别的局部的盘算绝对自力，还能够用盘算掩饰等战略进一步打消同步??的通讯开支。别的，对须要梯度积聚的情景，咱们还提出了缓存机制来累积各个积聚步统计的专家激活频率，使得盘算节点较少、只停止一次通讯到达的平衡范畴无限的情形下，也能逐步近似全局统计的激活频率。扩展平衡的范畴带来稳固的晋升咱们在三种参数范围（3.4B 激活 0.6B, 15B 激活 2.54B，43B 激活 6.6B）下分辨练习了 120B 跟 400B tokens，对照了差别的平衡范畴（Balance BSZ）对模子机能的影响。全部模子都应用了细粒度专家、共享专家及 dropless 战略（专家不会摈弃超越容量的tokens）。能够看到，将平衡范畴从个别框架实现的 4，8 或许 16 增年夜到 128 以上后模子在 Benchmark 指标跟 PPL 都有显明晋升。咱们在 3.4B 激活 0.6B 的模子练习 400B tokens 到设置长进一步对照了模子后果跟着平衡范畴的变更，能够看到 balance BSZ 从 2 到 128 模子的 PPL 在疾速下降，在 128 后逐步饱跟。现在主流 MoE 框架中即便是停止了机内通讯，对较年夜的模子 balance BSZ 也个别在 8 到 16 的，这进一步表现了咱们通讯方式的意思。剖析试验假设验证前文提到，这篇任务的动身点是在一个 micro-batch 中，数据的起源较为单一的，进而招致 MoE 模子须要将相似起源的数据平均调配到全部expert上，咱们改良了这一点进而失掉了晋升。但是，咱们也能够假设 global batch 是由于应用了更多的 token 来统计 expert 激活频率进而增加了方差，使得负载平衡丧失愈加稳固，进而晋升练习洗哦啊过。位了愈加谨严地对照这两种假设，咱们引入了一种对照的试验设置：Shffuled batch balance， 即咱们从global batch中随机抽取一个子集（这个子集的巨细即是micro batch的巨细）统计专家激活频率，进而盘算负载平衡丧失。Shuffled batch balance 跟 micro-batch balance领有雷同的token数量，跟 global-batch balance领有雷同的token散布。咱们发明，shuffled batch balance 跟 global batch balance 的表示多少乎分歧，都明显好于 micro batch balance。阐明，引入 global-batch 取得晋升的重要起因是在一个愈加通用、多样的 token 聚集上盘算丧失。进而验证了咱们的动身点跟假设。增加大批部分平衡丧失能进步模子效力只应用全局平衡会招致部分平衡状态有所下降，这会必定水平影响 MoE 的盘算效力。咱们进一步试验了在重要应用全局平衡的情形下，在练习进程中增加部分平衡（默许实现的 LBL，丧失权重为全局 LBL 的 1%）限度对模子机能跟效力的影响。能够看到，增加部分平衡能晋升模子的速率（每个更新步耗时从 1.64秒晋升到1.59秒），同时模子的后果也多少乎不受影响。同期相干任务以及探讨已有任务 GRIN 也提出了 Global Load Balance Loss Adaptations，但是更多将这一平衡方式作为练习框架只应用张量并行、不应用专家并行的上风。GRIN 中并不从 specialization 或是对模子 performance 影响等方面探讨应用 Global Load Balance 的念头，也不展现单一应用 Global Load Balance 的影响。Wang et al. 提出在基于MoE的年夜言语模子练习中，负载平衡丧失跟言语模子丧失犹如杠杆一样须要衡量，由于两者的优化目的并纷歧致。因而，他们提出了一种基于专家抉择频率更新的偏向项（bais term），在不转变路由分数的情形下均衡专家抉择，从而去失落了用来帮助练习的负载平衡丧失（auxiliary-loss free）。基于专家抉择频率更新的偏置项，以在不转变路由评分的情形下均衡专家抉择。然而，他们不比拟该方式在专家抉择频率是依据 micro-batch 盘算跟依据 global-batch 盘算时的机能差别。这项任务也被利用到 deepseek-v3 的练习中。deepseek-v3 的技巧讲演（同期任务）中夸大了这项技巧的专家抉择频率是基于 global-batch 停止盘算，并在小范围上探讨了基于global batch 应用 LBL 的成果，也发明这两种方式成果类似。而咱们的任务不只在年夜范围上体系验证了这种方式的无效性，还具体析了平衡范畴对机能的影响，并融化证实了 global-batch 是经由过程归入更多样化的范畴信息从而明显提机能。论断咱们回想了现在 MoE 练习框架中平衡丧失，发明现在的实现方法会将全部来自雷同范畴的部分输入都平均调配，限度了专家的分化。经由过程轻量的通讯将部分平衡抓紧为全局平衡，MoE 模子的机能跟专家特异性都失掉了明显的晋升。咱们以为这一停顿处理了现有MoE练习中的一个要害成绩，为MoE模子的优化供给了新的视角，并有助于构建愈加可说明的模子。只管咱们的试验重要会合在基于言语的义务上，咱们盼望咱们的任务可能为在差别范畴练习更年夜范围、更无效的 MoE 模子供给辅助。