【论文速读】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

重点内容：

• Transform