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

本文是LLM系列文章，针对《Mixture-of-Depths: Dynamically allocating compute in transformer-based language models》的翻译。

摘要

基于Transformer的语言模型将FLOP均匀地分布在输入序列中。在这项工作中，我们证明了Transformer可以学习将FLOP（或计算）动态分配到序列中的特定位置，从而优化沿序列在模型深度上的不同层的分配。我们的方法通过限制token数量来强制执行总计算预算(𝑘) 其可以参与在给定层的自注意力和MLP计算。要处理的token由网络使用顶部确定-𝑘 路由机制。自从𝑘 是先验定义的，与其他条件计算技术不同，这个简单的过程使用具有已知张量大小的静态计算图。然而，由于的身份𝑘 token是流动的，这种方法可以在时间和模型深度维度上非均匀地消耗FLOP。因此，计算支出在总和上是完全可预测的，但在token级别是动态的和上下文敏感的。以这种方式训练的模型不仅可以学习动态分配计算，而且可以高效地进行分配。这些模型与等效FLOPS和壁时钟训练时间的基线性能相匹配，但每次前向传球需要一小部分FLOP，并且在训练后采样期间可以加快50%以上。

1 引言

2 背景