Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.
review_content: summary: Paper proposes a remedy to the representation collapse issue of sparsely gated (activated) mixture of expert modes (sMoE) in Transformers. The hypothesis, and later supported with observations is that, the top-1 routing mechanism pushes the hidden representations clustering around expert centroids, introducing a tendency toward representation collapse - which might be hurting the model accuracy. 

To mitigate the representation collapse, deemed to be harmful, paper proposes adding two more operations to the routing procedure: (1) dimensionality reduction of the hidden representations that the routing decisions used to base on, and (2) L2 normalization of the activations (logits) of pre-softmax/sigmoid at the top-1 expert selection routine. 

Two enhancements combined, paper demonstrates improved upstream (pretraining perplexity) and downstream (accuracy after fine-tuning)  quality gains over the dense and regular top-1 sMoE models on various benchmarks. 
strengths_and_weaknesses: Strengths:

- Identification of "representation collapse" of MoE models as one of the potential issue of these family of models sheds light on the failure modes of them. 

- New metrics, like routing fluctuation ratio, and UMAP of token-expert assignment, are novel and great contributions to the community moving forward.

Weakness:

- Proposed method is focusing only on top-1 expert choice sMoE models which are not widely used, and might be attempting to solve a problem that only exist in that space. Which is okay, but the paper is not scoped, structured and narrated that way. Do we know, the widely used top-2 expert choice, where the first expert is chosen according to eq. (3) and the second is sampled randomly from the same distribution, has the same problem? 

- The assumption that the number of experts are going to be much smaller than the hidden dimension does not always hold (e.g., Lepikhin et al. 2020 use 2048 experts w/ 1024 hidden dimensions). Similar to the first weakness, scoping of the problem is one of the weaknesses.  

questions: Q1 - What does "conservative" mean in line 128?

Q2 - What are the initial and final values of \tau (in line 131) that are learned by the model - does the model learn higher or lower values than the one initially set?

Q3 - In line 142, you mention "we employ ... seq2seq for nmt" but there is no nmt experiments in the paper, is this intended?

Q4 - What is the alpha / temperature for the language mixing (section 4.1 - Pre-training Data)?

Q5 - You mention in 4.3 (last sentence) that the sigmoid gating has different behavior in upstream evaluation, but do not provide any data points / tables / figures - have I missed those?

Q6 - Representation collapse (and others in Figure 2 & 3) are using which set to collect data? Dev / test / train ? Please specify. 

Q7 - Why do you think routing consistency should compare the first and one of the last checkpoints to measure the stats? What is the hypothesis there? The model may change the chosen of expert throughout the course of training, which should be okay (please correct me if I'm wrong), rather than committing to an expert at the beginning of the training and sticking with it. Wouldn't it make more sense to take a look at the last n-checkpoints and computing the statistics over that? Please word your hypothesis why comparing first with the last checkpoint supports the argument for routing consistency. 

Q8 - All models seem to have trained for the same number of tokens, but some of them use extra parameters and extra priors (dim-reductions + L2 norm). Given these, wouldn't it make more sense to compare two models with the same upstream (pretraining) perplexity on their downstream quality? How do we make sure the #tokens that we set in the pretraining not biasing the proposed method over the others ?

Q9 - Follow up on Q8, have we adjusted the learning rate schedules for all the different approaches (a la chinchilla) ?

limitations: adequate
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper makes the following contributions:
1. Identify that common MoE routing mechanisms that rely on matching token representations to expert representations suffer from representation collapse - learnt token representations become too similar to expert representations collapsing to the subspace of expert representations.
2. Propose a simple fix to the representation collapse problem - utilizing a low rank projection and cosine similarity metric for routing.
3. Evaluate their approach and baseline routing algorithms on multilingual text MLM pre-training and downstream evaluation on the Xtreme benchmark.
4. Qualitative and quantitative analysis demonstrating that the approach presented in the paper improves the representation collapse in MoEs.
strengths_and_weaknesses: Strengths:
1. Identify a previously unknown representation collapse problem in Sparse Mixture-of-Experts models. Propose and evaluate a simple modification to the routing architecture that addresses the problem.
2. Careful experimentation evaluating the approach on multilingual text MLM pre-training and on downstream tasks from the Xtreme benchmark.
3. Intuitive qualitative and quantitative analysis of the learnt token representations that highlight the collapse and that the proposed modification addresses the problem.
4. The paper is well motivated, clearly written and makes it easy to understand the details of the experiments.

Weaknesses:
1. Limited downstream evaluations: The downstream evaluations on the Xtreme benchmark only evaluate on English-only fine-tuning and zero-shot evaluations. The zero-shot evaluation setting is a little brittle and can benefit from regularization effects and minor fine-tuning differences. It would strengthen the results to: (i) include language-level evaluation results (perhaps in the appendix). (ii) Conduct evaluations in non zero-shot settings perhaps using translated data to simulate multilingual fine-tuning (Translate-Train).
2. Non-standard upstream evaluation: The upstream evals evaluate MLM perplexity on a custom validation split. This is not a standard upstream evaluation setting and would be hard to replicate if there are minor differences in experimental setups or vocabularies or test split. I would advise the authors to repeat these experiments on a more standard benchmark; perhaps a causal LM with perplexity evaluations on standard test sets.
3. Limited experimental setting: The paper analyzes and evaluates MoEs trained with Top-1 routing. While theoretically you would expect generalized MoEs trained with Top-K routing to suffer from similar representational collapse, it would strengthen the paper to validate that their approach helps when MoEs are trained with Top-2 routing as well.


Overall I recommend weak accept. This paper identifies and addresses a problem with the sparse-MoE architecture which could be valuable to large model training for a variety of different tasks and applications. However, there are several limitations with their evaluation and the scope of experiments presented in the paper.

EDIT: Updated recommendation to accept following author rebuttal. The proposed algorithm seems to result in small but consistent improvements for MoE training on multilingual text pre-training and Machine Translation.
questions: Please see weaknesses above.
limitations: N/A
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper looks at the problem of representational collapse in Mixture Of Experts (MoE) based transformer layers. Specifically, they show that for a softmax based gating function (for combining the MoE layer output with the hidden representation from the previous layer) , the gradient updates to token representations consists of a term that is simply a linear combination of the $N$ expert vectors, and since $N$ is much less than the dimensionality of the token vectors themselves, this portion of the gradient update is “low-dimensional” (specifically, $<= N$). 

The paper argues that this could lead to representational collapse where token vectors receive updates that do not fully utilize the entire dimensionality of the vectors, leading to poorly learnt representations. In response, the paper proposes a simple (and somewhat effective) fix for this: Simply project the representations into a lower dimensional space before feeding it into the MoE layer. This ensures that the token vectors “interact” with the MoE layers (and hence with the corresponding gradients) via the low dimensional mapping. From experiments, they show that this simple fix results in representations with much lesser representational collapse (figure 2c) and in-turn modest gains on down-stream performance (Table-1)
strengths_and_weaknesses: Strengths: 
* Paper makes a clever observation regarding representational collapse in mixture of experts models, which to the best of my knowledge, has not been studied or pointed out previously. 
* Conditional computation is becoming increasingly important for scaling to ever larger model sizes and thus the overall theme of the paper is extremely relevant.
* The experiments on cross-lingual NLP show some modest imrprovements.

Weaknesses:
* Method: This work only shows representational collapse to be a problem for softmax gating, but this doesn’t seem to be the case for sigmoid? In particular, the contribution from the $J_2$ term only ends up pushing the token embedding towards $e_k$ (the chosen expert). So while one can say that all tokens that are assigned expert $k$ will receive a gradient which has a direction along $e_k$ (hence some amount of clustering), the $J_1$ term will ensure that they also get a gradient pointing in a different direction. (This second comment also applies to the case when you use a softmax gating).
* Experiments: 
    * Maybe I misunderstood the model, but it seems like only a single MoE layer is added after layer 6, instead of having several MoE layers (perhaps every other layer is an MoE layer). The reason why i ask this is because if there are more MoE layers, that increases the chances of representational collapse, and could probably show bigger gains for your approach. If there is only one MoE layer, then the representational collapse might not be that beneficial. 
    * On the flipside, it would be great to have a toy experiment on simple toy data (even linear regression) where you have only a single MoE layer with a softmax gating, where you can clearly demonstrate representational collapse.
    * It’s unclear whether the difference b/w Dense (no MoE) and MoE models is because of the MoE layers or just because of more parameters i.e. did you ensure that the Dense models have the same number of parameters? 
    * The UMAP plots don’t seem to suggest representational collapse: It’s really hard to conclude that 2a and 2b have different amounts of spread in the representations. It appears that for both of these, tokens that are mapped to the same expert are very close together. The plot in fig2c is nice but I would like really like to see more analysis / quantitative results on representational collapse. Some suggestions:
        * you could sample a large enough number of tokens and see if the subspace spanned by them is larger for your approach vs the baseline (presumably, for the baseline you’d expect the subspace to be much lower dimensional). 
        * you could also look at spread i.e. given a cluster of all token vectors that are assigned the same expert, what is thethe average L2 distance of each token vector from the mean?  

questions: * Method: Could the authors clarify the role of other technical ``bells and whistles’’ for preventing representational collapse: using cosine similarity instead of dot-products (L2 normalization), using a temperature parameter in the softmax? I can see these being applied to the baselines as well - was this also done?
* Experiments: Could the authors comment on why the MoE parameters were kept frozen while finetuning? I only somewhat buy the overfitting argument because pre-trained transformers have been finetuned on extremely small data without much “overfitting” and several tricks can be employed to prevent it (such as early stopping, using low learning rates etc).

Please also see the weaknesses section from above for a more comprehensive list of questions on experiments. 
limitations: Limitations adequately addressed.
review_title: null

Meta-review: metareview: This paper highlights that contemporary sparse Mixture-of-Experts (sMoE) models suffer from representation collapse in the gating mechanism. The paper then proposes a simple fix reducing the dimensionality of the gating representations and using cosine similarity. The paper shows qualitatively that the new representation suffers less from collapse and that model trained with the new algorithm exhibit small but consistent improvements when evaluated on cross-lingual language understanding and machine translation.

All of the reviewers had a concern that the method is tested only in the top-1 setting vs. the most common practice of using top-2 or more experts. The paper added experiments on top of GShard that alleviates this concern.

Another concern was limited scope of evaluation. The authors added MT results as well as in-language results, which showed the same trend of small, but consistent improvements.

As a result of the discussion, two of the three reviewers were happy of how their concerns were addressed and increased their score to recommend that the paper is accepted. The third reviewer has not been active in discussion. 
confidence: Certain
Decision: Accept