还记得谷歌大脑团队去年 6 月份发布的 43 页论文《Scaling Vision with Sparse Mixture of Experts》吗他们推出了史上最大规模的视觉模型 V—MoE，实现了接近 SOTA 的 Top—1 准确率如今，谷歌大脑开源了训练和微调模型的全部代码

在过去几十年里，深度学习的进步是由几个关键因素推动的:少量简单而灵活的机制，大型数据集，更专业的硬件配置，这些技术的进步使得神经网络在图像分类，机器翻译，蛋白质预测等任务中取得令人印象深刻的结果。

可是，大模型以及数据集的使用是以大量计算需求为代价的最近的研究表明，增强模型的泛化能力以及稳健性离不开大模型的支持，因此，在训练大模型的同时协调好与训练资源的限制是非常重要的一种可行的方法是利用条件计算，该方法不是为单个输入激活整个网络，而是根据不同的输入激活模型的不同部分这一范式已经在谷歌提出的 pathway愿景和最近的大型语言模型研究中得到了重视，但在计算机视觉中还没有得到很好的探索

稀疏门控混合专家网络 在自然语言处理中展示了出色的可扩展性可是，在计算机视觉中，几乎所有的高性能网络都是密集的，也就是说，每个输入都会转化为参数进行处理

去年 6 月，来自谷歌大脑的研究者提出了 V—MoE，这是一种基于专家稀疏混合的新视觉架构当应用于图像识别时，V—MoE 在推理时只需要一半的计算量，就能达到先进网络性能此外，该研究还提出了对路由算法的扩展，该算法可以在整个 batch 中对每个输入的子集进行优先级排序，从而实现自适应图像计算这允许 V—MoE 在测试时能够权衡性能和平滑计算最后，该研究展示了 V—MoE 扩展视觉模型的潜力，并训练了一个在 ImageNet 上达到 90.35% 的 150 亿参数模型

V—MoE

谷歌大脑在 ViT 的不同变体上构建 V—MoE:ViT—S，ViT—B，ViT—L 和 ViTH，其超参数如下:

ViT 已被证明在迁移学习设置中具有良好的扩展性，在较少的预训练计算下，比 CNN 获得更高的准确率ViT 将图像处理为一系列 patch，输入图像首先被分成大小相等的 patch，这些 patch 被线性投影到 Transformer 的隐藏层，在位置嵌入后，patch 嵌入由 Transformer 进行处理，该 Transformer 主要由交替的自注意力和 MLP 层组成MLP 有两个层和一个 GeLU 非线性

为了大规模扩展视觉模型，该研究将 ViT 架构中的一些密集前馈层 替换为独立 FFN 的稀疏混合可学习的路由层为每个独立的 token 选择对应的专家也就是说，来自同一图像的不同 token 可能会被路由到不同的专家在总共 E 位专家中，每个 token 最多只能路由到 K位专家这允许扩展模型的大小，同时保持每个 token 计算的恒定下图更详细地显示了 V—MoE 编码器块的结构

V—MoE Transformer 编码器块

实验结果

谷歌大脑首先在大型图像数据集 JFT—300M 上对模型进行一次预训练。

下图左展示了模型在所有大小时的预训练结果然后，使用一个新的 head将模型迁移至新的下游任务他们探索了两种迁移设置:在所有可用的新任务示例上微调整个模型或者冻结预训练网络并使用少量示例仅对新 head 调整

下图右总结了模型迁移至 ImageNet 的效果，其中每个图像类别仅在 5 张图像上训练。

左为 JFT—300M 数据集上的 Precision1 曲线图，右为 ImageNet 5—shot 的准确率曲线图。让我们来看一看，这个用于促进大规模DNN模型开发的高性能MoE库有哪些优势？其表现又如何？。

对于这两种情况，谷歌大脑发现，在给定训练计算量时，稀疏模型显著优于密集模型或者更快地实现相似性能为了探索视觉模型的极限，他们在 JFT—300M 扩展数据集上训练了一个具有 150 亿参数，24 个 MoE 层的模型这个迄今为止最大的视觉模型在 ImageNet 上实现了 90.35 的 Top—1 准确率

优先路由

在实践中，由于硬件限制，使用动态大小的缓冲区效率不高，因此模型通常为每个专家使用预定义的缓冲区容量一旦专家变「满」，超出此容量的分配 token 将被丢弃并不会被处理因此，更高的容量会产生更高的准确性，但它们的计算成本也更高

谷歌大脑利用这种实现约束来使 V—MoE 在推理时更快通过将总组合缓冲区容量降低到要处理的 token 数量以下，网络被迫跳过处理专家层中的一些 token该模型不是以某种任意方式选择要跳过的 token，而是学习根据重要性分数对它们进行排序这样可以保持高质量的预测，同时节省大量计算

在高容量下，Vanilla 和优先路由都可以很好地处理所有 patch但是，当减小缓冲区大小以节省计算时，Vanilla 路由选择处理任意 patch，通常导致预测不佳，BPR 智能地优先选择处理重要 patch，使得以更低的计算成本获得更佳的预测

事实证明，适当地删除 token 对于提供高质量和更有效的推理预测至关重要当专家容量减少时，Vanilla 路由机制的性能会迅速下降相反，BPR 对低容量更为稳健

总体而言，谷歌大脑观察发现，V—MoE 在推理时非常灵活:例如，可以减少每个 token 选择的专家数量以节省时间和计算，而无需对模型权重进行任何进一步的训练。，据微软亚洲研究院发布，作为目前唯一被证明能够将深度学习模型扩展到万亿以上参数的方法，MoE能让模型学习更多信息，并为计算机视觉，语音识别，自然语言处理和机器翻译系统等提供支持。近期，微软亚洲研究院发布了一个高性能MoE库——Tutel，并在8x80GBNVIDIAA100和8x200Gbps的InfiniBand网络的AzureNDmA100v4节点上进行了实验。

探索 V—MoE

由于关于稀疏网络的内部工作原理还有很多待发现，谷歌大脑还探索了 V—MoE 的路由模式一种假设是，路由器会根据某些语义背景学会区分并分配 token 给专家

为了测试这一点，他们在下面展示了两个不同 MoE 层的图，一个非常早期，另一个更靠近 headx 轴对应 32 个专家中的每一个，y 轴显示图像类别的 ID图中每个条目都显示了为与特定图像类对应的 token 选择专家的频率，颜色越深表示频率越高

结果显示，虽然在早期层几乎没有相关性，但在网络后期，每个专家只接收和处理来自少数几个类别的 token因此，可以得出结论，patch 的一些语义聚类出现在网络的更深层

更高的路由决策与图像类别相关。

谷歌大脑相信这只是计算机视觉大规模条件计算的开始异构专家架构和条件可变长度路由也是有潜力的研究方向稀疏模型尤其有益于数据丰富的领域，例如大规模视频建模他们希望开源的代码和模型能够吸引更多研究人员关注该领域

。