Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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. 

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.   

Strengths:

+ This paper is well-written and easy to follow;

+ 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);

+ 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;

+ The proposed method has promising results on different MTL problems, especially the large-scale real-world recommender system application;

Concerns: 

1. 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:

    - 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.

    - 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?

2. 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. 

3. 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.

4. Experimental Results:
   - 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].

   - 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?

   - 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.


[1] Max B. Paulus., et al. Gradient Estimation with Stochastic Softmax Tricks. NeurIPS 2020.

[2] Krzysztof Maziarz., et al. Gumbel-matrix routing for flexible multi-task learning. arXiv preprint arXiv:1910.04915.

[3] Ximeng Sun., et al. Adashare: Learning what to share for efficient deep multi-task learning. NeurIPS 2020.

[4] Nathan Silberman., et al. Indoor segmentation and support inference from rgbd images. ECCV 2012.

[5] Marius Cordts., et al. The cityscapes dataset for semantic urban scene understanding. CVPR 2016

[6] Amir R Zamir., et al. Taskonomy: Disentangling task transfer learning. CVPR 2018.

[7] Ozan Sener and Vladlen Koltun. Multi-Task Learning as Multi-Objective Optimization. NeurIPS 2018.
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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).
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.

* 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).

* Clarity: The paper is well-written.

* 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).

I have the following detailed questions:

* 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.

* 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].

* 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?).

* Q4: I notice that the submission history was "hidden as part of the resubmission bias experiment". Could the authors provide more information about that?

[1] Differentiable Top-k Operator with Optimal Transport, NeurIPS 20. https://arxiv.org/abs/2002.06504



limitations_and_societal_impact: The authors adequately addressed the limitations and potential negative societal impact of their work.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The submission studies how to design a differentiable module for selecting K experts out of N experts in a sparse MoE model.
The proposed method reparametrize the softmax distribution over N the experts as log(N) binary variables.
It then proposes a smoothed S-shaped continuous approximation of the hard binary function.
main_review: Pros:

- A novel differentiable N-select-K module that is built upon new ideas (the binary encoding and the S-shaped smoothed function).
Such 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).

- The writing is of high quality, and the source code is provided. 

- It presents extensive empirical analysis and consistent improvement over the existing method.

Cons:

- 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.

- 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.

- 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.


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_title: null

final_score: rating: 9: Top 15% of accepted NeurIPS papers, strong accept
review_content: summary: 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.
main_review: Originality: 
The proposed method is novel for mixture-of-experts models.  The differences from previous work are clearly described.


Quality: 
The proposed binary encoding scheme is reasonable, and the superiority of the differentiable sparse gate is supported by both the theoretical analysis the experimental results.


Clarity: 
The paper is well organized and easy to understand.
I would like to know how W is initialized to be 0 < S(Wx) <1 for Per-example Gating.


Significance: 
The 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.
If the experimental results on the same task as [32] are provided, the impact of this paper might be more significant.

limitations_and_societal_impact: No discussion about potential negative societal impact of their work are provided.  I cannot find one, too.
review_title: null

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)