# COMET: Learning Cardinality Constrained Mixture of Experts with Trees and Local Search

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

## Key facts

- Cited by 2 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 37 of its 73 references that are in the map.
- First posted on arXiv 5 Jun 2023.

## Placed in

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

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [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)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [SemEval-2017 Task 1: Semantic Textual Similarity - Multilingual and Cross-lingual Focused Evaluation](https://mlsanity.com/paper/1708.00055/) (2017)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Dynamic Routing Between Capsules](https://mlsanity.com/paper/1710.09829/) (2017)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [SkipNet: Learning Dynamic Routing in Convolutional Networks](https://mlsanity.com/paper/1711.09485/) (2017)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [Hash Layers For Large Sparse Models](https://mlsanity.com/paper/2106.04426/) (2021)
- [Distilling a Neural Network Into a Soft Decision Tree](https://mlsanity.com/paper/1711.09784/) (2017)
- [Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts](https://mlsanity.com/paper/2206.02770/) (2022)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Efficient Large Scale Language Modeling with Mixtures of Experts](https://mlsanity.com/paper/2112.10684/) (2021)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning](https://mlsanity.com/paper/2106.03760/) (2021)
- [Taming Sparsely Activated Transformer with Stochastic Experts](https://mlsanity.com/paper/2110.04260/) (2021)
- [Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference](https://mlsanity.com/paper/2110.03742/) (2021)
- [The Tree Ensemble Layer: Differentiability meets Conditional Computation](https://mlsanity.com/paper/2002.07772/) (2020)
- [Fast Best Subset Selection: Coordinate Descent and Local Combinatorial Optimization Algorithms](https://mlsanity.com/paper/1803.01454/) (2018)

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

- [A Survey on Inference Optimization Techniques for Mixture of Experts Models](https://mlsanity.com/paper/2412.14219/) (2024)
- [End-to-end Feature Selection Approach for Learning Skinny Trees](https://mlsanity.com/paper/2310.18542/) (2023)

## Abstract (arXiv)

The sparse Mixture-of-Experts (Sparse-MoE) framework efficiently scales up model capacity in various domains, such as natural language processing and vision. Sparse-MoEs select a subset of the "experts" (thus, only a portion of the overall network) for each input sample using a sparse, trainable gate. Existing sparse gates are prone to convergence and performance issues when training with first-order optimization methods. In this paper, we introduce two improvements to current MoE approaches. First, we propose a new sparse gate: COMET, which relies on a novel tree-based mechanism. COMET is differentiable, can exploit sparsity to speed up computation, and outperforms state-of-the-art gates. Second, due to the challenging combinatorial nature of sparse expert selection, first-order methods are typically prone to low-quality solutions. To deal with this challenge, we propose a novel, permutation-based local search method that can complement first-order methods in training any sparse gate, e.g., Hash routing, Top-k, DSelect-k, and COMET. We show that local search can help networks escape bad initializations or solutions. We performed large-scale experiments on various domains, including recommender systems, vision, and natural language processing. On standard vision and recommender systems benchmarks, COMET+ (COMET with local search) achieves up to 13% improvement in ROC AUC over popular gates, e.g., Hash routing and Top-k, and up to 9% over prior differentiable gates e.g., DSelect-k. When Top-k and Hash gates are combined with local search, we see up to 100× reduction in the budget needed for hyperparameter tuning. Moreover, for language modeling, our approach improves over the state-of-the-art MoEBERT model for distilling BERT on 5/7 GLUE benchmarks as well as SQuAD dataset.