DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning
Authors: Hussein Hazimeh
Zhe Zhao
Aakanksha Chowdhery
Maheswaran Sathiamoorthy
Yihua Chen
Rahul Mazumder
Lichan Hong
Ed Chi
Venue: NeurIPS 2021 Poster (2021)
Keywords: Mixture of Experts
Sparse Mixture of Experts
Sparsity
Sparse Gate
Subset Selection
Multi-task Learning

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 gradient-based 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-k.

DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning
Hussein Hazimeh', Zhe Zhao', Aakanksha Chowdhery', Maheswaran Sathiamoorthy'
Yihua Chen', Rahul Mazumder", Lichan Lichan Hong'. Ed H. Chi
GGoogle (hazimeh, zhezhao, chowdhery mlogn, ,hchen, lichan, Maasaacuusetts Institute of Technology, rahulmazomit.cm edu
Abstract
The Mixture-of-Experts (MoE) architecture showing promising results in improv- ing parameter sharing in multi-task learning (MTL) and in scaling capacity neural networks. of MoE models use trainable "sparse gate" select 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 gradient- based 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 DSelect-k on both synthetic and real MTL datasets with up to 128 tasks. Our experiments indicate that DSelect-k nificant improvements in prediction and expert selection over popular MoE gates. Notably, real-world, large-scale recommender system, DSelect-k achieves over 22% improvement in predictive performance compared to Top-k. We provide open-source implementation of DSelect-k
achieve statistically sig-
Introduction
The Mixture of Experts (MoE) [4] is the basis of many state-of-the-art deep learning models. For example, Moß-bbased layers are being used perform efficient computation in high-capacity neural networks and to improve parameter sharing in multi-task learning (MTL) 332 22 21. In its simplest form, MoE consists of set of experts (neural networks) and a trainable gate. The gate assigns weights to the experts on per-example basis, and the MoE outputs 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 (left), we show an example of simple MoE architecture that can be used as a standalone learner or as a layer in neural network.
Figure D
The literature on the MoE has traditionally focused on softmax-based gates, in which all experts are assigned nonzero weights [1]. To enhance the computational efficiency and interpretability of MoE models, recent works use sparse gates that assign nonzero weights only a small subset of the experts [ 32 28 21 Existing sparse gates are not differentiable, and reinforcement learning
algorithms are commonly used for training [ 28 In 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 our experiments).
htttp::/ggttiei t/git:bb//tttt... com/google--oo researche (aogg//aarreaadr/r////////aaaaae
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Task 1 NN
Task
Gate
E1
E2
E3
Gate
Gate 2
Input
Input
Figure 1: (Left)): An example of MoE that can be used as a standalone learner or layer in a neural network. Here "Ei" denotes the i-th expert. (Right): A multi-gate MoE for learning two tasks simultaneously. "Task NN" a neural network that generates the output Task i.
In this paper, we introduce DSelect-k: a continuously differentiable and sparse gate for MoE. Given user-specified parameter k; the gate selects most & out of the n experts. This explicit control over sparsity leads to a cardinality-constraiin challenging. To circumvent this challenge, we propose novel, unconstrained reformulation that is equivalent to the original problem. The reformulated problem uses 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 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 simple, closed-form expression. In contrast, state-of the-art differentiable methods for stochastic k-subset selection and Top-k relaxations, such as [2] 39 require solving an optimization subproblem (for each input example) compute the gate's output.
optimization problem, which is computationally
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
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 particular, can be useful. The goal of MTL. is to learn multiple tasks simultaneously by using 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 (rightt) shows an example of a multi-gate MoE, in the simple case two tasks. Here, each task has its own gate that adaptively controls the extent of parameter sharing. In our experiments, study the effectiivnnes of DSelect-k in the context of the multi-gate MoE.
Contributions: On a high-level, our main contribution is DSelect-k: new continuously differ- entiable 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) i out of the experts, where & is user-specified parameter. This leads to a challenging, cardinality-constr optimization problem. To deal with this challenge, we develop 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 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-an gates in terms of parameter sharing and predictive performance. (iv) We provide an open-source implementation of DSelect-k.
'These methods 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 [37 Recently, MoE-based models are showing success in deep learning. For example, [22 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] 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 stochastic s-subet 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 MoF
we consider deterministic subset that can improve interpretability and allows for efficient implementations [14 17 32] In contrast,
is common assumption in MoE models
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, will discuss in Sections [3d 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 sparse variants of the softmax function that can output sparse probability vectors, e.g., see [24] 26 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.
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' C TRP. In its simplest form, the MoF consists of a set of n experts (neural networks) f: : X' - IR"" i E (1, 2, ..., n), and gate g X' -> IR. that assigns weights to the experts. The gate' 's output is assumed to be probability vector, i.e., g(x) 2 O and Ei g(z); = 1, for any E X. Given an example x E X, the corresponding output of the MoE is weighted combination of the experts:
Σ (())(());
Next, we discuss two popular choices for the gate 9(.) that can be directly optimized using SGD,
Softmax Gate: A classical model for g(x) is the softmax gate: (Ax+ where (.) is the softmax function, A E Raxp is trainable weight matrix, and b E R is 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 4=00 0.
Top-k Gate: This 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 a (KeeрTTpK[AAz + 6)). where for any vector KeepTopK u); := v; if o; is in the top k elements of v, and KeepTopK(v); == -o otherwise") This gate is conceptually appealing since it allows for direct control over the number of experts to select
Too balance the load across experts,
B2 add noise and additional regularizers to the model.
3

---

Top-k
Dselect-k
Expert Weight
Expert Welght
03-
0.3
02-
Number of Epochs
Number of Epochs
Figure 2: Expert weights output by Top-k (left) and DSelect-k (right) during training synthetic data generated from MoE, under static gating. Each color represents a separate expert. Here DSelect-k recovers the true experts used by the data-generating model, whereas Top-k does not recover and exhibits oscillatory behavior. See Appendix C.2 for details on the data and setup.
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 done using gradient- based methods. To gain more insight, in Figure (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 attributed its discontinuous nature: 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 high-level comparison between DSelect-k and Top-k. Similar to Top-k, DSelect-k can select & out of the n experts and can be trained using gradient- optimization methods. A major advantage of DSelect- k over Top-k k is that it is continuously differentiable, which leads to more stable selection of experts during training see Figure (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, uses all the available experts, so conditional training is not possible. As we discuss in Section 3 after 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 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.
3 Differentiable and Sparse Gating
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 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' C RR, an output space v. and an associated loss function :: + R. We denote the set of N training examples by = {(x,, ys) E X' V)N We consider a learning model defined by the MoE in Equation (T). For simplicity, we assume that the experts are scalar-valued and belong to a class of continuous functions H. We assume that the number of experts n = 2m for some integer m-iin Appendix [.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, ....,..... i} by [i].
In Section 3.1 we develop DSelect-k for static gating setting. Then, in Section 3.2 we generalize it
the eer-example setting.
3.1 DSelect-k for Static Gating
Our goal here is to develop a static gate that selects convex combination of at most out of the n experts. The output of the gate can be thought of as probability vector w with at most & nonzero entries, where w; is the weight assigned to the expert f;. natural way to minimize the empirical risk of the MoE model is by solving the following problem:

---

min Juunna
Σ
{v..
Σ {((x)ws))
(2a)
N
(2.0)ED
1=1
s.t.
|w0 l0 S k
(2b)
1
Σ
w; w
(2c)
In the above, the Lo norm of w, ||llo. is equal to the number of nonzero entries in w. Thus, the cardinality constraint (2b) ensures that the gate selects most & experts. Problem (2) is
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, an equivalent unconstrained optimization problem, based on a binary encoding scheme. However, the
first transform Problem into
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 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 problem and discuss how SGD can be applied.
we smooth the unconstrained
3.1.1 Single Expert Selection using Binary Encoding
The single expert selector (selector, for short) a fundamental construct that we will later use to convert Problem (2) to unconstrained optimization problem. At high-level, the single expert selector chooses the index of out of the n. experts and returns 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 0j² Generally, the selector can choose any of the experts, and its choice is determined by a binary encoding variables, as we will describe next.
The selector is parameterized by m (recall that m = log, n)) binary variables, 21, 72. .. Zm. where view these variables collectively as binary number: 2m 21. The integer represented by the latter binary number determines which expert to select. More formally. let / be the integer represented by the binary number 2m ... 21. The selector is function r R m + {0, 13"h which maps 2 := [21, .2,..., ; Zm 1" to a one-hot encoding of the integer (1 + 1). For example, if all the 2,'s are 0, then the selector returns one-hot encoding of the integer 1. Next, define the selector r(2). For easier exposition, we start with the special case of experts and then generalize to n experts.
Special case of 4 experts: In this case, the selector uses two binary variables ø1 and ø2. Let be the integer represented by the binary number 2221. Then, the selector should return a one-hot encoding of the integer (1 + 1). To achieve this, we define the selector r(2) as follows:
r(=) = -, 22 12
_122]"
(3)
where E; : 1 - z.. By construction, exactly one entry in r(=) is 1 (specifically, r(2) 1) and the rest of the entries are zero. For example, if 21 = 2n then r(z); = 1 and r(=), = 0, (2, 3, 4).
General case of 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 I, define B(I) as the set of indices of the nonzero entries in the binary representation of (where we assume that the least significant bit indexed by 1). For example, B(0) 0. B(1) = (1). B(2) {2), and B(3)= = {1,2}. For every E [n], we define the i-th entry of r(2) as follows:
r(=), = II
II
(1-2;)
(4)
JEB(i11)
JE(m))B((11)
In the above, r(z); is product of m binary variables, which is equal to 1 iff the integer (i 1) is represented by the binary number 2m Fm-1 ..... 21. Therefore, r(z) returns a one-hot encoding of the 1..
index of the selected expert. Note that when n = 4, definitions and 4 are equivalent.
3.1.2 Multiple Expert Selection via Unconstrained Minimization
In this section, we develop combinatorial gate that allows for transforming Problem 2 into an unconstrained optimization problem. We design this gate by creating & instances of the single expert
5

---

selector r(.), and then taking a convex combination of these l instances. More formally, for every i E [k], let 20() E {0, 1 be a (learnable) binary vector, so that the output of the i-th instance of the selector is r(2()). Let Z be a k x matrix whose i-th row is 2(1), Moreover, let a E R be a vector of learnable parameters. We define the combinatorial gate q as follows:
a(a, (), Σ (a))(())))
where recall that o(.) is the softmax function. Since for every E [k], r(2()) is one-hot vector, we have (q( Z)llo S k. Moreover, since the weights of the selectors are obtained using a softmax, have g(a, 2) 0 and Σi=1 g(a, Z); 1. Thus, g(a, Z) has the same interpretation of in Problem (2. without requiring any constraints. Therefore, we propose replacing w in the objective of
Problem (2) with q(a, 2) and removing all the constraints. This replacement leads to equivalent unconstrained optimization problem, as we state in the next proposition.
Proposition Problem
2 equivalen to:
Σ
(V. Jv.Eno EMle)aaa,
min aa
N
(5)
(aN))D
E (0,1)" (E
The proof of Proposition]] 1 iis in Appendix B. Unlike Problem Problem ( does not involve any constraints, aside from requiring binary variables. However, these binary variables cannot 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).
3.1.3 Smooth Gating
In this section, we present procedure 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.
This is continuously differentiable and S-shaped function, similar in shape to the logistic function. However, unlike the logistic function, the smooth-step function can output and 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:
Smooth-step Function:
if < -2/2 if 7/- 7/2
S(t)=
2
36+ 2
if 1f2
The parameter controls the width of the fractional region (i.c., the region where the function is strictly between 0 and 1). Note that S(t) is continuously differentiable all points-this follows since at the boundary points 17/2, we have: (2(-)/2 (2(
0. This function has been recently used for conditional computation soft trees and is popular in the computer graphics literature 88 (29..
We obtain DSelect-k from the combinatorial gate q(a, Z) by (i) relaxing every binary variable in Z be continuous in the range (-oc, +oo), i.c., E kxm and (ii) applying the smooth-step function to Z element-wise. Formally, DSelect-k is function defined as follows:
Smoothing:
ū(a, :) = :(a. S(2))
Σ (())))((=)))))
6
where the matrix S(Z) is obtained by applying S(.) to Z element-wiss Note that ā(a, 2) is continuously differentiable so it is amenable to first-order methods. If S(Z) is binary, then q(a, Z)
'Equivalent means that the two problems have the same optimal objective, and given optimal solution for problem, construct optimal solution for the other.

---

selects at most & experts (this holds since ÿ(a, Z) g(a, S()), and from Section 3.1.2 selects at most & experts when its encoding matrix 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 binary matrix so that the cardinality constraint is enforced.
We propose using q(a, 2) in MoE, which leads the following optimization problem:
Σ (v. Σ M))) 1=1
min fnm.aa
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(a, 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 our proposal uses a smaller number of parameters. Next, we discuss how DSelect-k's 's parameters should be initialized in order ensure that it is trainable.
Initialization: By the definition of the smooth-step function, if ((zg)) is binary then S'(Z))) = 0. 8F_ O. This implies that, during optimizatio, if ((2;) becomes binary, the
and consequently aZ. variable Zi 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, This way, the Zıy's 's can have nonzero gradients at the start of optimization.
initialize each Zij that 0 < S(2,) < 1.
Accelerating Convergence to Binary Solutions: Recall that we need S(Z) to converge to a binary matrix, in order for the gate I to respect the cardinality constraint (i.e., to select at most & experts). Empirically, we observe that if the optimizer runs for a sufficiently large number of iterations, then (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 E Rmm a E R&, and Z E Rkxm r(S(z)) and q( Z) belong to the probability simplex.
The proof of the proposition is in Appendix B.1 Proposition implies that, during training, the output of each single expert selector used by a(, Z). i.c., r(S(2()) for i [k]. belongs to the probability simplex. Note that the entropy of each r()(2(0)) is minimized by any one-hot encoded vector. Thus, for each (f)), towards one-hot encoded vectors; equivalently, this encourages convergence towards a binary SZZ). Specifically, we solve the following regularized variant of Problem 7
add an entropy regularization term that encourages convergence
Σ
1 N (1,
Σ Maa)aaa, 2)) + AN(Z) 2=1
min
ft.... fu..
(R.ViED
where 0(2)= == Σ ((r(S(=(1)) and h(.) is the entropy function. The hyperparameter À is non-negative and controls how fast each selector converges one-hot encoding. In our experiments, we tume over a range of À values. When selecting the best hyperparameters from tuning, we disregard any À whose corresponding solution does not have 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 ((Z) separately. Another alternative is to anneal the parameter 7 of the smooth-step function towards zero.
Softmax based Alternative to Binary Encoding: Recall that our proposed selectors in 6 i.c., (S) E [k), learn one-hot vectors exactly (by using binary encoding). One practical alternative for learning one-hot vector is by using a softmax function with temperature annealing. Theoretically. this alternative cannot return one-hot vector, but after training, the softmax output can be transformed to one-hot vector using heuristic (e.g., by taking argmax). In Appendix C.I we perform an ablation study in which we replace the selectors in DSelect-k with softmax functions (along with temperature annealing entropy regularization).
7

---

3.2 DSelect-k for Per-example Gating
In this section, we generalize the static version of DSelect-k, a(a, 2), 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 E RExp, W E R E [k]. be a set of learnable parameters. Given an input example æ E IR, we set a = Gæ and z(0) = W a in ((a, Z) (to simplify the presentation, do not include bias terms). Thus, the per-example version of DSelect-k is function v defined as follows:
G(G, W, 2)
Σ ((C)))(((W))))
In the above, the term r(S(W(z)) represents the i-th single expert selector, whose output depends on the example v;; thus different examples are free to select different experts. The term o (Gz); deter- mines the input-dependent weight assigned to the i-th selector. The gate v(G, W,r) x) is continuously differentiable in the parameters G and W, we propose optimizing it using first-order methods. Similar to the case of static gating, if ((W())))) is binary for all E [k]. then each r (S(W((z)) will select exactly one expert, and the example x will be assigned to at most & experts.
To encourage S(W())), i E I:] become binary, introduce an entropy regularizer, similar in essence to that in static gating. However, unlike static gating, the regularizer here should be on per-example basis, so that each example respects the cardinality constraint. By Proposition 2 for any i [k), r (S((()))) belongs the probability simplex. Thus, for each example in the training data, we introduce a regularization term of the form: ((W, = EiEN h ()(S(W(()))) and minimize the following objective function:
1 (v. (V. (.).).....)) + AO(W.
Σ
(
where À is non-negative hyperparameter. Similar to the case of static gating, we tune over range of À values, and we only consider the choices of À that force the average number of selected experts per example to be less than equal to k. If the application requires that the cardinality constraint be satisfied strictly for every example (not only on average), then annealing y in the smooth-step function towards zero enforces this.
4 Experiments
We study the performance of DSelect-k in the context of MTI. and compare with state-of-the-art gates and baselines. the rest of this section, we present experiments on the following real MTI datasets: MovieLens, Multi-MNIST, Multi-Fashion MNIST, and 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 multi-gan 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 Gumbel- softmax based gate [33 -nllue 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 MoF! with softmax gate (which uses all the available experts). The second is a shared bottom model 44, where all tasks share the same bottom layers, which are in turn connected to task- task-specific neural nets.
Experimental Setup: All competing models were implemented in TensorFlow We used Adam 18 and Adagrad 77] 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
4.1 MovieLens
Dataset: MovieLens 10 is 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 binary classification problem for predicting whether the user will watch a particular movie. Task
8

---

is a regression problem to predict the user's rating (in {1, 55) for 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, re- spectively. We optimize a weighted average of the two losses, i.e., the final loss function is a (Loss of Task 1) + (1 a) (Loss of Task 2). and we report the results for a E {0.1, 0.5, 0.9}. The same loss function is also used for tuning and testing. The architecture consists of multi-gate MoE with 8 experts, where each of the experts and the task-specific networks is composed of ReLU- activated dense layers. For each a, 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.I.
Results: In Table 1 we report the test loss and the average number of selected experts. The results indicate that for all values of o, 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 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 a. 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).
a 0.1
a 0.5
a
0.9
Loss
Experts
Loss
Experts
Loss
Experts
DSSelct--k
4015 1
2.7
3804 1 3
1.5
3090 1 2
1.3
Static
Top-k Gumbel Softmax
4012 4
2.0
3818 =
2.0
3693 6
2.0
4171 3
2.7
3898 1
2.6
3688 1 4
3.6
DSelect-k
4006 6
1.5
3823 1
1.2
3679 1
1.1
Per-example Top-k
4027 + 8
2.0
3841
2.0
3741 3
2.0
Soltmax MoE
4090 1
3960
8.0
3847
1 10
8.0
Baselines
Shared Bottom
4037 + 2
3868
3687
Table 1: Test loss (with standard error) and average number of selected experts on MovieLens. The parameter is the weight of Task I's loss (see text for details). The test loss is multiplied by 10
4.2 Multi-MNIST and Multi-Fashion MNIST
Datasets: We consider two image classification datasets: Multi-MNIST and Multi-Fashin BI
which are multi-task variants of the MNIST (20 and Fashion MNIST 337 datasets. We construct the Multi-MNIST dataset similar to [11] 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 36 x 36 images with some overlap between the digits. The Multi-Fashion is constructed in a similar way by overlaying consider two classification tasks:
images from the Fashion MNIST dataset. For each dataset, Task 1 is to classify the top-left item and Task 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.
We use cross-entropy loss for each task and optimize the sum of the losses The model is multi-gate MoE with 8 experts, where each expert is convolutional neural network and each task-specific network is composed of a number of dense layers. We tune the optimization and gate-specifee 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]
Experimental Details:
Results: In Table 2 we report the test accuracy and the number of selected experts for the Multi- MNIST and Multi-Fashin datasets. On MuIti-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- Fastion, we again see that DSelect-k (static) performs best in terms of accuracy.
'Due the symmetry in the problem, assigning the two tasks equal weights is reasonable choice.
9

---

Muhi-MNIST
Mult-Fashion MNIST
Accuracy
Accuracy
Experts
Accuracy
Accuracy
Experts
DSelect-K
92.56 2.003
90.98 0.01
1.7
83,78 = 0.05
83.34 # 0.05
1.5
Statie
Top-k
91.93 +: 0.06
90.03 1: 0.08
4
83.44 = 0.07
82.06 +: 0.08
Gumbel Softmax
92.3 0.05
90.63 0.05
1.8
83.66 0.07
83.28 + 0.05
DSSllct-K
92.42 0.003
90.7 0.03
1.5
83.69 + 0.04
83.13 0.04
1.5
Per-example
Top-k
92.27 +: 0.03
0.04
83.66 0.04
83.15 + 0.01
Softmax MoF!
92.61 + 0.03
91.0 0.03
8
83.48 0.04
82.81 + 0.04
Baselines
Shared Bottom
91.3 3..34:0004
89.47 + 0.04
82.05 1: 0.05
81.37 1: 0.06
Table Test accuracy (with standard error) and number of selected experts on Multi- Salti-MMiST/Fashion.
Tasks Methods
DSelect-k
Top-k
Tegporment
E. Task (AUC)
0,8103
0.0002
0.748I + 0.0198
Ergugeriettsaak
E. Task (AUC)
0.8161
0.0002
0.7624 1. 0.0169
Segaoement tork
E. Task (RMSE)
0.2874
0.0002
0.3406 + 0.0180
Orgaperent task
E. Task (RMSE)
0.8781
0.0014
1.1213 + 0.0842
trangerentted S
E. Task (AUC)
0.7524
0.0003
0.5966 + 0.0529
S. Task 1 (AUC)
0.6133 1 0.0066
0,5080 1 0.0047
S. Task 2 (AUC)
0.8468 = 0.0289
0.5981
+ 0.0616
Sutulldienn teak
S. Task 3 (AUC)
0.9259 + 0.0008
0.6665 + 0.0091
Table 3: Average performance (AUC and RMSE) and standard error on a real-world recommender system with tasks: "E." on the recommender system. and "S." denote engagement und sutisfaction tasks, respectively.
Figure 3: Expert weights of the DSelect-k gates
4.3 A Large-scale Recommender System
We study the performance of DSelect-k and Top-k in a real- yeal-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 candidate generator followed by multi-task
ranking model, and it adopts framework similar to [40] (3). The ranking model makes predictions for 6 classification and 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 likes and dislikes). We construct the dataset from the system's '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 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 repetitions (using random initialization). For additional details, see Appendix D.3.
Results: In Table 3 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 we show a heatmap of the expert weights chosen by the DSelect- k gates. Notably, for DSelect-k, all engagement tasks share at least one expert, and two of the satisfaction tasks share the same expert.
Conclusion
We introduced DSelect-k: a continuously differentiable and sparse gate for MoE, which can be trained using first-order methods. Given user-specified parameter k, the gate selects at most & of the 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-ther -h-te-rr MoE gates and MTL baselines.
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 underly- ing applications. We do not see direct, negative socictal impacts from our proposal.
Societal Impact:
10

---

Acknowledgements:
The research was conducted while Hussein Hazimeh was at Google, and part of the writing was done during his time MIT. At MIT, Hussein Hazimeh and Rahul Mazumder acknowledge research funding from the Office of Naval Research [Grant ONR-N000141812200
References
[1] Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional com- putation in neural networks for faster models. CoRR, abs/1511 Abs//511.062971 1.06297, 2015.
URL http
//arxiv.
06297
[2] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 100:: 343..432, 2013.
[3] Mathieu Blondel, Andre Martins, and Vlad Niculae. Learning classifiers with fenchel-young losses: Generalized entropies, margins, and algorithms. In The 22nd International Conference on Artificial Intelligence and Statisties, pages 606 615. PMLR, 2019.
[4] Rich Caruana. Multitask learning. Machine learning, 28(1):41-755 1997.
[5] Jianbo Chen, Le Song. Martin Wainwright, and Michael Jordan. Learning to explain: An theeoeicio perspective on model interpretation. In International Conference Machine Learning, pages 883-892. PMLR, 2018.
[6] Gonçalo Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 2174-215 184, 2019.
[7] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[8] David S Ebert, Kenton Musgrave, Darwyn Peachey, Ken Perlin, and Steven Worley, Texturing & modeling: procedural approach. Morgan Kaufmann, 2003.
[9] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961 2021.
[10] Maxwell Harper and Joseph A Konstan. The movielens datasets: History and context. Acm transactions interactive intelligent systems (tiis) 5(4):1-19, 2015.
[11] Hussein Hazimeh, Natalia Ponomareva, Petros Mol, Zhenyu Tan, and Rahul Mazumder. The tree ensemble layer: Differentiability meets conditional computation. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 19 of Proceedings of Machine Learning Research, pages 4138-4148, Virtual, 13-18 Jul 2020. PMLR.
[12] Yani Ioannou, Duncan Robertson, Darko Zikic, Peter Kontschieder, Jamie Shotton, Matthew Brown, and Antonio Criminisi. Decision forests, convolutional networks and the models in-between. arXiv preprint arXiv: 012550 2016.
[13] Robert A Jacobs. Bias/variance analyses of mixtures-of-experts architectures. Neural computa- tion, (003669333,, 1997.
[14] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
[15] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with sofbel-softtax. arXiv preprint 1rXiv:1611.01144 2016.
[16] Wenxin Jiang and Martin A Tanner. On the identifiability of mixtures-of-exxerrs. Neural Networks, 1293:1253-1258 1999.
[17] Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2): 1994.

---

[18] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors. 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[19] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. urXiv:1312.0 14, 2013.
arXiv preprint
[20] Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lece lec.m.com//xx///mmsstt. 2, 2010.
[21] Dmitry Lepikhin, HyoukJoong Lee. Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang. Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
[22] Jiaqi Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H Chi. Modeling task relationships in multi-task learning with multi-gate mixture-of-expere In Proceedings of the 24th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pages 1930-1939, 2018.
[23] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv: 1,007 12, 2016.
[24] Andre Martins and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. In International Conference on Machine Learning, pages 1614 1623. PMLR, 2016.
[25] Max Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris Maddison. Gradient estimation with stochastic softmax tricks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33. pages 5691-5704. Curran Associates, Inc., 2020. URL bttps ioo.r neurips caaaper.dd..
cc httppppewebb0o////aaaaaffffaaaae
[26] Ben Peters, Vlad Niculae, and André FT Martins. Sparse sequence-to-sequee models. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1504-1519, 2019.
[27] Prajit Ramachandran and Quoe V Le. Diversity and depth in per-example routing models. In International Conference on Learning Representations, 2018.
[28] Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In International Conference Learning Representations, 2018
[29] Randi J Rost, Bill Linea-Kane, Dan Ginsburg, John Kessenich, Barthold Lichtenbelt, Hugh Malan, and Mike Weiblen. OpenGL shading language. Pearson Education, 2009.
[30] Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint arXiv:1706.05009. 2017.
[31] Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. In Advances neural information processing systems, pages 5856-3866 2017.
[32] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The pparsely-gated mixture-of- experts layer. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings, 2017.
[33] Ximeng Sun, Rameswar Panda, Rogerio Feris, and Kate Saenko. Adashare: Learning what to share for efficient deep multi-task learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 8728-8740. Curran Associates, Inc., 2020. URL / http://pp.......ooa4a..s.ss..srr Papeoe-ppppppppp pdf
12

---

[34] Hongyan Tang, Junning Liu, Ming Zhao, and Xudong Gong, Progressive layered extraction (ple): A novel multi-tast task learning (md) model for personalized recommendations. Fourteenth ACM Conference on Recommender Systems, pages 269-278, 2020.
[35] Xin Wang, Fisher Yu, Zi-Yi Dou, Trevor Darrell, and Joseph E Gonzalez. Skipnet: Learning dynamic routing in convolutional networks. In Proceedings of the European Conference on Computer Vision (ECCV), pages 409-424, 2018.
[36] Yuyan Wang, Zhe Zhao, Bo Dai, Christopher Fifty, Dong Lin, Lichan Hong, and Ed H Chi. Small towers make big differences. arXiv preprini arXiv:2008.0580 05808, 2020.
[37] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint 1XXiv:100.007220 0770800774, 20177.
[38] Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. In International Joint Conference on Artificial Intelligence, 2019.
[39] Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. Advances in Neural Information Processing Systems, 33, 2020.
[40] Zhe Zhao, Lichan Hong, Li Wei, Jilin Chen, Aniruddh Nath, Shawn Andrews, Aditee Kumthekar, Maheswaran Sathiamoorthy, Xinyang Yi, and Ed Chi. Recommending what video to watch multitask ranking system. In Proceedings of the 13th ACM Conference on Recommender Systems, pages 43-51, 2019.
13


Reviews:
[{"review_id": "rS2BcGqp9eO", "reviewer": "Reviewer_1DHt", "summary": "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).", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "* Originality:  MoE has been increasingly important in multi-task learning and large-scale model training. This work attempts to solve the discontinuous and non-smooth issue in the top-k operation in MoE by introducing a differentiable select-k method named DSeleck-k. The DSelect-k method is based on a novel binary encoding formulation. The reviewer want to point out another work about differentiable top-k [1], detailed comments can be found in Q2.\n\n* Quality: The submission is technically sound. The binary encoding formulation is novel and interesting to me. The proposed method is evaluated in the multi-task learning setting (in MovieLens and MNIST).\n\n* Clarity: The paper is well-written.\n\n* Significance:  This work improves the MoE training. It achieves comparable or better performance in the context of multi-task learning. Another important application of MoE, large-scale pre-training, was not covered (See detailed comments in Q1).\n\nI have the following detailed questions:\n\n* Q1: As the author mentioned in related work, one of the most recent applications of MoE is in language model pre-training (like GShard and Switch Transformer), where the FFN in the Transformer model is replaced by a MoE layer to enlarge the number of parameters. E.g., the Switch Transformer has 1.6 trillion parameters. Could the authors discuss the possibility of applying DSelect-k in training huge models like GShard and Switch Transformer. My main concern is that, the proposed DSelect-k has to forward all the available experts at the early training stage, however it is super expensive to fully forward a trillion-parameter model.\n\n* Q2: there has been earlier work [1] on designing differentiable algorithm for top-k operation. Different from this work, [1] views the top-k as the solution to an Optimal Transport problem. I am not an expert in this field, but could the authors discuss the difference between this work and [1].\n\n* Q3: Actually, the experiment on the  large-scale, real-world recommender system looks more interesting to me. I wonder why the authors put them in the appendix (perhaps because of the page limit?).\n\n* Q4: I notice that the submission history was \"hidden as part of the resubmission bias experiment\". Could the authors provide more information about that?\n\n[1] Differentiable Top-k Operator with Optimal Transport, NeurIPS 20. https://arxiv.org/abs/2002.06504\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "I don't think there are any ethical issues with this paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "fX0RHxWJSx", "reviewer": "Reviewer_XzFT", "summary": "This work proposes Dselect-k, a continuously differentiable gate, to select a set of at most k out of n experts for the Mixture-of-experts (MoE) architecture. The Dselect-k based MOE can be trained end-to-end by the first-order optimization method, and it has promising performances on synthetic and different real-world MTL problems. \n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have reviewed this work before for ICML 2021, and now act as an opt-in reviewer with a very tight deadline. So my comments will mainly focus on the improvements/changes this work has made, as well as the unsolved concerns in my previous review.   \n\nStrengths:\n\n+ This paper is well-written and easy to follow;\n\n+ It is glad to see this work now has a clear discussion with the highly related stochastic subset selection methods. But there are still some concerns on the comparison (see below);\n\n+ In addition to predictive performance, the obtained sparse gate can also lead to better model interpretability and computation savings on inference, which is desirable for real-world applications;\n\n+ The proposed method has promising results on different MTL problems, especially the large-scale real-world recommender system application;\n\nConcerns: \n\n1. Comparison with Stochastic Subset Selection: The proposed Dselect-k is closely related to the stochastic subset selection approaches [1] for selecting a k-subset with gradient-based methods. As pointed out in the paper (and the appendix), there are some important differences that make Dselect-k a better choice for MTL-MOE applications considered in this work, such as the deterministic property and better interpretability. I appreciate the thorough discussion provided in the current manuscript, but still have some concerns on the comparison:\n\n    - Experimental Comparison: This paper claims the deterministic/stochastic subset selection approaches are for different applications. However, as discussed in the appendix, there are closely related works that use the stochastic gumbel-softmax trick to select different components for an MTL neural network (as considered in this work) [2,3]. Indeed, this work had an experimental comparison with the method proposed in [3] in the earlier version. Why are those experimental results removed in this version? I still believe a careful experimental comparison with the stochastic subset selection is needed to truly show the pros/cons of the proposed method.\n\n    - Analysis on Determinism and Interpretability: This work has a valid claim that the deterministic approach would be more suitable for applications that require better interpretability. It is understandable that making the simple Gumbel-softmax gate [3] deterministic at inference will lead to dense solutions or possibly all zeros, since the gumbel-softmax gate method has independent predictions for each gate. But is it still the case for the deterministic version of stochastic subset selection [1] where the output is always a k-hot vector?\n\n2. Exact Control: It seems the proposed method can only guarantee selecting *at most* k items but not *precisely* k items as in Top-k and other k-subsect selection methods. \n\n3. Good Performance v.s. Sparsity: It is still unclear why sparsity in gating is also good for task performance. I think the dense gate approach could provide a more robust prediction (kind of ensemble) at the cost of higher computation cost.\n\n4. Experimental Results:\n   - The large-scale real-world recommender system application is now moved to the appendix, and the remaining experiments (MovieLens and Multi-MNIST) are quite toy-ish. I believe the industry-level recommender system application will not be open source. It would be much better to compare the proposed method on the open-source and widely-used MTL problems, such as NYUv2[4], CityScapes[5], and Taskonomy [6].\n\n   - In the MovieLens problem, the two tasks have different losses (classification and regression) in different scales. It is not suitable to only report the combined loss, especially when some of them are very close to each other (e.g., static DSelect-k and top-k). Why not report both the task 1/2 losses, or task 1 accuracy v.s. task 2 mse, as for MultiMNIST?\n\n   - DSelect-k typically selects 1 ~ 2 experts in all experiments. When it selects one expert, it simply reduce to the shared bottom (all tasks share the same expert) or the single-task baseline (each task has its own expert). It would be good to have a single-task baseline for all experiments, which usually has better performance than many MTL baselines in the cost of one model (1 expert) for each task. See results in [7] for MultiMNIST, and MMOE (Ma et al., 2018) for UCI Census-income.\n\n\n[1] Max B. Paulus., et al. Gradient Estimation with Stochastic Softmax Tricks. NeurIPS 2020.\n\n[2] Krzysztof Maziarz., et al. Gumbel-matrix routing for flexible multi-task learning. arXiv preprint arXiv:1910.04915.\n\n[3] Ximeng Sun., et al. Adashare: Learning what to share for efficient deep multi-task learning. NeurIPS 2020.\n\n[4] Nathan Silberman., et al. Indoor segmentation and support inference from rgbd images. ECCV 2012.\n\n[5] Marius Cordts., et al. The cityscapes dataset for semantic urban scene understanding. CVPR 2016\n\n[6] Amir R Zamir., et al. Taskonomy: Disentangling task transfer learning. CVPR 2018.\n\n[7] Ozan Sener and Vladlen Koltun. Multi-Task Learning as Multi-Objective Optimization. NeurIPS 2018.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "bgAfziNP8o3", "reviewer": "Reviewer_GzAm", "summary": "The authors propose a novel differentiable sparse gate for mixture-of-experts models. \nAt-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.\nThe proposed method outperforms conventional methods in multitask learning settings.", "questions": "", "limitations": "", "rating": 9, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: \nThe proposed method is novel for mixture-of-experts models.  The differences from previous work are clearly described.\n\n\nQuality: \nThe proposed binary encoding scheme is reasonable, and the superiority of the differentiable sparse gate is supported by both the theoretical analysis the experimental results.\n\n\nClarity: \nThe paper is well organized and easy to understand.\nI would like to know how W is initialized to be 0 < S(Wx) <1 for Per-example Gating.\n\n\nSignificance: \nThe proposed differentiable gating technique can not only be a baseline method to train mixture-of-experts models but also be applicable to other models, such as attention models.\nIf the experimental results on the same task as [32] are provided, the impact of this paper might be more significant.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No discussion about potential negative societal impact of their work are provided.  I cannot find one, too.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "Y05isTaLxE", "reviewer": "Reviewer_GSeM", "summary": "The submission studies how to design a differentiable module for selecting K experts out of N experts in a sparse MoE model.\nThe proposed method reparametrize the softmax distribution over N the experts as log(N) binary variables.\nIt then proposes a smoothed S-shaped continuous approximation of the hard binary function.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\n- A novel differentiable N-select-K module that is built upon new ideas (the binary encoding and the S-shaped smoothed function).\nSuch a radically different alternative to the existing methods (Top-K, Gumbel-softmax variants, etc) is welcome since the existing methods are still hard to tune (based on my personal experience with sparse MoE).\n\n- The writing is of high quality, and the source code is provided. \n\n- It presents extensive empirical analysis and consistent improvement over the existing method.\n\nCons:\n\n- The proposed smooth function S(t) may lead to dead experts (i.e., some experts will no longer be used after a certain number of training iterations) or may make the model converge to a bad local optimum, because the proposed function has no gradient once t is below -gamma/2 or above gamma/2.\n\n- The proposed module can select the same expert for more than one time. In some applications, we may want to select k DISTINCT experts, which is not supported by the proposed module.\n\n- As far as I know, most sparse MoE models suffer from the dead-expert issue (a significant number of experts are never selected after optimization). There is no analysis of how the proposed method performs when it comes to the dead-expert issue.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors are well aware that the proposed method may be slower than the Top-k method when it comes to training since it only supports conditional training partially due to the soft approximation. Please see Cons in Main Review for my extra concerns.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All four reviewers agreed that the paper is very well written, easy to read, the idea is novel, and the experimental results are convincing and consistent. In the discussion, there was no disagreement in accepting the paper. As the authors promised, I would like to encourage the authors to include comparisons against Gumbel softmax based apporachs or some other stochastic subset selection methods. It would be great to prepare text and figures to highlight the shortcomings of these approaches (e.g. impossibility of per-example gating etc.).
Decision: Accept (Poster)