# Mod-Squad: Designing Mixture of Experts As Modular Multi-Task Learners

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Placed in Mixture of LoRA experts, one of 12 directions in Mixture-of-experts.
- Builds on 34 of its 40 references that are in the map.
- First posted on arXiv 15 Dec 2022.

## Placed in

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

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Depth Map Prediction from a Single Image using a Multi-Scale Deep Network](https://mlsanity.com/paper/1406.2283/) (2014)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Once-for-All: Train One Network and Specialize it for Efficient Deployment](https://mlsanity.com/paper/1908.09791/) (2019)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Taskonomy: Disentangling Task Transfer Learning](https://mlsanity.com/paper/1804.08328/) (2018)
- [Rethinking the Value of Network Pruning](https://mlsanity.com/paper/1810.05270/) (2018)
- [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)
- [Rethinking Few-Shot Image Classification: a Good Embedding Is All You Need?](https://mlsanity.com/paper/2003.11539/) (2020)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Soft Filter Pruning for Accelerating Deep Convolutional Neural Networks](https://mlsanity.com/paper/1808.06866/) (2018)
- [Slimmable Neural Networks](https://mlsanity.com/paper/1812.08928/) (2018)
- [Comparing Rewinding and Fine-tuning in Neural Network Pruning](https://mlsanity.com/paper/2003.02389/) (2020)
- [Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts](https://mlsanity.com/paper/2206.02770/) (2022)
- [Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout](https://mlsanity.com/paper/2010.06808/) (2020)
- [UberNet: Training a `Universal' Convolutional Neural Network for Low-, Mid-, and High-Level Vision using Diverse Datasets and Limited Memory](https://mlsanity.com/paper/1609.02132/) (2016)
- [AdaShare: Learning What To Share For Efficient Deep Multi-Task Learning](https://mlsanity.com/paper/1911.12423/) (2019)
- [Attentive Single-Tasking of Multiple Tasks](https://mlsanity.com/paper/1904.08918/) (2019)
- [NDDR-CNN: Layerwise Feature Fusing in Multi-Task CNNs by Neural Discriminative Dimensionality Reduction](https://mlsanity.com/paper/1801.08297/) (2018)
- [Routing Networks: Adaptive Selection of Non-linear Functions for Multi-Task Learning](https://mlsanity.com/paper/1711.01239/) (2017)
- [M3ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design](https://mlsanity.com/paper/2210.14793/) (2022)
- [Branched Multi-Task Networks: Deciding What Layers To Share](https://mlsanity.com/paper/1904.02920/) (2019)
- [Mixture of Attention Heads: Selecting Attention Heads Per Token](https://mlsanity.com/paper/2210.05144/) (2022)

## Built on by (most cited first, 15 of 15)

- [Towards Modular LLMs by Building and Reusing a Library of LoRAs](https://mlsanity.com/paper/2405.11157/) (2024)
- [SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore](https://mlsanity.com/paper/2308.04430/) (2023)
- [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)
- [Sparse Universal Transformer](https://mlsanity.com/paper/2310.07096/) (2023)
- [Configurable Foundation Models: Building LLMs from a Modular Perspective](https://mlsanity.com/paper/2409.02877/) (2024)
- [DiffusionMTL: Learning Multi-Task Denoising Diffusion Model from Partially Annotated Data](https://mlsanity.com/paper/2403.15389/) (2024)
- [ModuleFormer: Modularity Emerges from Mixture-of-Experts](https://mlsanity.com/paper/2306.04640/) (2023)
- [Experts Weights Averaging: A New General Training Scheme for Vision Transformers](https://mlsanity.com/paper/2308.06093/) (2023)
- [Information Flow Control in Machine Learning through Modular Model Architecture](https://mlsanity.com/paper/2306.03235/) (2023)
- [Diffusion Models For Multi-Modal Generative Modeling](https://mlsanity.com/paper/2407.17571/) (2024)
- [Task-Based MoE for Multitask Multilingual Machine Translation](https://mlsanity.com/paper/2308.15772/) (2023)
- [Contrastive Learning and Mixture of Experts Enables Precise Vector Embeddings](https://mlsanity.com/paper/2401.15713/) (2024)
- [RobuMTL: Enhancing Multi-Task Learning Robustness Against Weather Conditions](https://mlsanity.com/paper/2601.10921/) (2026)
- [Learngene Search Across Multiple Datasets for Building Variable-Sized Models](https://mlsanity.com/paper/2605.08209/) (2026)

## Abstract (arXiv)

Optimization in multi-task learning (MTL) is more challenging than single-task learning (STL), as the gradient from different tasks can be contradictory. When tasks are related, it can be beneficial to share some parameters among them (cooperation). However, some tasks require additional parameters with expertise in a specific type of data or discrimination (specialization). To address the MTL challenge, we propose Mod-Squad, a new model that is Modularized into groups of experts (a 'Squad'). This structure allows us to formalize cooperation and specialization as the process of matching experts and tasks. We optimize this matching process during the training of a single model. Specifically, we incorporate mixture of experts (MoE) layers into a transformer model, with a new loss that incorporates the mutual dependence between tasks and experts. As a result, only a small set of experts are activated for each task. This prevents the sharing of the entire backbone model between all tasks, which strengthens the model, especially when the training set size and the number of tasks scale up. More interestingly, for each task, we can extract the small set of experts as a standalone model that maintains the same performance as the large model. Extensive experiments on the Taskonomy dataset with 13 vision tasks and the PASCAL-Context dataset with 5 vision tasks show the superiority of our approach.