负载均衡 is All You Need

来知乎几年了，宣传自己的文章倒是第一次。 这是我第二篇一作。第一篇是关于投机解码(Speculative Decoding) 的，改了改多个draft 之间如何选择的算法，做的实在太niche（小众）。而且实际应用上也没人真的会在batch inference之上再用multi-draft（多草稿）额外开几倍的并行，算法也就没了意义，也就必要跟大家讲。 25年尝试用RL 训练Generative Reward Model（生成式奖励模型）。想法很简单，如果模型生成的评价与人类一致，则（元）奖励为1，否则为0。结果因为太懒，又病了一个月，康复之后发现DeepSeek 发了篇一摸一样的，自己做的质量当然没法比。于是又做了一些补救的尝试，强行注入novelty，然而新算法RL 老是训崩，至今不知是VeRL 的问题还是我的问题，只好作罢。 失败了这么多次，也只好静下心来。我发觉虽然我high level 的想法很多，但是机器学习的基础一直没打好，编程的习惯也没有养成。幸得暑假的时候去公司实习有卡，我没有好好做老板的任务，倒是借这个机会认真补了补基础。回到学校后也一改往日的散漫，一点点学习架构，torch，通信，优化，终于做出了一篇自认为还可以的文章。 想做MoE 很久了。第一次动这个念头是24年看到Mixture of Depth（MoD），谷歌的一篇文章，讲的是可以训练一个路由器把transformer 中的MLP 和Attention 都实现可以跳过，实现了动态的计算量分配。我后来才知道，MoD 将动态算力分配作为其卖点，但其实这完全是expert choice（专家选择）路由的功劳。它唯一不同的是尝试吧Attention 也跳掉，但看起来做的并不对，毕竟跳过记忆的写入（即Key 和Value 的计算）和跳过记忆的读取（即Query 的计算）是两码事，并不一定可以用同一个路由器实现。 先大概讲一下MoE 的架构和路由设计吧。 简单来说，MoE 就是把模型里原来固定的MLP 改成了很多个专家。每次进来一个词（token），路由器就会给这个词和所有专家匹配，然后决定把它交给哪几个专家来处理。这就是路由。 现在市面上的主流方法都是 Token Choice（次元选择路由），即每个词自己挑匹配分最高的那几个专家。但如果所有词都喜欢某几个热门专家，冷门专家就没事干了。为了防止大家扎堆，大家就得强行给模型加各种限制和惩罚，强迫它们做到负载均衡，这就把事情搞得很复杂。 第二种叫 Expert Choice（专家选择路由）。思路反过来了，让每个专家主动去当前的这一大批词里，挑出得分最高的那些词。这样分配下来，每个专家的工作量绝对是一碗水端平的。所以它天然就能完美地做到负载均衡。同时，不同数量的专家可以挑选同一个词，也使得模型可以根据词的特性，动态地分配算力。 有了这两个优点，大家应该可以理解我为何如此的执着于专家选择路由了。 但是，为了挑词，专家必须得等这一大批词全齐了才能做决定。如果是做文本生成（也就是每次只能往外蹦一个词），模型没法提前拿到“未来”的词的路由值来做对比，所以它没法给大语言模型用。同时，模型还可以把路由作为未来信息泄露的通道，把未来的词的路由值泄露给当前的词，以此作弊，导致训练崩溃。这是我不希望看到的。 然而，EC 真的不能做成因果的吗？ 专家混合模型 (Mixture of Experts) 为了下文，我们先定义专家混合模型的各种符号和算法。 在Transformer里，MoE 层通常会把标准的密集前馈网络（FFN）替换为一个路由器（router）和 个专家（experts）。此处，G代表Granularity 颗粒度，即每个token 选择的专家数量。E代表Expansion rate 扩展率，即MLP 的总参数量与激活参数量的比率。假设我们有一个包含 个 token 的 batch，每个 token 的特征表示为 ，那么路由器首先会计算出一个分数矩阵 基于这组分数 ，路由规则会给出一个二值的分配矩阵 ，其中 表示第 个 token 激活了第 个专家，否则为 0。每个被选中的专家会计算出自己的输出 ，并乘上一个门控权重 。所以对于 token ，整个 MoE 层的输出就是所有被激活专家输出的加权和 这里决定矩阵 的路由规则非常关键，它同时控制着算力怎么分配，以及各个专家之间的负载均衡（load balance）。我们可以把 MoE 的路由过程形式化为一个带约束的优化问题。即在计算资源受限的情况下，寻找一个能让总路由分数最大化的分配矩阵 。因为分数越高，意味着 token 和专家越匹配，而且通过门控 ，这个专家对最终输出的贡献也会越大。 Token Choice (基于 Token 的路由) 标准 Token Choice（TC）路由的优化目标可以写成如下形式 在这里，“稀疏性约束”保证了每个 token 必须且只能选择 个专家；而“负载均衡约束”保证了每个专家必须正好处理 个 token。 但问题是，要精确求解这个带有两个硬约束的优化问题，需要用到像匈牙利匹配算法这样复杂度高达 的组合算法，这在实际训练中显然是不可能的。所以绝大多数 Token Choice 的方法（比如 Switch Transformer）都选择妥协。它们严格满足稀疏性约束（直接让每个 token 取 Top- 的专家），然后通过引入辅助损失（Auxiliary Loss）或一些无损的负载均衡策略，来尽力近似满足负载均衡约束。 Expert Choice (基于专家的路由) 虽然负载均衡约束对于防止路由坍塌（所有 token 都涌向极个别专家）至关重要，但“稀疏性约束”其实并没有什么实质好处。于是 Expert Choice (EC) 路由干脆掀桌子，完全去掉了稀疏性约束，只在 batch 内强制执行负载均衡。此时原始的优化问题就简化为了 这个被放宽的优化问题有一个极其简单的闭式解（closed-form solution），每个专家直接在当前的 batch 里挑选分数 Top- 的 token 就行了。 这种设计带来了两个巨大的好处。 完美的负载均衡，每个专家在设计上就强制处理恰好 个 token。 动态算力分配，有的 token 可能没被任何专家看上（跳过计算），有的 token 可能被多个专家同时相中。这就实现了基于 token 重要性自适应分配算力。 然而，EC 路由为了满足每个 batch 内的负载均衡，引入了一个致命的因果性问题（causality problem）。大家注意看，由于是专家挑 token，一个 token 能不能被选中（即 的值），不仅取决于它自己的分数，还取决于当前 batch 内所有其他 token 的分数。这其中自然也包括了在自回归生成时根本还不存在的“未来” token。虽然后续有一些工作（比如把它扩展到 batch 级别的 top-）能在一定程度上缓解这个问题，但只要路由决策还依赖于当前的 batch 构成，因果性问题就无法被彻底根除。 Expert Threshold (专家阈值路由) 正则化（regularization）是机器学习中常用的技术，用于防止塌缩。然而我们到底需要哪些正则化？哪些是无效的约束？ 回到前文。首先，Token Choice 的稀疏性约束是没有意义的。我们并不需要硬性规定每个token 必须使用多少计算量。这既然不有利于负载均衡，又限制了模型的灵活性。 Expert Choice 的负载均衡约束是有意义的。没有它，模型的路由就会坍塌，即所有token 都涌向极个别专家。然而，它要求每个批次都严格做到负载均衡，是否过于苛刻？毕竟，不同批次的数据分布是不同的，强制每个批次都做到负载均衡，反而会导致决策的依据不稳定。 因此，我们提出Expert Threshold (ET) 路由。我们既不要求每个token 激活固定数量的专家，也不要求每个批次都严格做到负载均衡。我们只要求在预期下，每个专家处理的token 比例相同。 换言之，相比于EC，我们不再提取每个批次中每个专家的Top- 次元，而是整个数据分布中每个专家的top 的次元。这使得我们不再需要依赖于当前的批次构成，也就解决了因果性问题。 具体方法，我们对每个批次的cutoff，即第 大的路由得分，计算一个指数移动平均（EMA），得到每个专家的阈值，以此来估计整个数据分布中每个专家的top 的大小。而对于训练和推理，我们只需要使用这个阈值来判断一个token 是否应该被激活。如果大于，说明这个token 属于前，则激活；否则，不激活。 这样，我们既解决了因果性问题，又保证了负载均衡。同时，我们也不需要依赖于当前的批次构成，也就解决了因果性问题。 ET 与 EC 的联系 概念上，ET 可以理解成在一个无限大的 batch 上做 Expert Choice。 EC 是每个专家从当前 batch 的 个 token 里挑 Top- 个。当 很小的时候，随便换一个 token 就可能改变分界线（cutoff）的位置。但当 趋于无穷，分界线收敛到路由分数分布上的一个固定分位点，每个 token 的路由决策就只取决于自己的分数了，跟 batch 里有谁无关。 ET 做的事情就是直接逼近这个极限。我们不等 batch 变无穷大，而是用 EMA 来估计那个全局的分位点阈值，然后拿它来做路由。 换个角度看，ET 和 EC 其实是在”稳定阈值”和”稳定负载”之间做了不同的取舍。EC 让阈值随 batch 变化来保证每个 batch 的负载完美均衡，代价是路由决策会随 batch 构成波动。ET 反过来，固定阈值来保证路由决策的稳定，代价是每个 batch 内的专家利用率会有一点小波动。 有趣的是，正因为这种等价关系，ET 的阈值路由可以直接作为 EC 训练模型的因果推理方案。只要在 EC 训练完之后，记录下每个专家的 cutoff EMA 作为阈值，推理时就可以无缝切换到 ET，不需要任何重训练。 实际上我最初的想法是从这里来的，即想办法把EC 的批次大小做得尽量大。当时我在飞机上阅读DiffMoE这篇文章，他们就强调了批次大小的重要性。很多相关的实验也都表明，批次大小越大，模型的性能越好。于是我干脆干到无限大，这样就解决了因果性问题。

添加图片注释，不超过 140 字（可选） 实验结果 实验是在 Nanochat（Karpathy 写的开源 GPT-2 训练代码）上做的，分了两个规模。小的 d12 模型有 575M 参数（195M 激活），大的 d20 模型有 2.4B 参数（561M 激活）。每个 MoE 层用了 16 个路由专家加 1 个共享专家，扩展率 ，也就是说平均每个 token 只激活 1 个路由专家加上那个共享专家。训练数据来自 FineWeb-Edu 100B，分别训了 10B 和 11.2B 个 token。 我对比了三大类路由方法，所有模型的架构和参数量完全相同，只有路由规则不同。 主要结论

添加图片注释，不超过 140 字（可选） ET 在两个规模上都稳定地优于 Token Choice，d12 上 CE loss 低 0.05，CORE benchmark 高 1.89；d20 上差距更大，CE loss 低 0.067，CORE 高 2.83。 EC 在批次足够大的时候（512k tokens），训练 loss 和 ET 基本持平，这验证了我们前面的理论。显式的大批次选择和 EMA 阈值估计，殊途同归。 在 d20 规模上，ET 以微弱优势超过了最好的 EC（CE loss 2.620 vs 2.621，CORE 25.14 vs 24.98）。

添加图片注释，不超过 140 字（可选） 一个值得注意的细节是，EC 训练时虽然不是因果的，但推理时我们统一使用 ET 的 cutoff EMA 作为阈值来做因果推理。所以上面所有 EC 和 ET 的 CORE 分数都是在因果条件下评测的，不存在信息泄露。 还有一个有意思的现象。EC 在小批次（2k）的时候表现很差，甚至不如没有任何负载均衡的 Token Choice。但随着批次增大，性能稳步提升，到 512k 的时候就和 ET 差不多了。这条曲线恰好印证了前面说的，ET 就是 EC 在无限大批次上的极限。

当使用ET来推理时，sequence EC（批次大小为2k 次元）会出现严重的训推不一致。 更多实验结果请读我的论文。 我这里有很多好康的（可解释性的图） 结论 这篇文章的核心观察其实很简单。Expert Choice 路由之所以不能因果，是因为它在每个批次内做 top- 选择，路由决策依赖于同 batch 里的其他 token。但如果把批次做到无限大，这个依赖就消失了，每个 token 的路由就只取决于它自己和一个全局阈值。ET 就是用 EMA 来逼近这个极限的。 说到底，我们并不需要每个批次都严格地负载均衡。只需要在整个数据分布上做到期望均衡就够了。放松了这一个约束，因果性问题就自然而然地解决了，而且性能不降反升。 希望这个工作能让大家重新审视 Expert Choice 路由。它一直被认为和自回归语言模型不兼容，但其实只差一步。 后记 1月底投完稿，我想着再完善完善，加些可视化可解释性的实验，改一改故事，preprint预印本就一直拖到现在。本以为，22年EC，三年多没人做后续，大概也不会有人抢。没想到上两周，苏剑林老师想出了几乎一样的想法，当然视角不同，算法核心是一样的。对数学好的人来说，同一个算法能讲出五六种不同的故事，但其实都是等价的。再次引一下，供各位参考比对。 附录 专家选择模型到底会泄露多少信息呢？ DeepSeek之前给了个trivial的结论，即它最多泄露N 选k 中路由组合的信息。 我这里给出一种构造性的通信方法，说明，假如无限精度，此上限是可以实现的。换言之哪怕批次大小干到无限也会有泄漏 幸好，现实世界都是有限精度，模型也没那么聪明。