Statistical Advantages of Perturbing Cosine Router in Mixture of Experts

2024-05-23

AI Summary

This work addresses the slow statistical convergence of cosine routers in Mixture-of-Experts (MoE) models. We first establish—via theoretical analysis—that parameter estimation under standard cosine routing exhibits only logarithmic convergence due to weak identifiability of routing directions and non-identifiability of input ℓ² norms. To resolve this, we propose a perturbed cosine router that injects controlled noise into the input ℓ² norm, breaking symmetry-induced degeneracy and restoring strong identifiability. We prove that this perturbation elevates both expert weight and function estimation convergence to polynomial rates (O(1/n^α), α > 0), thereby filling a critical gap in the statistical analysis of MoE routing mechanisms. Our approach integrates least-squares estimation, PDE-based modeling, and principled perturbation design. Empirical validation on synthetic data as well as image and language tasks demonstrates substantial mitigation of representation collapse and consistent improvements in generalization performance.

Technology Category

Machine Learning: Mixture of Experts (MoE)Search and Optimization: Learning to SearchPlanning, Routing, and Scheduling: Model-Based Reasoning

Application Category

User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systemsGraph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphsSecurity and Privacy: Large-scale security measurements

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^{ au}(n))$ where $ au>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 extit{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.

Problem

• Analyzes slow convergence rates in cosine router MoE.
• Proposes perturbed cosine router to improve estimation rates.
• Validates improved polynomial convergence rates empirically.

Innovation

• Perturbed cosine router enhances MoE performance
• Noise addition stabilizes cosine router convergence
• Improved polynomial estimation rates in MoE