Mixture-of-Depths：在基于 Transformer 的语言模型中动态分配计算量
机器学习
2024-04-04 v1 计算与语言
摘要
基于 Transformer 的语言模型在输入序列上均匀分配 FLOPs。在本工作中，我们证明了 Transformer 可以转而学习将 FLOPs（或计算量）动态分配给序列中的特定位置，从而在模型深度的不同层上优化沿序列的分配。我们的方法通过限制参与给定层自注意力和 MLP 计算的 token 数量（）来强制执行总计算预算。需要处理的 token 由网络使用 top- 路由机制确定。由于 是先验定义的，这一简单过程使用具有已知张量大小的静态计算图，不同于其他条件计算技术。然而，由于 个 token 的身份是流动的，该方法可以在时间和模型深度维度上非均匀地消耗 FLOPs。因此，计算消耗在总量上是完全可预测的，但在 token 级别上是动态且上下文敏感的。以这种方式训练的模型不仅学会了动态分配计算量，而且做到了高效。这些模型在等效 FLOPS 和训练挂钟时间下匹配基线性能，但每次前向传播仅需一小部分 FLOPs，并且在训练后采样期间步进速度可提高 50% 以上。
引用
@article{arxiv.2404.02258,
title = {Mixture-of-Depths: Dynamically allocating compute in transformer-based language models},
author = {David Raposo and Sam Ritter and Blake Richards and Timothy Lillicrap and Peter Conway Humphreys and Adam Santoro},
journal= {arXiv preprint arXiv:2404.02258},
year = {2024}
}