引入混合深度，DeepMind 新设计可大幅提升 Transformer 效率。
Transformer 的重要性无需多言，目前也有很多研究团队致力于改进这种变革性技术，其中一个重要的改进方向是提升 Transformer 的效率，比如让其具备自适应盘算能力，从而可以节省下不必要的盘算。
正如不久前 Transformer 架构的提出之一、NEAR Protocol 联合创始人 Illiya Polosukhin 在与黄仁勋的对话中说到的那样：「自适应盘算是接下来必须出现的。我们要关注，在特定问题上具体要花费多少盘算资源。」
其实人类就天生具备自适应盘算的能力 —— 人在解决各种不同的问题时，会自然地分配不同的时间和精力。
语言建模也应如此，为了得到准确的预测结果，并不需要为所有 token 和序列都投入同样的时间或资源。但是，Transformer 模型在一次前向传播中却会为每个 token 花费同等的盘算量。这不禁让人哀叹：大部分盘算都被浪费了！
理想情况下，如果可以不实行非必要的盘算，就可以降低 Transformer 的盘算估算。
条件式盘算这种技术可在需要实行盘算时才实行盘算，由此可以减少总盘算量。之前许多研究者已经提出了多种可以评估何时实行盘算以及运用多少盘算量的算法。
但是，对于这个颇具挑战性的问题，普遍运用的解决形式可能无法很好地应对现有的硬件限制，因为它们往往会引入动态盘算图。最有潜力的条件式盘算方法反而可能是那些能协调运用当前硬件栈的方法，其会优先运用静态盘算图和已知的张量大小（基于对硬件的最大利用而选取这个张量大小）。
近日，Google DeepMind 研究了这个问题，他们希望运用更低的盘算估算来缩减 Transformer 运用的盘算量。
论文标题：Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
论文地址：https://arxiv.org/pdf/2404.02258.pdf
他们设想：在每一层中，网络必须学会为每个 token 做决策，从而动态地分配可用盘算估算。在他们的具体兑现中，总盘算量由用户在训练之前设定并且不再更改，而非网络工作时实行决策的函数。这样一来，便可以提前预知并利用硬件效率收益（比如内存足迹减少量或每次前向传播的 FLOPs 减少量）。该团队的实验表明：可以在不损害网络整体性能的前提下获得这些收益。
DeepMind 的这个团队采用了类似于混合大师（MoE） Transformer 的方法，其中会在整个网络深度上实行动态 token 层面的路由决策。
而与 MoE 不同的是，这里他们的采用是：要么是将盘算应用于 token（和标准 Transformer 一样），要么就是通过一个残差连接绕过它（保持不变，节省盘算）。另一个与 MoE 的不同之处是：这里是将这种路由机制同时用在 MLP 和多头注意力上。因此，这也会影响网络处理的键值和查询，因此该路由不仅要决定更新哪些 token，还要决定哪些 token 可供关注。
DeepMind 将这一策略命名为 Mixture-of-Depths（MoD），以突显这一事实：各个 token 在 Transformer 深度上通过不同数量的层或模块。我们这里将其翻译成「混合深度」，见图 1。
MoD 支持运用者权衡考量性能与速度。一方面，运用者可以运用与通例 Transformer 同等的训练 FLOPs 来训练 MoD Transformer，这可为最终的对数概率训练目标带来多达 1.5% 的提升。另一方面，MoD Transformer 运用更少的盘算量就能达到与通例 Transformer 同样的训练损坏 —— 每一次前向传播的 FLOPs 可少最多 50%。
这些结果表明，MoD Transformer 可以学习智能地路由（即跳过不必要的盘算）。
兑现混合深度（MoD）Transformer
概况来说，其策略如下：
设定一个静态的盘算估算，该估算低于等价的通例 Transformer 所需的盘算量；做法是限制序列中可参与模块盘算（即自注意力模块和后续的 MLP）的 token 数量。举个例子，通例 Transformer 可能允许序列中的所有 token 都参与自注意力盘算，但 MoD Transformer 可限定仅运用序列中 50% 的 token。
针对每个 token，每个模块中都有一个路由算法给出一个标量权重；该权重表示路由对各个 token 的偏好 —— 是参与模块的盘算还是绕过去。
在每个模块中，找到最大的前 k 个标量权重，它们对应的 token 会参与到该模块的盘算中。由于必定只有 k 个 token 参与到该模块的盘算中，因此其盘算图和张量大小在训练过程中是静态的；这些 token 都是路由算法认定的动态且与上下文有关的 token。
路由规划
该团队考虑了两种学习到的路由规划（见图 2）：token 采用型和大师采用型。
在 token 采用型路由规划中，路由算法会跨盘算途径（比如跨 MoE Transformer 中的大师身份）生成针对每个 token 的概率分布。然后 token 会被传送到它们偏好的途径（即概率最高的途径），而辅助损坏可以确保所有 token 不会收敛到同一途径。token 采用型路由可能会有负载平衡问题，因为不能确保 token 在可能的途径之间划分适当。
大师采用型路由则是将上述规划反过来：不是让 token 采用它们偏好的途径，而是让每条途径基于 token 偏好采用前 k 个 token（top-k）。这能确保负载完美平衡，因为每条途径总是保证 k 个 token。但是，这也可能导致某些 token 被过处理或欠处理，因为某些 token 可能是多条途径的前 k 名，另一些 token 则可能没有相应途径。
DeepMind 的采用是运用大师采用型路由，原因有三。
第一，它无需辅助性的平衡损坏。
第二，由于选取前 k 名这一操作取决于路由权重的幅度，因此该路由规划允许运用相对路由权重，这有助于确定当前模块盘算最需要哪些 token；路由算法可以通过适当地设定权重来尽力确保最关键的 token 是在前 k 名之中 —— 这是 token 采用型路由规划无法做到的。在具体的用例中，有一条盘算途径本质上是 null 操作，因此应该避免将重要 token 路由到 null。
第三，由于路由只会经由两条途径，因此单次 top-k 操作就能高效地将 token 分成两个互斥的集合（每条盘算途径一个集合），这能应对上面提及的过处理或欠处理问题。
该路由规划的具体兑现请参看原论文。
采样
尽管大师采用型路由有很多优点，但它也有一个很明显的问题：top-k 操作是非因果式的。也就是说，一个给定 token 的路由权重是否在前 k 名取决于其之后的路由权重的值，但在实行自回归采样时，我们无法获得这些权重。
为了解决这个问题，该团队测试了两种方法。
第一种是引入一个简单的辅助损坏；实践证明，其对语言建模主目标的影响程度为 0.2%− 0.3%，但却能够让模型自回归地采样。他们运用了一个二元交叉熵损坏，其中路由算法的输出提供 logit，通过选取这些 logit 中的 top-k，就能提供目标（即，如果一个 token 在 top-k 中，就为 1，否则为 0）。
第二种方法是引入一个小的辅助 MLP 预测器（就像是又一个路由算法），其输入与路由算法的一样（具有 stop gradient），但其输出是一个预测结果：token 是否在序列的 top-k 中。该方法不会影响语言建模目标，实验表明也不会显著影响该步调的速度。
有了这些新方法，就可以通过采用路由到的 token 来实行自回归采样，也可以根据路由算法的输出绕过一个模块，这无需依赖任何未来 token 的信息。实验结果表明，这是一种相对简单辅助任务，可以很快兑现 99% 的准确度。
结果
训练，isoFLOP 比较
首先，该团队训练了一些 FLOP 估算相对较小（6e18）的模型，以确定最优的超参数（见下图 3）。
总体而言，可以看到 MoD Transformer 会将基准 isoFLOP 曲线向右下方拖动。也就是说，最优的 MoD Transformer 的损坏比最优的基准模型更低，同时参数也更多。这种效应带来了一个幸运的结果：存在一些和最优基准模型表现一样好甚至更好的 MoD 模型（同时步调速度更快），尽管它们本身在其超参数设置下并不是 isoFLOP 最优的。举个例子，一个 220M 参数量的 MoD 变体（图 3 中的 3 号模型）稍优于 isoFLOP 最优基准模型（参数量也是 220M，图 3 中的 1 号模型），但这个 MoD 变体在训练期间的步调速度快了 60% 以上。
下图 4 给出了总 FLOPs 为 6e18、2e19 和 1e20 时的 isoFLOP 分析。可以看到，当 FLOP 估算更大时，趋势依然继续。
下图 5 给出了一个运用交织的路由模块训练的 MoD Transformer 的路由决策。尽管其中存在大量绕过模块的情况，但这个 MoD Transformer 依然能兑现优于通例 Transformer 的性能。
自回归评估
他们也评估了 MoD 变体的自回归采样表现，结果见下图 6。这些结果表明 MoD Transformer 所带来的盘算节省不仅仅局限于训练设置。
混合深度与大师（MoDE）
MoD 技术可以自然地与 MoE 模型整合起来，组成所谓的 MoDE 模型。下图 7 展示了 MoDE 及其带来的提升。
MoDE 有两种变体：分阶段 MoDE 和集成式 MoDE。
其中分阶段 MoDE 是在自注意力步调之前进行路由绕过或到达 token 的操作；而集成式 MoDE 则是通过在通例 MLP 大师之间集成「无操作」大师来兑现 MoD 路由。前者的优势是允许 token 跳过自注意力步调，而后者的好处在于其路由机制很简单。
该团队注意到，以集成方式兑现 MoDE 明显优于直接降低大师的能力、依靠丢弃 token 来兑现残差路由的设计。