Beyond Parameter Count: Implicit Bias in Soft Mixture of Experts

1. The result in Theorem 1 is not really solid as the upper bound still hinges upon the input value. Moreover, since the authors choose a specific target function rather than consider an arbitrary convex function, I am quite skeptical about the generalization of the result. Additionally, the way this theorem is stated confuse me when I first read it. I think the writing of the statement could be improved.

2. The paper lacks the theoretical proof for the result that the soft mixture of multiple small experts is capable of representing simple target functions, which weakens the contribution of this paper considerably. Personally, the experiments on the synthetic data in Appendix B for empirically showing this result is not sufficient to convince me.

3. The principle of the proposed algorithm for choosing the best subset of important experts based on their weights in the matrix $C(X)$ seems to be reasonable. However, I believe that the claim that this algorithm can help reduce computation at inference would be more convincing if the experiments are conducted on bigger model and dataset.

• Clarity in Methodological Exposition: Certain aspects of the methodology, such as the approach for selecting the number of experts $k$ for prediction in real applications, would benefit from more detailed and clearer explanations.

• Clarity in Technical Proof Exposition: Including a proof sketch at the beginning of Appendix A would enhance readability and provide a clearer roadmap for understanding the technical details. Additionally, a more detailed explanation of why the equalities in lines 711-716 hold would further improve clarity and support the reader's comprehension of the argument. Clarification is needed regarding the notation for norms: are $|\cdot|$ and $\| \cdot \|_2$ intended to represent the same norm?

• Limited Theoretical Analysis for Multi-Expert Configurations: Although the paper offers theoretical insights into the limitations of a single expert, it does not extend the theoretical framework to cover multiple-expert $n$ or multiple-slot configurations. A more detailed analysis of how increasing the number of experts impacts scalability, convergence, or approximation accuracy would enhance the strength of the claims presented.

• Notation Section: The manuscript would benefit from a dedicated notation section to clarify certain symbols that may be unclear, such as the norm $|\cdot|$ or $\| \cdot \|_2$ for matrices and vectors. Providing a notation section would enhance readability and ensure that readers can follow the mathematical expressions consistently throughout the paper.

• More Accurate References: Lines 040 and 549-550: The correct citation should be: "Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma (2024). Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27), e2311878121." Lines 073 and 633-635: The correct citation should be: "Johan Obando-Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures of experts unlock parameter scaling for deep RL, 2024. Proceedings of the 41st International Conference on Machine Learning, PMLR 235:38520-38540, 2024." Lines 073 and 571-574: The correct citation should be: "Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y.Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1280–1297, Bangkok, Thailand. Association for Computational Linguistics."

• W1. SoftMoE Formulation: The SoftMoE formulation presented in Section 2 of this paper differs from that of the original SoftMoE paper [1]. In the original formulation, each expert processes $p$ slots, making $\Phi \in \mathbb{R}^{d \times (n \cdot p)}$ and $C(X) \in \mathbb{R}^{m \times (n \cdot p)}$. However, in this paper, the parameter $p$ is not included. Notably, setting $p=1$ in the original formulation results in an incomparability with other papers using SoftMoE, where $p$ is recognized as a significant hyperparameter. Furthermore, in line 243, the authors cite [2], claiming that this work demonstrates that even a single-expert SoftMoE can outperform traditional architectures. This comparison is inaccurate, as $p$ is not set to 1. Specifically, after reviewing the code, I observed that $p \approx m/n$; see lines 55-64 in this GitHub file: Link.

• W2. Theoretical Contribution: I don't find the results presented in Theorem 1 insightful.

  • W2.1. Trivial Result: Ignoring $p$ (the number of slots) in the SoftMoE formulation leads to a trivial outcome for $n=1$. Here, matrix $C(X)$ becomes an all-one vector of size $m$ instead of a matrix of size $m \times p$, resulting in identical output tokens and thereby ignoring any temporal information in the tokens, effectively replacing them with a fixed token. This trivializes the ineffectiveness of SoftMoE with $n=1$, which is not the case when $p > 1$.

  • W2.2. Notion of Approximation: The notion of approximation presented by the authors is non-standard. I encourage the authors to provide references or pointers to similar approaches in the literature to clarify this aspect.

  • W2.3. Proof Technique: After reviewing Appendix A, I noticed that the proof relies on a special case where a contradiction arises as matrix norms approach infinity. This is acknowledged by the authors in Section 3, where they mention that normalizing the input makes the results from Theorem 1 inapplicable.

• W3. Lack of comparison for Algorithm 1: The paper lacks a comparison with the sparse routing mechanism considered in the original SoftMoE paper [1] (cf. Section 3.2). (Correct me if I am wrong)

References:
[1] Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." arXiv preprint arXiv:2308.00951 (2023).
[2] Obando-Ceron, Johan, et al. "Mixtures of experts unlock parameter scaling for deep RL." arXiv preprint arXiv:2402.08609 (2024).

Section 3 on the representation failure of a single expert is quite trivial and obvious. Each expert essentially sees X only through a single d-dimensional linear projection, and hence any function of X that uses all of X in a non-linear manner clearly cannot be represented by a single expert regardless of how complex the expert is. Mutiple experts can potentially overcome this because they would each look at different projections.

Section 4 results on specialization are also quite weak. Choosing experts for each input separately can lead to an obviously absurd conclusion as follows. Having 10 experts each of whom predict a single class constantly would be considered highly specialised.

The proposed algorithm 1 is literally just picking the experts with the highest weight for each example, and is the first common sense thing that a practitioner looking to reduce inference complexity would do. The novelty of this approach is very minimal.