MoE 系列超详细解读 (二)｜ImageNet准确率超过90%？V-MoE：视觉领域的稀疏门控混合专家模型

本文目录

1 V-MoE：视觉领域的稀疏门控混合专家模型
(来自 Google Brain)
1 V-MoE 论文解读
1.1 背景：提高模型容量的新方法：稀疏门控混合专家模型
1.2 V-MoE 的贡献
1.3 V-MoE 的计算方法
1.4 专家的缓冲区容量
1.5 模型具体细节
1.6 通过 Batch Prioritized Routing 丢弃无用的 token
1.7 V-MoE 训练数据
1.8 上游任务实验结果
1.9 Few-Shot 实验结果和完全微调实验结果
1.10 放大 V-MoE

太长不看版

稀疏门控混合专家网络 (Sparsely-gated Mixture of Experts networks, MoE) 这种方法已经在自然语言处理领域中表现出了出色的可扩展性。但是在计算机视觉中，几乎所有性能网络都是 "密集 (Dense) 的"，也就是说，每个输入都由所有的参数来处理。

本文就提出了视觉领域经典的稀疏门控混合专家网络 Vision MoE (V-MoE)，它是 Vision Transformer 的稀疏版本，V-MoE 是一种可扩展的架构，其性能和最大的密集网络适配。在图像识别任务中，V-MoE 与 SOTA 的网络的性能相匹配，同时在推理时只需要不到一半的 FLOPs。作者还对路由算法进行了扩展，对整个批次中每个输入的子集进行优先级排序，从而对每幅图像实现自适应的计算，使得测试时的计算更加平滑。

1 V-MoE：视觉领域的稀疏门控混合专家模型

论文名称：Scaling Vision with Sparse Mixture of Experts

论文地址：

https//arxiv.org/pdf/2106.05974.pdf

• 1 V-MoE 论文解读：

1.1 背景：提高模型容量的新方法：稀疏门控混合专家模型

在深度学习的实践经验中，增加网络容量和数据集大小通常会提高模型的性能，这种方法在 NLP 领域已经取得了成功，比如各种各样的 Transformer。

但是，训练以及部署这样的模型需要很大的代价，十分昂贵。一方面的原因是这些深度神经网络通常是 "密集 (Dense)" 的，这个含义是：对于任何单个输入，都需要模型的全部参数来处理。因此，扩大模型的规模往往意味着高昂的计算代价。

那么，能不能有一种方法，可以在提高模型尺寸的同时，保持其训练和推理的成本大致恒定？

稀疏门控混合专家网络 (Sparsely-gated Mixture of Experts networks) 这种方法就是为了实现这个目的，其已经在自然语言处理领域中表现出了出色的可扩展性，能够在解锁万亿参数的模型时以更少的资源实现训练和推理。

本文把这种方法应用到了 CV 领域中，引入的模型称之为 V-MoE，是一种 Vision Transformer (ViT) 架构的稀疏变体。V-MoE 用稀疏的 MoE 层取代了 ViT 中密集的 MLP (又称 Expert)。每个输入图片被 "路由" 到不同 Expert 的子集中。但是，MoE 这种技术路线有其独特的劣势：即不可微性 (non-differentiability)，使得训练这样的模型本身就具有挑战性。

作者在本文探索了多种设计的范式和思路，并为 V-MoE 的预训练和微调提出了一种行之有效的策略，明显优于其他 Dense 的模型。而且，V-MoE 模型还很灵活，我们可以自由调节已经训练好的模型的稀疏度来自由改变其性能和推理成本的 trade-off。使用了 V-MoE 技术之后，我们可以把一个 ViT 视觉模型扩展到 15B 参数大小，这是迄今为止最大的视觉模型，其性能也和 SOTA 的 Dense 模型相当，同时需要更少的时间进行训练。

1.2 V-MoE 的贡献

V-MoE 的贡献可以概括为：

• 超大规模的视觉模型：

稀疏门控混合专家网络是一种用于视觉的分布式稀疏激活 Transformer 模型，包含 24 个 MoE 层，每层包含 32 个 Expert，总参数量 15B。这个模型可以稳定地训练，无缝微调。

• 性能和推理开销：

V-MoE 在上游、few-shot 和完整微调指标方面大大优于 Dense 的竞争者。此外，在推理时，V-MoE 模型可以调整为：1) 计算量或实际运行时间一半的情况下匹配最大的 Dense 模型性能，或者 2) 成本相当的情况下显著优于最大的 Dense 模型。

• 基于优先级的 Batch 路由算法 (Batch Prioritized Routing)：

V-MoE 提出一种新的基于优先级的路由算法，让 V-MoE 丢弃掉最没用的 Patch，从而对每个图像投入较少的计算。效果是 V-MoE 可以节约掉 20% 的计算开销，同时匹配最大的 Dense 模型性能。

• 一些路由决策可视化分析：

帮助启发更多该领域相关的路由决策算法研究。

1.3 V-MoE 的计算方法

一句话概括：不同的 Expert 模型负责不同的输入的部分。

旨在激活不同输入[5]的网络的不同子集。专家混合模型是一个特定的实例化，其中不同的模型“专家”负责输入空间的不同区域

V-MoE 的架构如下图1所示。V-MoE 把 ViT 的一部分 Block 里面的 MLP 换成了 Sparse MoE。MLP 的表达式为:

每个 Expert 的函数是: ，它们都有相同的架构 ，但是权重都不同 。

的路由函数 是:

式中， 。 函数的功能是除了最大的 个值之外，其余所有的值都置为 0 ， 独立采样 。

1.4 专家的缓冲区容量

定义：

专家的缓冲区容量 (每个 Expert 处理的 token (即图像 Patch) 数)：

Batch Size:

每张图片的 Patch 数:

每个 token 所选专家的数量：

Expert 的总数:

模型的容量:

V-MoE 控制每个 Expert 处理的 token (即图像 Patch) 数不变:

如果路由器为给定的 Expert 分配超过 个 token，则只处理其中的 。剩余的标记并不会完全 "丟失"，因为它们的信息由残差连接保留 (图 1 的顶部图)。此外，如果 ，几个 Expert 尝试处理一个 token。token 永远不会完全丟弃。如果 Expert 被分配少于 个 token，则其其余缓冲区被零填充。

1.5 模型具体细节

V-MoE 一共是来自5种尺寸的 ViT 模型：ViT-S(mall), ViT-B(ase), ViT-L(arge) 和 ViTH(uge)，具体如下图2所示。

有3个主要设计决策会影响模型的成本：

• MoE 层的数量：

本文尝试了两种：每一层都使用 MoE，通常命名为 Every xxx，或者只有最后几层使用 MoE，通常命名为 Last xxx。作者发现，尽管使用较少的 MoE 层减少了模型的参数数量，但它对质量的影响通常很小，并且可以显着加快模型，因为会产生更少的通信开销。

• Expert 的数量:

V-MoE 模型的成本不依赖于 Expert 的总数，而是每个 token 所选择的专家数量 。V-MoE 默认使用 ，而作者发现 Expert 总数 是实验设置中的最佳点。

• 模型容量 :

在上游训练期间，默认设置 以给出少量松弛，而不会显着增加成本。

1.6 通过 Batch Prioritized Routing 丢弃无用的 token

原始的 Routing 算法如下图3的 Algorithm1 所示, 设 为输入图片， 表示第 个 token 和第 个 Expert 的路由权重。当每个 token 所选择的专家数量 时循环 2 次。每次循环时遍历每个 Patch，通过 Router 选择出一个专家 和权重 。如果专家 的容量还没满，就把当前的 Patch 给这个专家处理；如果专家 的容量已经满了，就跳过当前的 Patch ，就没有专家处理它了。

Batch Prioritized Routing 算法如上图3的 Algorithm2 所示，思路就是优先处理最重要的 token (路由权重高的)。为了优先处理最重要的 token，作者提出对这 个 Patch 的重要性进行排序，说白了就是先为每个 Patch 计算出一个重要性分数 ，可以取 。然后按照这个重要性分数对这 个 Patch 做重新排序。再根据排序后的 个 Patch 按照原始的方法去做路由。

1.7 V-MoE 训练数据

V-MoE 在私有数据集 JFT-300M 上做预训练，它包含了约 305M 的训练集和 50000 的验证集。类别数 18291 (每张图像平均 1.89 个标签)。

对于 ImageNet 上的 Few-Shot 实验，作者只在每个类别使用 1, 5, 或者 10 shots 来调整上游模型，在验证集上评估得到的模型。

1.8 上游任务实验结果

V-MoE 的上游任务预训练数据集是 JFT-300M，它是一个多标签数据集，因此作者通过 precision@1 来衡量模型性能。如下图 4(a) 和图5所示是不同 V-MoE 和 ViT 变体相对于总训练时间的质量。图4的结果是 V-MoE last-n，就是最后的 个 Block 使用了 MoE 块 (V-MoE-H 的 ，其他情况为 )。图5是各种配置下的 V-MoE 的实验结果。

可以看到，专家模型在所有模型大小上都提供了显着的收益，FLOPs 只略有增加，但是 Pareto frontier 获得了提升。

1.9 Few-Shot 实验结果和完全微调实验结果

Few-Shot 实验结果：

作者评估了 Few-Shot 线性迁移学习，评估预训练学习到的表征的质量。给定来自新数据集的训练样本 ，使用预训练模型来提取每个图片的表征 ，并再单独训练一个线性回归模型把表征映射到标签 。实验结果如图 3(b) 所示, 是 5-shot 的 ImageNet 实验结果，V 学习到的表征的质量也优于 ViT。图4也有 1-shot，10-shot 的 ImageNet 实验结果。

完全微调实验结果：

完全微调就是在一个新任务上面重新初始化最后的分类头，并且完全微调所有的参数 (包括 Backbone 和分类头)。作者大致遵循 Vision Transformer 的标准微调做法，作者还在微调期间使用辅助损失函数。微调实验有2种配置，如下：

• 使用全部的训练数据微调

在训练策略方面，给每个专家使用 0.1 的 dropout rate，结果如下图6所示。V-MoE 的性能也优于 Dense 的模型。

• 小数据集微调

使用 VTAB 数据集，实验结果如下图7所示。虽然 V-MoE-H/14 的性能相似，但 V-MoE 在 ViT-L/16 级别的模型上表现出了显着的提升，这表明尽管这些模型的规模很大，但它们仍然可以通过少量数据进行微调，而且不需要额外的技巧。

1.10 放大 V-MoE

作为一种放大 ViT 模型的方法，作者自然希望验证下 V-MoE 方法对于超大 ViT 模型的泛化性能。为此，作者增加了模型的大小并使用更大的预训练数据集：JFT-3B 是 JFT-300M 的更大版本，它包含近 3B 图像，并且用 30k 类嘈杂注释。

除此之外，作者还遵循了[1]这个工作的方式，再对 Dense 模型放大的方法做了改进之后得到如下的放大 Sparse Mixture of Expert Layers 模型的方法：

• 低精度：使用 bfloat16 而不是 float32 来存储梯度移动平均。
• 学习率衰减：使用 inverse square root schedule 替换 linear schedule。
• 权重衰减：模型中 kernel weights 的权重衰减率为 0.03，bias 不使用，head kernel 的权重衰减为 3.0。
• 模型的 head：不使用 ViT 的 head，其第1个 token 被选择和使用，使用[2]中基于 Self-Attention 的 head。

作者训练了一个带有 48 个 MoE Block 的 V-MoE 模型，每个 Block 32 Experts，每个 token 被 2 个 Expert 处理，即 。模型一共一共 14.7B 参数，使用 V-MoE-15B 表示。

作者成功训练了 V-MoE-15B，这是迄今为止最大的视觉模型。它在 5-shot ImageNet 上的准确率为 ，令人印象深刻，在完全微调时的准确率更是达到了惊人的 。当时 ImageNet 上当前最先进的技术是 Meta Pseudo-Labelling (MPL)。MPL 使用 ImageNet 伪标记在未标记的 JFT-300M 上训练基于 EffecentNet 的模型，达到了 。

参考

1.Scaling Vision Transformers

2.Scaling Vision Transformers

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT