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

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

谷歌终于更新了 Transformer 架构。

最新发布的Mixture-of-Depths（MoD），改变了以往 Transformer 计算模式。

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

结果显示，在等效计算量和训练时间上，MoD 每次向前传播所需的计算量更小，而且后训练采样过程中步进速度提高 50%。

这一方法刚刚发布，就马上引发关注。

MoE 风头正盛，MoD 已经来后浪拍前浪了？

还有人开始 " 算账 "：

听说 GPT-4 Turbo 在 Blackwell 上提速 30 倍，再加上这个方法和其他各种加速，下一代生成模型可以走多远？

所以 MoD 如何实现？

迫使大模型关注真正重要信息

这项研究提出，现在的大模型训练和推理中，有很多计算是没必要的。

比如预测下一个句子很难，但是预测句子结束的标点符号很简单。如果给它们分配同样的计算资源，那么后者明显浪费了。

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

所以研究人员提出了MoD。

它在输入序列中的特定位置动态分配 FLOPs（运算次数或计算资源），优化不同层次的模型深度中的分配。

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

因为 token 数量是事先定义好的，所以这个过程使用一个已知张量大小的静态计算图，可以在时间和模型深度上动态扩展计算量。

下图右上图中的橙色部分，表示没有使用全部计算资源。

这种方法在节省计算资源的同时，还能提高效率。

这些模型在等效的 FLOPS 和训练时间上与基线性能相匹配，但每次前向传播所需的 FLOP 更少，并且在训练后采样时提速 50%。

对比来看，如果为每一个 token 生成一个概率分布，每个 token 根据最高概率被送去对应的 " 专家 "，可能会导致负载不平衡。

如果反过来，这能保障负载平衡，但是可能导致某些 token 被过度处理或处理不足。

最后来看论文中使用的 Expert-choice MoD，router 输出的权重被用于确定哪些 token 将使用 transformer 亏啊计算。权重较大的 token 将参与计算，权重较小的 token 将通过残差连接绕过计算，从而解决每次向前传播的 FLOPs。

最后，研究团队展示了 MoD 在不同实验中的性能表现。

首先，他们使用相对较小的 FLOP 预算（6e18），以确定最佳超参数配置。

通过这些实验，作者发现 MoD 方法能够 " 拉低并向右推移 "isoFLOP 基线曲线，这意味着最优的 MoD 方法在更低的损失水平上拥有更多的参数。

通过 isoFLOP 分析，比较 6e18、2e19 和 1e20 FLOPs 的总计算预算下的模型性能。

结果显示，在更多 FLOP 预算下，FLOP 最优的 MoD 仍然比基线模型有更多的参数。

存在一些 MoD 变体，在步骤速度上比 isoFLOP 最优基线模型更快，同时实现更低的损失。这表明在训练之外，MoD 的计算节省仍然有效。

同时，研究团队还探讨了 MoD 和 MoE 结合的可能性—— MoDE。

结果表明而这结合能提供更好的性能和更快的推理速度。

网友：联想到了 ResNet

MoD 推出后马上引发了不小关注。

有人感慨，MoE 还没有弄清楚呢，MoD 都已经来了！

这么高效的方法，让人马上联想到了 ResNet。

不过和 ResNet 不同，MoD 跳过连接是完全绕过层的。

还有人表示，希望这种方法是完全动态的，而不是每个层固定百分比。

这项研究由 DeepMind 和麦吉尔大学共同带来。

主要贡献者是 David Raposo 和 Adam Santoro。

他们二人都是 DeepMind 的研究科学家。此前共同带来了神作《Relational inductive biases, deep learning, and graph networks》。

这篇论文目前被引次数超过 3500 次，论文核心定义了 Inductive bias（归纳偏置）概念。

论文地址：

https://arxiv.org/abs/2404.02258

参考链接：

[ 1 ] https://twitter.com/TheSeaMouse/status/1775782800362242157

[ 2 ] https://twitter.com/_akhaliq/status/1775740222120087847

— 完 —

【 火热报名中】中国 AIGC 产业峰会

定档 4 月 17 日

峰会已经邀请到数位代表技术、产品、投资、用户等领域嘉宾，共论生成式 AI 产业最新变革趋势。

最新确认嘉宾包括：商汤科技杨帆、轻松集团高玉石、印象笔记唐毅、蚂蚁集团李建国等，。

峰会将全程线上下同步直播，欢迎预约直播 ⬇️

点这里关注我，记得标星噢

一键三连「分享」、「点赞」和「在看」

科技前沿进展日日相见 ~