深度混合:在基于Transformer的语言模型中动态分配计算资源Mixture-of-Depths: Dynamically allocating compute in transformer-based
language models
基于Transformer的语言模型将浮点运算(FLOPs)均匀分布于输入序列中。本文展示了一种新方法,即Transformer能够学习动态地将FLOPs(或计算资源)分配给序列中的特定位置,从而在模型深度的不同层级上优化序列计算资源的分配。我们的方法通过限制每一层中能参与自注意力和多层感知机(MLP)计算的token数量(k)来确保总计算预算。具体处理哪些token由网络通过top-k路由机制决定。由于k值是预先设定的,这一过程使用的是静态计算图,其张量大小已知,与其它条件计算技术不同。尽管如此,由于k个token的身份是动态变化的,该方法能够在时间和模型深度维度上非均匀地分配FLOPs。因此,计算总支出是完全可预测的,但在token级别上是动态且上下文敏感的。采用这种方式训练的模型不仅能动态分配计算资源,而且效率极高。这些模型在同等FLOPs和训练时间下,性能与基准模型相当,但在每次前向传播中所需的FLOPs大幅减少,且在训练后采样时速度可提升高达50%。