最近，专家混合模型MoE受到了学术界和工业界的广泛关注，其能够对任意输入来激活模型参数中的一小部分来将模型大小与推理效率分离，从而实现模型的轻量化设计。目前MoE已经在自然语言处理和计算机视觉进行了广泛的应用，本文介绍一篇来自Apple的最新工作，在这项工作中，苹果转而探索使用稀疏的MoE来缩小视觉Transformer模型（ViT）的参数规模，使其能够在移动端的推理芯片上更加流畅的运行。为此，本文提出了一种简化且适合移动设备的Mobile V-MoEs模型，将整个图像而不是单个patch路由输入给专家，并且提出了一种更加稳定的MoE训练范式，该范式可以使用超类信息来指导路由过程。作者团队通过大量的实验表明，与对应的密集ViT相比，本文提出的Mobile V-MoE可以在性能和效率之间实现更好的权衡，例如，对于 ViT-Tiny模型，Mobile V-MoE在ImageNet-1k上的性能比其密集模型提高了3.39%。对于推理成本仅为54M FLOPs的更小的ViT版本，本文方法实现了4.66%的改进。
文章链接：
https://arxiv.org/abs/2309.04354
一、引言
稀疏专家混合模型是一种可以将模型大小与推理效率解耦的神经网络加速手段，直观上理解，MoEs[1]是一种可以被划分为多个“专家”模块的神经网络，“专家”模块与一个路由模块联合训练，在MoEs中，每个输入仅由一小部分模型参数处理（又称条件计算）。相比之下，普通的密集模型则会激活每个与输入有关的参数，如下图所示（b）（c）所示，MoE首先使用路由模块从输入图像中选取一些patch，然后再将这些patch送入到专家模块中进行计算。
虽然目前在CV领域，Transformer架构代替CNN架构已经成为一种趋势，但是现有基于ViT架构的MoEs方法仍然无法像卷积结构一样很好的在移动端进行部署，因此，本文作者想使用条件计算来将注意力头的计算量进行缩减，此外提出了一种更加简化且更适合于移动设备的稀疏MoE设计，即首先使用路由模块将整个图像的表征（而不是图像块）直接分配给专家模块，作者还对这一结构设计了一套专门的训练范式，引入了语义超类的概念来指导路由器的训练来避免专家分配不平衡的问题。本文通过广泛的实验表明，所提出的稀疏MoE方法可以达到ViT模型性能与效率之间的全新平衡。
二、 本文方法
2.1 稀疏MoEs
2.2 适用于轻量级ViT的MoEs
先前用在大规模网络中的稀疏MoE使用的都是块级路由，即输入x 是已经划分好的图像patch，这种方式通常需要为每个图像块激活更多的专家参数。因此，每个patch的路由会增加整体路由机制的计算和内存开销，从而降低整体模型的效率。为了使MoEs在轻量级网络上正常工作，本文作者提出直接使用图像级的路由机制，即输入x 为整个图像，以此来减少激活专家参数的数量。
此外，在对稀疏MoE训练时存在一个很明显的缺陷，那就是如果从头开始训练网络端到端的学习路由机制，可能会导致训练过程不稳定和出现专家崩溃现象。所谓专家崩溃，就是由于训练数据集本身的数据分布等因素导致绝大多数的输入图像仅能路由到整体模型参数中的一小部分专家，而其他专家在训练期间都会被忽略。为了解决这一问题，现有的方法是设置一些额外的辅助损失来确保不同专家之间的负载平衡，但是这样又会增加训练过程的复杂性，因此本文作者提出了一种新型的超类路由机制，即先将数据集中的所有类别划分为一组超类来显示训练路由器，使每个专家参数专门研究一个超类。
三、 实验效果
本文的实验在ImageNet-1K数据集上进行，该数据集包含大约128万张训练图像，本文所有的对比方法和模型版本均在该训练集上从头端到端训练，然后在包含5万张图像的验证集上计算top-1识别准确率。
作者通过缩放Transformer总层数（12、9、6）和隐藏层特征维度（384、192、96和64）来控制Mobile V-MoEs与其对应的密集ViT的模型大小。上图展示了本文方法与其对应参数规模的ViT模型的识别准确率对比，可以看到本文提出的Mobile V-MoEs在所有的模型规模上都优于对应的ViT模型。从视觉ViT的基本范式出发，模型内部MLP的嵌入特征维度应是隐藏层特征维度的4倍。
四、总结
目前，在深度学习模型落地部署领域，正在经历着从CNN向视觉ViT过度的大潮流，基于CNN的移动端轻量级网络（如MobileNet）也亟待升级。本文介绍了一种移动端ViT轻量化的最新技术，作者将稀疏MoEs迁移到视觉ViT模型架构中，与其对应的密集ViT相比，稀疏MoE可以实现高效的性能与效率权衡，这使得将更多类型的视觉ViT模型部署到移动端计算设备上成为可能。此外本文作者展望到，如果能将稀疏MoEs技术应用到CNN和视觉ViT结合的算法模型上，可能会得到更好的推理效果。
参考
[1] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy
Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017.
[2] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herv´e J´egou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347–10357. PMLR, 2021.
作者：Seven_
Illustration by IconScout Store from IconScout
快抢沙发！
