Paper: MOESART: An Effective Sampling-based Router for Sparse Mixture of Experts (ICLR-2024)

Reviewer comment (Experiments):
Much of the paper compare MOESART against the standard (nondifferentiable) Top-k strategies. To give a better picture of this work's contribution, it is important to also compare MOESART against differentiable top-k strategies in both technical contributions, performances, and time complexity. The authors already mentioned the relevant baselines in Hazimeh et al., 2021; Ibrahim et al., 2023;  and Sander et al., 2023 but none are further discussed.

Author response:
- **Comparison with differentiable routers.**

  1. **Technical Contributions**. Differentiable routers e.g., DSelect-k (Hazimeh et al., 2021) and COMET (Ibrahim et al., 2023) consider optimization-based smooth formulations based on binary representations to make routing decisions. In contrast, we consider sampling-based routing. The differentiable routing strategies e.g., DSelect-k and COMET can not do sparse training and thus can not be used to scale up model training which is a major focus of this work. These routers require computing gradients with respect to a larger set of experts per-input example while training, making them significantly more expensive. In contrast, MOESART requires gradients with respect to exactly $k$ e.g., 2 experts per-input. Although, in principle, these differentiable routers may partially allow sparse training because of their dense-to-sparse nature, the open-source implementations of these routing approaches do not support such partially conditional training and perform fully dense training. In contrast, our proposed routing approach MOESART completely allows sparse training and can easily substitute existing sparse implementations of Top-$k$ style approaches.

  2. **Performance comparison** We have run new experiments to compare our routing approach with some state-of-the-art differentiable routers. These results are added in the revised draft of the paper in Supplement (Section S4). We compare across various datasets and report the performance metrics below. MOESART appears to be quite competitive with differentiable routers. Notably, MOESART even outperforms differentiable routers across many tasks.

      a. *Recommender systems*:

| Dataset | Router | Test Loss ($\times 10^{-2}$)  $\downarrow$  | Task-1 AUC  $\uparrow$      | Task-2 MSE $\downarrow$ |
|------------------------------|-----------------|---------------------------------|-------------------------------|-----------------------------------|
| Books (TW=(0.1,0.9))       | DSelect-$k$   | $241.64\pm0.14$                 | $59.05\pm0.14$                | $2.615\pm0.002$                   |
|                                          | COMET          | $\textbf{241.13}\pm0.11$    | $\textbf{66.13}\pm0.09$   | $\textbf{2.612}\pm0.001$      |
|                                          | MOESART     | $242.47\pm0.09$                 | $64.99\pm0.13$                | $2.626\pm0.001$                  |
| Books (TW=(0.9,0.1))        |  DSelect-$k$  |  $~~74.22\pm0.03$              | $77.61\pm0.02$                | $\textbf{2.632}\pm0.002$    |
|                                           | COMET          | $~~74.29\pm0.03$              | $77.30\pm0.02$                | $2.636\pm0.002$                  |
|                                           | MOESART     | $~~\textbf{73.68}\pm0.02$ | $\textbf{78.03}\pm0.03$   | $2.641\pm0.003$                  |
| MovieLens (TW=(0.1,0.9)) |  DSelect-$k$  | $~~73.81\pm0.03$              | $\textbf{87.79}\pm0.03$   | $0.7715\pm0.0003$              |
|                                           | COMET          | $~~74.19\pm0.03$              | $87.67\pm0.07$                | $0.7756\pm0.0004$               |
|                                           | MOESART     | $~~\textbf{73.60}\pm0.02$ | $87.33\pm0.03$                | $\textbf{0.7684}\pm0.0002$  |
| MovieLens (TW=(0.9,0.1)) | DSelect-$k$   | $~~41.40\pm0.03$              | $91.44\pm0.01$                | $0.7582\pm0.0005$                |
|                                           | COMET          | $~~41.25\pm0.04$              | $91.45\pm0.02$                | $0.7563\pm0.0008$                |
|                                           | MOESART     | $~~\textbf{40.89}\pm0.02$ | $\textbf{91.61}\pm0.01$   | $\textbf{0.7430}\pm0.0003$   |

      b. Image Datasets:
| Dataset | Router | Test Loss ($\times 10^{-2}$) $\downarrow$ | Task-1 Accuracy $\uparrow$ | Task-2 Accuracy $\uparrow$ |
|---------------|----------|--------------|--------------|---------------|
| Multi-MNIST | DSelect-$k$ | $~~6.93\pm0.06$ | $98.14\pm0.02$ | $97.68\pm0.03$ |
|          | COMET | $~~6.83\pm0.06$ | $98.21\pm0.02$ | $97.63\pm0.03$ |
|          | MOESART | $~~\textbf{5.86}\pm0.03$ | $\textbf{98.40}\pm0.02$  | $\textbf{97.92}\pm0.02$ |
| Multi-FashionMNIST | DSelect-$k$ | $~36.88\pm0.21$ | $87.37\pm0.07$  | $86.61\pm0.09$ |
|          | COMET | $~34.88\pm0.13$ | $87.97\pm0.06$   | $87.42\pm0.06$ |
|          | MOESART | $~\textbf{32.85}\pm0.11$   | $\textbf{88.56}\pm0.06$ | $\textbf{88.02}\pm0.07$ |

- References
  - Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, et al. DSelect-k: Differentiable selection in the mixture of experts with applications to multi-task learning. NeurIPS, 2021.
  - Shibal Ibrahim, Wenyu Chen, Hussein Hazimeh, et al. Comet: Learning cardinality constrained mixture of experts with trees and local search. KDD, 2023.

Reviewer comment (Theory):
The theoretical guarantee of MOESART is also rather weak. There are other alternatives such as REINFORCE, that is unbiased but are not widely used in this setting because of its high variance.

Author response:
- **Theoretic rigor of MOESART** 
  - Our work has a similar rigor as other Sparse MoE papers that have practically shown benefits of Sparse MoE e.g., (Shazeer et al., 2017; Fedus et al., 2022; Zoph et al., 2022; Zhou et al., 2022; Chi et al., 2022). The theoretical properties of Top-k routing were not shown for any of these methods in literature. Interestingly, theoretical properties of vanilla Top-k routing in Gaussian MoE have started to be investigated recently, e.g., see a concurrent ICLR submission (https://openreview.net/forum?id=jvtmdK69KQ). Hence, studying theoretic properties of sparse routing methods is an open research area and require independent investigations.
  - Our motivation to consider the MOESART objective stems from the following goal: We aim to sparsely approximate the dense-softmax-based MoE. In order to show evidence for this goal, we motivate in terms of smaller bias:
    - The loss distribution across samples of MOESART tends to be close to dense softmax-based classical MoE objective — this is shown in Figure 2. Given that the loss distribution is similar, this can suggest that the gradients have relatively low bias.
    - The adjusted router weights  $\tilde{g}$ (based on eq (3) in paper) after sampling have a smaller bias i.e., they are close in expectation to the dense router weights $g$,
$$ ||(E[\tilde{g}_1],\cdots,E[\tilde{g}_n])-(g_1,\cdots,g_n)||_2~~~\text{is small}$$ 
This is shown to be true with an ablation study in Figure S1 in Supplement Section S3.1. 

These explanations show good justification for why MOESART is effective for routing.

- **REINFORCE-style estimators** As reviewer notes, it has been pointed out in literature that REINFORCE-style estimators tend to have prohibitively high variance. It is especially true for networks that have other sources of randomness (i.e., dropout or other independent random variables). Despite attempts aimed at reducing variance of such estimators, such estimators have failed to work in MoE training — see Kool et al. (2021).

After the new experiments and clarifications, we would like to ask the reviewer to kindly revisit their evaluation and consider increasing the score.
- References
  - Wouter Kool, Chris J. Maddison, and Andriy Mnih. Unbiased gradient estimation with balanced assignments for mixtures of experts. In I (Still) Can’t Believe It’s Not Better! NeurIPS 2021 Workshop, 2021.

Reviewer comment (Experiments):
The NLP experiment in Table 2 should include X-MoE, which was originally proposed for this application.

Author response:
- **X-MoE Results for Table 2** We would like to highlight that authors of X-MoE (Chi et al., 2022) didn’t consider distillation of a pre-trained non-MoE model to an MoE model in their experiments. However, on reviewer’s request, we have run new experiments with X-MoE router in this setup and we report the numbers below on GLUE benchmarks. We observe some performance gain of X-MoE over Top-$k$ router on two tasks, however X-MoE doesn’t appear to consistently outperform Top-$k$, which appears to be consistent with our experiments on recommender systems and image datasets in Table 1. MOESART consistently outperforms X-MoE on these benchmarks as well.

|             | RTE         | CoLA         | MRPC       | SST-2       | QNLI     | QQP      | MNLI       |
|---------|------------|--------------|----------------|-------------|----------|-----------|------------|
|             | Acc $\uparrow$   | Mcc $\uparrow$ | F1 $\uparrow$    | Acc $\uparrow$ | Acc $\uparrow$   | F1 $\uparrow$ | m/mm $\uparrow$  |
| MoEBERT with Top-$k$     | 68.95    | 57.86  | 86.76    | 92.78    | 91.63  | 88.10   | 85.14  |
| MoEBERT with X-MoE       | 67.51     | 55.48   | 86.76   | 92.55    | 91.73   | 88.25    | 85.08  |
| MoEBERT with MOESART | $\textbf{70.40}$     | $\textbf{58.20}$ | $\textbf{88.24}$     | $\textbf{93.12}$ | $\textbf{91.95}$ | $\textbf{88.27}$ | $\textbf{85.21}$ |

Reviewer comment (Presentation):
The proposed method is presented quite poorly in Section 3. The Trimmed Lasso regularization is a general strategy but is presented in the middle of the method discussion.

Author response:
- **Presentation** We appreciate the reviewer for the suggestion to reorganize Section 3. Following the reviewer’s suggestion, we have deferred the discussion on Trimmed Lasso to a separate section (Section 3.2 in the revised draft) after the complete description of MOESART (sampling-based loss approximation and reweighting).

Reviewer comment (Novelty):
What are the contributions of MOESART compared to existing differentiable Top-k in terms of technical contributions, performances, and time complexity.

Author response:
- **Comparison with differentiable routers.**

  1. **Technical Contributions**. Differentiable routers e.g., DSelect-k (Hazimeh et al., 2021) and COMET (Ibrahim et al., 2023) consider optimization-based smooth formulations based on binary representations to make routing decisions. In contrast, we consider sampling-based routing. The differentiable routing strategies e.g., DSelect-k and COMET can not do sparse training and thus can not be used to scale up model training which is a major focus of this work. These routers require computing gradients with respect to a larger set of experts per-input example while training, making them significantly more expensive. In contrast, MOESART requires gradients with respect to exactly $k$ e.g., 2 experts per-input. Although, in principle, these differentiable routers may partially allow sparse training because of their dense-to-sparse nature, the open-source implementations of these routing approaches do not support such partially conditional training and perform fully dense training. In contrast, our proposed routing approach MOESART completely allows sparse training and can easily substitute existing sparse implementations of Top-$k$ style approaches.

  2. **Performance comparison** We have run new experiments to compare our routing approach with some state-of-the-art differentiable routers. These results are added in the revised draft of the paper in Supplement (Section S4). We compare across various datasets and report the performance metrics below. MOESART appears to be quite competitive with differentiable routers. Notably, MOESART even outperforms differentiable routers across many tasks.

      a. *Recommender systems*:

| Dataset | Router | Test Loss ($\times 10^{-2}$)  $\downarrow$  | Task-1 AUC  $\uparrow$      | Task-2 MSE $\downarrow$ |
|------------------------------|-----------------|---------------------------------|-------------------------------|-----------------------------------|
| Books (TW=(0.1,0.9))       | DSelect-$k$   | $241.64\pm0.14$                 | $59.05\pm0.14$                | $2.615\pm0.002$                   |
|                                          | COMET          | $\textbf{241.13}\pm0.11$    | $\textbf{66.13}\pm0.09$   | $\textbf{2.612}\pm0.001$      |
|                                          | MOESART     | $242.47\pm0.09$                 | $64.99\pm0.13$                | $2.626\pm0.001$                  |
| Books (TW=(0.9,0.1))        |  DSelect-$k$  |  $~~74.22\pm0.03$              | $77.61\pm0.02$                | $\textbf{2.632}\pm0.002$    |
|                                           | COMET          | $~~74.29\pm0.03$              | $77.30\pm0.02$                | $2.636\pm0.002$                  |
|                                           | MOESART     | $~~\textbf{73.68}\pm0.02$ | $\textbf{78.03}\pm0.03$   | $2.641\pm0.003$                  |
| MovieLens (TW=(0.1,0.9)) |  DSelect-$k$  | $~~73.81\pm0.03$              | $\textbf{87.79}\pm0.03$   | $0.7715\pm0.0003$              |
|                                           | COMET          | $~~74.19\pm0.03$              | $87.67\pm0.07$                | $0.7756\pm0.0004$               |
|                                           | MOESART     | $~~\textbf{73.60}\pm0.02$ | $87.33\pm0.03$                | $\textbf{0.7684}\pm0.0002$  |
| MovieLens (TW=(0.9,0.1)) | DSelect-$k$   | $~~41.40\pm0.03$              | $91.44\pm0.01$                | $0.7582\pm0.0005$                |
|                                           | COMET          | $~~41.25\pm0.04$              | $91.45\pm0.02$                | $0.7563\pm0.0008$                |
|                                           | MOESART     | $~~\textbf{40.89}\pm0.02$ | $\textbf{91.61}\pm0.01$   | $\textbf{0.7430}\pm0.0003$   |

      b. Image Datasets:
| Dataset | Router | Test Loss ($\times 10^{-2}$) $\downarrow$ | Task-1 Accuracy $\uparrow$ | Task-2 Accuracy $\uparrow$ |
|---------------|----------|--------------|--------------|---------------|
| Multi-MNIST | DSelect-$k$ | $~~6.93\pm0.06$ | $98.14\pm0.02$ | $97.68\pm0.03$ |
|          | COMET | $~~6.83\pm0.06$ | $98.21\pm0.02$ | $97.63\pm0.03$ |
|          | MOESART | $~~\textbf{5.86}\pm0.03$ | $\textbf{98.40}\pm0.02$  | $\textbf{97.92}\pm0.02$ |
| Multi-FashionMNIST | DSelect-$k$ | $~36.88\pm0.21$ | $87.37\pm0.07$  | $86.61\pm0.09$ |
|          | COMET | $~34.88\pm0.13$ | $87.97\pm0.06$   | $87.42\pm0.06$ |
|          | MOESART | $~\textbf{32.85}\pm0.11$   | $\textbf{88.56}\pm0.06$ | $\textbf{88.02}\pm0.07$ |

- References
  - Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, et al. DSelect-k: Differentiable selection in the mixture of experts with applications to multi-task learning. NeurIPS, 2021.
  - Shibal Ibrahim, Wenyu Chen, Hussein Hazimeh, et al. Comet: Learning cardinality constrained mixture of experts with trees and local search. KDD, 2023.

Reviewer comment (Experiments):
How much did the Trimmed Lasso regularization contribute to the performance gains?

Author response:
- **Contribution of the Trimmed Lasso regularization** Ablation study showing the performance of MOESART without and with trimmed lasso regularization is given in Table S2 in Supplement Section S3.3. In some cases, e.g., Books dataset, we observe substantial performance improvement with trimmed lasso regularization. In other cases e.g., MovieLens, we observe smaller improvement with trimmed lasso regularization. Recall that $g(x)$ is not sparse by construction in MOESART. Sampling injects sparsity in $\tilde{g}(x)$ during training. The goal of trimmed Lasso regularization is to encourage $g(x)$ to accumulate router probability mass per-input $x$ in the top $k$ elements, making the sampling process more deterministic per-input for inference by the end of training. Given that other Top-$k$ style routers are sparse by construction, it does not make sense to augment those routers with trimmed lasso regularization.

Reviewer comment (Experiments):
If possible, please include X-MoE in Table 2.

Author response:
- **X-MoE Results for Table 2** We would like to highlight that authors of X-MoE (Chi et al., 2022) didn’t consider distillation of a pre-trained non-MoE model to an MoE model in their experiments. However, on reviewer’s request, we have run new experiments with X-MoE router in this setup and we report the numbers below on GLUE benchmarks. We observe some performance gain of X-MoE over Top-$k$ router on two tasks, however X-MoE doesn’t appear to consistently outperform Top-$k$, which appears to be consistent with our experiments on recommender systems and image datasets in Table 1. MOESART consistently outperforms X-MoE on these benchmarks as well.

|             | RTE         | CoLA         | MRPC       | SST-2       | QNLI     | QQP      | MNLI       |
|---------|------------|--------------|----------------|-------------|----------|-----------|------------|
|             | Acc $\uparrow$   | Mcc $\uparrow$ | F1 $\uparrow$    | Acc $\uparrow$ | Acc $\uparrow$   | F1 $\uparrow$ | m/mm $\uparrow$  |
| MoEBERT with Top-$k$     | 68.95    | 57.86  | 86.76    | 92.78    | 91.63  | 88.10   | 85.14  |
| MoEBERT with X-MoE       | 67.51     | 55.48   | 86.76   | 92.55    | 91.73   | 88.25    | 85.08  |
| MoEBERT with MOESART | $\textbf{70.40}$     | $\textbf{58.20}$ | $\textbf{88.24}$     | $\textbf{93.12}$ | $\textbf{91.95}$ | $\textbf{88.27}$ | $\textbf{85.21}$ |

Reviewer comment (Experiments):
The paper uses top-k != 1 for all their experiments and doesn't show results for k=1.

Author response:
- **k!=1** Our motivation to consider k > 1 was based on earlier work (Du et al., 2022; Zhou et al., 2022; Zoph et al., 2022; Shazeer et al., 2017; Lepikhin et al., 2021; Ruiz et al., 2021; Hazimeh et al., 2021). They consider k > 1 for good trade-off between predictive performance and the training/serving efficiency of the model. Note that for k = 1, as is the case for Top-k (Shazeer et al., 2017; Ruiz et al., 2021), the gradient with respect to the router parameters is zero. Hence, we do not run any experiments for k = 1.

Reviewer comment (Novelty):
Is the paper saying it introduces "weighting expert outputs by output of router (aka sampling weight)"? I've seen this done in multiple MoE implementations.

Author response:
- **Clarification** It appears that there might be a confusion about “weighting expert outputs by output of router”. It is true that all routers have some router weights that combine expert outputs, which is a standard MoE learning paradigm. However, we propose a novel adjustment of the weights after we sample from $g$. This adjustment is discussed in eq (3). This adjustment reduces the bias of the adjusted weights $\tilde{g}$ in expectation $||{(E[\tilde{g}_1],\cdots,E[\tilde{g}_n])-(g_1,\cdots,g_n)}||_2$. The reduction in bias is shown in ablation study in Figure S1 in Supplement Section S3.1.

Reviewer comment (Evaluation):
It is unclear why it is desired to use a sampling-based expert index assignment. Figure 2 suggests that MOESART has a smaller gap between training and inference, but it is not clear whether such a smaller gap is a result of the introduced regularization.

Author response:
- **Clarification** Sampling based expert assignment can have a smaller bias as it gives more room for exploration amongst experts compared to Top-k routing. In Figure 2, the trimmed lasso regularization was set to 0 to particularly highlight the usefulness of the sampling based routing scheme without any regularization. We have clarified this in the figure caption in the revised draft.

Reviewer comment (Theory):
It is unclear why it is desired to update the scaling factor as proposed. As mentioned by the author, the gradient of the expert sampling is typically ignored by MoE training (but can be estimated with the policy gradient method), and the router training is conducted with only the scaling factor. Intuitively, modifying the scaling factor could lead to better gradient estimation and thus better MoE training, but it is unclear how this is achieved by the proposed method.

Author response:
- **Scaling**  We adjust/scale the weights after we sample from $g$. This adjustment, outlined in eq (3), can reduce the bias of the adjusted weights $\tilde{g}$, which can be measured in expectation as $||{(E[\tilde{g}_1],\cdots,E[\tilde{g}_n])-(g_1,\cdots,g_n)}||_2$.
We study this bias in ablation study in Section S3.1 in the Supplement. In this ablation study we consider the bias for eq (3) for different choices for $g$: uniform, random, etc. We show the plots in Fig. S1. For uniform setting, we observe the bias to be very similar to some alternative choices. However, for sufficiently non-uniform distributions, we can observe that there can be a significant reduction in bias with the proposed adjustment strategy in eq (3). In Sparse MoE, the distribution $g(x)$ can be very different across different inputs $x$, hence our proposed strategy is expected to have a lower bias overall across samples. A smaller bias can lead to improved performance. We observe this to be the case in Sparse MoE, which we show in Table S1 in the ablation study in Supplement Section S3.2.

- **Gradient Estimation with policy gradient method** It has been shown in Kool et al. (2021) that estimators based on policy gradient do not work well in Sparse MoE training.

Reviewer comment (Writing):
In the current draft, the proposed method is positioned as a sampling method. However, from the reviewer's understanding, the sampling process of the expert index remains unchanged from the softmax sampling, and the proposed method focuses instead on introducing additional regularization and changing expert output scaling. The reviewer suggests changing the wording to something like "carefully designed sampling".

Author response:
- **Change choice of words** Our use of the term “carefully designed” in the phrase “carefully designed sampling and weighting strategies” was primarily a reference to adjustment strategies. Hope this clarifies the term. We are happy to change the wording to be more precise in the revision.

Reviewer comment (Experiments):
It should be noted that all experiments in this work were conducted on small-sized datasets, whereas existing Mixture of Experts (MoE) structures typically focus on large-scale training and inference scenarios. This difference in dataset size should be taken into consideration when evaluating the applicability of the proposed method.

Author response:
**Medium/Relatively small training setups** The training setups we considered have been considered previously in some sparse mixture of experts literature (Zuo et al., 2022; Hazimeh et al., 2021; Ibrahim et al., 2023). Our experiments were run on academic-level compute resources. We are limited to a GPU cluster where our resources are capped at 4 V100 Tesla GPUs with 12-hour time-limits per job. Hence, we do not have the compute bandwidth to run pre-training on 100 billion parameter models, which are typically trained
on 100-1000s of GPUs for many days. Within the compute budget, we have still considered diverse models/datasets such as medium sized models e.g., MoEBERT on GLUE and SQuAD benchmarks. Interestingly, our work exposes that some of the advanced top-k
routing strategies do not generalize well across datasets/training setups/domains. Hence, we believe that these reasonably sized settings prove to be good regimes to explore novel routing strategies.

Reviewer comment (Evaluation):
The motivation behind MOESTART is to achieve results that closely resemble those of the top-k softmax gate. However, it is intriguing to observe that MOESTART consistently outperforms the softmax gate in all the conducted experiments. The authors should provide additional clarification or insights as to why MOESTART exhibits superior performance compared to the traditional softmax gate, despite its primary goal of approximation.

Author response:
**Performance difference between Softmax vs Sampling-based MOESART** The reviewer makes an interesting observation that MOESART tends to consistently outperform traditional softmax. We observe some differentiably sparse routers e.g., DSelect-k,
COMET can also outperform softmax routing — see Table S4 in Supplement Section S4 in the revised draft. Hence, sparsity can be beneficial on these datasets to improve generalization. Additionally, softmax routing with linear parameterization is less richer in terms of parameterization than tree-based COMET (Ibrahim et al., 2023) and may also be less richer in comparison to MOESART because sampling in MOESART may induce more nonlinearity due to sparsity. Dense nature of softmax activates all experts for all samples, which may affect the localized/specialization property as highlighted in Ramamurti & Ghosh (1998). Given that MOESART can outperform COMET, it is possible that sparsity in MOESART may induce more localized decision making, hence leading to improved performance.