大家还记得去年6月谷歌大脑团队发布的43页论文《Scaling Vision with Sparse Mixture of Experts》吗？他们推出了历史上最大的视觉模型V-MoE，并取得了接近SOTA的Top-1精度。今天，GoogleBrain开源了用于训练和微调模型的全部代码。在过去的几十年里，深度学习的进步是由几个关键因素驱动的：少量简单灵活的机制、庞大的数据集和更专业的硬件配置。在蛋白质预测等任务中取得了令人印象深刻的结果。然而，大型模型和数据集的使用是以大量计算需求为代价的。最近的研究表明，增强模型的泛化能力和鲁棒性离不开大模型的支持。因此，在训练大型模型时，协调好训练资源的局限性非常重要。一种可能的方法是利用条件计算，不是为单个输入激活整个网络，而是根据不同的输入激活模型的不同部分。这种范式在Google的pathway愿景（一种新的AI解决方案，可以克服现有系统的许多缺点，同时增强其优势）和最近的大规模语言模型研究中得到了强调。它在计算机视觉中还没有得到很好的探索。稀疏门控混合专家网络(MoE)在自然语言处理方面表现出出色的可扩展性。然而在计算机视觉中，几乎所有的高性能网络都是稠密的，即每一个输入都转化为一个参数进行处理。去年6月，谷歌大脑的研究人员提出了V-MoE（VisionMoE），一种基于专家稀疏混合的全新视觉架构。当应用于图像识别时，V-MoE通过在推理时只需要一半的计算来实现最先进的网络性能。此外，该研究提出了路由算法的扩展，该算法在整个批次中优先考虑每个输入的子集，从而实现自适应图像计算。这允许V-MoE在测试时权衡性能和平滑计算。最后，该研究展示了V-MoE扩展视觉模型和训练150亿参数模型的潜力，该模型在ImageNet上达到90.35%。论文地址：https://arxiv.org/pdf/2106.05974.pdf代码地址：https://github.com/google-research/vmoeV-MoE谷歌大脑在ViT的不同变体上构建V-MoE：ViT-S(mall),ViT-B(ase),ViT-L(arge),和ViTH(uge)，具有以下超参数：ViT已被证明在迁移学习设置中具有良好的扩展性，在较少的预训练下计算，更高的准确率比CNN获得。ViT将图像处理为一系列补丁。输入图像首先被分成大小相等的块。这些补丁被线性投影到Transformer的隐藏层。位置嵌入后，patchembedding（token）由Transformer处理。Transformer主要由交替的self-attention和MLP层组成。MLP有两层和GeLU非线性。对于VisionMoE，研究用MoE层替换其中的一个子集，其中每个专家都是一个MLP，如下图所示：为了大规模扩展视觉模型，研究在ViT中使用了一些密集的前馈层体系结构(FFN)被独立FFN（称为专家）的稀疏混合所取代。可学习的路由层为每个单独的令牌选择专家。也就是说，来自同一图像的不同令牌可能会被路由到不同的专家。在总共E位专家（E通常为32）中，每个代币最多只能路由到K（通常为1或2）位专家。这允许在保持每个令牌计算不变的同时缩放模型的大小。下图更详细地显示了V-MoE编码器块的结构。V-MoETransformerEncoderBlock实验结果GoogleBrain首先在大型图像数据集JFT-300M上对模型进行了一次预训练。下图左侧显示了模型在所有尺寸（从小的s/32到巨大的H/14）下的预训练结果。然后，使用新的头部（模型中的最后一层）将模型传输到新的下游任务（例如ImageNet）。他们探索了两种传输设置：在新任务的所有可用示例上微调整个模型，或冻结预训练网络并使用少量示例仅针对新头部进行调整（所谓的少镜头传输）。右下图总结了模型传输到ImageNet的效果，其中每个图像类仅在5张图像上进行训练（称为5-shot传输）。左边是JFT-300M数据集上的Precision@1曲线；右边是ImageNet5-shot的准确率曲线。对于这两种情况，GoogleBrain发现对于给定的训练计算量，稀疏模型明显优于密集模型或更快地达到类似性能。为了探索视觉模型的极限，他们在JFT-300M扩展数据集上训练了一个具有150亿个参数和24个MoE层（来自48个块）的模型。迄今为止最大的视觉模型在ImageNet上达到了90.35的Top-1精度。优先路由在实践中，由于硬件限制，使用动态大小的缓冲区效率不高，因此模型通常为每个专家使用预定义的缓冲区容量。一旦专家变得“满”，超过这个容量的分配代币将被丢弃而不被处理。因此，更高的容量会产生更高的精度，但它们的计算成本也更高。GoogleBrain利用此实现约束使V-MoE在推理时更快。通过将总组合缓冲区容量减少到要处理的令牌数量以下，网络被迫跳过处理专家层中的某些令牌。该模型不是以某种任意方式选择要跳过的标记（如之前的工作），而是学习根据它们的重要性分数对它们进行排名。这在节省大量计算的同时保持了高质量的预测。他们将这种方法称为批量优先级路由(BPR)，动态图如下所示：在高容量下，普通路由和优先级路由都适用于所有补丁。然而，在减少缓冲区大小以节省计算的同时，香草路由处理任意补丁，通常导致预测不佳；BPR智能地确定处理重要补丁的优先级，从而以更低的计算成本实现更好的预测。事实证明，正确删除标记对于提供高质量和更高效的推理预测至关重要。当专家的能力下降时，普通路由机制的性能会迅速下降。相比之下，BPR对低容量更稳健。总体而言，GoogleBrain观察到V-MoE在推理方面非常灵活：例如，可以减少每个标记选择的专家数量以节省时间和计算量，而无需对模型权重进行任何进一步训练。探索V-MoE由于稀疏网络的内部工作还有很多有待发现，GoogleBrain还探索了V-MoE的路由模式。一种假设是路由器学会根据某些语义上下文（例如“汽车”专家、“动物”专家等）来区分和分配令牌给专家。为了测试这一点，他们在下方显示了两个不同的MoE层的图，一个非常早，一个更靠近头部。x轴对应32位专家中的每一位，y轴显示图像类别的ID（从1到1000）。图表中的每个条目都显示了专家被选择为与特定图像类别对应的标记的频率，颜色越深表示频率越高。结果表明，虽然在早期层几乎没有相关性，但每个专家在网络的后期阶段只接收和处理来自几个类别的令牌。因此，可以得出结论，一些补丁的语义聚类发生在网络的更深层。更高的路由决策与图像类别相关联。谷歌大脑认为这只是计算机视觉大规模条件计算的开始。异构专家架构和条件变长路由也是潜在的研究方向。稀疏模型在数据丰富的领域特别有用，例如大规模视频建模。他们希望开源代码和模型能够吸引更多的研究人员进入该领域。
