【论文速读】MOD，《Mixture-of-Depths: Dynamically allocating compute in transformer-based language models》

论文链接：https://arxiv.org/html/2404.02258v1
最近看到有招聘要求中提到熟悉MoE、MoD。MoE前面学习过，赶紧学习一下MoD！

Transformer模型在前向传播中对每个token均匀分配计算资源，这在处理不同的语言建模问题时并不高效，因为并非所有token和序列都需要相同的计算量来进行准确预测。MoD模型旨在解决这一问题，通过动态分配计算资源，以更高效地使用计算预算。
文章的要点包括：

• 提出了一种新的Transformer模型变体，即MoD，它能够根据上下文动态地分配计算资源。
• 证明了MoD模型能够在保持性能的同时，减少每次前向传播所需的FLOPs，从而加快推理速度。
• 展示了MoD模型可以与现有的硬件堆栈（优先考虑静态计算图和已知张量大小）和谐工作，以提高硬件利用率。
• 通过实验验证了MoD模型在不同配置下的性能，以及与MoE模型结合使用时的性能提升（MoDE模型）。
• 提供了对MoD模型在训练和推理阶段的采样方法的见解，以及如何通过学习到的路由决策来提高模型效率。

MoD模型的方法要点：

1. 设置静态计算预算：通过限制序列中可以参与计算的token数量，来定义总的计算预算。
2. 使用per-block路由器：为每个token生成一个标量权重，表示路由器对该token参与计算的偏好。
3. 选择top-k token：根据路由器权重选择top-k个token参与计算，保持计算图和张量大小在训练期间静态，而token的参与是动态和上下文敏感的。
   

摘要

这篇文章介绍了一种名为Mixture-of-Depths（MoD）的新型Transformer模型，它能够动态地在序列的不同位置分配计算资源（FLOPs）。MoD模型通过限制每层可以参与自注意力和多层感知器（MLP）计算的token数量来执行这一动态分配。这种机制通过一个top-k路由方法来确定哪些token将被处理，其中k是预先定义的。与MoE（Mixture of Experts）模型不同，MoD模型可以选择对token应用计算或通过残差连接传递token（保持不变，节省计算）。MoD模型在训练时能够学习动态分配计算，并且这样做是高效的。在等效的FLOPs和训练时间下，MoD模型与基线性能相当，但在每次前向传播中需要的FLOPs更少，并且在训练后的采样中可以快达50%。

1. Introduction

重点内容：

• 并非所有问题都需要相同数量的时间或努力来解决。类似地，在语言建模中，并非所有token和序列都需要相同的时间或努力来准确进行预测。然而，现有的Transformer模型在前向传播中对每个token分配相同数量的计算资源。
• 理想的Transformer模型应该通过不必要地消耗计算资源来使用更小的总计算预算。
• 条件计算（Conditional computation）是一种尝试通过仅在需要时消耗计算资源来减少总计算的技术。不同的算法提供了关于何时以及使用多少计算资源的解决方案。
• 现有的硬件限制可能导致一般形式的条件计算问题难以解决，因为它们倾向于引入动态计算图。与现有硬件堆栈协调一致的条件计算方法可能更有前景，这些方法优先考虑静态计算图和已知的张量大小以最大化硬件利用率。
• 本文考虑了使用静态计算预算进行语言建模的问题，该预算可以小于普通Transformer所使用的预算。网络必须学习如何通过每个token的决策动态分配可用的计算资源。
• 利用了类似于专家混合（Mixture of Experts，MoE）Transformer的方法，其中在网络深度上对token级别的路由决策进行动态选择。与MoE不同的是，MoD可以选择对token应用计算（就像标准Transformer一样），或者通过残差连接传递它（保持不变并节省计算）。
• MoD技术还允许在性能和速度之间进行权衡。一方面，可以训练MoD Transformer，在等效训练FLOPs的情况下，最终对数概率训练目标比普通Transformer提高多达1.5%，并且在训练时所需的实际时间相同。另一方面，可以训练MoD Transformer，在等效FLOPs的情况下实现与普通Transformer相当的训练损失，但每次前向传播使用的FLOPs更少，因此在步骤上更快。这些结果表明MoD Transformers能够智能地路由（即跳过不必要的计算），因为尽管每次前向传播的FLOP足迹更小，但它们仍然能够实现相同或更好的序列对数概率。

2. Background

重点内容：

• Transformer架构已经成为实际人工智能革命的主力，带来了前所未有的能力，但代价是昂贵的训练运行和提供程序。这激发了人们对使Transformer架构更高效的巨大兴趣。
• 条件计算是减少总计算量的一种有希望的方法，通过学习机制来决定何时以及如何消耗计算资源。这个术语由Bengio在2013年引入，并在接下来的几年中进一步探讨。
• 最近的工作已经为Transformer开发了条件计算方法。一些工作集中在“提前退出”上，即学习决定何时结束对给定token的计算，允许token跳过退出决策后的所有剩余Transformer层。
• 其他工作开发了迭代Transformer层的方法，使用共享权重进行自适应数量的步骤。还有工作通过构建在适配器方法上的条件计算，在微调设置中使用条件计算，以学习跳过冻结预训练权重的块，只运行一个小的微调适配器。
• 一种成功的条件计算公式是“专家混合”层（MoE），最初在LSTMs的背景下开发，后来的工作显示了MoE与Transformer结合使用的引人注目的经验结果。与尝试节省或消耗额外计算的其他条件计算方法不同，MoE Transformers使用条件逻辑将token路由到多个专家MLP中的一个，同时保持总计算支出恒定。
• MoD方法可以被看作是使用MoE Transformer的路由逻辑，但不是有多个专家，MoD部署了一个可以动态跳过的单一专家。
  
  图1展示了Mixture-of-Depths（MoD）变换器的概念。MoD Transformer与专家混合（Mixture-of-Experts，MoE）Transformer类似，使用路由器来选择可能的计算路径。但在MoE Transformer中，选择通常是在标准块的计算（例如自注意力和多层感知器，MLP）之间进行。而在MoD Transformer中，可能的选择包括一个标准块的计算或一个残差连接（即不进行任何操作，从而节省计算资源）。
  图1中特别指出，在MoD中，一些token选择了第二条路径，即通过残差连接直接传递，而不是经过自注意力和MLP计算。这意味着MoD的总FLOP（浮点运算次数）比常规Transformer或MoE Transformer要小。
  图的右上角展示了一个训练好的模型对一个短序列（为了可视化目的截断为64个token）的路由决策。通过检查这些选择，可以发现有些token是由后面的块的层处理的，它们在整个模型的深度中只通过了相对较少的块。
  图的右下角分别展示了常规Transformer与基于“早退”的条件计算的计算路径。
  图1强调了MoD如何在不同层或块中动态地分配token，以及如何通过智能路由来减少计算量，从而在保持性能的同时提高效率。

3. Implementing Mixture-of-Depths Transformers

实现总体策略如下：

• 设置静态计算预算：通过限制序列中可以参与单个块计算的token数量（例如，参与自注意力和随后的MLP），来少于等效常规Transformer的计算预算。例如，常规Transformer可能允许序列中的所有token参与自注意力，MoD可能限制为序列中的50%的token。见第3.1节。
• 使用每个块的路由器：为每个token发出一个标量权重，该权重表示路由器对该token参与块计算或绕过它的偏好。见第3.2节。
• 识别 top-k 标量权重（每个序列，每个块）以选择将参与块计算的那些token。由于精确地有 k 个token 将参与块的计算，计算图和张量大小在训练期间保持静态；只是token的参与是动态的和上下文敏感的，由路由器决定。见第3.3节。

3.1 定义计算预算
为了执行每次前向传播的总计算预算，定义容量的概念，它定义了给定计算的输入token的总数（例如，参与自注意力的token，MoE中的某个专家等）。例如，每个常规Transformer块中的自注意力和MLP的容量是 T —— 序列和批次中的token总数。另一方面，MoE Transformer每个专家MLP的容量小于 T，以便更均匀地将总计算分配给每个专家。但是，由于他们每个块使用多个专家，他们的总容量大约等于常规Transformer的容量。

通常，token容量决定了使用条件计算的变换器的总FLOPs，而不是任何路由决策的结果。这是因为静态图实现考虑了最坏情况决策；例如，即使相对较少的token最终路由到它，计算的输入也将被填充到它的容量，和/或如果容量超出，token将从计算中丢弃。

可以通过降低计算的容量来实现与常规Transformer相比每次前向传播使用更小的计算预算的目标。然而，随意使用较小的计算预算将导致性能下降。假设某些token可能不需要像其他token那样多的处理，而这些token可以通过学习来识别。因此，如果网络学会选择正确的token来填充其容量，那么它可能保留其性能。在以下部分中，描述了可以为此目的使用的路由方案。

3.2 绕过变换器块的路由
考虑的设置是将token路由到两个计算路径之一：(1) 自注意力和MLP块，(2) 残差连接。后者在计算上很便宜，并且产生的块输出完全由其输入值决定。前者在计算上很昂贵。

如果将路径(1)的容量设置为小于 T（序列和批次中的token总数），那么每次前向传播的总FLOPs将少于普通变换器。例如，如果将一个块的容量设置为 T/2（即普通变换器中token数量的一半），那么自注意力期间的query-times-key矩阵乘法在FLOP强度上仅为普通变换器的25%（(T/2)^2 vs T^2）。MLP的FLOP节省也可以通过类似计算确定。

直观地说，每次前向传播的总FLOPs（以及完成前向传播所需的时间）会随着缩减块容量的激进程度而减少。然而，下游性能也会受到缩减块容量的激进程度以及我们实现的路由算法的影响。

在一个极端情况下，如果将每个块的容量保持在 T，并将每个token路由到（而不是绕过）每个块，那么就恢复了一个常规Transformer。在另一个极端情况下，如果我们将每个块的容量设置为0，并将所有token绕过每个块，那么我们就留下了一个非常快的模型，它不与Transformer的大部分参数进行交互，并且无疑性能很差。假设在这两者之间的某个地方存在一个最优模型，它比常规Transformer更快，并且表现一样好，甚至更好，同时还更快地进行步骤。

3.3 路由方案
简单地，可以利用随机性来路由token，类似于层或块的“dropout”。将这种路由方案作为对照，并会展示它与常规Transformer相比性能明显不足。

再假设学习到的路由是更好的。直观地说，网络应该能够学习哪些token需要或多或少的处理。如果正确地认为Transformer经常消耗比他们需要做出预测更多的计算，那么一个实证问题就是能多激进地缩减每个块的容量，以及能够负担得起绕过每个块多少token。

图2：路由方案。当使用token-choice路由（左）时，token被引导到他们选择的计算路径。如果给定路径超过其容量（例如，在这个例子中超过两个token），那么多余的token必须被丢弃（紫色token）。最终丢弃哪个token取决于底层代码的确切实现。例如，通常给序列或批次顺序中较早的token优先权。在expert-choice路由（中）中，每个路径精确地选择k个token，基于token的路由器权重进行top-k机制。在这里，如果token不是任何给定路径的top-k，它们将被丢弃（橙色token），一些token甚至可能被引导到多个路径（黄色token）。在这项工作中，采用了expert-choice路由（右）。但是，因为只通过两条路径路由，一个单一的top-k操作可以高效地将token分成两个互斥的集合，每个计算路径一个，防止了上述的过度或不足处理问题。

考虑了两种学习到的路由方案（见图2）：token-choice和expert-choice。在token-choice路由中，路由器为每个token在计算路径（例如，在MoE Transformers中的专家身份）上产生概率分布。然后，token被引导到他们最喜欢的路径——即具有最高概率的路径——辅助损失确保所有token不会汇聚到同一路径。token-choice路由可能会有负载平衡问题，因为没有保证token在可能的路径之间适当地分配自己。“expert-choice路由”将这个配方颠倒过来：不是让token选择他们喜欢的路径，而是每个路径根据token的偏好选择top-k个token。这确保了完美的负载平衡，因为每个路径都保证有k个token被引导到。然而，这可能会导致一些token的处理过度或不足，因为一些token可能在多个路径的top-k中，或者一个也不在。

选择利用expert-choice路由有几个原因。首先，它消除了辅助平衡损失的需要。其次，由于top-k操作取决于路由器权重的大小，这种路由方案允许相对路由权重帮助确定哪些token最需要块的计算；路由器可以通过适当设置权重来尝试确保最关键的token在top-k中，这在token-choice路由方案中是不可能的。对于特定用例，其中一个计算路径本质上是一个空操作，关键是将重要token从空操作路由出去可能至关重要。第三，因为只通过两条路径路由，一个单一的top-k操作可以高效地将token分成两个互斥的集合，每个计算路径一个，防止了上述的过度或不足处理问题。

3.4 路由实现
作为高层次直觉的提醒，每个token由路由器处理以产生一个标量权重，然后使用top-k权重来选择将通过Transformer块的token身份，该块包括自注意力和随后的MLP。

假设有一个长度为 S 的token嵌入序列，对于给定层 l；即 𝑋_𝑙={𝑥𝑖_𝑙∣𝑖 是一个整数, 1≤𝑖≤𝑆}.
对于给定层 l 的每个token嵌入，路由器权重是一个标量，由线性投影产生，𝑟𝑖_𝑙=𝑤𝜃^𝑇 * 𝑥𝑖^𝑙

目标是使用这些路由器权重来确定每个token的块计算的输出。假设 𝑃𝛽(𝑅_𝑙) 是路由器权重集合 𝑅_𝑙的第 𝛽 百分数，其中 𝛽=1−𝐶/𝑆，C 是用户定义的每个批次元素的容量（一个小于 S 的整数，定义了序列中的token数量，这些token将由给定函数处理）。给定token的块输出是

这里，路由器权重大于阈值，就执行f函数（包括自注意力和随后的MLP）。注意，给定token的输出 可能依赖于其他token。与基线相比，MoD可以节省计算量，因为块的计算𝑓包括的token数量少于通常（C < S），使得自注意力和MLP更便宜。

值得注意的是，函数 𝑓 的输出乘以路由器权重。这将路由器权重沿着“梯度路径”放置，因此在语言建模任务的过程中使它们受到梯度下降的影响（文章试验过在通过块计算绕过的token的计算路径中也包括路由器权重的版本，但似乎足够了——并且在实现上更简单——只包括那些不绕过块计算的token的计算路径上的路由器权重）。

3.5 采样
尽管专家选择路由有许多优点，但它有一个明显的问题：top-k操作是非因果的。这意味着一个给定token的路由权重是否在序列的top-k中，取决于它之后token的路由权重的值，这在自回归采样时是无法访问的。

为了解决这个问题，测试了两种方法。第一种是引入一个简单的辅助损失，它在实证上对主要语言建模目标的影响大约是0.2-0.3%，但允许从模型中自回归采样。使用的是二元交叉熵损失，其中路由器的输出提供logits，这些logits的top-k选择提供目标（即，如果一个token在top-k中则为1，如果不在则为0）。直观地说，这个损失将路由器输出的sigmoid中心设置在0.5左右；被选为top-k的token被推动产生高于0.5的路由器输出，而不在top-k中的token则被推动产生低于0.5的路由器输出。第二种方法引入了一个小型的辅助MLP预测器（类似于第二个路由器），它接收与路由器相同的输入（带有停止梯度），但其输出是预测那个token是否会在序列的top-k中。这种方法不影响语言建模目标，并且在实证上对步骤速度没有显著影响。

有了这些新方法，可以通过基于路由器的输出选择将token路由到或绕过块来进行自回归采样，这并不依赖于来自未来token的任何信息。我实证证据表明，这是一个相对容易的辅助任务，很快就能达到99%的准确率。

3.6 训练方法
所有模型都使用相同的基本超参数配置（例如，余弦调度等于训练步骤的1倍，批量大小为128，序列长度为2048），除了在isoFLOP分析中改变层数、头数和嵌入大小以产生不同大小的模型。

4 结果

4.1 训练，isoFLOP比较

图3：MoD超参数调整。MoD的变体被训练了6e18 FLOPs，以确定进一步isoFLOP分析的最佳超参数。在左侧的图表中，灰色框表示性能优于isoFLOP最优基线的模型。最佳的MoD变体是可以选择每隔一个块进行路由的，并且使用top-k为256（也就是说，256个，或者说序列的token的12.5%由自注意力和随后的MLP处理，而1792个，或者说87.5%的序列的token绕过块）。右侧显示了选定模型的学习曲线。值得注意的是，模型#3达到了与isoFLOP最优基线相同的性能，但由于每次前向传播需要较少的FLOPs，其步骤速度提高了66%。

首先用相对较小的FLOP预算（6e18）训练模型以确定最佳超参数（见图3）。总的来说，MoD变换器将基线isoFLOP曲线“向下和向右”拖动。也就是说，最优的MoD变换器实现了比最优基线更低的损失，并且有更多的参数。这种效应的一个幸运结果是，存在较小的MoD模型，虽然它们在超参数设置上本身不是isoFLOP最优的，但仍然与最优基线模型性能相当或更好，同时步骤速度更快。例如，一个220M参数的MoD（图3中的模型#3）稍微优于isoFLOP最优基线（也是220M，图3中的模型#1），但在训练期间步骤速度提高了60%以上。至关重要的是，当在等效硬件上运行时，这两种模型变体的墙钟训练时间大致相同（见图3）。

图4展示了在不同FLOPs预算（6e18、2e19和1e20）下，使用12.5%容量的MoD（Mixture-of-Depths）变换器变体进行的isoFLOP分析。在这个分析中，训练了大小从60M到3B参数不等的模型。图4的右侧描绘了每次前向传播的相对FLOPs（以isoFLOP最优基线为标准化）。

从图中可以观察到以下几点：
FLOPs优化趋势：MoD变换器在FLOPs优化方面具有比基线更多的参数，这一趋势在更大的FLOPs预算下仍然持续。
性能与速度：存在MoD变体，它们在步骤速度上比isoFLOP最优基线快（以在等效硬件上训练时的步骤每秒计），同时实现了更低的损失。图4中展示的是标准化的每次前向传播的FLOPs，而不是实际的墙钟步骤时间，但根据实验，这两者是紧密相关的。可以制作一个类似的图表来展示相对的墙钟步骤时间，基本趋势是相同的。
步骤速度增益来源：MoD变换器的步骤速度增益有两个来源。首先，由于部分token绕过块，MoD变换器的参数每次前向传播的FLOPs比基线少。其次，由于isoFLOP最优的MoD变换器比isoFLOP最优基线更大且损失更低，存在较小的MoD变体，它们的性能与isoFLOP最优基线一样好或更好，并且由于模型更小，它们的步骤速度更快。
最优MoD：图4还揭示了另一个重要发现：最优的MoD变换器是那些每次前向传播使用与isoFLOP最优基线相同数量FLOPs的模型。这一发现允许人们直接预测哪种尺寸的MoD变换器将在给定的isoFLOP训练预算下表现最佳：只需为给定的MoD配置（即容量和路由频率）调整模型大小，使其每次前向传播使用的FLOPs与isoFLOP最优基线一样多，就能得到该配置下表现最优的MoD变体。
模型尺寸调整：实证上发现，增加模型的深度比增加宽度更优，当向模型添加FLOPs时。
总的来说，图4说明了MoD 如何在保持或提高性能的同时，通过减少每次前向传播所需的FLOPs来加快模型的推理速度，并且展示了不同模型尺寸对于达到isoFLOP最优性能的影响。

图5展示了一个经过训练的MoD（Mixture-of-Depths）的路由决策，采用了间隔的12.5%容量路由块与全注意力块的交错训练方式。具体来说，图5揭示了以下几点：
路由选择性：在路由块中，选择通过自注意力和多层感知器（MLP）计算进行路由的token数量是稀疏的。这意味着在这些块中，只有少数token会参与到计算密集型的自注意力和MLP操作，而大部分token则会选择绕过这些操作。
路由偏好：尽管网络有时会优先选择某些token在不同深度的块中进行路由，但整体上，网络倾向于在路由块中让尽可能少的token参与计算，从而节省计算资源。
路由决策可视化：左侧的图表展示了路由决策，其中可以看到序列末尾处有一个垂直的深蓝色条带，这表明在序列的这一部分，有更多的token选择绕过计算块。
路由器权重分布：右侧的直方图显示了路由器权重的分布情况，根据辅助损失的设计，大约有12.5%的权重高于0.5，而87.5%的权重低于0.5。这表明网络已经学会了如何根据token的重要性来分配路由权重。
性能与计算效率：尽管进行了激进的路由绕过，但变换器仍然能够相对于基线实现性能提升。这表明MoD模型在减少计算量的同时，仍然能够有效地捕捉语言数据中的关键信息。
token参与度与预测难度：初步分析表明，更频繁地参与块计算的token与具有更高熵的输出预测相关，这可能对应于更难做出的预测。这暗示了网络可能已经学会根据预测的难度来动态调整token的路由策略。
总的来说，图5提供了对MoD 如何在不同深度的块中动态地对token进行路由的直观理解，以及这种路由策略如何影响模型的性能和计算效率。

4.3 Mixture-of-Depths-and-Experts (MoDE)

图7展示了Mixture-of-Depths-and-Experts (MoDE)模型的概念和效果。MoDE模型是将MoD（Mixture-of-Depths）技术与MoE（Mixture of Experts）模型结合起来的一种新型结构。
图7中展示了以下内容：
MoDE模型的两种实现方式：介绍了两种将MoD技术与MoE结合的方法。第一种是分阶段的MoDE（Staged MoDE），它先实施MoD机制，然后再进行MoE机制。第二种是集成的MoDE（Integrated MoDE），它使用单一的路由操作将token分配给专家或者无操作（no-op）。
性能提升：图7展示了MoD和MoE各自以及结合使用时的性能提升。这表明MoD提供的改进可以与MoE的改进相加，从而可能实现比单独使用任一技术更好的性能。