[论文解读] Mobile V-MoEs: Scaling Down Vision Transformers via Sparse Mixture-of-Experts

Erik Daxberger, Floris Weers+8 more|arXiv (Cornell University)|Sep 8, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用 3

一句话总结

本文提出 Mobile V-MoEs，一种稀疏的专家混合（MoE）架构，该架构将整个图像而非单个图像块路由至视觉Transformer（ViTs）中的专用专家，从而实现高效的条件计算。通过使用语义超类来稳定路由器训练，该方法在性能与效率的权衡上达到当前最优，相较于密集ViTs，在ImageNet-1k上将ViT-Tiny的性能提升了3.39%，将参数量更小的54M FLOP模型性能提升了4.66%。

ABSTRACT

Sparse Mixture-of-Experts models (MoEs) have recently gained popularity due to their ability to decouple model size from inference efficiency by only activating a small subset of the model parameters for any given input token. As such, sparse MoEs have enabled unprecedented scalability, resulting in tremendous successes across domains such as natural language processing and computer vision. In this work, we instead explore the use of sparse MoEs to scale-down Vision Transformers (ViTs) to make them more attractive for resource-constrained vision applications. To this end, we propose a simplified and mobile-friendly MoE design where entire images rather than individual patches are routed to the experts. We also propose a stable MoE training procedure that uses super-class information to guide the router. We empirically show that our sparse Mobile Vision MoEs (V-MoEs) can achieve a better trade-off between performance and efficiency than the corresponding dense ViTs. For example, for the ViT-Tiny model, our Mobile V-MoE outperforms its dense counterpart by 3.39% on ImageNet-1k. For an even smaller ViT variant with only 54M FLOPs inference cost, our MoE achieves an improvement of 4.66%.

研究动机与目标

• 通过缩小视觉Transformer（ViTs）的规模，改进资源受限视觉应用中的性能与效率权衡。
• 解决现有稀疏MoE设计中对单个图像块进行路由所导致的效率低下问题，此类设计在推理时需要加载大量专家。
• 通过利用语义超类来引导路由器学习，从而在小规模模型中稳定MoE训练并防止专家崩溃。
• 证明在小型ViT模型中，图像级路由相比图像块级路由能实现更高的效率与准确性。

提出的方法

• 提出一种面向移动设备的MoE设计，其中单个路由器将整个输入图像分配给专家，而非对单个图像块进行路由。
• 采用基于softmax的路由函数并结合top-k选择机制，每张图像仅激活E个专家中的k个。
• 提出一种稳定的训练流程，利用语义超类作为监督信号来引导路由器，降低训练不稳定性。
• 采用混合架构，包含L个MoE-ViT层后接密集ViT层，实现在保持模型容量的同时实现选择性计算。
• 通过在超类预测上施加辅助损失来实现负载均衡，防止训练过程中专家分布不均。
• 端到端联合训练路由器与模型其余部分，但引入超类监督以提升收敛性并防止专家崩溃。

Figure 1 : Accuracy vs. FLOPs for ViTs of different sizes. Labels (e.g. 12 $\times$ 192, which is ViT-Tiny) refer to the number of ViT layers (e.g. 12) and the hidden embedding dimension (e.g. 192). The sparse MoEs outperform their corresponding dense baselines across different model scales. Fig. 3(

实验结果

研究问题

• RQ1 稀疏MoE能否有效用于在资源受限的移动视觉应用中缩小ViTs规模？
• RQ2 在小型ViT模型中，图像级路由是否在推理效率与准确性方面优于图像块级路由？
• RQ3 语义超类监督能否稳定小规模模型中的MoE训练并防止专家崩溃？
• RQ4 小型ViT模型中MoE超参数（如专家数量、MoE层深度、top-k路由）的最佳配置是什么？
• RQ5 Mobile V-MoEs在性能与效率权衡方面相较于密集ViTs和标准MoE-ViTs的表现如何？

主要发现

• 对于ViT-Tiny（12×192），Mobile V-MoE在ImageNet-1k上的top-1准确率相比密集ViT基线提升了3.39%。
• 对于仅含54M FLOPs的更小ViT变体，Mobile V-MoE相比其密集基线实现了4.66%的准确率提升。
• 最优专家数量E为10，超过此数量后性能趋于停滞，且由于分类难度增加，路由器准确率下降。
• 当L=2个MoE层时达到最佳模型性能，因为更深的MoE使用会因特征信息减少而导致路由器准确率下降。
• 结合语义超类的图像级路由优于端到端学习的路由和随机超类基线，尤其在FLOPs效率方面优势显著。
• 每张图像使用k=1或k=2个专家时，性能增益与FLOPs的比值最佳，k值更高则收益递减。

(a) Dense ViT