David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro
基于 Transformer 的语言模型通常将 FLOPs 均匀地分配到输入序列的各个位置。本文证明，Transformer 可以学习动态地将 FLOPs（或计算资源）分配到序列中的特定位置，并沿着序列长度和模型深度对分配进行优化。该方法通过限制每层中参与自注意力和 MLP 计算的 token 数量（$k$ 个）来强制执行总计算预算。待处理的 token 由网络使用 top-$k$ 路由机制确定。由于 $k$ 是预先定义的，该方法具有静态计算图，张量大小已知，这与其他条件计算技术不同。然而，由于这 $k$ 个 token 的身份是动态变化的，该方法可以在时间维度和模型深度维度上非均匀地消耗 FLOPs。因此，计算总开销完全可预测，但在 token 级别上是动态且上下文敏感的。经过这样训练的模型不仅学会了动态分配计算资源，而且效率很高。这些模型在同等 FLOPs 和训练时间下达到了与基线模型相当的性能，但每次前向传播所需的 FLOPs 更少，并且在训练后采样阶段，步进速度可提升高达 50%。