引入混合深度,DeepMind 新设计可大幅提升 Transformer 效率。
论文标题:Mixture-of-Depths: Dynamically allocating compute in transformer-based language models 论文地址:https://arxiv.org/pdf/2404.02258.pdf
设定一个静态的计算预算,该预算低于等价的常规 Transformer 所需的计算量;做法是限制序列中可参与模块计算(即自注意力模块和后续的 MLP)的 token 数量。举个例子,常规 Transformer 可能允许序列中的所有 token 都参与自注意力计算,但 MoD Transformer 可限定仅使用序列中 50% 的 token。 针对每个 token,每个模块中都有一个路由算法给出一个标量权重;该权重表示路由对各个 token 的偏好 —— 是参与模块的计算还是绕过去。 在每个模块中,找到最大的前 k 个标量权重,它们对应的 token 会参与到该模块的计算中。由于必定只有 k 个 token 参与到该模块的计算中,因此其计算图和张量大小在训练过程中是静态的;这些 token 都是路由算法认定的动态且与上下文有关的 token。