Scaling Vision with Sparse Mixture of Experts

By Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, Neil Houlsby. 2021

1,110 citations recorded by the source corpus.

Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are"dense", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-of-the-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Radiology, Nuclear Medicine and Imaging, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Radiology, Nuclear Medicine and Imaging

In Model Architectures

In Efficiency & Compression

Citations in the corpus

A Review of Sparse Expert Models in Deep Learning (2022, background): “2021) introduces an improvement for top-k routing named Batch Prioritized Routing (BPR) for ViT image classiﬁcation models.”
The Efficiency Misnomer (2021, background): “For instance, there seems to be a rising trend towards sparse models (Fedus et al., 2021; Riquelme et al., 2021) which usually have an incredibly large number of trainable parameters but maintain the FLOPs and speed of dense models.”
VATT: Transformers for Multimodal Self-Supervised Learning from Raw Video, Audio and Text (2021, positive): “Such observation encourages further studies on the possibility of utilizing Mixture-of-Experts [79, 28, 76] to increase the model’s capacity for simultaneous multimodal perception.”
Swin Transformer V2: Scaling Up Capacity and Resolution (2021, methodology): “To address the data hungry problem, previous large vision models typically utilize huge labelled data such as JFT-3B [17, 56, 80].”
Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts (2022, methodology): “Figure 3: What necessitates entropy losses? Classic refers to the standard formulation (importance + load losses [1]).”
MVSS-Net: Multi-View Multi-Scale Supervised Networks for Image Manipulation Detection (2021, background): “Compared to ResNet, big vision models with billions of trainable parameters have shown superior performance in natural image classification [52].”
Sparse-MLP: A Fully-MLP Architecture with Conditional Computation (2021, methodology): “According to existing MoE-based models (Riquelme et al. 2021; Xue et al. 2021), the performance is insensitive to λ.”
Time-MoE: Billion-Scale Time Series Foundation Models with Mixture of Experts (2024, positive): “However, sparse networks, such as mixture-of-experts models (Jacobs et al., 1991), which dynamically route inputs to specialized expert networks, have shown comparable or even superior generalization to dense models while being more efficient (Fedus et al., 2022; Riquelme et al., 2021).”
A Survey of Visual Transformers (2021, background): “Nowadays, the most recent supervised Transformers are exploring both the structural combination [40], [59] and scaling laws [39], [211].”
From Sparse to Soft Mixtures of Experts (2023, methodology): “Batch Priority Routing (BPR) (Riquelme et al., 2021) significantly improves the model performance, especially in the case of K = 1 (Appendix F, Table 7).”
Taming Sparsely Activated Transformer with Stochastic Experts (2021, background): “There are other works in different research ﬁelds (e.g., computer vision) that draw different conclusions than ours (Riquelme et al., 2021).”
SimMIM: a Simple Framework for Masked Image Modeling (2021, background): “Such claim is also partly supported by recent works of sparse inference [19] that the recognition accuracy has very little drop after throwing a large portion of image features [23,41,51].”
SalUn: Empowering Machine Unlearning via Gradient-based Weight Saliency in Both Image Classification and Generation (2023, background): “In particular, weight sparsity (Frankle & Carbin, 2018) has gained recognition as an important driver of modularity, leading to improvements in various aspects of ML, including efficiency (Riquelme et al., 2021), interpretability (Wong et al., 2021), and robustness (Chen et al., 2022b).”
Residual Mixture of Experts (2022, methodology): “Thus, to finetune a pretrained non-MoE transformer into an MoE transformer, we need to replace several MLP layers to MoE layers [18, 41].”
Tutel: Adaptive Mixture-of-Experts at Scale (2022, methodology): “It can be seen that 32 and 64 perform the best, which is consistent with that in previous works [9, 36].”
DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (2024, background): “Then, inspired by Riquelme et al. (2021), we drop tokens with the lowest affinity scores on each device until reaching the computational budget.”