Statistical Advantages of Perturbing Cosine Router in Mixture of Experts

TL;DR

This paper analyzes the statistical properties of the cosine router in Mixture of Experts, revealing slow convergence rates that can be improved by adding noise, and validates these findings through extensive simulations.

Contribution

It provides the first comprehensive theoretical analysis of the cosine router in MoE, demonstrating how noise perturbation improves estimation rates.

Findings

Without perturbation, estimation rates are as slow as O(1/ log^τ(n)).

Adding noise to the cosine router improves convergence to polynomial rates.

Simulation studies confirm the theoretical improvements in both synthetic and real data.

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 \textit{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.