谷歌更新Transformer架构，更节省计算资源！50%性能提升

AIGC动态3个月前发布 QbitAI

290 0 0

文章摘要

【关键词】 谷歌更新、Transformer架构、MoD技术、计算资源分配、训练效率提升

谷歌最近更新了其Transformer架构，推出了一种名为Mixture-of-Depths（MoD）的新技术。

这项技术通过动态分配大模型中的计算资源，跳过一些不必要的计算，从而显著提高了训练效率和推理速度。

MoD的核心思想是在输入序列中的特定位置动态分配FLOPs（浮点运算次数），以优化模型不同层次深度中的资源分配。

这种方法不仅节省了计算资源，还提高了模型的效率，使得每次前向传播所需的FLOP更少，并且在训练后采样时步进速度提高了50%。

MoD的提出基于这样一个观察：在大模型的训练和推理中，很多计算是不必要的。

例如，预测下一个句子可能很难，而预测句子结束的标点符号则相对简单。

如果给这两种任务分配相同的计算资源，显然对后者来说是一种浪费。

理想情况下，模型应该只给那些需要准确预测的token分配更多的计算资源。

MoD正是基于这样的理念，通过限制给定层的自注意力和MLP计算的token数量，迫使神经网络学会主要关注真正重要的信息。

这种方法使用一个已知张量大小的静态计算图，在时间和模型深度上动态扩展计算量。

MoD的另一个优势在于它能够保障负载平衡，避免了传统方法中可能出现的某些token被过度处理或处理不足的问题。

在MoD中，router输出的权重被用于确定哪些token将使用transformer计算，权重较大的token将参与计算，而权重较小的token则通过残差连接绕过计算。

研究团队通过实验展示了MoD在不同计算预算下的性能表现。

他们发现，MoD方法能够在更低的损失水平上拥有更多的参数，并且在训练之外，MoD的计算节省仍然有效。

此外，研究团队还探讨了MoD与MoE（Mixture-of-Experts）结合的可能性，这种结合被称为MoDE，结果表明这种结合能提供更好的性能和更快的推理速度。

MoD的推出引发了业界的广泛关注，有人将其与ResNet相比较，但MoD的跳过连接是完全绕过层的，与ResNet有所不同。

这项研究由DeepMind和麦吉尔大学共同进行，主要贡献者是David Raposo和Adam Santoro，两位都是DeepMind的研究科学家，此前共同发表了引用次数超过3500次的论文《Relational inductive biases, deep learning, and graph networks》，定义了Inductive bias（归纳偏置）的概念。

MoD的研究成果已经在论文中公布，链接为：https://arxiv.org/abs/2404.02258。