arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2404.02258cs.LGcs.CL
收藏

Mixture-of-Depths：在基于 transformer 的语言模型中动态分配计算量

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

Google DeepMind（谷歌DeepMind）；McGill University（麦吉尔大学）；Mila（米拉研究所）

David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro

阅读 PDF
arXiv 页面 ↗HTML 原文 ↗

AI总结：

本文提出 Mixture-of-Depths，通过 top-$k$ 路由限制每层参与计算的 token，使 transformer 在固定总计算预算下实现 token 级动态分配，并以更低前向 FLOPs 保持基线性能、最高加速 50% 以上。

摘要：

基于 Transformer 的语言模型会在整个输入序列上均匀分配 FLOPs。在本工作中，我们证明 transformer 可以转而学习将 FLOPs（或计算量）动态分配给序列中的特定位置，并沿序列为模型不同深度的各层优化这种分配。我们的方法通过限制在给定层中能够参与自注意力和 MLP 计算的 token 数量（$k$）来强制执行总计算预算。待处理的 token 由网络使用 top-$k$ 路由机制决定。由于 $k$ 是预先定义的，与其他条件计算技术不同，这一简单过程使用具有已知张量尺寸的静态计算图。尽管如此，由于这 $k$ 个 token 的身份是流动的，该方法能够在时间和模型深度两个维度上不均匀地消耗 FLOPs。因此，计算支出在总量上完全可预测，但在 token 层面是动态且上下文敏感的。以这种方式训练的模型不仅学会了动态分配计算，而且执行得很高效。在等效 FLOPs 和训练挂钟时间下，这些模型达到了与基线相当的性能，但每次前向传播只需要一小部分 FLOPs，并且在训练后采样期间的单步速度最高可提升 50% 以上。

英文摘要：

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-$k$ routing mechanism. Since $k$ is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the $k$ tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50\% faster to step during post-training sampling.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。

↑