# M3ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design

arXiv: https://arxiv.org/abs/2210.14793 · HTML: https://mlsanity.com/paper/2210.14793/

## Key facts

- Cited by 98 papers in the map (data through 14 Oct 2026).
- Placed in Vision mixture-of-experts, one of 12 directions in Mixture-of-experts.
- Builds on 71 of its 117 references that are in the map.
- First posted on arXiv 26 Oct 2022.

## Placed in

- Mixture-of-experts / [Vision mixture-of-experts](https://mlsanity.com/moe/vision-mixture-of-experts/)

## Builds on (most cited first, 30 of 71)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://mlsanity.com/paper/2102.12122/) (2021)
- [Vision Transformers for Dense Prediction](https://mlsanity.com/paper/2103.13413/) (2021)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [LXMERT: Learning Cross-Modality Encoder Representations from Transformers](https://mlsanity.com/paper/1908.07490/) (2019)
- [Point Transformer](https://mlsanity.com/paper/2011.00931/) (2020)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Florence: A New Foundation Model for Computer Vision](https://mlsanity.com/paper/2111.11432/) (2021)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [Image Transformer](https://mlsanity.com/paper/1802.05751/) (2018)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [PCT: Point cloud transformer](https://mlsanity.com/paper/2012.09688/) (2020)
- [Taskonomy: Disentangling Task Transfer Learning](https://mlsanity.com/paper/1804.08328/) (2018)
- [Transformer in Transformer](https://mlsanity.com/paper/2103.00112/) (2021)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [Cross-stitch Networks for Multi-task Learning](https://mlsanity.com/paper/1604.03539/) (2016)

## Built on by (most cited first, 30 of 98)

- [Accelerating Distributed MoE Training and Inference with Lina](https://mlsanity.com/paper/2210.17223/) (2022)
- [Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy](https://mlsanity.com/paper/2310.01334/) (2023)
- [Mixtures of Experts Unlock Parameter Scaling for Deep RL](https://mlsanity.com/paper/2402.08609/) (2024)
- [Graph Mixture of Experts: Learning on Large-Scale Graphs with Explicit Diversity Modeling](https://mlsanity.com/paper/2304.02806/) (2023)
- [Multi-Task Dense Prediction via Mixture of Low-Rank Experts](https://mlsanity.com/paper/2403.17749/) (2024)
- [Sparse Diffusion Policy: A Sparse, Reusable, and Flexible Policy for Robot Learning](https://mlsanity.com/paper/2407.01531/) (2024)
- [A Survey on Inference Optimization Techniques for Mixture of Experts Models](https://mlsanity.com/paper/2412.14219/) (2024)
- [Bridging Remote Sensors with Multisensor Geospatial Foundation Models](https://mlsanity.com/paper/2404.01260/) (2024)
- [TaskExpert: Dynamically Assembling Multi-Task Representations with Memorial Mixture-of-Experts](https://mlsanity.com/paper/2307.15324/) (2023)
- [Mixture of Experts Meets Prompt-Based Continual Learning](https://mlsanity.com/paper/2405.14124/) (2024)
- [SiDA-MoE: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Models](https://mlsanity.com/paper/2310.18859/) (2023)
- [Demystifying Softmax Gating Function in Gaussian Mixture of Experts](https://mlsanity.com/paper/2305.03288/) (2023)
- [Statistical Perspective of Top-K Sparse Softmax Gating Mixture of Experts](https://mlsanity.com/paper/2309.13850/) (2023)
- [Enhancing NeRF akin to Enhancing LLMs: Generalizable NeRF Transformer with Mixture-of-View-Experts](https://mlsanity.com/paper/2308.11793/) (2023)
- [SwapMoE: Serving Off-the-shelf MoE-based Large Language Models with Tunable Memory Budget](https://mlsanity.com/paper/2308.15030/) (2023)
- [Omni-SMoLA: Boosting Generalist Multimodal Models with Soft Mixture of Low-rank Experts](https://mlsanity.com/paper/2312.00968/) (2023)
- [Ten Lessons We Have Learned in the New "Sparseland": A Short Handbook for Sparse Neural Network Researchers](https://mlsanity.com/paper/2302.02596/) (2023)
- [Sigmoid Gating is More Sample Efficient than Softmax Gating in Mixture of Experts](https://mlsanity.com/paper/2405.13997/) (2024)
- [Vision Transformers on the Edge: A Comprehensive Survey of Model Compression and Acceleration Strategies](https://mlsanity.com/paper/2503.02891/) (2025)
- [Edge-MoE: Memory-Efficient Multi-Task Vision Transformer Architecture with Task-level Sparsity via Mixture-of-Experts](https://mlsanity.com/paper/2305.18691/) (2023)
- [A General Theory for Softmax Gating Multinomial Logistic Mixture of Experts](https://mlsanity.com/paper/2310.14188/) (2023)
- [On Least Square Estimation in Softmax Gating Mixture of Experts](https://mlsanity.com/paper/2402.02952/) (2024)
- [A Survey: Collaborative Hardware and Software Design in the Era of Large Language Models](https://mlsanity.com/paper/2410.07265/) (2024)
- [Psychometry: An Omnifit Model for Image Reconstruction from Human Brain Activity](https://mlsanity.com/paper/2403.20022/) (2024)
- [Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design](https://mlsanity.com/paper/2410.19123/) (2024)
- [From Google Gemini to OpenAI Q* (Q-Star): A Survey of Reshaping the Generative Artificial Intelligence (AI) Research Landscape](https://mlsanity.com/paper/2312.10868/) (2023)
- [DiffusionMTL: Learning Multi-Task Denoising Diffusion Model from Partially Annotated Data](https://mlsanity.com/paper/2403.15389/) (2024)
- [Mixture of Experts (MoE): A Big Data Perspective](https://mlsanity.com/paper/2501.16352/) (2025)
- [Selective Task Group Updates for Multi-Task Optimization](https://mlsanity.com/paper/2502.11986/) (2025)
- [Statistical Advantages of Perturbing Cosine Router in Mixture of Experts](https://mlsanity.com/paper/2405.14131/) (2024)

## Abstract (arXiv)

Multi-task learning (MTL) encapsulates multiple learned tasks in a single model and often lets those tasks learn better jointly. However, when deploying MTL onto those real-world systems that are often resource-constrained or latency-sensitive, two prominent challenges arise: (i) during training, simultaneously optimizing all tasks is often difficult due to gradient conflicts across tasks; (ii) at inference, current MTL regimes have to activate nearly the entire model even to just execute a single task. Yet most real systems demand only one or two tasks at each moment, and switch between tasks as needed: therefore such all tasks activated inference is also highly inefficient and non-scalable. In this paper, we present a model-accelerator co-design framework to enable efficient on-device MTL. Our framework, dubbed M3ViT, customizes mixture-of-experts (MoE) layers into a vision transformer (ViT) backbone for MTL, and sparsely activates task-specific experts during training. Then at inference with any task of interest, the same design allows for activating only the task-corresponding sparse expert pathway, instead of the full model. Our new model design is further enhanced by hardware-level innovations, in particular, a novel computation reordering scheme tailored for memory-constrained MTL that achieves zero-overhead switching between tasks and can scale to any number of experts. When executing single-task inference, M3ViT achieves higher accuracies than encoder-focused MTL methods, while significantly reducing 88% inference FLOPs. When implemented on a hardware platform of one Xilinx ZCU104 FPGA, our co-design framework reduces the memory requirement by 2.4 times, while achieving energy efficiency up to 9.23 times higher than a comparable FPGA baseline. Code is available at: https://github.com/VITA-Group/M3ViT.