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

By Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, M. Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong, Ed H. Chi. 2021

204 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Information Systems, Management Science and Operations Research, Model Architectures

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 Information Systems

In Management Science and Operations Research

In Model Architectures

Citations in the corpus

Mixtral of Experts (2024, methodology): “There are multiple alternative ways of implementing G ( x ) [6, 15, 35], but a simple and performant one is implemented by taking the softmax over the Top-K logits of a linear layer [28].”
A Review of Sparse Expert Models in Deep Learning (2022, background): “Hazimeh et al. (2021) proposes DSelect-k, which is a smooth version of the top-k routing algorithm that improves over standard top-k routing.”
Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts (2022, background): “Many works aim to improve the gating mechanism itself, by making it differentiable [38], reformulating as a linear assignment task [39] or even swapping it out for a simple hashing algorithm [40].”
Dense-to-Sparse Gate for Mixture-of-Experts (2021, methodology): “DSelectK Hazimeh et al. (2021) also utilizes the dense-to-sparse techniques to train the multi-gates MoE in the multi-task learning scenario.”
A Survey on Mixture of Experts in Large Language Models (2024, unknown): “[73]introduceDSelect-k, whichisasmoothversionofthetop-k gating algorithm that improves over standard top-k gating.”
UniPELT: A Unified Framework for Parameter-Efficient Language Model Tuning (2021, background): “One notable example is Mixtureof-Experts (MoE) (Shazeer et al., 2017; Hazimeh et al., 2021), which maintains a set of experts (neural networks) and one or more trainable gates to select a combination of the experts that is specific to each input example.”
ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, background): “Hazimeh et al. (2021) proposed a continuously differentiable sparse gate with demonstrated improvements over vanilla top-k gating.”
Designing Effective Sparse Expert Models (2022, background): “Hazimeh et al. (2021) proposed a continuously differentiable sparse gate with demonstrated improvements over vanilla top-k gating.”
Mixture-of-Experts in Remote Sensing: A Survey (2026, unknown): “DSelect-k [44] replaces the hard top-k operator with a continuous relaxation based on a differentiable selection mechanism, enabling end-to-end gradient-based learning while still producing sparse expert assignments.”