ImageNet准确率超过90%！谷歌大脑开源V-MoE，用稀疏条件计算来训练目前最大的视觉模型！（NeurIPS 2021）

▊ 写在前面

稀疏的专家混合网络 (Sparsely-gated Mixture of Experts networks (MoEs)) 在自然语言处理中显示出出色的可扩展性。然而，在计算机视觉中，几乎所有SOTA网络都是 “密集的”，也就是说，对于每个输入，每个参数都要参与计算 。

在本文中，作者提出了一种视觉MoE (V-MoE)，这是视觉Transformer的稀疏版本，它是可扩展的，并且可以与最大的密集网络性能相近。当应用于图像识别时，V-MoE与最新网络的性能相近，同时在推理时只需要一半的计算开销。

此外，作者提出了路由算法的扩展，该算法可以对整个batch中的每个输入的子集进行优先级排序，从而实现每幅图像的自适应计算。这使得V-MOE能够在测试时能够更好的权衡性能和计算量。最后，作者展示了V-MoE在视觉模型缩放方面的潜力，并在ImageNet上训练了一个15B参数的模型，达到了准确率90.35% 。

▊ 1. 论文和代码地址

Scaling Vision with Sparse Mixture of Experts

论文地址：https://arxiv.org/abs/2106.05974

代码地址：https://github.com/google-research/vmoe

▊ 2. Motivation

深度学习的一系列研究表明，增加网络容量和数据集大小通常会提高性能。在计算机视觉中，在大型数据集上预训练的大型模型通常会达到SOTA的水平。这种方法在自然语言处理(NLP)中取得了更大的成功，在自然语言处理中，大型的预训练模型无处不在，并且在许多任务上表现得非常好。

Text Transformers是迄今为止最大的模型，有些模型的参数超过100B。然而，这种模型的训练和推理成本都是非常昂贵的。这是因为这些深度网络通常是 “密集的”—每个样本都使用每个参数进行处理，因此具有很高的计算成本。

相比之下，条件计算旨在通过仅将参数子集应用于每个样本来增加模型容量，同时保持训练和推理成本大致恒定 。在NLP中，稀疏混合专家(MOE) 越来越受欢迎，能够用更少的资源进行训练和推理，同时解锁万亿个参数的模型。

在这项工作中，作者探索了大规模视觉模型的条件计算，提出了视觉MoE (V-MoE) ，这是最近的ViT结构的稀疏变体，用于图像分类。V-MoE用稀疏MoE层代替ViT中的密集前馈层的子集，其中每个图像patch被 “路由” 到 “专家” 的子集。

由于路由模式的独特性和不可微性，在深度稀疏模型中进行路由是具有挑战性的。作者探索了各种设计选择，并为V-MoE的预训练和迁移学习提供了一个有效的方法，且性能显著优于其密集的网络。作者进一步证明了V-MoE模型是非常灵活的，通过调节输入或模型权重的稀疏水平，可以在推理期间对已经训练好的模型进行性能和计算开销的平衡。

使用V-MoE，作者可以扩展到15B参数的模型大小，这是迄今为止最大的视觉模型，该模型与最先进的密集模型的性能相匹配，同时只需要更少的时间来训练。V-MOE可以与VIT的成本相匹配，同时实现更好的性能。为了控制性能和计算的权衡，作者提出了批量优先路由算法（Batch Prioritized Routing） ，该算法重新利用模型稀疏性来跳过某些patch的计算，从而减少了对无信息图像区域的计算。

▊ 3. 方法

3.1. The Vision Mixture of Experts

3.1.1. Conditional Computation with MoEs

条件计算旨在针对不同的输入，激活网络的不同子集 。专家混合模型是一种特定的实例化，其中不同的模型 “专家” 负责输入空间的不同区域 。MoE的计算可以表示为，其中是输入的数据，是专家计算的函数，是 “路由” 函数，用来决定每个专家的权重。和都由神经网络参数化。

但是按照这个定义，这仍然是一个密集的网络。然而，如果是稀疏的，即模型被限制为仅分配个非零权重，则不需要计算未使用的专家，这样就在训练和推理时降低了模型的计算量。

3.1.2. MoEs for Vision

作者在视觉Transformer (ViT) 的背景下探索稀疏性在视觉中的应用。事实证明，ViT在迁移学习设置中可以很好地扩展，能够在更少的预训练计算下达到比CNN具有更高的准确性。ViT将图像处理为patch序列，输入图像首先被分为大小相等的patch，然后被线性投影为token。添加positional embeddings后，token由Transformer处理，该Transformer主要由交替的自注意和MLP层组成。

MLP具有两层FC结构和GeLU非线性函数：。对于视觉MoE，作者用MoE层替换其中的一个子集，其中每个专家都是MLP。如上图所示。所有专家具有相同的结构，但权重不同。

3.1.3. Routing

对于V-MOE中的每个MoE层，作者使用路由函数，其中是除了最大k个值的元素，将向量的其他元素设置为零的操作；是在上的采样结果。

实验中，作者使用k=1或k=2。在视觉Transformer的上下文中，是网络某一层上的图像token的表示。因此，V-MOE路由patch表示，而不是整个图像。

作者在专家权重的Softmax之后用。这允许模型能够在k=1的情况下进行训练，并且在k>1的情况下执行得更好。最后，作者在激活值中添加了少量标准差为的噪声，从而进一步提高了性能。

3.1.4. Expert’s Buffer Capacity

在训练期间，稀疏模型可能只训练一小部分专家 。这种模式会导致两个问题。首先，统计效率低下 ：在模型专注于单个专家的学习时，该模型并不比密集模型更强大。其次，计算效率低下 : 向专家分配数据的不平衡可能会导致硬件利用率低下。

为了消除不平衡并简化实现，作者将每个专家的缓冲容量（即每个专家处理的token数量）固定，并使用辅助损失函数训练模型，以保持负载平衡。

作者将专家的缓冲容量 ()定义在batch处理中的图像数量 ()，每个图像的token数量 ()，每个token的选定专家数量 ()，专家总数 (和容量比率 ()之上，即：Ž。

如果路由器向给定专家分配了多于个token，则仅处理其中的个token。剩余的token并不完全 “丢掉”，因为它们的信息由剩余的连接保留。另外，如果，则有多位专家会处理每个token，token永远不会完全丢弃。如果为专家分配的token少于个token时，则其缓冲区的其余部分就用零填充。

作者使用容量比率来调整专家的能力。当C>1时，将添加一个空闲容量，以考虑潜在的路由不平衡。当新数据可能来自与上游训练中非常不同的分布时，这通常有助于微调。在C <1的情况下，路由器会被迫忽略一些分配。

3.2. Skipping Tokens with Batch Prioritized Routing

为了提高路由效率，作者提出了一种新的路由算法（BPR），允许模型对重要token进行优先排序。通过同时减少每个专家的容量，丢弃最没用的token。直观地，并非每个patch对于给定图像进行分类都同样重要的，例如，可以删除大多数背景patch，使模型仅关注具有相关实体的patch。

Vanilla Routing

路由函数逐行应用于一个batch的输入。一个batch处理包含N个图像，每个图像由P个token组成；的每一行对应于图像的特定token的D维表示。相应地，表示第t个token和第i个专家的路由权重。在所有路由算法中，对于i <j，每个TOP-i分配都比任何TOP-j分配具有更高的优先级。路由器首先尝试分配所有第i个专家选项，然后再分配第j个选项。普通路由算法如下所示：

Batch Prioritized Routing (BPR)

为了关注“最重要”的token，作者提出了计算每个token的优先级分数，并在分配之前对进行相应排序。根据token的最大路由权重对token进行排序，即。前k个权重的总和用同样的方式计算，即。

作者将路由器输出重用为分配优先级的代理。上图可视化了容量越来越小的Batch Prioritized Routing (BPR)算法的token优先级结果。由于batch中所有图像上的所有token彼此竞争，因此不同的图像可以接收不同的计算量。BPR的算法如下所示：

▊ 4.实验

4.1. Upstream results

从上表可以看出，本文的模型在各个数据集上都表现不错，并且在ImageNet上基于fine-tuning设置达到了90.35%的准确率。

上图显示了不同V-MOE和ViT变体的总训练计算量和时间。

4.2. Linear few-shot results

上图显示，在5-shot ImageNet下的计算量和训练时间的结果。

上表展示了在1，5，10-shot下，本文方法的性能优于其他的方法。

4.3. Full fine-tuning results

上图展示了在ImageNet上进行fine-tuning的结果。

▊ 5. 总结

在本文中，作者使用稀疏条件计算来训练一些迄今为止最大的视觉模型 ，显示了在表示学习和迁移学习方面的显著改进。除了V-MoE之外，作者还提出了批量优先路由（Batch Prioritized Routing ，BPR） ，这利用了模型的稀疏性来引入输入的稀疏性，可以在不进一步调整模型的情况下完成此操作，从而允许重新使用具有稀疏条件计算的训练模型。

作者认为，这只是视觉条件计算的开始；工作的扩展可以包括扩大专家数量，减少对数据的依赖以及改善稀疏模型产生特征的迁移。