COMET: Learning Cardinality Constrained Mixture of Experts with Trees
  and Local Search
Shibal Ibrahim Wenyu Chen Hussein Hazimeh Natalia Ponomareva Zhe Zhao Rahul Mazumder
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\times$ 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.