Multi-Head Mixture-of-Experts

Summary

The work presents an advanced MoE model architecture to further improve the expert utilization. Authors conduct comprehensive experiments on three pre-training tasks, English-focused language modeling, multi-lingual language modeling and masked multi-modality modeling with models at different scales (300M to 7B).

Strengths

1) Important research problem. Better expert utilization has great potential to improve the MoE scalability. 2) Comprehensive experiments. From 300M to 7B. From NLP to CV. 3) Cool visualization. Figure 1 is very helpful in understanding why and how this model works.

Weaknesses

I only have one but important concern: the end-to-end training and inference throughput. I understand that the theoretical computation and communication cost is not high, but since we are conducting a more complicated/fine-grained routing decision, the routing process may take longer in practice. Considering that MoE models are sometimes trained and deployed at very large scale, the intra-node all2all is sometimes required. Will the proposed approach slow down this? I can increase my score to positive if this concern can be solved.

Questions

See weakness.

Rating

4

Confidence

5

Soundness

2

Presentation

3

Contribution

3

Limitations

NA

Summary

This paper proposes Multi-Head Mixture-of-Experts (MH-MoE), a simple yet effective routing strategy that splits each input token into multiple sub-tokens for expert routing. This operation significantly enhances the ratio of activated experts for each token, enabling more fine-grained assigning of tokens. Through extensive experimental results across different parameter scales (300M to 7B) and three 15 pre-training tasks, the paper validates the effectiveness of the proposed MH-MoE strategy.

Strengths

1. The paper is well-written and easy to follow. 2. The proposed MH-MoE is simple yet effective. It is easy to be integrated into existing frameworks with MoE optimization. This improves the feasibility of the method. 3. The experiments cover both CV and NLP modeling scenarios and the results are consistently positive.

Weaknesses

1. My major concern is about the experimental settings. Implementing MH-MoE based on the X-MoE instead of the vanilla SMoE is strange. The improvements shown by the experiments may not be generalizable to other MoE structures, which makes the proposed method less convincing. 2. Lack of experiments on scaling computations (activated experts). Though the paper conducts scaling experiments on scaling total parameters (number of experts), the computation remains static. The information conveyed by this kind of scaling is limited as the computation is also important in affecting the performance of MoE. I’m not expecting additional experiments as the time is limited during rebuttal. However, a computational scaling experiment would make the method more convincing.

Questions

Can you provide the balance loss curves for your method and the baselines? I’m curious about the effect of your method on load balancing as it introduces more sub-tokens for routing.

Rating

5

Confidence

5

Soundness

2

Presentation

4

Contribution

3

Limitations

None

Summary

This paper introduces Multi-Head Mixture-of-Experts (MH-MoE), a method for training MoE models with enhanced expert activation. In particular, MH-MoE splits each input token into multiple sub-tokens, which are processed by a set of experts in parallel, and then merges them back into their original token form. The results show improved performance over dense and X-MoE models in parameter-matching settings.

Strengths

The authors propose a simple yet effective idea of splitting tokens into sub-tokens to boost expert activation. The paper is clear and well-written. The authors verify their approach on a variety of tasks, including English-focused language modeling, multi-lingual language modeling, and masked multi-modality modeling, across different parameter scales (ranging from 300M to 7B by scaling the number of experts).

Weaknesses

1. Experiments are performed on a limited number of models and baselines, making the results less convincing. - The experiments are limited by a single model architecture. The authors apply their approach by modifying the previously released X-MoE model. - The authors compare their performance with the mentioned X-MoE model, whereas there are a variety of MoE approaches that have been proposed and released recently. - The ablation studies are limited. The performance of SMoE heavily depends on the choice of hyper-parameters, one of the most important being the number of experts to be activated (referred to as top-k). How does the activation map change with an increase in k? Would MH-MoE still perform better, or is there a saturation point? In other words, does MH-MoE seem superior to X-MoE only because a small k=2 was chosen? 2. There is a lack of discussion on limitations. The authors claim that their approach does not have any limitations. I would urge the authors to reflect on the guidelines and reconsider their answer. For example, the weaknesses mentioned above (but not limited to) should be considered.

Questions

Please see the reasons to reject.

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

No. authors did not disclose any limitations of their method. See "Weaknesses" section for more details.

Summary

The paper proposed a new architectural extension to large pre-trained models (e.g. LLMs, or Multi-modal models). Specifically, the focus is on Mixture of Experts (MoE) models and the paper proposes a new method by introducing subtokens and routing each subtoken to experts. The resulting architectural change performs better on down-stream tasks and demonstrates a higher use of a variety of experts. The author's also conduct a variety of ablation studies to test various aspects and changes of their proposed architectural change.

Strengths

Strengths: - Achieves strong performance on down-stream tasks and seems to out-perform the baseline methods presented in the paper. - Paper claims that they achieve better results with same cost - Good reporting on experiments for reproducibility - Decent amount of experiments and evaluations

Weaknesses

Weaknesses: 1. No actual comparison in terms of "cost" is given except for parameter count and complexity calculations in the appendix. It would be interesting to see actual evaluation of cost (e.g. computational cost of training & computational cost of inference vs. MoE alone). 2. Evaluation and comparison is not done on publicly available and popular (& stat-of-the-art) models of relevant sizes (e.g. Mistral, Aya, Gemma ...). Therefore, it is not clear if this method actually performs better or just on the models and training runs presented in this paper. 3. Clarity a. some parts are unclear, e.g. equations 2,3,6,7 speak about the new additions of the method, however, it is not clear at all that there is MLP layer present (only later in the paper it is mentioned), however, it is "still" not clear whether and how it applied during equations 2,3,6,7. b. some mistakes in the paper e.g. Paper checklist 7, the answer is N/A. However, N/A means no experiments were run, the correct answer should be No.) 4. The core claim that the problem with MoE is: "low experts activation issue" has a few problems: a. It is never clearly explained why this is an issue (and no evidence is given that this is a real issue). b. On one hand the paper presents that the method increases activations, however, the paper also shows that increasing the number of heads (and therefore activations) seems to also decrease performance after a specific number of heads [Figure 5, Section 4.4]. (Great for reporting this, however, this also challenges the introduction / proposition of the method). I.e. it seems that "low experts activation" is not always a problem, or only a problem to a certain degree. => Therefore, it would be interesting to revise this hypothesis, and also to provide a counter hypothesis (with evidence). 5. Reporting on sharding of the models on the DGX-2 system is not quite described. (Which is quite essential to run it sensibly.)

Questions

Questions: 1. How does your model compare against some popular & state-of-the-art models (e.g. MoE LLMs like Mistral) in terms of performance? 2. Referring to "Weakness 4.b", what is your explanation of why your method actually works better? 3. What is the actual cost of running standard MoE vs. yours (in terms of practical training or inference time) - given the same computation budget.

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

Limitations: 1. Limitations are not really considered in the paper. 2. It would be really interesting to understand what does NOT work in the method in practice? a. What did not work during developing of the method? b. When is the method worse than previous other methods? c. Computational limitations? d. etc. etc.

Thank you for your response

-> A1: Some recent SOTA MoE models include SMEAR [1], BTX [2], DeepSeekMoE [3], DEMix [4], etc There were two points raise in the question: 1) building on top of other models; 2) baselining against SOTA models. I understand that it is impossible to compare performance with all the variety of MoE models that where developed over the last two years (since 2022 for XMoE and 2023 for AdaMV), but it is always beneficial to select those most recent models of similar architectures that 1) were evaluated on the same datasets used in the paper to understand where SOTA is; 2) architectures that aimed to solve same issue (activation sparsity). Comparison against two selected baselines creates false premise that activation problem has not been challenged, and doesn't give enough information about where the proposed method lies against SOTA results (even though it is not necessary that the proposed method must improve them). 1. Soft Merging of Experts with Adaptive Routing 2. Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM 3. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models 4. DEMix Layers: Disentangling Domains for Modular Language Modeling -> A2 I understand that in the limited time of rebuttal, it is impossible to perform convincing ablation studies to answer the raised questions. Initial results shared are promising, and I encourage authors to continue investigations.