150亿参数，google开源了史上最大视觉模型V-MoE的全部代码

在过去几十年里，深度学习的进步是由几个关键因素推动的：少量简单而灵活的机制、大型数据集、更专业的硬件配置，这些技术的进步使得神经搜集在图象分类、机器翻译、蛋白质预测等任务中取得令人印象深刻的结果。然而，大模型以及数据集的使用是以大量估计需求为代价的。最近的研讨表明，增强模型的泛化能力以及稳健性离不开大模型的支持，因此，在训练大模型的同时协调好与训练资源的限制是非常重要的。一种可行的方法是利用条件估计，该方法不是为单个输出激活整个搜集，而是根据不同的输出激活模型的不同部分。这一范式已经在google提出的 pathway（一种全新的 AI 解决思路，它可以克服现有系统的许多缺点，同时又能强化其优势）愿景和最近的大型语言模型研讨中得到了重视，但在估计机视觉中还没有得到很好的摸索。希罕门控混合大师搜集 (MoE) 在自然语言处置惩罚中展示了出色的可扩大性。然而，在估计机视觉中，几乎所有的高性能搜集都是密集的，也就是说，每一个输出都会转化为参数进行处置惩罚。去年 6 月，来自google大脑的研讨者提出了 V-MoE（Vision MoE ），这是一种基于大师希罕混合的新视觉架构。当应用于图象识别时，V-MoE 在推理时只需要一半的估计量，就能达到先进搜集性能。此外，该研讨还提出了对路由算法的扩大，该算法可以在整个 batch 中对每一个输出的子集进行优先级排序，从而实现自适应图象估计。这允许 V-MoE 在测试时能够权衡性能和平滑估计。最后，该研讨展示了 V-MoE 扩大视觉模型的潜力，并训练了一个在 ImageNet 上达到 90.35% 的 150 亿参数模型。论文地址：https://arxiv.org/pdf/2106.05974.pdf代码地址：https://github.com/google-research/vmoeV-MoEgoogle大脑在 ViT 的不同变体上构建 V-MoE：ViT-S(mall)、ViT-B(ase)、ViT-L(arge) 和 ViTH(uge)，其超参数如下：ViT 已被证明在转嫁学习设置中具有良好的扩大性，在较少的预训练估计下，比 CNN 获得更高的准确率。ViT 将图象处置惩罚为一系列 patch，输出图象首先被分成大小相等的 patch，这些 patch 被线性投影到 Transformer 的隐藏层，在位置嵌入后，patch 嵌入（token）由 Transformer 进行处置惩罚，该 Transformer 主要由交替的自注意力和 MLP 层组成。MLP 有两个层和一个 GeLU 非线性。对于 Vision MoE，该研讨用 MoE 层替换其中的一个子集，其中每一个大师都是一个 MLP，如下图所示：为了大规模扩大视觉模型，该研讨将 ViT 架构中的一些密集前馈层 (FFN) 替换为独立 FFN 的希罕混合（称之为大师）。可学习的路由层为每一个独立的 token 选择对应的大师。也就是说，来自同一图象的不同 token 可能会被路由到不同的大师。在总共 E 位大师（E 通常为 32）中，每一个 token 最多只能路由到 K（通常为 1 或 2）位大师。这允许扩大模型的大小，同时保持每一个 token 估计的恒定。下图更详细地显示了 V-MoE 编码器块的结构。

V-MoE Transformer 编码器块实验结果google大脑首先在大型图象数据集 JFT-300M 上对模型进行一次预训练。下图左展示了模型在所有大小（从 small s/32 到 huge H/14）时的预训练结果。然后，使用一个新的 head（一个模型中的最后一层）将模型转嫁至新的下游任务（如 ImageNet）。他们摸索了两种转嫁设置：在所有可用的新任务示例上微调整个模型或者冻结预训练搜集并使用少量示例仅对新 head 调整（即所谓的小样本转嫁）。下图右总结了模型转嫁至 ImageNet 的效果，其中每一个图象种别仅在 5 张图象上训练（叫做 5-shot transfer）。

左为 JFT-300M 数据集上的 Precision@1 曲线图；右为 ImageNet 5-shot 的准确率曲线图。对于这两种情况，google大脑发现，在给定训练估计量时，希罕模型显著优于密集模型或者更快地实现相似性能。为了摸索视觉模型的极限，他们在 JFT-300M 扩大数据集上训练了一个具有 150 亿参数、24 个 MoE 层（出自 48 个块）的模型。这个迄今为止最大的视觉模型在 ImageNet 上实现了 90.35 的 Top-1 准确率。优先路由在实践中，由于硬件限制，使用动态大小的缓冲区（buffer）效率不高，因此模型通常为每一个大师使用预定义的缓冲区容量。一旦大师变「满」，超出此容量的分配 token 将被丢弃并不会被处置惩罚。因此，更高的容量会产生更高的准确性，但它们的估计成本也更高。google大脑利用这种实现约束来使 V-MoE 在推理时更快。通过将总组合缓冲区容量降低到要处置惩罚的 token 数量以下，搜集被迫跳过处置惩罚大师层中的一些 token。该模型不是以某种任意方式选择要跳过的 token（就像以前的工作那样），而是学习根据重要性分数对它们进行排序。这样可以保持高质量的预测，同时节省大量估计。他们将这种方法称为批量优先级路由（Batch Priority Routing, BPR） ，动态示意图如下所示：

在高容量下，Vanilla 和优先路由都可以很好地处置惩罚所有 patch。但是，当减小缓冲区大小以节省估计时，Vanilla 路由选择处置惩罚任意 patch，通常导致预测不佳；BPR 智能地优先选择处置惩罚重要 patch，使得以更低的估计成本获得更佳的预测。事实证明，适当地删除 token 对于提供高质量和更有效的推理预测至关重要。当大师容量减少时，Vanilla 路由机制的性能会迅速下降。相反，BPR 对低容量更为稳健。总体而言，google大脑观察发现，V-MoE 在推理时非常灵活：例如，可以减少每一个 token 选择的大师数量以节省时间和估计，而无需对模型权重进行任何进一步的训练。摸索 V-MoE由于关于希罕搜集的内部工作原理还有很多待发现，google大脑还摸索了 V-MoE 的路由模式。一种假设是，路由器会根据某些语义背景（如「汽车」大师、「动物」大师等）学会区分并分配 token 给大师。 为了测试这一点，他们在下面展示了两个不同 MoE 层的图，一个非常早期（very early-on），另一个更靠近 head。x 轴对应 32 个大师中的每一个，y 轴显示图象种别的 ID（从 1 到 1000）。图中每一个条目都显示了为与特定图象类对应的 token 选择大师的频率，颜色越深表示频率越高。结果显示，虽然在早期层几乎没有相关性，但在搜集后期，每一个大师只接收和处置惩罚来自少数几个种别的 token。因此，可以得出结论，patch 的一些语义聚类出现在搜集的更深层。

更高的路由决策与图象种别相关。google大脑相信这只是估计机视觉大规模条件估计的开始。异构大师架构和条件可变长度路由也是有潜力的研讨方向。希罕模型尤其有益于数据丰富的领域，例如大规模视频建模。他们希望开源的代码和模型能够吸引更多研讨人员关注该领域。原文链接：https://ai.googleblog.com/2022/01/scaling-vision-with-sparse-mixture-of.html?continueFlag=b96fa8ed72dfc82b777e51b7e954c7dc