1: \begin{abstract}

2: 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: \modelaname, which relies on a novel tree-based mechanism. \modelaname~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  \textit{any} sparse gate, e.g., Hash routing, Top-k, DSelect-k, and \modelaname. 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, \modelbname~(\modelaname~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.

3: \end{abstract}

4: