Scaling Vision with Sparse Mixture of Experts

Abstract

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.

arXiv comment44 pages, 38 figures

References in corpus (14)

• Gaussian Error Linear Units (GELUs)
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• High-Performance Large-Scale Image Recognition Without Normalization
• CondConv: Conditionally Parameterized Convolutions for Efficient Inference
• A Large-scale Study of Representation Learning with the Visual Task Adaptation Benchmark
• Conditional Computation in Neural Networks for faster models
• Learning Factored Representations in a Deep Mixture of Experts
• Carbon Emissions and Large Neural Network Training
• Deep Mixture of Experts via Shallow Embedding
• Routing Networks and the Challenges of Modular and Compositional Computation
• Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning
• Biased Mixtures Of Experts: Enabling Computer Vision Inference Under Data Transfer Limitations

Cited by in corpus (11)

• Swin Transformer V2: Scaling Up Capacity and Resolution
• A Survey of Visual Transformers
• SimMIM: A Simple Framework for Masked Image Modeling
• Advancing Robust Underwater Acoustic Target Recognition through Multi-task Learning and Multi-Gate Mixture-of-Experts
• Swin2-MoSE: A New Single Image Super-Resolution Model for Remote Sensing
• INTERN: A New Learning Paradigm Towards General Vision
• DS-Net++: Dynamic Weight Slicing for Efficient Inference in CNNs and Transformers
• Go Wider Instead of Deeper
• Sparse MoEs meet Efficient Ensembles
• The Efficiency Misnomer
• Unbiased Gradient Estimation with Balanced Assignments for Mixtures of Experts