Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization

Summary

In the paper, the author introduces $\mu$MoE layers, which factorize large weight tensors to facilitate implicit computation, thereby accelerating the computation process. By increasing the number of experts, the model enhances its specialization in vision tasks. Further, the author conducted experiments with the GPT2 and MLP-Mixer models, demonstrating that these models maintain high accuracy and exhibit increased specialization. Both qualitative and quantitative evaluations conducted by the author were commendably thorough and well-executed. The experiment is comprehensive.

Strengths

1. **Innovative Metric:** The introduction of a quantitative metric for "class-level polysemanticity" is a novel approach to evaluating the specialization of experts within the model, providing a clear measure of the impact of increasing the number of experts. 2. **Comprehensive Evaluation:** The paper undertakes a broad range of experiments to validate the effectiveness of $\mu$MoE layers. The inclusion of both qualitative and quantitative analyses strengthens the argument for the utility of µMoE layers. 3. **Competitive Performance with Added Benefits:** The ability of $\mu$MoE layers to compete with larger models like GPT-2 and MLP-Mixer, coupled with the advantages of additional interpretability, is highlighted as a significant strength of the study.

Weaknesses

1. **Limited Scope:** The research primarily focuses on vision models and specific datasets, which may limit the generalizability of the findings. Expanding the validation to more diverse models (LLaMA, gemma, and $e.t.c$) and datasets could provide a more comprehensive view of the $\mu$MoE layers' effectiveness. 2. **Lack of Robustness Analysis:** The paper would benefit from including experiments that assess the robustness and performance of µMoE layers on out-of-distribution data, offering insights into the model's reliability under varied conditions. 3. **Need for Implementation Details:** There is a notable gap in the disclosure of practical implementation details and computational requirements for µMoE layers. Addressing this could aid in the reproducibility and application of the findings.

Questions

1. Could you specify the computational hardware and configuration utilized for these experiments? 2. Is it feasible to apply this methodology to large-scale MoE Language Models with a substantial number of experts, similar to the extensive parameter count found in Mixtral 8$\times$7B?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

The authors adequately discuss the limitations and do not have the potential negative societal impact.

Summary

The paper proposes the use of (factorized) Multilinear Mixture of Experts as an alternative to Sparse MoEs, which are prone to training issues due to the sparse top-k activation. Several factorization options are described and compared, which result in models of better accuracy in vision tasks, when matching the parameter count, compared to models using vanilla dense MLPs. The paper also includes experiments on a language task, where the proposed method is almost as good as the baseline.

Strengths

- The proposed approach (using either the CP structure, or Tensor Ring structure, or both) achieves better accuracy than MLPs on image tasks, on two different architectures: MLP-Mixer S-16 and CLIP B-32, with models of similar parameter count. - The discussion about different factorization strategies for the MLPs in the appendix is a delight. In fact, all the subsections in the appendix contain many useful details. I enjoyed reading several of them. - The main paper contains a thorough and well written description of the proposed method, as well as the experiments. Any of the details and ablations that are missing from the main paper can be found in the appendix.

Weaknesses

- The results on GPT-2 are close to the dense MLPs, but not better than them. This is problematic, since this architecture is the biggest one trained (with 124M params) and probably the one trained on the richest data source. This begs the question of whether the proposed approach can scale well to bigger backbone architectures and datasets. - Appendix C.2 discusses why the low-rank assumption in MLPs is reasonable. In particular, a trained MLP-Mixer is taken and the rank of the matrices is reduced with truncated SVD. This shows that the low rank approximation is reasonable once the model is trained, but completely ignores the effect of this approximation _during training_, especially when the architecture size and training data scales. - Most importantly: there is not quality comparison between the proposed approach and Sparse MoEs. The recipes for training Sparse MoEs have improved significantly, and they are present in most state-of-the-art language models, and multi-modal models (for instance, take a look at the [Mixtral report](https://arxiv.org/abs/2401.04088), [DeepSeekMoE](https://arxiv.org/abs/2401.06066), [Gemini 1.5 report](https://arxiv.org/abs/2403.05530), [Databrick's DBRX](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm) which is available on HuggingFace, and many others). In vision tasks, Soft MoEs in ["From Sparse to Soft Mixture of Experts"](https://arxiv.org/abs/2308.00951) have shown also much better results than vanilla dense transformers and Sparse MoEs. Thus, either of these should be considered as a stronger baseline than the use of vanilla MLPs. - A discussion about the speed of the different methods is also missing from the main paper, only parameter count is considered, which may not be the only (or main) limiting factor. ----------- During the rebuttal the authors satisfactorily addressed two of my weaknesses. They also pointed out that the method is not an alternative to Sparse MoEs, but to provide a better parametrization of standard (dense) models. I still believe that Sparse MoEs would be a better baseline than Dense MoEs (formerly referred to as Soft MoEs in the paper), as I mentioned in my comment to the author's rebuttal. I also find concerning the fact that when the method is applied on relatively large models, it provides little benefit, although I acknowledge that it may be more interpretable. Given this, I'm slightly increasing my initial rating from "Reject" to "Borderline accept".

Questions

This is a suggestion, not a question: The term "Soft MoE" used to refer to the seminal work of Robert A. Jacobs et al. from 1991 can be confusing with the recent work ["From Sparse to Soft Mixture of Experts"](https://arxiv.org/abs/2308.00951) by Puigcerver et al, from last year. Since the term Soft is only used once in the 1991 work, I would suggest using "Dense MoE" to refer to it, which is also more opposed to "Sparse MoE".

Rating

5

Confidence

5

Soundness

2

Presentation

4

Contribution

1

Limitations

No negative societal impacts are specific to this work, in my opinion.

Summary

This paper proposes the Multilinear Mixture of Experts (µMoE) layer for scalable expert specialization by performing an implicit computation on prohibitively large weight tensors entirely in factorized form. Both qualitative and quantitative results show that scaling µMoE layers when fine-tuning foundation models for vision tasks leads to specialized experts at the class level.

Strengths

* The problem studied in this paper is well-motivated, specialization of experts is one of the most important goals in MoE. * The formulated illustration of factorization is clear. * Figures and tables are clear and easy to read.

Weaknesses

* It would be better if more results on language models could be included.

Questions

* The description of the dense µMoE layer in Section 3.1 is a little bit confusing to me, for example, what is the difference between it and the soft MoE layer? * According to the methodology description, e.g. in Section 3.1.1, there is no “top-k” routing in the µMoE layer. Nevertheless, how are the “top-activating” experts selected in Figures. 4 & 5?

Rating

5

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes.

Summary

The paper proposes a new way of computation in a mixture of expert (MoE) layer. Unlike SparseMoE, instead of using the top-K operator and then doing full computation only for that expert, this paper do not do topK, but does a linear combination of all experts, but each expert is now a low rank matrix, which saves compute/parameters. The authors do show some properties that performing MoE computation in such a way can help in interpretability of the experts, their task specialization and so on. This can further help in certain cases like to mitigate biases.

Strengths

- The idea seems to be simple enough to be easily implemented in existing MoE setups. - Anecdotal and some statistical results suggests that it does help in mitigating bias in certain cases in a constrained setting.

Weaknesses

- Figure 2 can be made much stronger by showing statistical differences between the two different cases -- like entropy of experts, overlap among experts, distribution, etc, and also comparing it with similar parameter and/or compute sized sparse/soft MoEs. - Again, for Figure 3, there should be a way to compare it at least with sparse MoE. - Comparison with Sparse and SoftMoE in vision and LLMs are missing -- keeping the same parameter size and/or compute fixed across algorithms, as well as how does performance scale with more parameters in the MoE part. These are essential, as increasing the capacity of the network by increasing parameter count, while preserving the same compute is one of the main motivations behind MoEs, other than the interpretability lens which the works focuses mostly on. - Several ablations seems to be missing -- performance improvements when you increase number of experts, placement of the MoE layers, experiments on different model sizes, larger datasets, effect of different non-linearities in the gating function, etc.

Questions

- As mentioned above, I think it is quite important to show how \muMoE perform compared to SoftMoE and SparseMoE. - Making Figure 2 more statistically significant rather than anecdotes. - More experiments needs to be conducted to showcase the efficacy of the interpretability of the experts and its use cases. And how are those benefits compared to Soft and SparseMoE. Suggestions: - I think the notation used in Eqn 1 is a bit difficult to follow, specifically with the colons used without brackets.

Rating

6

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

The authors have identified the limitations of the work.

Summary

The paper presents an interesting and novel approach with the µMoE layer, offering a method for scalable expert specialization. Extensive experiments in multi-dimensional have been conducted to show the effectiveness of the proposed method. However, the marginal improvements in results, coupled with several other weaknesses, including experimental variability, novelty of factorization, clarity of metrics, and scalability, limit the overall impact of the work.

Strengths

1. Introducing the Multilinear Mixture of Experts (µMoE) layer is novel and addresses the computational cost challenges associated with scaling the number of experts for fine-grained specialization. The mathematical formulation of µMoE layers is thorough, and the factorization approach is well-explained and technically sound. 2. Practicality: The proposed µMoE models demonstrate a more parameter-efficient alternative to traditional MoEs, an important consideration for practical applications. 3. The paper includes extensive qualitative and quantitative evidence to support the claims. Experiments on vision and language models illustrate the versatility and effectiveness of the proposed method. 4. This paper is well-written and easy to follow.

Weaknesses

1. The idea of factorization is not new, and it seems that the authors simply apply factorization in the MoE context. The authors should specify and address any new challenges that arise when applying factorization in the MoE context. 2. Marginal Improvements and lack of multiple runs: The experimental results in Table 2 show only marginal improvements over existing methods, which might not be compelling enough to justify the additional complexity introduced by the µMoE layers. All experiments were only run once without reporting mean and standard deviation. It is recommended that the authors add results from at least three runs to provide a better understanding of the variability and reliability of their results. 3. The metrics used in the paper are not clearly explained. The authors should include a paragraph either in the experiments section or in the appendix to explain the metrics used. 2. The evaluation captures expert behavior only on the test set, not on out-of-distribution data, which is crucial for understanding the robustness and generalizability of the model. 4. While the paper qualitatively demonstrates interpretability, it lacks objective, quantitative measures of the interpretability gains compared to other methods. 8. It is unclear how the proposed methods perform across different model architectures and sizes. The authors should discuss whether the method scales well and provide evidence to support their claims. More Comments: 1. The paper is generally well-written and clearly presents the proposed methodology and results. However, some sections, particularly those dealing with mathematical formulations, could benefit from additional explanations or visual aids to enhance comprehension. 2. Addressing the limitations and adding some discussions, such as out-of-distribution performance and quantitative measures for interpretability, would significantly strengthen the paper. Further research should also focus on demonstrating more significant performance gains and improving the ease of implementation to enhance the practical utility of the proposed method.

Questions

N/A

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors addressed limitations and the direction for future research.

Discussion summary

Dear Area Chair(s) and Reviewers, We thank the AC kindly for their effort in handling our submission, and for the help facilitating the discussion with the reviewers. Throughout the discussions, reviewers praised the paper for the well-motivated approach to the important goal of scalable expert specialization (**Reviewer XreH**), its technical soundness and `"extensive qualitative and quantitative evidence to support the claims"` (**Reviewer Su68**), and the innovative new metric introduced to quantify expert specialization (**Reviewer 1sqG**). **All three of the reviewers to respond to our rebuttals (**qC51**, **XreH**, and **YTQ3**) raised their score following our new experimental results**: **Reviewer qC51** increased their score to a 6 following the new comparisons to Soft and Sparse MoEs, where muMoEs demonstrated Pareto-optimal improvements to accuracy for CLIP fine-tuning. **Reviewer YTQ3** also increased their score after these experiments and additional clarifications about the paper's aims. We also performed new experiments on steering GPT-2 models for **Reviewer XreH**, who raised their support for the paper too. Two reviewers are yet to respond to our rebuttals. However, **our rebuttals address all concerns of the remaining reviewers**, which we summarize below (and describe how the key concerns are addressed): - **Reviewer-1sqG**: the reviewer already leaves a favorable score of 7, and we propose to acknowledge additional limitations in the camera-ready (regarding robustness to OOD data and experiments on very large models). - **Reviewer-Su68**: we clarify in our rebuttal that **the fairness results are indeed already over multiple runs**, and highlight where detail of the metrics can be located in the appendix. Given the lack of responses from the remaining reviewers, we kindly request that you consider our rebuttal's comprehensive resolution of the remaining concerns when making your decision about the paper. We will incorporate all changes suggested by the reviewers and clarifications that emerged from the insightful technical discussions into our camera-ready version. We are grateful to the reviewers for helping improve the paper significantly. Sincerely, Submission #3407 Authors.