Statistical Advantages of Perturbing Cosine Router in Mixture of Experts

Authors: Huy Nguyen, Pedram Akbarian, Huyen Trang Pham, Thien Trang Nguyen Vu, Shujian Zhang, Nhat Ho
Primary area: learning theory
Keywords: mixture of experts, cosine router, perturbation

Abstract:
The cosine router in Mixture of Experts (MoE) has recently emerged as an attractive alternative to the conventional linear router. Indeed, the cosine router demonstrates favorable performance in image and language tasks and exhibits better ability to mitigate the representation collapse issue, which often leads to parameter redundancy and limited representation potentials. Despite its empirical success, a comprehensive analysis of the cosine router in MoE has been lacking. Considering the least square estimation of the cosine routing MoE, we demonstrate that due to the intrinsic interaction of the model parameters in the cosine router via some partial differential equations, regardless of the structures of the experts, the estimation rates of experts and model parameters can be as slow as $\mathcal{O}(1/\log^{\tau}(n))$ where $\tau > 0$ is some constant and $n$ is the sample size. Surprisingly, these pessimistic non-polynomial convergence rates can be circumvented by the widely used technique in practice to stabilize the cosine router --- simply adding noises to the $\ell^2$-norms in the cosine router, which we refer to as *perturbed cosine router*. Under the strongly identifiable settings of the expert functions, we prove that the estimation rates for both the experts and model parameters under the perturbed cosine routing MoE are significantly improved to polynomial rates. Finally, we conduct extensive simulation studies in both synthetic and real data settings to empirically validate our theoretical results.

Decision: accept (poster)
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730431684305

thread_type:
official_review

<thread_object>
creator: reviewer_vrcc

<content>
summary:
The paper analyzes the generalization bounds of the vanilla cosine router used in Mixture of Experts (MoE) architectures and proposes the addition of noise to enhance these bounds when the experts are Feed-Forward Networks (FFNs) or polynomial functions. Experimental results validate the effectiveness of the proposed method and the theoretical findings.

strengths:
The paper is clearly and logically presented, making it easy to follow the methodology and findings. The experimental results validate the theoretical analysis and demonstrate the effectiveness of the proposed method, providing solid support for the authors' claims.

weaknesses:
As I am not familiar with the statistical techniques used in this field, I found it challenging to fully assess the novelty of this paper compared to prior work. Specifically, in my literature review, I noted that [1] appears closely related to this study. Could the authors clarify any key differences or advancements offered by their approach compared to [1]?
Additionally, as noted by [2], the sparse Mixture of Experts (MoE) approach appears to offer improved generalization capabilities.

[1] Nguyen, Huy, et al. "Statistical Advantages of Perturbing Cosine Router in Sparse Mixture of Experts." arXiv preprint arXiv:2405.14131 (2024).

[2] Li, Bo, et al. "Sparse Mixture-of-Experts are Domain Generalizable Learners." The Eleventh International Conference on Learning Representations.

questions:
1. The intuition behind Equation (4) is unclear. Could the authors elaborate on how the partial differential equation (PDE) influences generalization performance? Additionally, the benefit of the condition $\beta_1\top \frac{\partial \tilde{H}}{\partial \beta_1}(x, \beta_1m \eta)\neq 0$ is not clearly explained.
2. The types of experts used in the experiments should be specified more clearly in the Experimental section. For example, are the experts linear models, ReLU-based feedforward networks (FFNs), or other types?
3. Given that multi-layer perceptrons (MLPs) are commonly used in MoE architectures, how does the perturbed cosine router compare to a linear router with the same generalization bound? What specific advantages does it offer?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730578665129

thread_type:
official_review

<thread_object>
creator: reviewer_hemt

<content>
summary:
This paper studies the cosine router from a statistical perspective. The main theoretical result is that the cosine router cannot be learned with any polynomial rate, whereas a slight modification (the perturbed cosine router) can. This is then backed up with several experiments showing that the perturbed cosine router outperforms the cosine router.

strengths:
Routing schemes such as the cosine router are very important components in modern machine learning models, but there has been little theoretical study. This paper presents a compelling theoretical picture that a perturbed cosine router should have better rates than the standard cosine router. The experiments are also extensive and seem quite convincing. Finally, the paper is generally quite well-written.

weaknesses:
My main concern is that some of the most important theoretical conclusions of the paper are not formally stated, let alone proven: the claims directly after Theorem 2 about how the theorem implies slow rates for parameter estimation and expert estimation. I (roughly speaking) believe the intuition for (i), though I have some questions -- see below -- but I completely do not understand the intuition for (ii). The claim is that since the parameters \eta are hard to estimate, and h(\eta) is Lipschitz in \eta, then h(\eta) should be hard to estimate as well. But doesn't the inequality go in the wrong direction? e.g. if h(\eta) is constant in \eta, then it's trivial to estimate, and also optimally Lipschitz?

A more minor point -- see questions below -- is that the claim of 1/polylog(n) rate lower bound seems unsubstantiated, since all that is proven is that the rate cannot be any polynomial.

It's possible I am missing something, so I will be happy to increase my score if these concerns are addressed.

questions:
- In e.g. Theorem 2, it seems like an abuse of notation to write \bar{G}_n as in G_k(\Theta) when actually it's an estimator, i.e. a mapping from n samples to G_k(\Theta)? Or am I misunderstanding the statement?
- In Lemma 2, why isn't it enough for the limit to be finite? Why does it have to be 0? It seems like if it's a constant then the application on line 931 would still work, since it seems like C_1 doesn't have to be super small.
- I don't understand the claimed implications (i) and (ii) of Theorem 2. For (i), what if the rate of estimating \beta*_0 is bad? What if only one of {\beta*_1, \eta*} has bad rate and the other has good rate? Finally, just because the rate is slower than any polynomial doesn't mean that it could be 1/polylog(n)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730738272859

thread_type:
official_review

<thread_object>
creator: reviewer_idq2

<content>
summary:
This paper provides a theoretical understanding of the advantage of perturbing cosine router in Mixture-of-Experts (MoE) compared to the vanilla cosine router in terms of sample efficiency. In regression setting where the ground truth itself is generated by an MoE with cosine router and if the expert function satisfies a certain condition ("strong identifiability"), the paper shows that, by adding noise to the L^2 norm of the router weights and the tokens, the intrinsic interactions between the elements of the router parameters can be avoided and thus a better sample complexity can be achieved

strengths:
Originality: The paper provides the first theoretical study of the cosine router MoE and it's perturbed version and confirms the theoretical advantage of the later one, which can be considered as novel. 

Quality: The author addresses some technical challenges to complete the theoretical study. For example, the normalization of the cosine router introduces more sophisticated parameter interactions among the elements of the router parameters which needed to be addressed.

Clarity: The paper is well-written and the implication of the results is easily understandable

Significance: The results are significant from the sense that, it provides the theoretical validity of using perturbation in cosine router

weaknesses:
1. The scope of the theoretical result is limited as it required that, the ground truth is itself generated by a cosine routing MoE. It is not clear if practical datasets meet this assumption. Could you please discuss the implications of your results for real-world datasets that may not perfectly match the theoretical assumptions? Any intuition for how their findings might generalize to more realistic settings.

2. In the numerical experiments at section 5.1, the ground truth is generated by an MoE with perturbed cosine router rather than a clean cosine router. Also, the weights are initialized for SGD near the ground truth parameters (which are from the perturbed router). Therefore, for a fair comparison, the ground truth should be generated by the clean router to see whether the perturbed version converge faster (in terms of number of samples required) compared to the clean one. Is it possible to run an additional experiment with the ground truth is generated by a clean cosine router, to more directly test the theoretical claims about improved convergence of the perturbed version?

3. In practical MoE models, all the experts are FFN with ReLU activation with at least one hidden layer (i.e. two-layer FFN). However, in the paper, the exampled FFN version is one-layer. Can the analysis be extended to two-layer case? In other words, does the two-layer FFN with ReLU activation satisfy the strong identifiable condition? Could you please add some discussion on whether and how the analysis might extend to multi-layer FFNs, which are more common in practice?

questions:
The theoretical results are training algorithm agnostic. I'm curious about the analysis under SGD. Here, it has been claimed that adding perturbation (i.e. noise at the router) remove the parameter interaction among the elements of the router parameter. As SGD also add noises, can the SGD based algorithm also or remove the interaction (even for a linear router)? Can you comment on this?


---------------------------
After the response.
Thank you for your response to my questions. However, I decided to decrease my score to 5. My decision is based on the following reasons (related to weakness 1 and weakness 3 of my official review):

1.	The strong identifiability condition is not satisfied by a two-layer ReLU expert function. However, two-layer ReLU is the most common expert function used in practice. Moreover, there exists theoretical works analyzed the two-layer ReLU as expert functions of MoE (e.g., [1], [2])
2.	The authors claimed that the theoretical results can be extended to the more general setting where the ground truth is generated by an arbitrary regression function instead of the perturbed cosine router. A formal discussion has been provided about the generalization at section 4 of the paper. In that case, the formal results of the paper (i.e. Theorem 1-4) should be presented in the more general setting instead the current restrictive setting.

[1] Chen, Z., Deng, Y., Wu, Y., Gu, Q., & Li, Y. (2022). Towards understanding mixture of experts in deep learning. arXiv preprint arXiv:2208.02813.

[2] Chowdhury, M. N. R., Zhang, S., Wang, M., Liu, S., & Chen, P. Y. (2023, July). Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks. In International Conference on Machine Learning (pp. 6074-6114). PMLR.

After the author response again, I decided to change the score back to 6. The reason is in the comment.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731881667605

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear Area Chairs and Reviewers,

We would like to express our gratitude for your 
valuable feedback and constructive questions, which have greatly contributed to the improvement of our manuscript. Additionally, we really appreciate the following endorsement from the reviewers for our work:

**(i) Contributions:** The theoretical study of the cosine router MoE and its perturbed version is **novel** *(Reviewer idQ2)*. This paper presents a **compelling** theoretical picture that a perturbed cosine router should have better rates than the standard cosine router *(Reviewer hEMT)*.


**(ii) Soundness:**  The experiments are also **extensive** and seem quite **convincing** *(Reviewer hEMT)*. The experimental results validate the theoretical analysis and demonstrate the effectiveness of the proposed method, providing **solid** support for the authors' claims *(Reviewer VrcC)*.

**(iii) Presentation:** The paper is **well-written** *(Reviewer idQ2, Reviewer hEMT)* and the implication of the results is **easily understandable** *(Reviewer idQ2)*. The paper is **clearly and logically presented**, making it **easy to follow** the methodology and findings *(Reviewer VrcC)*.

Subsequently, we will address the concerns of the reviewers separately in our responses to their reviews and include according changes in the revision of our paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces the perturbed cosine router in Mixture of Experts (MoE) models, showing that it improves sample efficiency by adding noise to router parameters. The authors provide theoretical guarantees and empirical validation, demonstrating that the perturbed router outperforms the standard cosine router.

Reviewers appreciated the novelty and theoretical contribution of the paper, noting the solid empirical validation and the clear explanation of the method’s advantages. The results were seen as significant, addressing an important problem in MoE models.