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

1 Introduction . The Mixture of Experts ( MoE ) [ 14 ] is the basis of many state-of-the-art deep learning models . For example , MoE-based layers are being used to perform efficient computation in high-capacity neural networks and to improve parameter sharing in multi-task learning ( MTL ) [ 32 , 22 , 21 ] . In its simplest form , a MoE consists of a set of experts ( neural networks ) and a trainable gate . The gate assigns weights to the experts on a per-example basis , and the MoE outputs a weighted combination of the experts . This per-example weighting mechanism allows experts to specialize in different partitions of the input space , which has the potential to improve predictive performance and interpretability . In Figure 1 ( left ) , we show an example of a simple MoE architecture that can be used as a standalone learner or as a layer in a neural network . The literature on the MoE has traditionally focused on softmax-based gates , in which all experts are assigned nonzero weights [ 17 ] . To enhance the computational efficiency and interpretability of MoE models , recent works use sparse gates that assign nonzero weights to only a small subset of the experts [ 1 , 32 , 28 , 21 ] . Existing sparse gates are not differentiable , and reinforcement learning algorithms are commonly used for training [ 1 , 28 ] . In an exciting work , [ 32 ] introduced a new sparse gate ( Top-k gate ) and proposed training it using stochastic gradient descent ( SGD ) . The ability to train the gate using SGD is appealing because it enables end-to-end training . However , the Top-k gate is not continuous , which can lead to convergence issues in SGD that affect statistical performance ( as we demonstrate in our experiments ) . 1https : //github.com/google-research/google-research/tree/master/dselect_k_moe 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we introduce DSelect-k : a continuously differentiable and sparse gate for MoE . Given a user-specified parameter k , the gate selects at most k out of the n experts . This explicit control over sparsity leads to a cardinality-constrained optimization problem , which is computationally challenging . To circumvent this challenge , we propose a novel , unconstrained reformulation that is equivalent to the original problem . The reformulated problem uses a binary encoding scheme to implicitly enforce the cardinality constraint . We demonstrate that by carefully smoothing the binary encoding variables , the reformulated problem can be effectively optimized using first-order methods such as SGD . DSelect-k has a unique advantage over existing methods in terms of compactness and computational efficiency . The number of parameters used by DSelect-k is logarithmic in the number of experts , as opposed to linear in existing gates such as Top-k . Moreover , DSelect-k ’ s output can be computed efficiently via a simple , closed-form expression . In contrast , state-of-the-art differentiable methods for stochastic k-subset selection and Top-k relaxations , such as [ 25 , 39 ] 2 , require solving an optimization subproblem ( for each input example ) to compute the gate ’ s output . DSelect-k supports two gating mechanisms : per-example and static . Per-example gating is the classical gating technique used in MoE models , in which the weights assigned to the experts are a function of the input example [ 14 , 32 ] . In static gating , a subset of experts is selected and the corresponding weights do not depend on the input [ 28 ] . Based on our experiments , each gating mechanism can outperform the other in certain settings . Thus , we study both mechanisms and advocate for experimenting with each . MTL is an important area where MoE models in general , and our gate in particular , can be useful . The goal of MTL is to learn multiple tasks simultaneously by using a shared model . Compared to the usual single task learning , MTL can achieve better generalization performance through exploiting task relationships [ 4 ] . One key problem in MTL is how to share model parameters between tasks [ 30 ] . For instance , sharing parameters between unrelated tasks can potentially degrade performance . The multi-gate MoE [ 22 ] is a flexible architecture that allows for learning what to share between tasks . Figure 1 ( right ) shows an example of a multi-gate MoE , in the simple case of two tasks . Here , each task has its own gate that adaptively controls the extent of parameter sharing . In our experiments , we study the effectiveness of DSelect-k in the context of the multi-gate MoE . Contributions : On a high-level , our main contribution is DSelect-k : a new continuously differentiable and sparse gate for MoE , which can be directly trained using first-order methods . Our technical contributions can be summarized as follows . ( i ) The gate selects ( at most ) k out of the n experts , where k is a user-specified parameter . This leads to a challenging , cardinality-constrained optimization problem . To deal with this challenge , we develop a novel , unconstrained reformulation , and we prove that it is equivalent to the original problem . The reformulation uses a binary encoding scheme that implicitly imposes the cardinality constraint using learnable binary codes . ( ii ) To make the unconstrained reformulation smooth , we relax and smooth the binary variables . We demonstrate that , with careful initialization and regularization , the resulting problem can be optimized with first-order methods such as SGD . ( iii ) We carry out a series of experiments on synthetic and real MTL datasets , which show that our gate is competitive with state-of-the-art gates in terms of parameter sharing and predictive performance . ( iv ) We provide an open-source implementation of DSelect-k. 2These methods were not designed specifically for the MoE . 1.1 Related Work . MoE and Conditional Computation : Since MoE was introduced by [ 14 ] , an exciting body of work has extended and studied this model , e.g. , see [ 17 , 13 , 16 ] . Recently , MoE-based models are showing success in deep learning . For example , [ 32 ] introduced the sparse Top-k gate for MoE and showed significant computational improvements on machine translation tasks ; we discuss exact connections to this gate in Section 2 . The Top-k gate has also been utilized in several state-of-the-art deep learning models that considered MTL tasks , e.g. , [ 21 , 27 , 9 ] . Our work is also related to the conditional computation models that activate parts of the neural network based on the input [ 2 , 1 , 32 , 12 , 35 ] . Unlike DSelect-k , these works are based on non-differentiable models , or heuristics where the training and inference models are different . Stochastic k-Subset Selection and Top-k Relaxation : A related line of work focuses on stochastic k-subset selection in neural networks , e.g. , see [ 25 , 5 , 38 ] and the references therein . Specifically , these works propose differentiable methods for sampling k-subsets from a categorical distribution , based on extensions or generalizations of the Gumbel-softmax trick [ 23 , 15 ] . However , in the MoE we consider deterministic subset selection—determinism is a common assumption in MoE models that can improve interpretability and allows for efficient implementations [ 14 , 17 , 32 ] . In contrast , the stochastic approaches described above are suitable in applications where there is an underlying sampling distribution , such as in variational inference [ 19 ] . Another related work is the differentiable relaxation of the Top-k operator proposed by [ 39 ] . All the aforementioned works perform dense training ( i.e. , the gradients of all experts , even if not selected , will have to be computed during backpropagation ) , whereas DSelect-k can ( to an extent ) exploit sparsity to speed up training , as we will discuss in Sections 2 and 3 . Moreover , the stochastic k-subset selection framework in [ 25 ] ( which encompasses several previous works ) and the Top-k relaxation in [ 39 ] require solving an optimization subproblem to compute the gate output–each example will require solving a separate subproblem in the per-example gating setting , which can be computationally prohibitive . In contrast , DSelect-k ’ s output is computed efficiently via a closed-form expression . Sparse Transformations to the Simplex : These are sparse variants of the softmax function that can output sparse probability vectors , e.g. , see [ 24 , 26 , 6 , 3 ] . While similar to our work in that they output sparse probability vectors , these transformations can not control the sparsity level precisely as DSelect-k does ( through a cardinality constraint ) . Thus , these transformations may assign some examples or tasks sparse combinations and others dense combinations . MTL : In Appendix A , we review related literature on MTL . 2 Gating in the Mixture of Experts . In this section , we first review the MoE architecture and popular gates , and then discuss how these gates compare to our proposal . We will assume that the inputs to the MoE belong to a space X ⊂ Rp . In its simplest form , the MoE consists of a set of n experts ( neural networks ) fi : X → Ru , i ∈ { 1 , 2 , . . . , n } , and a gate g : X → Rn that assigns weights to the experts . The gate ’ s output is assumed to be a probability vector , i.e. , g ( x ) ≥ 0 and ∑n i=1 g ( x ) i = 1 , for any x ∈ X . Given an example x ∈ X , the corresponding output of the MoE is a weighted combination of the experts : n∑ i=1 fi ( x ) g ( x ) i . ( 1 ) Next , we discuss two popular choices for the gate g ( . ) that can be directly optimized using SGD . Softmax Gate : A classical model for g ( x ) is the softmax gate : σ ( Ax+b ) , where σ ( . ) is the softmax function , A ∈ Rn×p is a trainable weight matrix , and b ∈ Rn is a bias vector [ 17 ] . This gate is dense , in the sense that all experts are assigned nonzero probabilities . Note that static gating ( i.e. , gating which does not depend on the input example ) can be obtained by setting A = 0 . Top-k Gate : This is a sparse variant of the softmax gate that returns a probability vector with only k nonzero entries [ 32 ] . The Top-k gate is defined by σ ( KeepTopK ( Ax+ b ) ) , where for any vector v , KeepTopK ( v ) i : = vi if vi is in the top k elements of v , and KeepTopK ( v ) i : = −∞ otherwise3 . This gate is conceptually appealing since it allows for direct control over the number of experts to select 3To balance the load across experts , [ 32 ] add noise and additional regularizers to the model . and is trained using SGD . Moreover , the Top-k gate supports conditional training : in backpropagation , for each input example , only the gradients of the loss w.r.t . top k elements need to be computed . With a careful implementation , conditional training can lead to computational savings . However , the Top-k gate is not continuous , which implies that the gradient does not exist at certain inputs . This can be problematic when training is done using gradient-based methods . To gain more insight , in Figure 2 ( left ) , we plot the expert weights chosen by the Top-k gate during training with SGD . The results indicate an oscillatory behavior in the output of the Top-k gate , which can be attributed to its discontinuous nature : a small change in the input can lead to “ jumps ” in the output . Comparison with DSelect-k : We develop DSelect-k in Section 3 . Here we present a high-level comparison between DSelect-k and Top-k . Similar to Top-k , DSelect-k can select k out of the n experts and can be trained using gradient-based optimization methods . A major advantage of DSelectk over Top-k is that it is continuously differentiable , which leads to more stable selection of experts during training—see Figure 2 ( right ) . During inference , DSelect-k only needs to evaluate a subset of the experts , which can lead to computational savings . However , DSelect-k supports conditional training only partially . At the start of training , it uses all the available experts , so conditional training is not possible . As we discuss in Section 3 , after a certain point during training , DSelect-k converges to a small subset of the experts , and then conditional training becomes possible . Our experiments indicate that DSelect-k can have a significant edge over Top-k in terms of prediction and expert selection performance , so the full support for conditional training in Top-k seems to come at the expense of statistical performance .

Summary: The submission studies how to design a differentiable module for selecting K experts out of N experts in a sparse MoE model. The proposed method reparametrize the softmax distribution over N the experts as log(N) binary variables. It then proposes a smoothed S-shaped continuous approximation of the hard binary function.