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

Abstract:
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 gradientbased 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-k1.

heading: 1 Introduction
text: 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.

heading: 1.1 Related Work
text: 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 cannot 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.

heading: 2 Gating in the Mixture of Experts
text: 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.

heading: 3 Differentiable and Sparse Gating
text: In this section, we develop DSelect-k, for both the static and per-example gating settings. First, we introduce the problem setup and notation. To simplify the presentation, we will develop the gate for a single supervised learning task, and we note that the same gate can be used in MTL models. We assume that the task has an input space X ⊂ Rp, an output space Y , and an associated loss function ` : Y × R → R. We denote the set of N training examples by D = {(xi, yi) ∈ X × Y}Ni=1. We consider a learning model defined by the MoE in Equation (1). For simplicity, we assume that the experts are scalar-valued and belong to a class of continuous functionsH. We assume that the number of experts n = 2m for some integer m—in Appendix B.2, we discuss how the gate can be extended to arbitrary n. For convenience, given a non-negative integer i, we denote the set {1, 2, . . . , i} by [i]. In Section 3.1, we develop DSelect-k for static gating setting. Then, in Section 3.2, we generalize it to the per-example setting.

heading: 3.1 DSelect-k for Static Gating
text: Our goal here is to develop a static gate that selects a convex combination of at most k out of the n experts. The output of the gate can be thought of as a probability vector w with at most k nonzero entries, where wi is the weight assigned to the expert fi. A natural way to minimize the empirical risk of the MoE model is by solving the following problem:
min f1,...fn,w
1
N ∑ (x,y)∈D ` ( y, n∑ i=1 fi(x)wi ) (2a)
s.t. ‖w‖0 ≤ k (2b) n∑ i=1 wi = 1, w ≥ 0. (2c)
In the above, the L0 norm of w, ‖w‖0, is equal to the number of nonzero entries in w. Thus, the cardinality constraint (2b) ensures that the gate selects at most k experts. Problem (2) is a combinatorial optimization problem that is not amenable to SGD due to the cardinality constraint (2b) and the simplex constraints in (2c). In what follows of this section, we first transform Problem (2) into an equivalent unconstrained optimization problem, based on a binary encoding scheme. However, the unconstrained problem cannot be directly handled using SGD due to the presence of binary variables. Thus, in a second transformation, we smooth the binary variables, which leads to an optimization problem that is amenable to SGD.
Road map: In Section 3.1.1, we introduce the single expert selector: a construct for choosing 1 out of n experts by using binary encoding. In Section 3.1.2, we leverage the single expert selector to transform Problem (2) into an unconstrained one. Then, in Section 3.1.3, we smooth the unconstrained problem and discuss how SGD can be applied.

heading: 3.1.1 Single Expert Selection using Binary Encoding
text: The single expert selector (selector, for short) is a fundamental construct that we will later use to convert Problem (2) to an unconstrained optimization problem. At a high-level, the single expert selector chooses the index of 1 out of the n experts and returns a one-hot encoding of the choice. For example, in the case of 4 experts, the selector can choose the first expert by returning the binary vector [1 0 0 0]T . Generally, the selector can choose any of the experts, and its choice is determined by a set of binary encoding variables, as we will describe next.
The selector is parameterized by m (recall that m = log2 n) binary variables, z1, z2, . . . , zm, where we view these variables collectively as a binary number: zmzm−1 . . . z1. The integer represented by the latter binary number determines which expert to select. More formally, let l be the integer represented by the binary number zmzm−1 . . . z1. The selector is a function r : Rm → {0, 1}n which maps z := [z1, z2, . . . , zm]T to a one-hot encoding of the integer (l + 1). For example, if all the zi’s are 0, then the selector returns a one-hot encoding of the integer 1. Next, we define the selector r(z). For easier exposition, we start with the special case of 4 experts and then generalize to n experts.
Special case of 4 experts: In this case, the selector uses two binary variables z1 and z2. Let l be the integer represented by the binary number z2z1. Then, the selector should return a one-hot encoding of the integer (l + 1). To achieve this, we define the selector r(z) as follows:
r(z) = [z̄1z̄2, z1z̄2, z̄1z2, z1z2] T (3)
where z̄i := 1− zi. By construction, exactly one entry in r(z) is 1 (specifically, r(z)l+1 = 1) and the rest of the entries are zero. For example, if z1 = z2 = 0, then r(z)1 = 1 and r(z)i = 0, i ∈ {2, 3, 4}. General case of n experts: Here we generalize the selector r(z) to the case of n experts. To aid in the presentation, we make the following definition. For any non-negative integer l, we define B(l) as the set of indices of the nonzero entries in the binary representation of l (where we assume that the least significant bit is indexed by 1). For example, B(0) = ∅, B(1) = {1}, B(2) = {2}, and B(3) = {1, 2}. For every i ∈ [n], we define the i-th entry of r(z) as follows:
r(z)i = ∏
j∈B(i−1)
(zj) ∏
j∈[m]\B(i−1)
(1− zj) (4)
In the above, r(z)i is a product of m binary variables, which is equal to 1 iff the integer (i− 1) is represented by the binary number zmzm−1 . . . z1. Therefore, r(z) returns a one-hot encoding of the index of the selected expert. Note that when n = 4, definitions (3) and (4) are equivalent.

heading: 3.1.2 Multiple Expert Selection via Unconstrained Minimization
text: In this section, we develop a combinatorial gate that allows for transforming Problem (2) into an unconstrained optimization problem. We design this gate by creating k instances of the single expert
selector r(.), and then taking a convex combination of these k instances. More formally, for every i ∈ [k], let z(i) ∈ {0, 1}m be a (learnable) binary vector, so that the output of the i-th instance of the selector is r(z(i)). Let Z be a k ×m matrix whose i-th row is z(i). Moreover, let α ∈ Rk be a vector of learnable parameters. We define the combinatorial gate q as follows:
q(α,Z) = k∑ i=1 σ(α)ir(z (i)),
where we recall that σ(.) is the softmax function. Since for every i ∈ [k], r(z(i)) is a one-hot vector, we have ‖q(α,Z)‖0 ≤ k. Moreover, since the weights of the selectors are obtained using a softmax, we have q(α,Z) ≥ 0 and ∑n i=1 q(α,Z)i = 1. Thus, q(α,Z) has the same interpretation of w in Problem (2), without requiring any constraints. Therefore, we propose replacing w in the objective of Problem (2) with q(α,Z) and removing all the constraints. This replacement leads to an equivalent unconstrained optimization problem, as we state in the next proposition. Proposition 1. Problem (2) is equivalent4 to:
min f1,...fn,α,Z
1
N ∑ (x,y)∈D ` ( y, n∑ i=1 fi(x)q(α,Z)i ) z(i) ∈ {0, 1}m, i ∈ [k]
(5)
The proof of Proposition 1 is in Appendix B.1. Unlike Problem (2), Problem (5) does not involve any constraints, aside from requiring binary variables. However, these binary variables cannot be directly handled using first-order methods. Next, we discuss how to smooth the binary variables in order to obtain a continuous relaxation of Problem (5).

heading: 3.1.3 Smooth Gating
text: In this section, we present a procedure to smooth the binary variables in Problem (5) and discuss how the resulting problem can be optimized using first-order methods. The procedure relies on the smooth-step function, which we define next.
Smooth-step Function: This is a continuously differentiable and S-shaped function, similar in shape to the logistic function. However, unlike the logistic function, the smooth-step function can output 0 and 1 exactly for sufficiently large magnitudes of the input. The smooth-step and logistic functions are depicted in Appendix B.3. More formally, given a non-negative scaling parameter γ, the smooth-step function, S : R→ R, is a cubic piecewise polynomial defined as follows:
S(t) =  0 if t ≤ −γ/2 − 2γ3 t 3 + 32γ t+ 1 2 if − γ/2 ≤ t ≤ γ/2
1 if t ≥ γ/2
The parameter γ controls the width of the fractional region (i.e., the region where the function is strictly between 0 and 1). Note that S(t) is continuously differentiable at all points—this follows since at the boundary points ±γ/2, we have: S′(−γ/2) = S′(γ/2) = 0. This function has been recently used for conditional computation in soft trees [11] and is popular in the computer graphics literature [8, 29].
Smoothing: We obtain DSelect-k from the combinatorial gate q(α,Z) by (i) relaxing every binary variable in Z to be continuous in the range (−∞,+∞), i.e., Z ∈ Rk×m, and (ii) applying the smooth-step function to Z element-wise. Formally, DSelect-k is a function q̃ defined as follows:
q̃(α,Z) := q(α, S(Z)) = k∑ i=1 σ(α)ir ( S(z(i)) ) , (6)
where the matrix S(Z) is obtained by applying S(·) to Z element-wise. Note that q̃(α,Z) is continuously differentiable so it is amenable to first-order methods. If S(Z) is binary, then q̃(α,Z)
4Equivalent means that the two problems have the same optimal objective, and given an optimal solution for one problem, we can construct an optimal solution for the other.
selects at most k experts (this holds since q̃(α,Z) = q(α, S(Z)), and from Section 3.1.2, q selects at most k experts when its encoding matrix is binary). However, when S(Z) has any non-binary entries, then more than k experts can be potentially selected, meaning that the cardinality constraint will not be respected. In what follows, we discuss how the gate can be optimized using first-order methods, while ensuring that S(Z) converges to a binary matrix so that the cardinality constraint is enforced.
We propose using q̃(α,Z) in MoE, which leads to the following optimization problem:
min f1,...fn,α,Z
1
N ∑ (x,y)∈D ` ( y, n∑ i=1 fi(x)q̃(α,Z)i ) . (7)
Problem (7) can be viewed as a continuous relaxation of Problem (5). If the experts are differentiable, then the objective of Problem (7) is differentiable. Thus, we propose optimizing MoE end-to-end using first-order methods. We note that q̃(α,Z) uses (k + k log n) learnable parameters. In contrast, the Top-k and softmax gates (discussed in Section 2) use n parameters. Thus, for relatively small k, our proposal uses a smaller number of parameters. Next, we discuss how DSelect-k’s parameters should be initialized in order to ensure that it is trainable.
Initialization: By the definition of the smooth-step function, if S(Zij) is binary then S′(Zij) = 0, and consequently ∂`∂Zij = 0. This implies that, during optimization, if S(Zij) becomes binary, the variable Zij will not be updated in any subsequent iteration. Thus, we have to be careful about the initialization of Z. For example, if Z is initialized so that S(Z) is a binary matrix then the gate will not be trained. To ensure that the gate is trainable, we initialize each Zij so that 0 < S(Zij) < 1. This way, the Zij’s can have nonzero gradients at the start of optimization.
Accelerating Convergence to Binary Solutions: Recall that we need S(Z) to converge to a binary matrix, in order for the gate q̃ to respect the cardinality constraint (i.e., to select at most k experts). Empirically, we observe that if the optimizer runs for a sufficiently large number of iterations, then S(Z) typically converges to a binary matrix. However, early stopping of the optimizer can be desired in practice for computational and statistical considerations, and this can prevent S(Z) from converging. To encourage faster convergence towards a binary S(Z), we will add an entropy regularizer to Problem (7). The following proposition is needed before we introduce the regularizer. Proposition 2. For any z ∈ Rm, α ∈ Rk, and Z ∈ Rk×m, r(S(z)) and q̃(α,Z) belong to the probability simplex.
The proof of the proposition is in Appendix B.1. Proposition 2 implies that, during training, the output of each single expert selector used by q̃(α,Z), i.e., r(S(z(i))) for i ∈ [k], belongs to the probability simplex. Note that the entropy of each r(S(z(i))) is minimized by any one-hot encoded vector. Thus, for each r(S(z(i))), we add an entropy regularization term that encourages convergence towards one-hot encoded vectors; equivalently, this encourages convergence towards a binary S(Z). Specifically, we solve the following regularized variant of Problem (7):
min f1,...fn,α,Z ∑ (x,y)∈D 1 N ` ( y, n∑ i=1 fi(x)q̃(α,Z)i ) + λΩ(Z)
where Ω(Z) := ∑k i=1 h ( r(S(z(i))) ) and h(.) is the entropy function. The hyperparameter λ is non-negative and controls how fast each selector converges to a one-hot encoding. In our experiments, we tune over a range of λ values. When selecting the best hyperparameters from tuning, we disregard any λ whose corresponding solution does not have a binary S(Z). In Appendix C.3, we report the number of training steps required for S(Z) to converge to a binary matrix, on several real datasets. Other alternatives to ensure that S(Z) converges to a binary matrix are also possible. One alternative is to regularize the entropy of each entry in S(Z) separately. Another alternative is to anneal the parameter γ of the smooth-step function towards zero. Softmax-based Alternative to Binary Encoding: Recall that our proposed selectors in (6), i.e., r ( S(z(i)) ) , i ∈ [k], learn one-hot vectors exactly (by using binary encoding). One practical alternative for learning a one-hot vector is by using a softmax function with temperature annealing. Theoretically, this alternative cannot return a one-hot vector, but after training, the softmax output can be transformed to a one-hot vector using a heuristic (e.g., by taking an argmax). In Appendix C.1, we perform an ablation study in which we replace the selectors in DSelect-k with softmax functions (along with temperature annealing or entropy regularization).

heading: 3.2 DSelect-k for Per-example Gating
text: In this section, we generalize the static version of DSelect-k, q̃(α,Z), to the per-example gating setting. The key idea is to make the gate’s parameters α and Z functions of the input, so that the gate can make decisions on a per-example basis. Note that many functional forms are possible for these parameters. For simplicity and based on our experiments, we choose to make α and Z linear functions of the input example. More formally, let G ∈ Rk×p, W (i) ∈ Rm×p, i ∈ [k], be a set of learnable parameters. Given an input example x ∈ Rp, we set α = Gx and z(i) = W (i)x in q̃(α,Z) (to simplify the presentation, we do not include bias terms). Thus, the per-example version of DSelect-k is a function v defined as follows:
v(G,W, x) = k∑ i=1 σ(Gx)ir ( S(W (i)x) ) .
In the above, the term r ( S(W (i)x) ) represents the i-th single expert selector, whose output depends
on the example x; thus different examples are free to select different experts. The term σ(Gx)i determines the input-dependent weight assigned to the i-th selector. The gate v(G,W, x) is continuously differentiable in the parameters G and W , so we propose optimizing it using first-order methods. Similar to the case of static gating, if S(W (i)x) is binary for all i ∈ [k], then each r ( S(W (i)x) ) will select exactly one expert, and the example x will be assigned to at most k experts.
To encourage S(W (i)x), i ∈ [k] to become binary, we introduce an entropy regularizer, similar in essence to that in static gating. However, unlike static gating, the regularizer here should be on a per-example basis, so that each example respects the cardinality constraint. By Proposition 2, for any i ∈ [k], r ( S(W (i)x) ) belongs to the probability simplex. Thus, for each example x in the training
data, we introduce a regularization term of the form: Ω(W,x) := ∑ i∈[k] h ( r ( S(W (i)x) )) , and minimize the following objective function:∑ (x,y)∈D ( 1 N ` ( y, n∑ i=1 fi(x)v(G,W, x)i ) + λΩ(W,x) ) ,
where λ is a non-negative hyperparameter. Similar to the case of static gating, we tune over a range of λ values, and we only consider the choices of λ that force the average number of selected experts per example to be less than or equal to k. If the application requires that the cardinality constraint be satisfied strictly for every example (not only on average), then annealing γ in the smooth-step function towards zero enforces this.

heading: 4 Experiments
text: We study the performance of DSelect-k in the context of MTL and compare with state-of-the-art gates and baselines. In the rest of this section, we present experiments on the following real MTL datasets: MovieLens, Multi-MNIST, Multi-Fashion MNIST, and on a real-world, large-scale recommender system. Moreover, in Appendix C, we present an additional experiment on synthetic data (with up to 128 tasks), in which we study statistical performance and perform ablation studies.
Competing Methods: We focus on a multi-gate MoE, and study the DSelect-k and Top-k gates in both the static and per-example gating settings. For static gating, we also consider a Gumbelsoftmax based gate [33]–unlike DSelect-k this gate cannot control the sparsity level explicitly (see the supplementary for details). In addition, we consider two MTL baselines. The first baseline is a MoE with a softmax gate (which uses all the available experts). The second is a shared bottom model [4], where all tasks share the same bottom layers, which are in turn connected to task-specific neural nets.
Experimental Setup: All competing models were implemented in TensorFlow 2. We used Adam [18] and Adagrad [7] for optimization, and we tuned the key hyperparameters using random grid search (with an average of 5 trials per grid point). Full details on the setup are in Appendix D.

heading: 4.1 MovieLens
text: Dataset: MovieLens [10] is a movie recommendation dataset containing records for 4,000 movies and 6,000 users. Following [36], for every user-movie pair, we construct two tasks. Task 1 is a binary classification problem for predicting whether the user will watch a particular movie. Task 2
is a regression problem to predict the user’s rating (in {1, 2, . . . , 5}) for a given movie. We use 1.6 million examples for training and 200, 000 for each of the validation and testing sets.
Experimental Details: We use the cross-entropy and squared error losses for tasks 1 and 2, respectively. We optimize a weighted average of the two losses, i.e., the final loss function is α(Loss of Task 1) + (1 − α)(Loss of Task 2), and we report the results for α ∈ {0.1, 0.5, 0.9}. The same loss function is also used for tuning and testing. The architecture consists of a multi-gate MoE with 8 experts, where each of the experts and the task-specific networks is composed of ReLUactivated dense layers. For each α, we tune over the optimization and gate-specific hyperparameters, including the number of experts to select (i.e., k in DSelect-k and Top-k). After tuning, we train each model for 100 repetitions (using random initialization) and report the averaged results. For full details, see Appendix D.1.
Results: In Table 1, we report the test loss and the average number of selected experts. The results indicate that for all values of α, either one of our DSelect-k gates (static or per-example) outperforms the competing methods, in terms of both the test loss and the number of selected experts. In the static gating setting, there does not seem to be a clear winner among the three competing methods (Top-k, DSelect-k, and Gumbel Softmax), but we note that DSelect-k outperforms both Top-k and Gumbel Softmax for two out of the three choices of α. Notably, the softmax MoE is uniformly outperformed by the DSelect-k and Top-k gates, so sparsity in gating seems to be beneficial on this dataset. Our hypothesis is that softmax MoE is overfitting and the sparse gating methods are mitigating this issue. In Table C.2 in the appendix, we additionally report the individual task metrics (loss and accuracy).

heading: 4.2 Multi-MNIST and Multi-Fashion MNIST
text: Datasets: We consider two image classification datasets: Multi-MNIST and Multi-Fashion [31], which are multi-task variants of the MNIST [20] and Fashion MNIST [37] datasets. We construct the Multi-MNIST dataset similar to [31]: uniformly sample two images from MNIST and overlay them on top of each other, and (ii) shift one digit towards the top-left corner and the other digit towards the bottom-right corner (by 4 pixels in each direction). This procedure leads to 36× 36 images with some overlap between the digits. The Multi-Fashion is constructed in a similar way by overlaying images from the Fashion MNIST dataset. For each dataset, we consider two classification tasks: Task 1 is to classify the top-left item and Task 2 is to classify the bottom-right item. We use 100,000 examples for training, and 20,000 examples for each of the validation and testing sets.
Experimental Details: We use cross-entropy loss for each task and optimize the sum of the losses5. The model is a multi-gate MoE with 8 experts, where each expert is a convolutional neural network and each task-specific network is composed of a number of dense layers. We tune the optimization and gate-specific hyperparameters, including the number of experts to select, and use the average of the task accuracies as the tuning metric. After tuning, we train each model for 100 repetitions (using random initialization) and report the averaged results. For full details, see Appendix D.2.
Results: In Table 2, we report the test accuracy and the number of selected experts for the MultiMNIST and Multi-Fashion datasets. On Multi-MNIST, DSelect-k (static) outperforms Top-k and Gumbel Softmax, in terms of both task accuracies and number of selected experts. For example, it achieves over 1% improvement in Task 2’s accuracy compared to Top-k (static). DSelect-k (static) comes close to the performance of the Softmax MoE, but uses less experts (1.7 vs. 8 experts). Here DSelect-k (per-example) does not offer improvement over the static variant (unlike the MovieLens dataset). On Multi-Fashion, we again see that DSelect-k (static) performs best in terms of accuracy.
5Due to the symmetry in the problem, assigning the two tasks equal weights is a reasonable choice.

heading: 4.3 A Large-scale Recommender System
text: We study the performance of DSelect-k and Top-k in a real-word, large-scale content recommendation system. The system encompasses hundreds of millions of unique items and billions of users.
Architecture and Dataset: The system consists of a candidate generator followed by a multi-task ranking model, and it adopts a framework similar to [40, 34]. The ranking model makes predictions for 6 classification and 2 regression tasks. These can be classified into two categories: (i) engagement tasks (e.g., predicting user clicks, bad clicks, engagement time), and (ii) satisfaction tasks (e.g., predicting user satisfaction behaviors such as likes and dislikes). We construct the dataset from the system’s user logs (which contain historical information about the user and labels for the 8 tasks). The dataset consists of billions of examples (we do not report the exact number for confidentiality). We use a random 90/10 split for the training and evaluation sets.
Experimental Details: We use the cross-entropy and squared error losses for the classification and regression tasks, respectively. The ranking model is based on a multi-gate MoE, in which each task uses a separate static gate. The MoE uses 8 experts, each composed of dense layers. For both the DSelect-k and Top-k based models, we tune the learning rate and the experts’ architecture. Then, using the best hyperparameters, we train the final models for 5 repetitions (using random initialization). For additional details, see Appendix D.3.
Results: In Table 3, we report the out-of-sample performance metrics for the 8 tasks. The results indicate that DSelect-k outperforms Top-k on all tasks, with the improvements being most pronounced on the satisfaction tasks. In Figure 3, we show a heatmap of the expert weights chosen by the DSelectk gates. Notably, for DSelect-k, all engagement tasks share at least one expert, and two of the satisfaction tasks share the same expert.

heading: 5 Conclusion
text: We introduced DSelect-k: a continuously differentiable and sparse gate for MoE, which can be trained using first-order methods. Given a user-specified parameter k, the gate selects at most k of the n experts. Such direct control over the sparsity level is typically handled in the literature by adding a cardinality constraint to the optimization problem. One of the key ideas we introduced is a binary encoding scheme that allows for selecting k experts, without requiring any constraints in the optimization problem. We studied the performance of DSelect-k in MTL settings, on both synthetic and real datasets. Our experiments indicate that DSelect-k can achieve significant improvements in prediction and expert selection, compared to state-of-the-art MoE gates and MTL baselines.
Societal Impact: MoE models are used in various applications (as discussed in the introduction). DSelect-k can improve the interpretability and efficiency of MoE models, thus benefiting the underlying applications. We do not see direct, negative societal impacts from our proposal.
Acknowledgements: The research was conducted while Hussein Hazimeh was at Google, and part of the writing was done during his time at MIT. At MIT, Hussein Hazimeh and Rahul Mazumder acknowledge research funding from the Office of Naval Research [Grant ONR-N000141812298].

Review summaries:
This paper presents DSelect-k, a differentiable selection method in MoE. The DSelect-k is designed based on based on a novel binary encoding formulation, it requires to uses all the available experts at the early training stage, and is able to converge to a small subset of the experts so computation can be saved. The DSelect-k is mainly evaluated in the multi-task learning setting on datasets (e.g., in recommendation system).

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.

The authors propose a novel differentiable sparse gate for mixture-of-experts models.  At-most-k selection problem, a cardinality-constrained problem, is reformulated as an unconstrained problem using a binary encoding scheme.  A continuous relaxation of the binary variable can be optimized with first-order methods with careful initialization and regularization. The proposed method outperforms conventional methods in multitask learning settings.