Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

这篇文章展示了基于Transformer的语言模型可以学习动态地将FLOPs（浮点运算数）分配给序列中的特定位置，以优化模型深度中不同层次的序列分配。通过限制每个层次中可以参与自注意力和MLP计算的令牌数量（$k$），我们的方法强制执行了总的计算预算。网络使用一种名为top-$k$路由机制来确定要处理的令牌。由于$k$是预先定义的，因此这个简单的过程使用已知张量大小的静态计算图，与其他条件计算技术不同。尽管如此，由于$k$令牌的身份是不确定的，因此该方法可以在时间和模型深度维度上非均匀地扩展FLOPs。因此，计算支出在总体上是完全可预测的，但在令牌级别上是动态和上下文敏感的。这种训练方式不仅使模型学会了动态地分配计算，而且效率很高。这些模型与基准性能相匹配，具有相同的FLOPs和训练时间，但每个前向传递需要的FLOPs却只有一小部分，并且在训练后的采样过程中可以快50\%以上。

本论文试图通过动态分配FLOPs的方式解决transformer-based language models中FLOPs分配不均匀的问题，从而提高模型的效率。

论文提出了一种基于top-k路由机制的方法，通过限制每层参与自注意力和MLP计算的token数量，并根据网络动态地选择参与计算的token，从而实现FLOPs的动态分配。

该方法可以使模型在保持与基准模型相同FLOPs和训练时间的情况下，每个前向传递需要的FLOPs数量大大减少，且在后训练采样时速度提高50%以上。实验结果表明该方法可以有效地提高模型的效率和性能。

在相关研究方面，最近也有一些关于transformer模型FLOPs优化的研究，如《Reducing Transformer Depth on Demand with Structured Dropout》和《Dynamic Computation Allocation for Transformer-based Language Model》等。