# DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning

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

## Key facts

- Cited by 115 papers in the map (data through 14 Oct 2026).
- Placed in MoE theory and gating functions, one of 12 directions in Mixture-of-experts.
- Builds on 25 of its 45 references that are in the map.
- First posted on arXiv 7 Jun 2021.

## Placed in

- Mixture-of-experts / [MoE theory and gating functions](https://mlsanity.com/moe/moe-theory-and-gating-functions/)

## Builds on (most cited first, 25 of 25)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [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)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [Dynamic Routing Between Capsules](https://mlsanity.com/paper/1710.09829/) (2017)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [SkipNet: Learning Dynamic Routing in Convolutional Networks](https://mlsanity.com/paper/1711.09485/) (2017)
- [Learning to Explain: An Information-Theoretic Perspective on Model Interpretation](https://mlsanity.com/paper/1802.07814/) (2018)
- [Sparse Sequence-to-Sequence Models](https://mlsanity.com/paper/1905.05702/) (2019)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [AdaShare: Learning What To Share For Efficient Deep Multi-Task Learning](https://mlsanity.com/paper/1911.12423/) (2019)
- [Adaptively Sparse Transformers](https://mlsanity.com/paper/1909.00015/) (2019)
- [Routing Networks: Adaptive Selection of Non-linear Functions for Multi-Task Learning](https://mlsanity.com/paper/1711.01239/) (2017)
- [Reparameterizable Subset Sampling via Continuous Relaxations](https://mlsanity.com/paper/1901.10517/) (2019)
- [Gradient Estimation with Stochastic Softmax Tricks](https://mlsanity.com/paper/2006.08063/) (2020)
- [The Tree Ensemble Layer: Differentiability meets Conditional Computation](https://mlsanity.com/paper/2002.07772/) (2020)
- [Learning Classifiers with Fenchel-Young Losses: Generalized Entropies, Margins, and Algorithms](https://mlsanity.com/paper/1805.09717/) (2018)
- [Decision Forests, Convolutional Networks and the Models in-Between](https://mlsanity.com/paper/1603.01250/) (2016)
- [Flexible Multi-task Networks by Learning Parameter Allocation](https://mlsanity.com/paper/1910.04915/) (2019)
- [Small Towers Make Big Differences](https://mlsanity.com/paper/2008.05808/) (2020)

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

- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [AdaMerging: Adaptive Model Merging for Multi-Task Learning](https://mlsanity.com/paper/2310.02575/) (2023)
- [Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts](https://mlsanity.com/paper/2206.02770/) (2022)
- [UniPELT: A Unified Framework for Parameter-Efficient Language Model Tuning](https://mlsanity.com/paper/2110.07577/) (2021)
- [Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning](https://mlsanity.com/paper/2309.05444/) (2023)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [M3ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design](https://mlsanity.com/paper/2210.14793/) (2022)
- [ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models](https://mlsanity.com/paper/2310.04564/) (2023)
- [Soft Merging of Experts with Adaptive Routing](https://mlsanity.com/paper/2306.03745/) (2023)
- [Scaling Vision-Language Models with Sparse Mixture of Experts](https://mlsanity.com/paper/2303.07226/) (2023)
- [Modular Deep Learning](https://mlsanity.com/paper/2302.11529/) (2023)
- [Mixtures of Experts Unlock Parameter Scaling for Deep RL](https://mlsanity.com/paper/2402.08609/) (2024)
- [Mixture-of-Experts Meets Instruction Tuning:A Winning Combination for Large Language Models](https://mlsanity.com/paper/2305.14705/) (2023)
- [Model Compression and Efficient Inference for Large Language Models: A Survey](https://mlsanity.com/paper/2402.09748/) (2024)
- [Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers](https://mlsanity.com/paper/2303.01610/) (2023)
- [MetaGPT: Merging Large Language Models Using Model Exclusive Task Arithmetic](https://mlsanity.com/paper/2406.11385/) (2024)
- [ZoomNeXt: A Unified Collaborative Pyramid Network for Camouflaged Object Detection](https://mlsanity.com/paper/2310.20208/) (2023)
- [Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning](https://mlsanity.com/paper/2412.12953/) (2024)
- [Mixture of Experts Meets Prompt-Based Continual Learning](https://mlsanity.com/paper/2405.14124/) (2024)
- [CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models](https://mlsanity.com/paper/2404.08763/) (2024)
- [Towards MoE Deployment: Mitigating Inefficiencies in Mixture-of-Expert (MoE) Inference](https://mlsanity.com/paper/2303.06182/) (2023)
- [Emergent Modularity in Pre-trained Transformers](https://mlsanity.com/paper/2305.18390/) (2023)
- [Demystifying Softmax Gating Function in Gaussian Mixture of Experts](https://mlsanity.com/paper/2305.03288/) (2023)
- [Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts](https://mlsanity.com/paper/2403.09176/) (2024)
- [EvoMoE: An Evolutional Mixture-of-Experts Training Framework via Dense-To-Sparse Gate](https://mlsanity.com/paper/2112.14397/) (2021)
- [Statistical Perspective of Top-K Sparse Softmax Gating Mixture of Experts](https://mlsanity.com/paper/2309.13850/) (2023)
- [Dual-Balancing for Multi-Task Learning](https://mlsanity.com/paper/2308.12029/) (2023)
- [MoE-FFD: Mixture of Experts for Generalized and Parameter-Efficient Face Forgery Detection](https://mlsanity.com/paper/2404.08452/) (2024)

## Abstract (arXiv)

The Mixture-of-Experts (MoE) architecture is showing promising results in improving parameter sharing in multi-task learning (MTL) and in scaling high-capacity neural networks. State-of-the-art MoE models use a trainable sparse gate to select a subset of the experts for each input example. While conceptually appealing, existing sparse gates, such as Top-k, are not smooth. The lack of smoothness can lead to convergence and statistical performance issues when training with gradient-based methods. In this paper, we develop DSelect-k: a continuously differentiable and sparse gate for MoE, based on a novel binary encoding formulation. The gate can be trained using first-order methods, such as stochastic gradient descent, and offers explicit control over the number of experts to select. We demonstrate the effectiveness of DSelect-k on both synthetic and real MTL datasets with up to 128 tasks. Our experiments indicate that DSelect-k can achieve statistically significant improvements in prediction and expert selection over popular MoE gates. Notably, on a real-world, large-scale recommender system, DSelect-k achieves over 22% improvement in predictive performance compared to Top-k. We provide an open-source implementation of DSelect-k.