论文阅读：From Sparse to Soft Mixtures of Experts

From Sparse to Soft Mixtures of Experts
从稀疏到软专家混合
摘要
MOE的问题：训练不稳定、令牌丢失、无法缩放专家数量或无效微调。
提出：soft MOE，一个完全可微的稀疏Transformer，解决这些挑战，同时保持MoE的好处。soft MOE通过将所有输入令牌的不同加权组合传递给每个专家来执行隐式软分配。与其他MOE作品一样，Soft MoE中的专家仅处理（组合）令牌的子集，从而以较低的推理成本实现更大的模型容量。在视觉识别方面，Soft MOE大大优于标准Transformer（ViTs）和流行的MoE变体（Tokens Choice和Experts Choice）。降低推理成本的同时性能匹配。以较少的推理时间成本换取了较大的参数增长
1 intro
较大的transformer以增加的计算成本提高性能。模型大小和训练数据必须一起缩放，以最佳地使用任何给定的训练计算预算（大模型需要大训练样本支撑）。稀疏MOE可以扩张参数，但并非所有参数都会在一次样本中全部用到。稀疏MOE已经应用在语言，视觉，多模态方面。
稀疏MOE transformer 的核心在于离散优化问题：决定哪些模块应该应用于每个输入令牌。这些模块通常被称为专家，通常是MLP。现存方法：线性程序、强化学习、确定性固定规则（Roller等人，2021）、最佳运输（Liu等人，2022）、每个令牌的贪婪前k个专家（Shazeer等人，2017），或每个专家的贪婪前k个令牌（Zhou等人，2022年）（目前最常见的门控网络属于每个专家的贪婪前k个令牌）在许多情况下，需要启发式辅助损失来平衡专家的利用率并最小化未分配的令牌。小的推理批量、新的输入或迁移学习可能会加剧分布不平衡的问题。
软MOES不是使用稀疏和离散的路由来试图在令牌和专家之间找到一个好的硬分配，而是通过混合令牌来执行软分配。 特别是，我们计算所有令牌的几个加权平均值–权重取决于令牌和专家–然后由对应的专家处理每个加权平均值。
软MOE模型避免了上面提到的大部分由稀疏MOE核心的离散过程引起的挑战。 流行的稀疏MOE算法学习一些路由器参数，梯度的来源通常是两个：专家输出与选定的路由分数的右乘，以及可强制执行某些所需的行为的辅助损失，也取决于路由得分。（路由很重要但路由不可靠）
据观察，这些机制通常并不比随机固定路由机制更好。软 MoE 则避免了这一问题，因为每个路由（或混合）参数都是根据每个输入标记直接更新的。软路由可以在训练路由时提供稳定性；观察到，在训练过程中，大量输入标记会同时改变网络中的离散路由（batch太大的弊端？？），从而导致训练难题。此外，硬路由对many专家来说是一个挑战，大多数研究都是在几十个专家的情况下进行训练的。我们的研究表明，Soft MoE 可以扩展到数千名专家，而且在结构上是平衡的。最后，在推理过程中不会出现批量效应，即一个输入会影响路由（由于专家能力有限），从而影响对其他 输入的预测。

结果：与VIT比较（为什么不比Token MOE）,参数多但训练快，效果好，上有指标优秀，扩展到其他任务效果好，跨模态有优势。
2 soft MOE
m×d 表示一个序列的输入令牌，其中 m 是令牌的数量，d 是它们的维度。
每层有n个export,每个export有p个插槽（即容量），moe层的参数为Φd×(n·p)
图2：软 MoE 路由算法。Soft MoE 首先根据每个插槽的一些可学习参数，为每对输入令牌和插槽计算分数或对数。然后对每个插槽（列）的对数进行归一化，每个插槽根据这些权重（绿色）计算所有输入标记的线性组合。然后，每个专家（本作品中为 MLP）处理其插槽（例如，本图中每个专家处理 2 个插槽）。最后，对每个标记（即按行）进行相同的原始对数归一化处理，并用于合并每个输入标记的所有插槽输出（蓝色）。虚线框代表可学习参数。
相当于加了线性组合将每个原始样本组合后构成了新的插槽，避免对原始样本的处理，保证不会有原始样本被完全舍弃。
（行列是怎么定义的，行是插槽，列是原始样本batch）
看算法1：
D(调度参数)用于m维度
C（组合参数）用于n*p维度
最后按列组合得到输出
模型结构上，用软 MoE 块替换transformer的 MLP 块子集。特别是，我们通常会替换 MLP 块的后半部分。插槽总数是软 MoE 层的一个关键超参数，因为时间复杂度取决于插槽数量而非专家数量。
2.2 软 MoE 的特性以及与稀疏 MoE 的联系
完全可微
所有稀疏 MoE 算法的核心都是令牌和专家之间的分配问题，通常会受到一些特定的容量和平衡约束。
硬路由解法：Top-k 或 "令牌选择 "路由器会为每个令牌选择分最高的专家if该专家还有空位（即该专家尚未填满其容量）。专家选择 "路由器为每个专家选择top-capacity-scored最高的token**（会出现一个token被分配多次的情况）。其他研究还提出了更先进（通常也更昂贵）的算法来计算分配，例如基于线性规划算法、最优传输或强化学习。
不丢弃令牌和专家不平衡
前述方法都存在丢弃令牌和专家不平衡（一些专家获得的令牌远远多于其他专家）的问题，影响性能。softmoe不受影响，因为每个插槽都是由所有token的加权平均值填充的，且由于softmax，所有权重都是正值。
快
插槽总数是决定软 MoE 层成本的主要超参数。每个输入都会应用相同数量的 MLP。专家总数在计算中并不重要（时间层面）：如果插槽总数相同，则多专家&&每个专家插槽少和少专家&&每个专家插槽多的成本是相匹配的。我们必须满足的唯一限制条件是，插槽数必须大于或等于专家数（因为每位专家必须至少处理一个插槽）。Soft MoE 的主要优势在于完全避免了排序或 top-k** 操作，排序操作速度慢，且不适合硬件加速器。所以送分题 MOE快于大多数稀疏MOE。

稀疏模式与密集模式的特性
稀疏模式的稀疏性来自于专家参数只应用于输入词组的一个子集。然而，从技术上讲，软 MoE 并不稀疏，因为每个槽都是所有输入标记的加权平均值。每个输入标记都会部分激活所有模型参数。同样，所有输出标记也部分依赖于所有插槽（和专家）。最后，还要注意的是，"软模型 "并不是 “密集模型”，在 "密集模型 "中，每个专家都会处理所有输入标记，因为每个专家都只处理插槽的一个子集。

???每序列决定论
在容量受限的情况下，所有 Sparse MoE 方法都会将token按固定大小分组，并强制（或鼓励）实现组内平衡。
当分组中包含来自不同序列或输入的令牌时，这些令牌往往会相互竞争专家缓冲区中的可用位置。因此，模型在序列层面不再具有确定性，而只能在批次层面具有确定性，因为某些输入序列可能会影响对其他输入的最终预测。使用较大分组的模型往往为路由算法提供了更大的自由度，通常性能更好，但计算成本也更高。另一方面，当分组包含来自单个序列的标记时，模型不得不在每个输入序列上使用每个专家。这可能会导致更多的通用专家。此外，在训练和推理之间改变组的大小也会产生问题，因为标记到专家的分配可能会发生分布上的变化。
Soft MoE 巧妙地回避了所有这些挑战。由于它结合了每个输入序列中的所有标记，因此我们只需将组的大小设置为单个序列。每个专家都会处理来自每个输入的标记，这或许在一定程度上限制了高级专业化的数量。然而，这也意味着它对每个实例都具有确定性和快速性。
（不能通过原始token的路由做可解释性工作了）
2.3 实现
时间复杂度
归一化：MOE通常用于替换encoder中的前馈，所以输入是层归一化的。
分布式模型：当专家数量大幅增加时，单个设备的内存不可能容纳整个模型，尤其是在训练过程中或在大型模型骨干上使用 MoE 时。在这种情况下，我们采用标准技术将模型分布到许多设备上。分布模型通常会增加模型的成本开销，而我们上面基于 FLOPs 得出的时间复杂度分析并没有考虑到这一点。为了考虑这一差异，我们在所有实验中不仅测量了 FLOPs，还测量了以 TPUv3 芯片小时为单位的壁钟时间。
2.4 与其他方法的联系
现有的许多研究都通过合并、混合或融合输入标记来减少输入序列长度，通常使用具有固定密钥的类似注意力的加权平均法，试图减轻self-atti相对于序列长度的二次成本。虽然我们的调度和组合权重的计算方式与这些方法类似，但我们的目标并不是减少序列长度（虽然这是有可能的），实际上，我们是在每个软 MoE 层结束时，用组合权重对专家输出进行加权后，恢复原始序列长度。
多头注意力除了在加权平均中使用 softmax 外，还与软 MoE 有一些相似之处：h 个不同的头可以解释为不同的（线性）专家。区别在于，如果 m 是序列长度，每个输入标记的维度为 d，那么 h 个头中的每个头都会处理大小为 d/h 的 m 个向量。对 m′ 个输出标记中的每个标记使用不同的权重（即注意力权重）将 m 个结果向量组合起来。我们的专家是非线性的，他们会在这些专家的输入和输出端合并大小为 d 的向量。
有一些MOE方法时对专家参数加权组合，这些方法也是可微的，但复杂度较高，1）它们必须对专家参数进行平均，当使用参数较多的专家时，这可能会成为时间和/或内存的瓶颈；2）它们不能像软（和稀疏）MoEs 那样广泛地利用矢量化运算，因为每个输入都使用了不同的参数加权组合。我们推荐使用（Muqeeth et al, 2023）中的 "计算成本 "讨论来解决这些问题。
2.5 目前的局限性
自回归解码
Soft MoE 的关键方面之一在于巧妙地合并输入中的所有标记。这使得在自回归解码器中使用软 MoE 变得困难，因为在训练过程中必须保留过去和未来标记之间的因果关系。尽管可以使用注意力层中使用的因果掩码，但必须小心不要在令牌索引和槽索引之间引入任何相关性，因为这会导致每个专家接受训练的令牌索引产生偏差。在自回归解码器中使用软 MoE 是一个有前途的研究途径，我们为未来的工作留出空间。
懒专家和内存消耗
每个专家一个插槽往往是最佳选择。从性能的角度来看，使用两名专家各有一个插槽比为一名专家提供两个插槽（或token混合）更有效。我们假设相同的专家槽往往会在一定程度上对齐并提供较小的信息增益，而单个专家可能缺乏适应非常不同的槽预测的灵活性。因此，软 MoE 倾向于利用大量专家，尽管其成本仍然与密集主干网相似，但模型的内存需求可能会变得很大。
3 图像分类实验
三方面实验
1 训练的帕累托边界
2 推理时间优化
3 消融实验
3.2 稀疏路由算法
token level:每个token都会选择该token路由得分最高的前 K 个专家（。增加 K 通常会带来更好的性能，但代价是额外的计算成本。批量优先级路由(BPR)显著提高模型性能。
expert level：专家可以根据路由分数选择前 C token。C 是缓冲区大小，我们通常设置 E·C = c·T，其中 E 是专家数量，T 是组中令牌总数，c 是容量乘数。会出现一些令牌被多个专家同时选择，而其他一些令牌根本没有被选择。
3.3 训练的帕累托边界
vit各种规模在令牌选择、专家选择和软 MoE扩展方法下的结果。softmoe在flop和时间预算上都远远优于其他。
3.4 过训练
除了较短的运行和消融之外，我们还训练了许多模型更长的时间（几百万步），以在更大的计算规模上测试 Soft MoE 的能力。我们提出两个实验。
在第 3.4.1 节中，我们训练了从小到大不同大小的 ViT 和 Soft MoE，步长为 4M。图 4 和表 2 显示了结果。
对于给定的计算预算，Soft MoE 模型的性能远远优于 Vision Transformer 模型。例如，Soft MoE S/16 在 JFT 和 10-shot ImageNet 上的性能优于 ViT B/16，并且它还提高了完整 ImageNet 数据的微调分数，尽管其训练（和推理）成本要小得多。同样，Soft MoE B/16 的性能优于 ViT L/16 上游，微调后仅落后 0.5，同时速度快 3 倍，所需的 FLOP 数减少近 4 倍。最后，Soft MoE L/16 模型的性能优于密集的 H/14 模型，同时训练和推理服务的速度再次提高了 3 倍左右。
总结：在节省训练及推理时间的前提下完整效果优于密集模型，微调效果下降不显著（有提升）
其次，在第 3.4.2 节中，我们在超出标准密集缩放定律建议的最佳点的时间下，继续训练一些先前的 Soft MoE 模型。稀疏模型可以利用额外的容量来稳步提高其性能，从而产生非常强大的 Soft MoE 模型，并且推理成本特别低。由于具有较小主干的软 MoE 可以与较大视觉 Transformer 的质量相匹配，我们继续训练小型主干，以非常低的推理成本获得更高质量的模型。即使在额外（过度）训练之后，相对于较大 ViT 模型的总体训练时间也是相当的。对于这些长时间运行，我们观察到较长的冷却时间对于 Soft MoE 效果很好。因此，我们将冷却时间从 50k 步骤增加到最多 500k 步骤。
过度训练的 Soft MoE B/16 优于最佳 ViT 模型 (H/14)同时减少 10 倍的计算量（5.7 倍时间）
3 模型消融 附录比较全，正文中的结果较为简略
超参数探索：每个专家的最佳槽数。每个专家一到两个位置效果最好。
最佳专家数量。当每个专家使用一个槽位时，与输入令牌数量大致相同的专家效果最佳。
架构/算法消融。为了理清好处的来源，我们将 Soft MoE 与许多消融版本进行比较：确定性地将令牌 i 路由到专家 i、固定统一调度/组合权重等。
MoE 层放置。
3.5.1专家数量和每个专家的空位
当将 Soft MoE 应用于给定架构和输入序列长度时，必须决定使用多少个专家以及每个专家使用多少个插槽。插槽总数决定了 MoE 层中应用的工作量 (FLOP)（忽略较小的路由成本）。如果槽总数大于输入令牌的数量，则模型将比密集 Transformer 需要更多的 FLOP：MoE 层将处理更多“令牌”。相反，如果插槽数量低于原始token数量，Soft MoE 将相对于密集 Transformer 节省一些计算量。
每个专家的最佳槽数。
在这个实验中，计算总量是固定的，我们比较不同的配置。具体来说，我们将槽总数固定为 4096，并用不同数量的专家来训练模型。 MoE 算法通常无法很好地扩展到大量专家（超过 100 名）。
在专家选择和令牌选择 MoE 的情况下，所有专家缓冲区的并集大小也是每个输入图像 4096。对于 Soft MoE，我们只是改变专家的数量，保持专家联盟处理的token总数不变。对于稀疏 MoE（专家/token选择），有一个实现细节：“组大小”是一起路由的批次的子集。
一组中的所有token都竞争由每位专家选择。范围可以从一个图像/组到整个批次/组；后者更灵活，但增加了路由的计算开销.

最佳专家数量。在这里，模型的成本不匹配：更多的专家会增加成本（通过更多的槽位）。图 7 显示，对于软 MoE 和专家选择，专家越多效果越好（最多 1024 名）。

每个专家更多插槽：附录

3.5.2 算法消融：身份和统一路由
Soft MoE 依赖于学习如何为每个专家混合token。为了了解找到有用的输入标记线性组合的影响，我们通过测试一些自然选择来消除这一方面：
身份路由：令牌不混合，第一个令牌分配给第一个专家，第二个令牌分配给第二个专家，依此类推。
均匀混合：每个槽以相同的方式混合所有输入令牌：通过统一平均它们，用于调度和组合。在这种情况下，我们必须独立且随机地将每个专家初始化。

4 对比学习实验
我们测试学习到的表示在用于其他任务时是否也明显更好。在本节中，我们将探讨一种流行的范式：图像语言对比学习。
其中图像塔在图像分类任务上进行预训练，然后在图像-文本对数据集上训练文本编码器时冻结。

5 模型检验
将了解模型学习的路由的各个方面
token对插槽的贡献：虽然softmoe没有drop,但如果某些token的 logit 远低于其他token的 logit，则仍有可能对所有插槽贡献很小.我们想看看某些token是否以不成比例的方式对插槽做出贡献。
我们看到有一个沉重的token尾部，为插槽提供了更强的总贡献，并且跨层的形状有些相似。大约 2-5% 的token提供的总权重高于 2。此外，15% 到 20% 的token仅贡献总权重的 0.25。

即使是有softmax,仍然会有一些token相对而言不重要。

expert对output的贡献：有多少不同的槽最终对输出token有贡献。我们关注每个专家一个位置的情况。我们可以通过对批次中所有输出标记的线性组合中的相应系数进行平均来估算每个专家（相当于插槽）的总贡献。根据层的不同，一些专家对输出token的影响可能是其他专家的 3 倍到 14 倍。
每个槽的输入令牌数：对于每个槽，图 9（右）显示了需要多少个输入令牌才能在其线性组合中实现一定的累积权重。不同时段的分布差异很大。对于少数插槽，前 20-25 个token占插槽权重的 90%，而对于其他插槽，分布更加均匀，并且许多token有助于填充插槽。一般来说，我们看到插槽倾向于混合大量令牌，这与标准稀疏 MoE 不同
视觉评估：为了提供一些关于槽如何平均输入标记的直观信息，图 10 以图形方式显示了图 1 中所示图像的 8 个不同槽的线性组合。我们对补丁进行阴影处理，与它们在槽中的权重成反比；请注意，在传递给专家之前，所有标记表示最终都会组合成一个标记（具有隐藏维度 h）（与我们的图中不同，它们以通常的方式排列）。这些图对应于具有 128 位专家的 Soft MoE S/16，每个专家一个槽位，我们从 128 个槽位中精心挑选了 8 个槽位，以突出不同的槽位如何倾向于关注图像的不同元素。
（所以其实是有可能产生对应专家&&做可解释性的。另一方面，产生领域专家就意味着可以做可解释性吗？怎么理解）
6 讨论
稀疏模型可能面临基础设施方面的挑战，这可能会减缓其广泛采用的速度。由于这些模型最初是为了解锁海量模型尺寸而设计的，因此它们往往是分布式的，并且大多数路由算法需要额外的通信成本：额外的激活、梯度或专家参数跨设备发送。对于软MOE来说也是如此，专家也可能是分散的。
然而，现代密集模型现在足够大，它们也是分布式的，从而缩小了该轴上的差距。此外，稀疏性的好处在小模型尺度上表现出色，无论是在之前的工作还是 Soft MoE 中，都符合行业当前对更快推理的需求。
我们提出了 Soft MoE，一种新的稀疏 Transformer 架构，它避免了稀疏混合专家模型中常见的离散令牌到专家分配问题。通过将输入标记合并成线性组合，然后将其分派给专家，我们能够训练一个快速且完全可微的模型。
这些实验表明，Soft MoE 非常有效，并且远远优于其他方法，同时计算成本通常更低。如何处理语言解码器的因果屏蔽是未来工作中一个令人兴奋且有影响力的研究方向