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
Venue: ICLR 2025 Poster (2025)
Keywords: mixture of experts
cosine router
perturbation
Primary area: learning theory

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.

Published as a conference paper at ICLR 2025
STATISTICAL ADVANTAGES OF PERTURBING COSINE ROUTER IN MIXTURE OF EXPERTS
Huy Nguyen'
Pedram Akbarian"* !
Trang Pham :
Trang Nguyen
Shuijan Zhang!
Nhat Ho!
1 The University of Texas at Austin
2 Qualcomm AI Research
ABSTRACT
The cosine router in Mixture of Experts (MoE) has recently emerged tractive 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 (1/ log" (n)) where T > O is some constant and v is the sample size. Surprisingly, these pes- simistic non poolyoomiilll convergence rates can be circumvented by the widely used technique in practice to stabilize the cosine router - simply adding noises the f--oomms in the cosine router, which refer 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 empirically validate our theoretical results.
INTRODUCTION
Proposed by Jacobs et al. (1991) and Jordan & Jacobs (1994). Mixture of Experts (MoE) has been known as effective statistical method to incorporate the capabilities of various specialized models called experts. Different from conventional mixture models (Lindsay 1993) in which the mixture weights scalars, the MoE rather utilizes routing mechanism determine set of weights depending on an input token. In particular, the router first computes the similarity scores between each token and experts and then assigns more weights to the more relevant experts determined based on those scores. To further improve the scalability of the MoE, Shazeer et al. (2017) has recently introduced a sparse variant of this model, which routes each input to only subset of experts. This sparse MoE model allows to increase the number of learnable parameters with nearly constant computational overhead. As a consequence, the sparse MoF has been leveraged in several applications, including large language models (Jiang et al.) 2024 Puigcerver al. 2024 Zhou et al.
2023b Dai et al. 2024; Pham et al. 2024), computer vision (Riquelme et al. 2021 Liang et al.
2022), speech recognition (You et al. 2021 Gulati et al. 2020 Peng et al. 1996) continual learning
(Le et al. 2024. Li et al. (025). ), multi- task learning (Hazimeh et al. 2021 1). and other applications
(Han et al. 20044 Chow et al. 2023 Le al. 2025
In the above applications, practitioners often use a linear router which calculates the similarity score by taking the inner product of token hidden representation and expert embedding. Nevertheless, Chi let al. 2022) discovered that utilizing the linear router might induce the representation collapse issue. This phenomenon occurs when a fraction of experts govern the docision-making process, leading to the redundancy of other experts. In response, Chi et al. (2022) proposed an alternative known as
"qqual Contribution
QQullcomm Vietnam Company Limited

---

Published as a conference paper at ICLR 2025
cosine router. In partieular, this router begins with projecting the token hidden representation into a diw-demeniional space, followed by applying L normalization to both the token representations and expert embeddings. By doing so, the similarity scores become more stable, circumventing the dominance of certain experts. The efficacy of the cosine routing MoE has been experimentally demonstrated in language modeling (Chi et al.] 2022), and domain generalization (Li et al. 2023)
On the other hand, a comprchensive theoretical study of the cosine router has remained lacking.
In the literature, there have been some attempts to understand the MoE models with different types of gating functions whose outputs are the composition of some functions and the routing scores. First of all, considering the classification problem with cluster structures, Chen et al. (2022) demonstrated that the router operated by a neural network could learn the cluster- features. which helped divide a complex problem into simpler classification sub-problems that individual expert could handle. Next, (Ho et al. (2022) studied the expert estimation under an input-free gating Gaussian MoE model, showing that the rates for estimating experts depend on the algebraic structures among experts. Subsequently, the Gaussian MoE model with softmax gating was explored in (Nguyen al. 2023 2024a) which pointed out that interactions among model parameters via some partial differential equations (PDE) did harm the expert estimation rates. Saying that the setting of Gaussian MoE was far from practice, Nguyen et al. (2024c) rather took into account a regression framework with the regression function being linear router MoE model. They verified the benefits of formulating experts as feed-forward networks with popular activation functions like ReLU and sigmoid from the perspective of the expert estimation problem.
It is worth noting that the expert estimation problem allows us to capture how fast an expert specializes in specific task, which one of the most important problems in the MoE literature known as expert specialization (Dai et al. 2024; Krishnamurthy et al. (023). Furthermore, from the convergence analysis of expert estimation, we can gain several insights for designing the router and expert networks (see Section 4). Therefore, we will investigate the effects of the cosine router on the convergence of expert estimation in this paper. For that sake, let us now present the problem setting formally.
Problem setting, We assume that (X1, Y1), (X2, Y2), 12),,..,,(X.. (Xn, Yn) Rd R is an i.i.d sample of size n generated according to the following model
Y; = fa. (X,) + i= 1, 1.,,..,,
(1)
where regression function fa. (-) takes the following form:
k-
Softmax (91. Tx |18 ||=||| + Bos ((x,,)))
(x) =
(2)
Here, the function h(,, 17 is known the expert function, which assumed to be of parametric
form. Meanwhile, (36 fi, 17 ) are true yet unknown parameters in the parameter space + c
R x R44 x Rd's and G. == Ein ( denotes the associated mixing measure, i.e. a weighted sum of Dirac measures 5. Additionally, we define for any vector o = (11, ,,,,,,,,,, vk.) in Rt that Softmax(v::) := xxp(v::/_jt. exp(r;). In the cosine router in equation (2). we omit the step of reducing the dimension of the input token and assume that it has already been in a low-dimensional space for simplicity. Furthermore, we assume that X1, X2...., : Xn are i.i.d. samples from some probability distribution μs. Lastly, &1, ..,..., En are independent Gaussian noise variables such that [X, = and Var(s,|X,) = o² for all < i S n. Notably, the Gaussian assumption is just for the simplicity of the proof argument.
.......,,,
Least squares estimation (LSE).. To estimate the true parameters (19 Pi n/ ) or, equivalently, to estimate the true mixing measure Gs, we leverage the popular least squares method (see, e.g., van de Geer 2000). Formally, the mixing measure G., is approximated by
a
G, = arg
arg min G
Σ (y, (xi-
Under the exact-specified setting, i.e., when the true number of expert k. is known, the minimum in the above equation is subject to the set of all mixing measures with k, atoms, denoted by (O) == {G = Σ cxp([8))6( over-specified setting, i.e., when k. is unknown and the true model ( is over-specified by a mixture
(Boi, 31s. n.) E {}. On the other hand, under the
2

---

Published as a conference paper at ICLR 2025
of k experts where k > k., the minimum is subject to the set of all mixing measures with most A: atoms, i.c., (x))) == {G = Σ exp(36))0(8...)) : 1 S S k, (1., Th)) E }}.
Universal assumptions. In the sequel, we implicitly impose four following mild assumptions on the model parameters, which were widely used in previous works (Nguyen et al. 2024c), unless stating otherwise:
(A.1) Convergence of LSE: The parameter space θ S R x Rd x Id is compact, while the input space X C Rd is bounded. This helps ensure the convergence of least squares estimation.
(A.2) Distinct experts: The true parameters ....... ME. are distinct so that the experts are different from each other. Furthermore, the expert function h(.,)) is Lipschitz continuous w.r.t its parameters and bounded.
h(.nn)).. "nii....,((.....................
(A.3) Identifiability of the MoE: In order that the cosine routing MoE is identifiable, i.c., fa(œ) = fc. (x) for almost every implies that G G., we let Puk. O.
(A.4) depuddde router: To ensure that the router is input dependeeenn,nt, we assume that at least one among the parameters Si....... ................................ Six is non-zero.
Technical challenges. The normalization of parameters in the cosine router leads to a fundamental challenge in theory. In particular, to establish parameter and expert estimation rates based on the convergence rate of the regression function, we rely on decomposing the discrepancy fa., (x)- (a) into a combination of linearly independent terms by applying Taylor expansions to the product of the softmax's 's numerator and the expert function, i.e. (z, $1. n1) == exp [+]]]]]]]+]]]
3 h((,, n)) However,
the normalization of 81 in the cosine router leads an intrinsic interaction among the elements of the parameter 31 via the following PDE:
0H 8T
03 (I, 11, =
Although parameter interactions expressed in the language PDEs have been observed in Nguyen et al. (2024c). the structure of the above interaction is much more sophisticated (even hold for the first-order derivatives while those in Nguuen et al. (2024c) occurs only when taking the second-order derivatives). Thus, this PDE induces several linearly dependent terms in the Taylor expansion, and we have to aggregate their coefficients in order to form the desired combination of linearly independent terms. Then, the resulting coefficients become complex, thereby negatively affecting the convergence of expert estimation. To the best of our knowledge, such a phenomenon with the cosine router has never been observed in previous works.
Main contributions. In this work, we develop a comprehensive theoretical analysis of regression function estimation well as parameter and expert estimations under the cosine router MoE model (I). Our contributions are two-fold and can be summarized as follows (see also Table
1. Cosine router:
Equipped with the cosine router, we demonstrate that under both the exact-specified and the over-specified settings, the rates for estimating ground-truth parameters Bü Si, and 1/ are slower than any polynomial rates and, therefore, could be as slow as Op(1/ log" (n)), where is some constant. These slow rates are attributed to the internal interaction among router parameters expressed by the PDE in equation (4). As a result, the estimation rates for experts h(, mz are also negatively affected, and could be of order (r(1/ log" (n)).
Perturbed cosine router: In response, we propose novel router called perturbed cosine router in which we add noises to the L norms of the token representations and the expert embeddings.
This not only helps stabilize the router but also eliminates the intrinsic interaction in equation (4). Additionally, we also establish identifiability conditions to characterize expert functions that have faster estimation rates than others under the exact-specified and over-specified settings, respectively. Those conditions indicate that the rates for estimating experts, which are formulated as networks (FFNs) with widely used activation functions such as ReLU and GeLU, are significantly improved, ranging from Op(
(atanooppnn)))))) (/log(n)/n) to Op( (p(Vlog(n)) log(n)/n).
Outline. In Section 2 establish the convergence rates of parameter and expert estimations under the setting of the cosine router MoE, Then, we derive these rates when the cosine router is replaced by the perturbed cosine router in Section Based on these theoretical results, we provide few practical
3

---

Published as a conference paper at ICLR 2025
Table Summary of worst possible estimation rates for linear experts, polynomial experts and FFN experts equipped with the ReLU activation function under the MoE with linear router (Nguyen et al. 2024c), cosine router (ours) and perturbed cosine routers (ours).
Routers/ Experts
Linear: a x b
Polynomial: x + b)P,
ReLU FFN
Linear
1/ log' (n)
//log' (n)
n 174
Cosine
1/log" (n)
1/log" (n)
1/log (n)
Perturbed cosine
1/ log' (n)
n 174
1/4
Next, we empirically verify the (theoretical) benefits of the perturbed cosine router over the cosine router under both the synthetic and real data settings in Section before concluding the paper in Section 6 Finally, proofs and additional details of the experiments are deferred to the Appendices.
implications in Section 4
Notations. We let [n] stand for the set f1, nn for any n E N. Next, for any set S. we denote [S] as its cardinality. For any vector v R. and == (a1 .,,,,........................... aa) N we let v = 2 2/2 ...ug,, |o] == 24 + U22 +... + Vu and a! == aylaz!... agl, while ||a|| stands for its - value. Lastly, for any two positive sequences (a,)n21 and (bn m2 we write a, = (b,,)
or an S bn if there exists C > 0 such that on < Cb, for all n E N. Meanwhile, the notation an = (bb)) indicates that an/b, br is stochastically bounded.
2 COSINE ROUTER MIXTURE OF EXPERTS
In this section, characterize the parameter and expert estimation rates under the over-specified setting of the cosine router MoE. We first start with the convergence rate of the regression function estimation fa. to the true regression function fa. under the L norm in the following theorem:
Theorem 1. Given the least-square estimator Gn defined equation B.. the regression estimator fe_ (.) converges to the true regression function fc. (-) at the following rate:
Iifa fc. (c = 10 Orr(Vlo((n)/m)
The proof of Theorem 1 is in Appendix B. The result of Theorem I indicates that the re- gression estimation rate parametric. Therefore, as long as we can establish the lower bound Ilaa, fa. [[)(()) 2 [(G,, G.) where L is some loss function among parameters, we arrive the parameter estimation rate L(G,,,G, Op( /lo(()//n). This approach is the key component of the convergence rates of parameter and expert estimations under the cosine router MoE. In the sequel, we will consider the over-specified setting of the cosine router, while the results for the exact-specified setting will be presented in Appendix A.2
Recall that under the over-specified setting, the true number of experts is unknown. Then, based on the notion of Voronoi cells (Manole & Ho 2022). we will construct Voronoi loss function among parameters tailored to this setting.
Voronoi loss. Let G be mixing measure with k'' atoms w; == (81r, n:). Then, we distribute these atoms to the Voronoi cells generated by the atoms wj = (8i n1)) of G., which are defined as
A = A,(G) (i E [k'] ||wr will s ||ا; w 11, AA + #
(5)
Then, the Voronoi loss of interest is given by
h.
(G, = Σ Σ
exp(36e) -xx((3x))
exp( || + [1
1=1
EEA,
1=1 LEA,
where r 2 is some constant, A&1 := 31. Bi, and Anyy := -15
Note that, due to the parameter interaction inside the cosine router captured by the PDE the lower bound Ilfe_ -Galll2()) L,, (G.,, G.) does not hold true, and thus, we cannot achieve
4

---

Published as a conference paper at ICLR 2025
the desired bound (G,, G.) = (pr((Vlo(nn///)) mentioned at the beginning of Section 2 By contrast, we show in Appendix B33 an opposed result to the previous lower bound, saying that
lim 6 6-00 GEEL (e)) (G)..))
inf
IIfe Sc. (//) (i, [o. G.) - =,,
for any r 2 1. This result implies the following minimax lower bound of parameter estimation:
Theorem 2. Under the over-specified setting, the following minimax lower bound of estimating G.,
inf (e) GEG,(
sup
(Gn, G)] 2 n n-1/2
(())
holds true for any 2 1, where Efg indicates the expectation taken w.tt the product measure with fü and the infinum over all estimators taking values 9 (e).
See Appendix for for the proof of Theorem 2 There are two main implications of the above result:
(i) Parameter estimation rates. The minimax lower bound together with the formulation of Cir indicate that at least one among the rates for estimating parameters 8853 Pös- 'i is slower than any polynomial rates Op(n~1/24 and, thus, could be of order Op log" (n)), for some constant
0.
(ii) Router estimation rates: When the estimation rate of either 3 or S6 is slower than any polynomial rates, since the softmax function is Lipschitz w.r.t the Euclidean norm (Gao & Pavel 2018). we deduce that the worst possible rate of estimating the cosine router or the mixture weights in equation (2) could also be slower than any polynomial rates. In practice, the router and the expert networks are trained simultancously (see Section 1.2 in (Shazeer et al. 2017)). Thus, the slow convergence of the router might decelerate the model convergence.
(iii) Expert estimation rates. Assume that Gn E exp (Boz ( Since the expert h((,)) Lipschitz continuous, it follows that
sup h(x, 77t) h(z, n1)) 1,7 - nníl--
(6)
x
Consequently, the worst possible rate for estimating the expert h(xz, n;) is identical to the worst possible rate for estimating the parameter 17. For instance, if the expert function takes the polynomial form of hxx,, 1) := (mz)) (considered in (Mendes & Jiang 2012 Nguyen et al. 2021 2024c)), where assume E IR for simplicity, then we have [h(x, 17) h(
7-' zz) As result, the expert estimation rate is exactly the parameter estimation rate. Thus, when the estimation rate of n3 is slower than any polynomial rates, the worst possible estimation rates for the experts h(- n))) could also be slower than any polynomial rates and be as slow Op(1/ log"(n)). This indicates that the cosine router is even less sample efficient than the linear router (see also Table [T]. Note that by employing techniques of partitioning the input space in (Nguyen et al., 2024b), we can show that these results still hold when using the Top-K sparse softmax gating function. Therefore, in order to improve the sample efficiency while preserving the robustness to the representation collapse. we need to slightly modify the structure of the cosine router in the next section.
1) "
3 PERTURBED COSINE ROUTER MIXTURE OF EXPERTS
In this section, we demonstrate that the pessimistic non-polynonin convergence rates of parameter and expert estimation under the cosine router can be casily circumvented by the widely used technique in practice to stabilize the cosine router: adding noises to the L2 norm in the cosine router. Although this perturbation technique and the cosine router have been well studied in the literature, like to emphasize that to the best our knowledge, our work is the first to be aware of combining these two methods together which refer to the perturbed cosine router MoE. We now present the formulation of a MoE with the perturbed cosine router under the regression setting.
would
Problem setup for the perturbed cosine router MoE model. We assume that i.i.d. sample of size n: (X1, Y1), (X2, Y2), Yz),..., (x,, Yn) E dR x is generated according to the model
= ga. (X₂) + == i
where regression function gG. (-) takes the following form:
Tx (113)) + T1) [| + T2)
9c. (x) =
Softmax
+ 8
(8)
5

---

Published as a conference paper at ICLR 2025
Here, T1, 72 O are two noise hyper- syperr-prammeers. The main difference between the two regression functions fc, and 9G, is the noise syper-parameters T1, 72 that we add to the norms of the expert embeddings S4, and the token input which leads to the perturbed cosine router. By doing so, the pa- rameter interaction inside the router as in equation (4) does not occur. More specifically, let us denote
(2,, (1, 17) == exp (| +
3,7
h(x,, ri), then it can verified that 31 aMl ap, (x, B1, (1) + O.
Least squares estimation. Similar to the cosine router setting, we can estimate the unknown ground-truth parameters (3ha Big ) using the least-square estimator, which is given by:
G == arg min
(Xi-9o(X.))
G
In the following theory, we provide a convergence rate of regression function estimation under the perturbed cosine router MoE model.
Theorem 3. Given a least squares estimator G, defined in equation D the regression function estimation 9ū. (-) admits the following convergence rate:
l1a.. 9G. (() = Op( Op(Vlog(n)/n).
(10)
Proof of Theorem 3 is in Appendix C.I The result of Theorem proves that the regression function
estimation rate Op( /log(n//n) under the perturbed cosine router MoE is of the same order as that with the vanilla cosine router Theorem Following the similar proof strategy in the cosine router MoE in Section 2 for capturing the convergence rates of parameter and expert estimations under the perturbed cosine router MoE model, sufficient to establish the lower bound between the difference of regression functions and the difference of parameters under both the exact-specified and over-specified settings.
In this section, we study the over-specified setting of the perturbed cosine router. The results for the exact-specified setting of the perturbed cosine router is in Appendix A.2
We now derive condition called strong identifiability on the expert function h(-,.1) to identify which experts exhibit faster estimation rates than others under the over-specified setting.
Definition (Strong identifiability). An expert function x h(x,, 11) is called strongly identifiable if it is twice differentiable with respect to its parameter and the set of fanctions x:
@laul+|aal H 05g Ompa (x, B1s, n.) aq No. 02 No, 05 |aal la2l $
linearly independent for almost every x. for any 2 and distinct parameters mp.... Tek where we denote H (ar, 81, m)) = exp( [ )( 1).
87
Example. For experts formulated as neural networks, i.e. h(x, (a, b)) p(a x b), if the activation é is selected as 2(2) = ReLU(z), 6(2) tanh(-) ) or Φ(2) = 2n for 2 2, then they are strongly
identifiable. Conversely, a linear expert ((x,, (a, b)) = x fails to meet the strong identifiability the condition.
To capture the convergence behavior expert estimation rate under the over-specified setting in Theorem we will use the Voronoi loss defined as follows:
he
(, =
Σ exp(3oi) - exp( EEA,
Σ
Σ (o:( | +
|nall
1=1
jE(K.HAA,
Σ
Σ exp(Bor) |1
112 +1|
(11)
JE(N.]A SIVE,,
Theorem 4. Suppose that the expert function h(x, 11) satisfies the strong identifiability condition in Definition then the following 1.2 -lower bound holds for any mixing measure G (x(:)::
llga gG. (l & C2(G,Go).
Furthermore, this bound and the result in Theorem 3 imply that Le L2(Gn, G.) Or( Or((Tog(n)/n))

---

Published as a conference paper at ICLR 2025
The proof of Theorem 4 is in Appendix C+ A few comments regarding this theorem are in order:
(i) Parameter estimation rates. Under the over-specified setting, parameters $i nz which are fitted by one atom, i.e. [A;(G.)) = 1, share the same estimation rate of order Op( po(Vyogg /lg(n)//n) Meanwhile, those for parameters fitted by more than one atom, i.e. [Ay(G,)) > 1, are slightly slower, standing at order Op( (or(lloog /log(n)/n).
(ii) Expert estimation rates. Given the above parameter estimation rates and the inequality (6 we observe that the rates for estimating strongly identifiable experts k()) range from Op( V (llo(n)/m) to (p((/lo((n)nn) Notably, these rates apply even for polynomial experts of degree least two, i.e. h(x, (a, b)) (a x + 6)" with p 2 2. they satisfy the strong identifiability condition. By contrast, the estimation rates for polynomial experts when using the vanilla cosine router (see Theorem ( and the linear router (see Theorem 4.6, (Nguyen al. 2024c)) are significantly slower and could be of order Op(1 (1/ (1/log" (n)), where 7 > 0 is some constant (see also Table This observation highlights that our proposed perturbed cosine router is more sample efficient than both the linear router and the cosine router.
4 PRACTICAL IMPLICATIONS
We now discuss two important practical implications from the theoretical results of the paper.
Router and expert design: From the benefits of the perturbed cosine router for the expert estimation of MoE models, our theories suggest that when using the cosine router to avoid the repre- sentation collapse, practitioners should add noises to L norms of the token hidden representations and the expert embeddings to achieve favorable performance. Additionally, the strong identifiability condition also verifies the advantages of using non-linear expert networks over linear ones.
2. Misspecified settings. Thus far in the paper, we have only considered well-specified settings, namely, the data are assumed to be sampled from the (perturbed) cosine router MoE. Although it may look restrictive, the results under this setting lay an important foundation for more realistic misspec- ified setting where the data are not necessarily generated from those models. Under that misspecified setting, we assume that the data are generated from regression framework as in equation (T) but with an arbitrary regression function q(-), which is not (perturbed) cosine router MoE. Then, we can demonstrate that the LSE G, converges where (-)-) is regression function taking the form of the (perturbed) cosine router MoE. Further- more, the optimal mixing measure will be in the boundary of the parameter space GR (9), namely, G has A atoms. Thus, becomes sufficiently large, G, also has k atoms. The insights from our theories for the well-specified setting indicate that the Voronoi losses can be used to obtain the estimation rates of individual parameters of the LSE G, to those of G and therefore, achieve the following expert estimation rates under the misspecified settings, which will be empirically validated via numerical experiments in Appendix F
mixing measure T E arg mingee. |1 fo [.)))))
(2.1) Cosine router MoE: the worst expert estimation rate could be as slow Op(1/ log' (n)) for some T >> 0. It indicates that we still need an exponential number of data (roughly exp(1/e") where e the desired upproximation error) to estimate the experts as well as select important experts.
(2.2) Perturbed cosine router MoE: the slowest expert estimation rate is of order Or(n-1/4). Thus, we only need a polynomial number of data (roughly e - to estimate the experts. This explains why the perturbed cosine router is solution the parameter estimation problem, more generally, the expert estimation problem of the MoE models.
However, the convergence analysis under the misspecified setting suffers from the challenges of understanding the universal approximation power of the (perturbed) cosine router, which have remained clusive in the literature. However, since this is beyond the scope of our paper, we leave for future development.
EXPERIMENTS
In this section, we first conduct numerical experiments on synthetic data (cf. Section 5.1). and then carry out experiments with real data language modeling (cf. Section 5.2) and domain generalization (cf. Section 5.3) tasks. Our main goal is to empirically demonstrate the efficacy of the perturbed cosine router over the vanilla cosine router and the linear router in MoE models.
7

---

Published as a conference paper at ICLR 2025
Vantil (r 000
Vantll (r- 0.00 903
C2 0.)
Pestathed (r-
Perrohal (r M.
15
Semple
Sample
(a) Exact-specified setting with k 8 experts
(b) Over-specified setting with lc = K. +
Figure I: Logarithmic plots displaying empirical convergence rates. Subfigures la and lb depict the empirical averages of the Voronoi losses Ly(G,,, G.) (cf. equation (T2) and L (Gn. G.) ) (cf. equation (TT)) for the exact and over-specified settings, respectively. The blue lines depict the Voronoi loss associated with the perturbed router, whereas the green lines are indicative of the Voronoi loss
associated with the standard cosine router. The red dash-dotted lines are used to illustrate the fitted lines for determining the empirical convergence rate.
5.1 NUMERICAL EXPERIMENTS
We first perform numerical experiments on synthetic data empirically verify the theoretical convergence rates of the least squares estimation for both perturbed and vanilla cosine router MoE models. We generate synthetic data based on the model described in equation ( Specifically, we generate {(X,, YJ) C TRd x R by first sampling X;
( 1J4) for i = 1, 1...., n. Then, we generate Y, according to the model: Y, = (X,) + for i E [n]. where the regression function (,7)))=
fc. (+) is defined as: fa. (z) == Σi= 1= Softmax ([[s)]]]+) + (ix)) . ((a)) T x b;) The
input data dimension set at 32. We employ k, experts of the form & ((a;) T x + bt). where the activation function 6 is set to be the ReLU function. The details of the values of the parameters as well as the training procedure are in Appendix E.
Two experimental settings are examined: (1) Exact-specified, and (2) Over-specified. In the exact-specified setting, the model is fitted with the same number of experts as the data generation model, specifically lt k, 8. In the over-speddfeed setting, the model includes one additional expert, totaling k: = k, + 1 = 9 experts. In each setting, experiments are conducted using both the standard and the perturbed cosine routers, with set to zero for the standard router and 0.1 for the perturbed router. For each experiment, we calculate the Voronoi losses for every model and report the mean values for each sample size in Figure Error bars representing two standard deviations are also shown. In Figure Ta the empirical convergence rates of both the standard and perturbed routers are analyzed under the exact-specified setting. The perturbed router shows a rapid convergence rate of On P(n_,,)), while the standard vanilla router has noticeably slower rate of O(n (0."1...). Similarly, in Figure the convergence rates are assessed for the same routers under the over-speeified setting. Here, the perturbed router again shows a faster convergence rate of O0-47.77 compared to the cosine router's slower rate of On -0.05).
Results.
0.5),
Misspecified settings. In Appendix F we will also conduct numerical experiments for comparing the sample efficiency of the cosine router and its perturbed variant under the setting where the data generated from the regression framework with the same regression function.
5.2 LANGUAGE MODELING
In this section, we focus on the language modeling task (Bahl et al. 1983). a fundamental challenge in natural language processing that involves predicting the next word or character in sequence to evaluate model's ability to generate and understand text. To assess how different routers influence the model's ability to capture linguistic structures and enhance performance across varying levels of textual granularity, we compare the performance of perturbed and vanilla cosine router Mixture of Experts (MoEs) on both character-level (Graves) (2033) and word-level (Bengio et al.) 2000) tasks.
Datasets. We evaluate the model's pre-training capabilities character- language modeling using Enwik8 and Text8 datasets (Mahoney) (01T). and assess its word-level language modeling
performance on Wikitext-103 (Merity et al. 2016).
8

---

Published as a conference paper at ICLR 2025
Table 2: Performance of vanilla and perturbed cosine routers on language modeling tasks.
Enwik8 (BPC L)
Text8 (BPC 4)
Wikitext- 103 (PPL
Router/Experts
Small
Medium
Small
Medium
Small
Medium
Cosine Perturbed cosine
1.213
1.161
1,310
1.271
90.070
38.018
1.197
1.147
1.303
1.251
89.910
37.859
Metries. To quantify the performance of our perturbed cosine router relative to the vanilla cosine one, we utilize the Bit Per Character (BPC) metric (Graves (013) for character-level language modeling and Perplexity (PPL) (Jelinek al. 1977) for word-level language modeling tasks.
Architecture and training procedure. In order to alleviate the representation collapse issue associ- ated with estimating routing scores in the original space, follow Chi et al. (2022) first project input representations on lower-dimensional space and parameterize experts with corresponding lower- dimensional embeddings. Subsequently, we calculate the routing scores of inputs and embeddings in this reduced-dimensiona space using our proposed perturbed cosine router. Our experiments adopt the Switch Transformer (Fedus et al. (02T), which is fundamentally a sparse variant of the T5 encoder-decoder (Raffel et al. (020),, with MoË layers replacing the MLPs. Detailed information regarding the datasets, metrics, training setup, and hyperparameters for this task is provided in Appendix
Results. The empirical advantage of our proposed cosine router over the vamilla version when applied to language modeling tasks is demonstrated in Table The results indicate that the perturbed cosine router enhances the performance of the original cosine router in all datasets across both small and medium configurations. It notably improves results for the Enwik8 and Text8 datasets at various scales and slightly outperforms the original cosine router for the Wikitext- 1-0 dataset.
5.3 DOMAIN GENERALIZATION
We conduct experiments the applications of MoE models in domain generalization. Our objective is to empirically demonstrate the efficacy of our proposed perturbed cosine router over the vanilla cosine router in this field. Domain generalization (Zhou et al. 2023a) aims generalize model's performance to unseen test domains with distributions different from those encountered during training. Specifically, in domain generalization. a model expected to leverage multiple training datasets gathered from various domains and exhibit robustness to domain shifts during testing. Such ability of out-o--distiibution generalization largely hinges on the model's s capability to incorporate invariances across multiple domains (Li et al.] 2023). Given that distribution shifts in data correspond to distribution shifts (visual) attributes (Wiles al. 2022). capturing these diverse attributes and aligning them with invariant correlations is crucial. Mixture of Experts emerges as powerful tool for efficiently capturing these visual attributes, and has been proven effective in enhancing performance in domain generalization (Li et al. 2023). Therefore, we further justify the effectiveness of our perturbed cosine router in domain generalization.
Datasets. We followed the experimental setting of et al. 2023) and evaluated our method using 5 benchmark datasets in DomainBed: PACS, VLCS, OfficeHome,, TerraIncognita, and DomainNet. Each dataset is comprised of images for classification tasks from different domains.
Architecture.
Following Guirajani & Lopez-Paz (2021). we conduct experiments on ViT-S/16, which has an input patch size of 16 x 16, comprising 6 heads in multi-head attention layers, and a total of 12 transformer blocks. We adopt last-hwo two-layer configuration, where each MoE block comprises 6 experts. The router selects the top out of experts for each image patch.
Training procedure and result. We follow the training domain validation procedure outlined in (Li al.) 2023 Guirajani & Lapez-Paz) 202I), where each training domain is split into training and validation subsets. The final overall validation set consists of the validation subsets from all training domains. Subsequently, we select the model with the highest performance on the overall validation set. To ensure fair comparisons, the results are averaged over three runs.
Table 3 summarizes the experimental results. For each dataset, we report the average results across test domains. The results demonstrate that our perturbed cosine router consistently outperforms the linear and vanilla cosine router across all datasets, thereby convincingly justifying the effectiveness of adding noise cosine routers. Detailed performances for each domain reported in Tables and

---

Published as a conference paper at ICLR 2025
Table 3: Average oul-of-distribution test accuracies.
Router/Experts
PACS
VLCS
OfficeHome
Terralncognita
DomainNet
Avg.
Lincar
86.33
78.15
73.02
41.30
48.19
65.40
Cosine
87.22
78.99
73.27
45.55
48.45
66.70
Perturbed cosine
89.36
80.01
74.09
49.87
48.51
68.37
Table 4: Per-domain performance of PACS, VLCS, OfficeHome, Terralncognita.
Router/Experts
A
C
P
S
PACS
Linear
87.29
81.20
98.50
78.34
Cosine
89.24
86.11
97.60
75.92
Perturbed cosine
89.87
86.97
97.90
82.68
Router/Experts
C
L
S
VLCS
Linear
97.53
63.65
74.09
77.33
Cosine
98.59
67.42
70.88
79.07
Perturbed cosine
98.59
67,80
74.70
78.95
Router/Experts
A
C
P
R
Officeloome
Linear
72.99
57.27
79.03
82.78
Cosine
73.40
57.27
78.69
83.70
Perturbed cosine
74.64
57.85
79.59
84.27
Router/Experts
L100
L30
L43
L46
TerraIncognita
Linear
45.99
28.51
54.66
36.05
Cosine Perturbed cosine
50.00
37.49
53.02
11.67
57.59
43.30
56.93
11.67
Table 5: Per-domain performance of DomainNet.
Router/Experts
clipart
infograph
painting
quickdraw
sketch
DomainNet
Lincar
69.11
24.95
54.81
16.88
68.95
54.41
Cosine
68.05
24.48
55.75
17.39
69.41
55.59
Perturbed
68.31
24.52
55.03
17.90
69.46
55.83
6 CONCLUSION
In this paper, we investigate the impacts of the cosine router on the convergence rates of least squares estimation in MoE models. We figure out that owing to the parameter interaction inside the cosine router expressed by PDE, the rates for estimating parameters and experts are slower than any polynomial rates and therefore, could be as slow as Pp (n)). response this issue, we propose using the perturbed cosine router where we add noises to the L norms of the token representations and the expert embeddings in cosine router order eliminate previous parameter interaction. Equipped with this novel router, demonstrate that if the expert function satisfies the strong identifiability condition, then the parameter and expert estimation rates significantly improved to be of polynomial orders. Finally, conduct several experiments both synthetic and real-world data to empirically justify the theoretical results.
There are few limitations in current analysis. First of all, the assumption that the data sampled from the (perturbed) cosine router MoE is often violated in real-world settings. However, discussed in Section 4 our theories can totally be extended setting where the data are not necessarily generated from those models, which we leave for future development, Second, since the ground -trth parameters are implicitly assumed be independent of the sample size n. the parameter and expert estimation rates presented in this work are point-wise rather than uniform. To deal with this problem, we can utilize the techniques for characterizing the uniform parameter estimation rates in traditional mixture models (see (Heinrich & Kahn 2018 Do et al., 2023 Yan et al., 2025)). Nevertheless, since the adaptation of those techniques to the setting of the (perturbed) cosine router MoE is still challenging due to the complex structures of the (perturbed) cosine router, we believe that further technical tools need to be developed to achieve the desired uniform estimation rates.
more realistic misspecified
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
NH acknowledges support from the NSF IFML. 2019844 and the NSF Al Institute for Foundations of Machine Learning.
REPRODUCIBILITY STATEMENT
To facilitate the reproduction of our empirical results, present detailed descriptions of the data and the experimental setup in Section 5 and Appendix We will release our code upon the acceptance of submission. All datasets used this study are publicly available, enabling full replication of experiments.
REFERENCES
Lalit R. Bahl, Frederick Jelinek, and Robert L. Mercer. A maximum likelihood approach to continuous speech recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, PAMI-5(2): 179-190, 1983. doi: 10.1 109/TAAI.... CTT.MM..983.40000 4767777
Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. In Leen, Dietterich, and V. Tresp (eds.), Advances in Neural Information Processing Sys- lems, volume 13. MIT Press, 2000. URL fi
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-expee eeperttsssssssseeesooooooo layer in deep learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 25049-23062. Curran Associates, Inc., 2022.
Zewen Chi, Li Dong, Shaohan Huang. Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https //openreview net / i YCCCE3
Yinlam Chow, Azamat Tulepbergenov, Ofir Nachum, Dhawal Gupta, Moonkyung Ryu, Mohammad Ghavamzadeh, and Craig Boutilier, A Mixture-of Expert Approach to RL-based Dialogue Man- agement. The Eleventh International Conference on Learning Representations, 2023. URL https:/ http:///ppeeeeeeeeeeeeeeee net Forum? (otoormm11=FF0011oe222
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-cxperts 2024.
language models. arXiv preprint 24Xiv:2401.0440
Dat Do, Huy Nguyen, Khai Nguyen, and Nhat Ho. Minimax optimal rate for parameter estimation in multivariate deviated models. In Advances in Neural Information Processing Systems, volume 36, pp. 30096-30133. Curran Associates, Inc., 2023.
Nikoli Dryden and Torsten Hoefler. Spatial mixture-of-experts. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.) Advances A in 2022.
Neural Information Processing Systems,
William Fedus, Barret Zoph, and Noam M. Shazeer. Switch transformers. Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23: 12::111222 2021. URL hittps://ape www.::isssmmmtiicccccecee
Bolin Gao and Lacra Pavel. On the properties of the softmax function with application in game theory and reinforcement learning. arXiv preprint arXiv:
2018.
Alex Graves. Generating sequences with recurrent neural networks. arXiv preprint 1rXiv:1308.0850 2013.
11

---

Published as a conference paper at ICLR 2025
Anmol Gulati, James Qin, Chung- Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, and Ruoming Pang. Conformer: Convolution-augmennn Transformer for Speech Recognition. In Proc. Interspeech 2020, pp. 5036-5040, 2020. doi:
10.21457/Inerrspecch.200 20tp:/sppeeecc22001555555 2020-3015.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021. htt.::s//ppeereeeee....
id=1 [d==lodxeDDoRII
Xing Han, Huy Nguyen, Carl Harris, Nhat Ho, and Suchi Saria. Fusemoe: Mixture-of-cxperts transformers for fleximodal fusion.. In Advances in Neural Information Processing Systems, 2024.
Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, Maheswaran Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong, and Ed Chi. DSelect-k: Differentiable Selec- tion in the Mixture of Experts with Applications zato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 29335 29347.2947. Curran Asso- ciates, Inc.. 2021. URL utt.atpss//pppreeiiis...eooopppp papaper 20007//E1e//Bc222ccoere118882333
Multi-Task Learning. In M. Ran-
P. Heinrich and J. Kahn. Strong identifiability and optimal minimax rates for finite mixture estimation. The Annals of Statistics, 46(6);2844-2776) 2018.
Nhat Ho, Chiao-Yu Yang, and Michael I. Jordan. Convergence rates for Gaussian mixtures of experts. Journal of Machine Learning Research, 23(323):1-81, 2022.
Robert Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton, Adaptive mixtures of local experts. Neural Computation, 3, 1991.
Fred Jelinck, Robert Mereer, Lalit R Bahl, and James K Baker. Perplexity measure of the difficulty of speech recognition tasks. The Journal of the Acoustical Society of America, 62(S1): S63-S63, 1977.
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel. Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang. Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts. arxiv preprint arxiv 2401.04088, 2024.
M. I. Jordan and R. A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Computation,
1994.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Yamuna Krishnamurthy, Chris Watkins, and Thomas Gaertner. Improving expert specialization in mixture of experts. arXiv preprint arXiv: /4703, 2023.
Minh Le, An Nguyen, Huy Nguyen, Trang Nguyen, Trang Pham, Linh Van Ngo, and Nhat Ho. Mixture of experts meets prompt-based continual learning. In Advances in Neural Information Processing Systems, volume 38, 2024.
Minh Le, Chau Nguyen, Huy Nguyen, Quyen Tran, Trung Le, and Nhat Ho. Revisiting prefix- tuning: Statistical benefits of reparameterization among prompts. In The Thirteenth International Conference on Learning Representations, 2025. URLhttps::/ http:///pppereeeeeee net/ forum?
XI-QiTSaFX9255
Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang. Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of of-eeertss domain generalizable learners. In The Eleventh International
Conference on Learning Representations, 2023. URL https: http::/:////peeveee....... net ffrtffformm
Hd==ecc2nn990
12

---

Published as a conference paper at ICLR 2025
Hongbo Li, Sen Lin, Lingjie Duan, Yingbin Liang, and Ness Shroff. Theory mixture-of-exppe in continual learning, In The Thirteenth International Conference on Learning Representations. 2025.
Hanxue Liang. Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng. Cong Hao, and Zhangyang Wang. ViT: Mixtureoo--Experts Vision Transformer for Efficient Multi-task Learning with Mocel-Accelerator Co-design. In NeurlPS, 2022. URL htRphhttp///papeee nips. cc wwwpppppppp....c fi /ttpleeeppppee/////aaaaaaaa
740214 a html
Bruce G. Lindsay. Mixture models: Theory, geometry and applications. In NSF-CBMS Regional Conference Series in Probability and Statisties. IMS, Hayward, CA., 1995.
Matt Mahoney. Large text compression benchmark. http: / /www//wwwwk///...////// www.mattmaaoney..ee. net
text. htmll 2011.
Tudor Manole and Nhat Ho. Refined convergence rates for maximum likelihood estimation under finite mixture models. In Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 14979-15006. PMLR, 17-23 Jul 2022.
Eduardo F. Mendes and Wenxin Jiang. On convergence rates of mixtures of polynomial experts. Neural Computation, 24(11): 3)3025-3001, 2012. doi: 10.1 201162//EEC_aa0000 __35444
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Huy Nguyen, TrungTin Nguyen, and Nhat Ho. Demystifying softmax gating function in Gaussian mixture of experts. Advances in Neural Information Processing Systems, 2023.
Huy Nguyen, Pedram Akbarian, and Nhat Ho. Is temperature sample efficient for softmax Gaussian mixture of experts? In Proceedings of the ICML, 2024a.
Huy Nguyen, Pedram Akbarian, Fanqi Yan, and Nhat Ho. Statistical perspective of top-k sparse softmax gating mixture of experts. In International Conference on Learning Representations, 2024b.
Huy Nguyen, Nhat Ho, and Alessandro Rinaldo. On least square estimation in softmax gating mixture of experts. In Proceedings of the ICMI., 2024c.
TrungTin Nguyen, Faicel Chamroukhi, Hien Duy Nguyen, and Florence Forbes. Non-asymptotic model selection in dlock-diagonal mixture of polynomial experts models. arXiv preprint arXiv:2104.08952 2021.
Fengchun Peng, Robert A. Jacobs, and Martin A. Tanner. Bayesian Inference in Mfxtures-of-Exxee of-eextrts and Hierarchical Mixtures Experts Models With an Application Speech Recognition. Journal of the American Statistical Association, 91(435):953 1996. ISSN 01621459.
Quang Pham, Giang Do, Huy Nguyen, TrungTin Nguyen, Chenghao Liu, Mina Sartipi. Binh T. Nguyen, Savitha Ramasamy, Xiaoli Li, Steven Hoi, and Nhat Ho. Competesmoe effective training of sparse mixture of experts via competition. arXiv preprint arXiv:4402.025266 2024.
Joan Puigcerver, Carlosx Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In The Twelfth International Conference Learning Representations, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Mach. Learn. Res., 21(1), jan 2020. ISSN 1532-4435.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Advances in Neural Information Processing Systems, volume 34, pp. 8583-8595. Curran Associates, Inc., 2021.
13

---

Published as a conference paper at ICLR 2025
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- gated mixture-of-experss layer. In International Conference Learning Representations (ICLR), 2017.
Sara van de Geer. Empirical processes in M-estimation. Cambridge University Press, 2000.
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre-Alvise Rebuffi, Ira Ktena, Krishnamurthy Dj Dvijotham, and Ali Taylan Cemgil. A fine-grained analysis on distribution shift. In International Conference on Learning Representations, 2022. URL http:: /ttttppe///peeeeeeee. net/ forum?
id-D14LetuLdy.
Fanqi Yan, Huy Nguyen, Dung Le, Pedram Akbarian, and Nhat Ho. Understanding expert structures on minimax parameter estimation in contaminated mixture of experts. In Proceedings of The 28th International Conference Artificial Intelligence and Statisties, 2025.
Zhao You, Shulin Feng, Dan Su. and Dong Yu. Speechmoe: Scaling large acoustic models with dynamic routing mixture of experts. In Interspeech, 2021.
Bin Yu. Assouad, Fano, and Le Cam. Festschrift for Lucien Le Cam, pp. 423-435, 1997.
Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(4))4396 0411443336441,,,,,0 2023a. doi: 10.1 10.119//PAML22222.11555
Yanqi Zhou, Nan Du, Yanping Huang, Daiyi Peng, Chang Lan, Da Huang, Siamak Shakeri, David So, Andrew M Dai, Yifeng Lu, et al. Brainformers: Trading simplicity for efficiency. In International Conference on Machine Learning, pp. 42531 -25442. PMLR, 2023b.
14

---

Published as a conference paper at ICLR 2025
Supplementary Material for "Statistical Advantages of Perturbing Cosine Router in Mixture of Experts"
In this supplementary material, we first explore the exact-specified settings of the (perturbed) cosine router MoE model in Appendix A Next, provide proofs of theoretical results associated with the cosine router MoE and its perturbed counterpart in Appendix [ and Appendix [ respectively. Those proof's are partially supported by auxiliary results presented in Appendix [ Subsequently, in Appendix [
we specify the details for the experiments performed in Section Finally, we conduct further numerical experiments on the convergence of least squares estimation under the misspecified settings in Appendix F
A ADDITIONAL RESULTS
In this appendix, we provide the convergence analysis of parameter and expert estimation under the exact-specified settings of the cosine router MoE and its perturbed variant in Appendix [ and Appendix A.2 respectively.
A.1 EXACT-SPECIFED SETTING OF THE COSINE ROUTER MoE
Firstly, we start with the exact-specified setting of the cosine router MoE.
Recall that under the exact-specified setting, the true number of experts k, is known. According to the proof technique for deriving parameter estimation rates under the exact-specified setting in the literature (gguyen et al.) 2023). a key step is to apply the first-order Taylor expansions to the product of the softmax's ' numerator and the expert function, i.c.
7) = exp |
| h(2, n). 03, (x, 31, 7) = 0,
However, since the parameter interaction via the PDE in equation (4. i.c. BI
0H
holds even for the first derivatives of the function H, the convergence of LSE under the exact-specified setting is illustrated by minimax lower bound for estimating G. Theorem
Theorem Under the exact-specified setting, the following minimax lower bound of estimating G.
inf sup G.EE,, (9) GEEx. (0)
EfclCi, G)] 11/2
holds true for any r 2 where Eje indicates the expectation taken w.r. the product measure with fö and the infimum is over all estimators taking values in &s. (9).
Proof of Theorem 5 is deferred to Appendix B.2, It can be seen that the convergence behavior of parameter and expert estimation under the exact-specified setting is analogous that under the over-specified setting. That is, the rates for estimating parameters Bi, and 1/ as well experts Mc mz) are slower than any polynomial rates and thus, could be as slow (p(1/ log' (n)), where T is some constant.
A.2 EXACT-SPECIFEED SETTING OF THE PERTURBED COSINE ROUTER MoE
We now consider the exact-specified setting of the perturbed cosine router MoE model 7 To begin with, we introduce a condition called weak identifiability on the expert function /t,, 1) to characterize which experts have faster estimation rates than others under this setting.
Definition (Weak identifiability). An expert function I h(,, n) is said be weakly identifiable if it is differentiable w.r.1 its parameter and the set of functions in x
glarl+ 035 Oypa (x. BIr. n6) ag No. 1,,002 N 05 laal laz $1
linearly independent for almost every x. for any and pair-wise distinct parameters m..... Ikk where denote H (x, 31,n) exp(
8. x
){{I,)).
[10,
15

---

Published as a conference paper at ICLR 2025
Recall from the "Technical challenges" paragraph in Section | that key step to establish the expert estimation rates to decompose the difference fč. (z) - fc. (2) into a combination of linearly independent terms via Taylor expansions to the function H (-, 81, 11) Therefore, the purpose of the weak identifiability condition is avoid all potential parameter interactions as in equation (4), which
may lead to undesirable linearly dependent terms.
Example. For simplicity, we consider experts formulated as neural networks, i.e. h(z, (a. b)) ó(a + b). It can be validated that if the function ((-) is either popular activation such as ReLU(()) and tanh(-) or a polynomial d(z) z², for any p E N, then the expert h(z, (o, b)) is weakly identifiable. On the other hand, constant expert h(,, identifiability the condition.
constant fails to satisfy the weak
Next, we will use the Voronoi loss function La(G, G.) defined below to determine the estimation
rates for weakly identifiable experts Theorem 6
whose proof can be found in Appendix C.2
h. Ls(G, G.) j=1
k-
Σ
exp(3us) exp(965)
+ Σ
exp(Bor)
| + |
(12)
TEA;
j=LiEA,
Theorem 6. Assume that b().n) is a weakly identifiable expert function, then the following lower bound holds true for any G E Ek. (9):
llga - gG. ga.ll 2 G.)
Furthermore, this bound and the result Theorem imply that L
The bound Ca(G., G.) = Op( (0g(/log()//)) all the parameters 9i. ni enjoy the same parametric es- Furthermore, by employing the argument in equa- tion (6). we deduce that the rates for estimating experts h(. 12,) are also of order Op( Vlog(n)/m) Those rates substantially faster than their counterparts when using the vanilla cosine router, which could be as slow as (1( log" (n)) (see Theorem 5). This comparison highlights the benefits of our proposed perturb cosine router the vanilla cosine router.
timation rates, standing at order Op( v
B PROOF OF RESULTS FOR COSINE ROUTER MOE
In this appendix, we provide proofs for the theoretical results regarding the cosine router in stated in Section 2 including Theorem Theorem 5 and Theorem 2 in that order.
B. PROOF OF THEOREM 1
First of all, let us introduce the definitions of some necessary concepts for the proof, namely an e-bracket, bracketing number, a bracketing entropy, an 8-cover and covering number. In particular, let (R, || []) be the space of real-valued functions f X' + R. Then, the aforementioned concepts are defined as follows:
Definition 3 (e-bracket). Given two functions I(-) and U(), the bracket [L,U] is the set of all functions f E R such that T(x) S f(x) S (z) for all x X, and IUU T.]] S &.
Definition& 4 (Bracketing number). The bracketing number Ny(e, R. -11) is the minimum number of e-brackets needed to cover
Definition 5 (Bracketing entropy). The bracketing entropy 11 B R. || II) is the logarithm of the bracketing number N1 R. -1).
Definition 6 (E-cover). An E-cover of the set R under some norm || is set ( "a....,,yy. N,) such that for any f E R, there exists some E [N] such that || kill S E. Definition 7 (Covering number). The E-covering number (E, R, || II) is the minimum number of balls (T;; $) = [f E R : ||f - T|| < }} need to cover R.
Subsequently, we denote by Rr (+) the set of regression functions w.r.t mixing measures in G& (+). that is, (+) := (fc( : G E (0)). Additionally, for each 5 0, the T,² ball centered around the regression function fc. and intersected with the set (+(O) is defined as
0,) {f E (e) [lf fc. (1(3))) - 6).
16

---

Published as a conference paper at ICLR 2025
In order measure the size of the above set. van de Geer (2000) suggest using the following quantity:
(x(6))) = / (1, (o,0,, IIL dt 6,
(13)
82/2113
where (t, (6,t,, -l112)))))) stands for the bracketing entropy (van de Geer 2000) of (O, u)
under the Lo-norm, and è V 0 max 6). By using the similar proof argument of Theorem 7.4 and Theorem 9.2 in (van de Geer 2000) with notations being adapted to this work, we obtain the following lemma:
8)) that satisfies (8) is non-increen non-increasing function of Then, for some universal constant e and for some sequence (5n) such that vno, 2 e.t (8,), we achieve that
Lemma 1. Take ((5) 2 JB Rx(0,8)
P Pllfa. -fc. > 26) S cexp
nd²
c2
for all ó 2 8n-
General picture. We first show that when the expert functions are Lipschitz continuous, the following bound holds for any < E 1/4:
HB(E, R& (4,,, S log(1/s).
(14)
Given this bound, it follows that
[
HBR (t. RR (O,t,, (lx)))) dt Vt f log(1/t)dt v 6. (15) 82
JB(6, (o(6,R
/2
Let (()) = 8- [log( (1)()//)))// then (6)/82 is non-increasing function of 8. Furthermore, equa- tion [1] indicates that (8) (6, (O, 5)). In addition, let 5, (lo nn, then get that Vno 2 (8,) for some universal constant Finally, by applying Lemma I we achieve the desired conclusion of the theorem. As a consequence, suffices to demonstrate the bound (14).
Proof for the bound (4). In order prove the bracketing entropy bound in equation (4). leverage the proof arguments for the convergence of regression estimation in (Nguyen et al, 2024c).
Since the expert functions are Lipschitz continuous, then for any function fc Rx( have that fa(x) < M for all where M is some constant.
Let and (Tı,..., nN} be the 7-cover under the I norm of the Rx where N N(T, (+,),, || [ is the r-covering number of the metric space (R& || ||.=). Then, we construct the brackets of the form (1), (x)) for all i [N] follows:
(x(x) max mi (2) -T, 0}.
U.(2) == max n. (x) + r, M..
From the above formulation, it can be checked that Re(O) c: (M1) and U,(z) L,(x) < min M). Additionally, we get that
IU: - -illla (aat::0W((/// -((1i(c)- 1/2 dp(x)) 2T.
By definition of the bracketing entropy, we achieve that
H2(2r, R+) ))^)()) ) = log N1 (2r, Rr (9). [1)()))
< log N log ((r, (+), |l I1r.).
(16)
Therefore, is necessary to provide an upper bound for the covering number N. Indeed, let us denote A : {B8, 81) E R x gd (Bo, m) E } and Ω == {n E gd= (Bo. 31, E }). Since θ is compact set, A and Ω are also compact. Therefore, we can find 7-covers A and Ω2 for A and Ω respectively. Furthermore, it can be validated that
A < Op(r -d.+1)))
12 S Op(r -d(k-
17

---

Published as a conference paper at ICLR 2025
For each mixing measure G Ei= (xp() Gx((),, consider two other mixing measures G' and G defined as
G²
exp(Bun)((B...)))
-= exp (e (8..))))
Here, T7, S2 such that 77, is the closest Ti in that set, while (Boar (10) Dr is the closest to (Bur, 31:) in that set. Now, we aim to upper bound the term Ilfe far 2o In particular, we have
(3
[|311|1 ||c||| + Boo) [M(a,, mz) hax, 17.)
sup flf for
Softmax
k
(316) || ||=|||
sup Softmax ZEA
+ Sos
1/(2, md) (((x,
sup |,(2, h(,,))) zEX
i=1
sup || - 7ill ||=||| S
zEX
Above, the second inequality holds as the softmax weight is bounded by I, and the third inequality is due to the fact that the expert h(x,,- Lipschitz function w.r.t 7) and the input space X' is bounded, i.e. ||0||| S B for any X' for some constant B > 0.
Next, we demonstrate that ll follo 2 follows:
B 3
IlSor Sall
sup xEX
Σ
Softmax
+ Box)
Softmax
+
Bou)
||3rel1 ||0|||
1
||,|||
b
B Biat
3
sup dEX
Softmax
1|31a|1 -1|=/1
+ (b..) - Softmax
||3 ||==|||
+ 9ox)
[1(2,
sup [1] - Brell #EX
||0||| + 1800 Boell
sup [ B ST, FE
By the triangle inequality, we have
llfe fallo s llfa - fo + llfa - fallo S T.
By definition of the covering number, we deduce that
N(T, (()
Ar |
10r(n-(d:+1k O(n d2k)
-dak)
< OOnn(dnn1 (((n((41++4)))) (d++)))))
(17)
Putting the results in equations
(6) and 17 together, we achieve that
(2T, R. Rx (0), )Iх)()) log(1/7).
By setting 7 = e/2, we achieve that
(x(), $ log(1/s),
which completes the proof.
18

---

Published as a conference paper at ICLR 2025
B.2 PROOF OF THEOREM
Lemma 2. If the following holds for any r 2 1:
[l - (o-Jc.llum 0,
lim G-EG (G..L...G)))
inf
(18)
GG,
then we obtain that
inf G,,ER, (0) GEE, (0)
sup
Esclei, 2 1/2
(19)
Proof of Lemma 2 Indeed, from the Gaussian assumption the noise variables 6ç, we obtain that
Y, ~ [ff. (X,). (X,). and a fixed constant C1 0 which will choose later, we can find a mixing
for all E [n]. Next, the assumption in equation (18) indicates for
sufficiently small measure G", E &k. (Q) such that L1, (G', G.) = 2s and || - -f__lla ||I I Cie. From Le Cam's
lemma Yu 1997), the Voronoi loss function Li,r satisfies the weak triangle inequality, obtain that
sup Es G₂ G,EE, GEEN.
inf
G)]
12 Li,r G.)
2 E exp( - fc. (1))
2 E (nnn=))
(20)
where the second inequality is due to the fact that
(N) (XI, (So; (X) fc. (X))²
20
By choosing n -12 obtain that e exp( Cins?) n 1-2 exp(epppoc C1). As consequence,
we achieve the desired minimax lower bound in equation 19).
Main proof. It sufficient show that the following limit holds true for any r 2 1:
llfa fc. (la 0.
lim inf + GEE,, (0): To this end, we need construct a sequence of mixing measures Gn Ex. (9) that satisfies fi,r (Gn, G.) + 0 and
(21)
(G,
||a., Sc. (LP(() + 0, (G,, G.)
as nl + co. Next, let us take into account the sequence Gn = Ek= exp(3%))) ((pp.))) in which
exp(3%) = exp( for any 1 1 < k.;
3 = (1 + Bi and I7 = 37: for any 2 < i < ka;
n = n for any i k..
Consequently, it can be verified that when 0o, we have
(G, G.) = exp(X6)) [108 Bi [1]]]]]] exp(96i) (v))== 0,
Next, we demonstrate that |lfc. fa. wc.lwe SC, quantity
(G, G.) 0. For that purpose, we consider the
(Bi,) Tx ||st,||| ||a|||
Qx(
exp
[Sa (x) Sc. (x)),
(22)
19

---

Published as a conference paper at ICLR 2025
which can be decomposed follows:
(B1,) Tx
en(x)
exp(
exp
|| -||ao| h(
m/ h(z,n)))
exp
[| ||00|
h(2)1z)
j=l iE
(P91) Tx
(1i1)
(3p)
exp
11
||a||]
fa,
(1) (x) exp
1188,11 -|1211|
Son (cn((x))
1
(Bi,)
exp( exp(3li) exp((is)) exp
[(xx,])) fa. (x))
|| ||=|||
SEA,
== A (z) - (x)+ + Cz().
Since exp(X%) (xx( for all i [k.], we deduce that (z 0. Additionally, from the
choices of P and r7,", we can rewrite (r) as
(B1) x
(g x h(z,ni).
||]||] )
(x (x1()
exp
||-||
- exp
|8
B.T
Let us denote F(r, = exp
By applying the Taylor expansion of order r, we have
|-:
glo
(x
(0R
811)
(z. + R(x)
|a|=1 al
059
[+ 2) lol (B11)
ala F
= exp( /h m
(r, (11) + R(x),
a!
08g
|a|=1
where R(x) a Taylor remainder such that R(x//L115 (G, G.) as n + co. It implied from Lemma 3 (see Appendix D) that
Σ |a|=t
(3 d
alai
(x, B11 0.
0
for any 1 1 t S r, follows that (x) = R(x). This result indicates that A, (z) G.) + O + co. By arguing similarly, we also obtain that Br(x)/Li,, G.) + 0 + 00. Combine the previous results together, we achieve that
(G. (GG,G.G,........ 0.
Since the input space X' and the parameter space θ are both bounded, the term (Tt,) |-, Bbg also bounded. This result together with the formulation of Qn(x) in equation (22) suggests that [JG., (x) Sc. Sa. (x))) G.) - 0 for almost every E X. As consequence, we get that
Σe exp
llfa. - fa. (l1 +
(G,, G.)
as 00, and hence, achieve the result in equation 21
B.3 PROOF OF THEOREM 2
Similar to the proof of Theorem 5 Appendix B.2 we only need to demonstrate that the following limit holds true for any 2 1:
1 - fa.llal)))))) 0.
lim Ges (o) (G...)) LG,(G,
inf
(23)
For that purpose, it suffices to build sequence of mixing measures G, E Ga( that satisfies L1, (G., G.) + 0 and
[l,c. -SalleP(W) +0 0, Li, Gx)
oc. Let consider the sequence G, - (exp(()))))))eeeee))))p ((p in which
20

---

Published as a conference paper at ICLR 2025
cxp((81) exp( exp(((1), and exp(3%) exp($6(t-1)) for any 3 S +1+
5n = (1 1. 87.. 3/22 = (1+ # Bii and أüi = -1 for any 3 3 k. +1:
off n/2 ./i and ry.' . for any 3 S i I k. + 1.
Consequently, it can be verified that when + c0, we have
(Gn, G.) - exp(3(n) [1f (1.9
Pial'" + 115 - Biill"] exp(201) (vá) ) + 0,
Next, we demonstrate that llfa. - fa. fa.lllauu// Lir G.. +.
To this end, we consider the quantity
]
k-
(3i,) Tx
Qx(x)== =
Σ
exp
|| || + Sūg
[fa. (f.)
(24)
which can be decomposed as follows:
A. Qn(x)= = Σ
(B
(Bi,)
exp(
exp
|1 ||-=||
h(2,n,)) -exp
||, (ari|l-laa) ||ج|||
h(x, n5)
j=1 iEA;
k.
(B%) Ta
(
exp(
exp
|
||a:||
Sa. (x) - -exp
[1,8,11| -|ael|
San (x))
j=1 iEA;
k.
(Bi,)
exp( - exp((6,) exp
(2,,n)) Sc.
|| ||-|||
1=1
iEA,
= An (x) (x) (z).
From the choices of S7: and sff we can rewrite An (I) as
1 Ax(x) 2 (x)((%))h(a,,)
(BT:)Tx || ||
)
| || ||>||| )
( TT
exp
exp
Let us denote F(x, exp
8 x
By applying the Taylor expansion of order r, we have
d 0.99
Ax (((,)
M1(, 71)
(8)
811 a.
(z,
Bit) R(x)
la
(1+ (لح)))) lal (811)
glal F
(0p(Bi1)((1 7i)
Σ
(x, P11) + R(x),
n
03
where R(x) is a Taylor remainder such that R(x)/ (Gn, G.) + 0 as n + xc follows from Lemma 3 (see Appendix DD that
(1))) (al
alal
Σ la|=t
(3
0:90 (x, (1,
a!
n
(1 (-D)) n
glalF
(81 o
(x, 311) 0.
0
for any 1 1 t S r. Thus, we get that (,) = R(x). which implies that A, Li,r G.) +0 as 0o. By arguing similarly, also obtain that (x) Lr (n, G.) 0 as co.
21

---

Published as a conference paper at ICLR 2025
Furthermore, we have
2 exp( exp(((i)) exp ( (811) Txx
Cn(x) =
(((znn)) Sc. (1r)
||
||2:||
($i,) Tx
exp((0(+t)) exp(36))) exp
[M (z) - San (1))
[1:1 ||
0.
Putting the previous results together, we achieve that
Qn (n( G.) 0.
Since the input space X and the parameter space θ are both bounded, the term
Σ exp
| || + Púj also bounded. This result together with the formulation of (z)
(o1,)"a
equation (24) suggests that [Sa,, (x) Sc. (o.()))/R1,,(GG) G.) - 0 for almost every x E X. As consequence, get that
[a.. Sa. ( 0,
GG,s G.)
as oo, and hence, achieve the result in equation 23
C PROOF OF RESULTS FOR PERTURBED COSINE ROUTER MoE
In this appendix, we provide proofs for the theoretical results regarding the perturbed cosine router, namely Theorem 3 Theorem and Theorem in that order.
C.I PROOF OF THEOREM 3
Since the proof of Theorem can be done similar fashion to that of Theorem it is omitted.
C.2 PROOF OF THEOREM 6
In this proof, we aim to establish the following inequality:
inf |19G 9G. L G.) > 0.
(25)
GEE,, (8)
For that purpose, we divide the proof of the above inequality into local and global parts in the sequel.
Local part: In this part, we demonstrate that
lim GEU (O) 119a 9G. (c.ll1 0.
inf
(26)
Assume by contrary that the above inequality does not hold true, then there exists a sequence of mixing measures G,, = Et= exp( ( G (+) such that Lon = La(Gn, G.) + 0 and
l19G. - 9G. WlG.lll..//// Con + O,
(27)
as n. + oc. Let us denote by A == A, (G,) a Voronoi cell G, generated by the j-th components of G.. Since our arguments are asymptotic, we may assume that those Voronoi cells do not depend on the sample size, i.e. A₁ = Ay. Moreover, recall that under the exact-specified setting, each Voronoi cell has only one element. Therefore, we may assume WLOG that A, = {i}. and
Lan
exp( exp(37t) +
exp(
|
+1|
where we denote
A3% 8 Big and Anp
1//2 m
22

---

Published as a conference paper at ICLR 2025
Since Can - O, we get that (P%) 77,)) - (5 n7.' and exp(5%) - cxp()6;) as - for any ? [k+]]. Now, we divide the proof of the local part into three steps as follows:
Step 1: Taylor expansion. In this step, decompose the term
( (1 + m)) [| r2) + (6)) (x),(() gG_ (x)]
Qn(z) =
Σ
exp
into combination of linearly independent elements using Taylor expansion. In particular, let us then we have
denote P(x,8)) == exp
0
(x(x)
exp(
F(x, (R)(,,)) Flx, Bi))(z,) m])
exp((%t) F(a, - F(x, (7:) (x)
exp( exp(3 F(x, Bi)),,, (t) Far, $i.)a.. (x)
1=1
== A (x) - (x) + (z)
(28)
By means the first-order Taylor expansion, we have
A.(x) exp( ( glail F gaah h 0
(z, Bit) Ömpe m) + RI(x), (29)
|a|=1
0 cc. Similarly, also get that where R1(x) (I) is Taylor remainder such that Ri (x)/
exp( Dz
ahlF
Bx(z) Σ
(A
037 (1,9is)a,, (z) + R2(x),
where R2(x) is a Taylor remainder such that R2(2)/ Can - 0 as n - co. As a result, we deduce that
k.
Jlail 03
Joaslh
Qn(x)=
The
(x, (7)) Ompa (x, n²) + (x)
002
0
k.
Ohl F
S
(x, (i.99, (x) - R2(x),
(30)
05
hl=0
where we define
exp(
T Tan.a
(A ( )a (Am,")
ad
exp( (A (Apls))" )7
S
for any (a1, (a1)))) # (0a, 0) and # Og. Otherwise, T = S t,0. : exp(3gg) - exp((3)))
Step 2: Non-vanishing coefficients. this step, we show that not the ratios T: . Csn. and
/ Lan Lsn converge to zero. Indeed, assume by contrary that all of them converge to zero, i.c.
T Th Lyn
S Lan
+ 0,
+ 0
Then, it follows that
Z
Σ
exp( - exp(36t)
Zan
1==1
S,0,
+ 0:
23

---

Published as a conference paper at ICLR 2025
Ei= exp(3%))1A 1/1
di Zan Eis u==
Teny 00
0;
L3n
exp(
Zzm 2i=1
T T. + O.
Z3n
Due to the topological equivalence of norm-l and norm-2, we deduce that
exp(3 - 0.
exp(+%))) - 0.
Lan
C3n
As result, obtain that
Lan
exp((%)
exp((6t)
k. +
Σ exp( 1m1
(1A
[IASTI] + || -0.
Lan
Linn
which is a contradiction. Thus, at least one among the ratios T ag/ Can, and / C3n must not go zero as + oc.
Step 3: Application of Fatou's lemma. In this step, we demonstrate a result opposed to that in Step 2, i.e. the ratios o Snn. and //3n all converge to zero.
In particular, let denote by m, the maximum of the absolute values of agg Can. and S C3n- Since least one among those ratios must not approach zero as - Jc, we get that 1 1/m,, A oo as + CO.
Recall from the hypothesis in equation (7) that Igg.. - gG. (w) + n + oo, which indicates that l|9G., - 9G. ( L3n + due the equivalence between I.¹ (n)-norm and [2) norm. By means of the Fatou's 's lemma, we have
1la.. - 9G. 9G. 1IAA
/
Iga., (x) - 9G. (G. (x)
0 lim oc
lim inf n=
du(x) 2 0.
Mnn
my Con
This result implies that [ga,, (x) ga. (xn))/mm/(a) [3n] + for almost every x.
Let us denote
The To - 1i,01.02:
9tn//mLLsn 81,7
with note that at least one among the limits ti Si, is non-zero. Then, from the decomposition equation (30).
deduce that
glan F
glaoth Oypa
F (I, (i,gi,),., (x) 0,
(s,, Bis)
(x, m
Σ 81,y
0:900
037
0
for almost every x. Note that the expert function ht,, m) satisfies the condition in Definition then the above equation implies that fr .as
Si,7 0, for any E [ks,], 0n E N 02 No2 and
Y N4 such that Jayl + |az) hrl < 2. This contradicts the fact that at least one among the limits fi, Si,y is different from zero.
Hence, we obtain the local inequality in equation 26 Thus, we can find an e' such that
inf
11 sG. [llewn/ > 0.
GEE, (GC..,
Global part: Given the above result, it suffices to demonstrate that
inf GEE, (G):s,,,GG)))
|19c SG. /ll/ G.) > 0.
(31)
Assume by contrary that the inequality (31) does not hold true, then we can find a sequence mixing measures G² Ek. (0) such that La (G,, G-) > and
I 9G- [l
lim n-000
Ls G.)
24

---

Published as a conference paper at ICLR 2025
which indicates that llac, -9G- (1)) 0 as oc. Recall that θ is compact set, therefore, we can replace the sequence G", by one of its subsequences that converge to mixing measure G' E &x. (52). Since La(G', G.) > e', we deduce that L3(G", G.) e'.
Next, by invoking the Fatou's 's lemma, we have that
0 lim l99, 9G. 9G. (1it.() 2 lim inf gGy, (x) 9G. (x) 2 du(x)). n-
Thus, we get that 90 (x) = 9a. (x) for almost every From Proposition I we deduce that G' = G.. Consequently, follows that La(G", G.) O, contradicting the fact that La G.) E' >
Hence, the proof is completed.
C.3 PROOF OF THEOREM 4
In this proof, it sufficient to demonstrate the following inequality:
inf (E)
ll9c - Sc.lleeww/C((GGGG > 0.
(32)
This can be done by deriving its local part and the global part as in Appendix C.2 Since the global
part can be argued similarly, our main goal prove the local part:
lim 5-0 (E)
inf
llga - 9G. (leeele) G.)
(33)
Assume by contrary that the above inequality does not hold true, then there exists a sequence of mixing measures G, = Em= exp( (087.)))) in Gx (8) such that Len == C2(G,, G.) + and
llgc. (L2 +
(34)
as n + oc. Let us denote by A == A, a Voronoi cell Gn generated by the j-th components of G., Since our arguments are asymptotic, may assume that those Voronoi cells do not depend on the sample size, i.e., Aj- = Ay- Therefore, we may assume WLOG that
h-
Lan
exp( - exp( +
Σ
Σ
exp
||
+ || [??]
1=1 IEA,
JE
AEA,
Σ
Σ
exp (3)
|
+ ||
jE
where we denote A87 == PR Bi, and Anj :- 75
Now, we divide the proof of the local part into three steps follows:
Step 1: Taylor expansion. In this step, we decompose the term
(Bi (11 (| + ph.) 19c, (x) - 9G. (x)
(=)
exp
into a combination of linearly independent elements using Taylor expansion. In particular, let us then we have
denote F(x, = exp
Ae Qa(x)
exp(33r)
(aM,, - F(a, (a,, n]]
/=1 EEA,
exp(36i)
81,) -
- F (z,
go. (x)
j=liEA,
exp( (i, -xp(((()))))))))
exp(3t,))
- F(x, ((x,8i,))c, (2)
j=1
iEA;
(x) Ba(x) Cx()).
(35)
25

---

Published as a conference paper at ICLR 2025
Next, we continue to separate the term An(2) into two parts as
An(x) =
Σ
(xp())i)) ((x,,3))((z,, Fi, )a(a, nj) ni)]
A =1 sE
Σ
exp( ((x,BB,))(x m/) - F(x,, (1,
|E[A.]:|A,I211 E
== 1 (α) (z))
Similar equation 29 by applying the first-order and the second-order Taylor expansions to 1 (2) and (2)) respectively, we have
(n,1
E
exp (5
(A
(An)i)
alai F 0.
(I, (1, 01j02
@laa h
(x, n1j) + Ri (z),
a!
1E(R.)AA
|a|=1
2
exp (5
ala
alaalh
An,2(z)
E
(A )an (Anlj)
(x, (i, 01joz
(x, n7) + R2(x),
a!
for {1,2). Analogously, where (1) is Taylor remainder such that (x)x /2) 0 -0, we also get that B, (2) = Bn.1 + (z) where
a exp( (A
(x,(
E
(x. (e,Bij))e, (x) + Rs(x),
?
0.7
jElk.|:lA/1=1 1x
exp(
ahl (i,.9,) (x) + Ra(x),
(2)((z)=
Σ
(A
?!
a37
hn|1
in which (2) is Taylor remainder such that Rx(x)) Lzn + 0 as co, for E {3, As result, we deduce that
glail F
glazAh h (i, Ompe (x, R1) R2(x)
Qn(x)
T
0.5mm
|a|=0
0 F
Σ
Siy
(o,, Bi 19,) (x) Ra(x) Ra(x),
(36)
037
1=1
h|=0
where we define
cxp(5m)
(5m) (A
The jan,ag
(An)))
iE.A, a!
exp(3nn) (A )7
Sh =
?!
iEAs
for any (2) # (Oa. 0) and + Og. Otherwise, To = Sy,0, := EiEA, exp(3 exp()6))
Step 2: Non-vanishing coefficients. In this step, we show that not all the ratios T Ty Lzn, and
/ / Can converge to zero. Indeed, assume by contrary that all of them converge to zero, i.c.
T Han,0 Lzn
0,
0
Lzn
as n + 00. Then, it follows that
Lan Σ}= EiEA, exp(3%) exp(365)
Ei= 1 ++;; 0; k.
Lan
Lzn Σ}= Eier, exp( ||
Σ}= Edi [Tee,, 0,
+ 0:
Dzn
26

---

Published as a conference paper at ICLR 2025
Lan Ej= EIEA, exp(si))lAn))
||
Lon 2x=11 - T
jo . - 0;
Zzn Ljz LiEA, exp(
Din 2j=1 2d1 E EiEA, exp ( Din j==
3.2e,, 00 +
Lin
Due to the topological equivalence of the norm-l and norm-2, we deduce that
1 LE C2n jE,
O. ΣΣ exp((6))111)))) exp | 0, Czn 3=1iEA;
Thus, by taking the summation of the above limits, we obtain that 1 = Lzn/ C2n - as 12 - oc. which is contradiction. Consequently, at least among the ratios Ta az L2n and S7,7// L2n must not go to zero as -> oc.
Step 3: Application of Fatou's lemma. In this step, demonstrate result opposed to that in Step 2, i.e. the ratios Ta /a./ L2n, and / C2n all converge to zero.
In particular, let us denote by m, the maximum of the absolute values of T oz C2n, and Sy L2n- Since at least one among those ratios must not approach zero as n - Jo, we get that 1/m,, A oo as 00.
Recall from the hypothesis in equation 344 that lgg.. - 9a. (IL Len - O as n + oo, which indicates that |19G., 9G. (n) + 0 due to the equivalence between L (u))norm and L²(()- norm. By means of the Fatou's lemma, have
0 lim IlaG, - 9G. -9G. -GG.ll-t
lim I9a, (x) - gG. (x)) du(x)) 2 0. n-+00
MnL2n
My mmn2n
This result implies that [9c., (x) 9c. (x))// [m,, L2nl - for almost every x.
Let us denote
Th Tan maooe//naa - 1j,a,,a9
P5n/mnfe + 8j.7
with note that at least one among the limits faana 83n4 is non-zero. Then, from the decomposition equation (36). we deduce that
glas
glaoth Oype (x, n1)
OhiF (xc, Big (iy),,, (x) 0, 05
(s,, (ig
Sjn
0,900
f=1 |a|=0
j=1 1n|=0
for almost every z. Note that the expert function ht,, m7) satisfies the condition in Definition then the above equation implies that tj,as,,0 0z
$j,y 0, for any j [k,], a1 No, 02 No and
? Ná such that |ail + |ava hy < 2. This contradicts the fact that at least one among the limits "j,a,aas $J.y is different from zero.
Hence, we obtain the local inequality in equation 33 and complete the proof.
D AUXILIARY RESULTS
In this appendix, we present three additional results to facilitate the proof' in Appendix B and Appendix C
Proposition (Identifiabilityy. If (a)) = 9G. (x) holds true for almost every x, then it follows that G'.
Proof of Proposition
For almost every since go(z) = gG. (z), then we have
(Bix) Tu
Sofumax
(11Biell + m.) +2) Bos) 11(2,
(Bia)
Softmax
+ (6.)
(37)
(118))) + TI) {||=5| + T2)
27

---

Published as a conference paper at ICLR 2025
Due to the identifiability of the expert function h(,, n)), the set (h(z, 7//) : [K]], where 7/ ....... Te. are pair-wise different vectors for some k' N, is linearly independent for almost every æ.
Additionally, note that if k * k.. then there exists some index i E [k,] such that Çi ni for any (91.)
j [k.]. This result implies that Softmax (T which is contradiction. Thus, must have le k4. As result, it follows that
+ Box
for almost every æ.
Softmax (B1s) x + Box) : E [Ac]
Softmax (B1,) + Pos) [1:x]]
for almost every x. WLOG, we may assume that
(B1;) Tx (118rel + m) (||=]] + 12)
(01e) Tx
Sofumax
+ 9os Sofumax (1 m) (|1=0 72)
+ (ob
(38)
for almost every т, for any i E [K+], As the Softmax function is invariant to translations, then the equation (38) indicates that
(B1z)
(g1,)
(1ll +T1) (1|=1 T2)
(11 + TI) (1|=||+
Bor = fox + uo:
for some up E IR. The first equation implies that B1 = Big while the second equation together with the assumption Box -o lead Sp: 9%. for any E [k.]]
Let us consider æ E Xe, where E [g] such that (f...... (x) = {1,..., [h....,K.. K}. Then, the equation (37 can be rewritten
k.
(Bt,) (11Ø%) |l + T1) ([|=|| + T2)
exp(Pad) exp
h(2, Gi)
(gix) (1| + T2)
Σ exp( exp
Mzz,n)) nt);
(39)
(11
+ TT)
for almost every x E Xr. Next, we denote Pı, Pa,..., .., Pm as a partition of the index set [ki], where m S k; such that exp((3n) = exp(3rr) for any i E Pj and [k.]. On the other hand, when i and i' do not belong to the same set Pj. we exp(Soi) * exp(3or). Thus, we can represent equation (39) as
(Bi,) (11S%) + T1) [||=| + T2)
Σ
exp(Sox) exp
h(o,Ci)
(Bi:)
exp(86i) exp
h(x, n²),
(11 + +T1) (||==| + T2)
J=1 EP
for almost every x E Xz. Recall that we have Bu = Bi, and Bor = Pür for any E [k..]. then the above result leads
{n E P3) (n% :i P3).
for any j E [m]. As consequence, obtain that
G Σ
exp(Sox) =
exp( (Bon (xp()Bo)0(()).1))) = G.,
j=lier KEP
1=1 1
Hence, we reach the conclusion this proposition.
BT
Lemma 3. Let F(x, B1) = exp
For any vector 31 E Rd and t N, we have
|| ||
||x]||
d
a'F
Σ
B(u1))
g(me))
(I,
(40)
(g((m)
0
28

---

Published as a conference paper at ICLR 2025
Proof of Lemma 3
We will prove the above result by using the induction method. In particular, first show that it holds for t = 1. By taking the first derivative of F w.r.t 31. we have
OF
x. ||3: ll
||-|||
P.T.
91
||=|||
(x, 51)
Fxx, B1).
0
|
1|
Then, follows that
0F
31 11
||>|||
8 81 11
||||||
181 x
B1
(x. B1)
0,
0/1
|
or equivalently,
d1 3(m))
aF
(I, 31) = O.
as("))
21
Subsequently, assume that the equation
440
holds for t 1, i.e.
0--1F as(m)) 0g(un-1) 1) (1) 31) 0,
we will demonstrate that also holds for t. Note that the above left-hand side can be decomposed as
0--1p
0 F F
(uu 1)
(x, =
Σ
gfmu)
(u 1)
03 (1n) 08 uv 1)
03
m1
03 (
=1
g-1F
Σ
gluz)
3
gla glme)
(I, B1)
08(w2)
ue- 1) 08(mm-11
)g((u)
81-1F
01-1
Σ
8f 1) (8((u)))
(x, B1) + (8(w))ye-1
(x, B1);
03(4 0 0(3(")))022
a(8(m)))
We- ##
where u, is some index in [d). By taking the derivatives of both sides w.r.t g(u)), get
Σ
g(wn)
g(vx-1)
O'F
03(mm) a
F-
(wz)
9 g(u-1))
08 (m)) (1) B1)
00(w2))
05 t
M ph
1 1
Σ
(w2) g(uu 1) s)g(m))
FF 0
(2,, 31)
0860
1) ((8((m))
u2 22...,
+
g1-IF
+
(t- - 2)
Σ
(((L-))
(8(w))
0g (v 1) ((0(()) (I,
t 2
11-1
F'F
+
t- (
Σ
g(Me 1) (:3 ((u1))(12
03( 0 (1)) ) (x. (1, B1)
1u
0-1FF
g
+ (t 1) ((
(x, + (3(m))
(x, 31)
0(3(m))31= 3-1
Σ
di
B(m1)
gru 1)
(3((1)) (2) a (I,
F
US 1=
di
(u g--1F F
+ (1)
gluz) 3
(I, B1).
05(w2)) 0
29

---

Published as a conference paper at ICLR 2025
Therefore, it follows that
uel
0 Σ
((v))) ... olu-1) g(a))
O'F
31)
03 (1 ... 05(a))
0r 01-F
+ (t 1)
$ u2)
81 ( g((e)
(z, 81).
03(u2)) 0s7u)
It is worth noting that
0- F
Σ
B g(uz)
g(wr-1) glwe)
(1,
08(40) 03(4)
Σ
(w))
g
g(ux-1)
(x, B1) = 0.
00(wn) 03
Consequently, deduce that
d1
FºF
B(u1)
(I, = O.
( ... ...05(m) as(n))
Hence, we reach the conclusion in equation (40
E EXPERIMENTAL DETAILS
In this appendix, we provide the details for the numerical experiments on synthetic data, and the experiments with real data on language modeling conducted in Section S
E. EXPERIMENTAL DETAILS FOR SYNTHETIC DATA
Model details. We now provide the details for the model parameters in model
(8i))
(c)
Σ Softmax
((a)) (aa))(++()))
(41)
(11 + +T) (1+] +T)
The variance of Gaussian noise is specified as x² = 0.01. For simplicity. the perturbations for both ||ж||| and || || are considered identical, denoted by Ti T2 The true parameters for the router, (8i, (or) Rd x IR, are drawn independently from isotropic Gaussian distribution with zero mean and variance o² = 0.01/a /d for 1 < S and otherwise are set to zero. Similarly, the true parameters of the experts, (at b,) Rd R. drawn independently of isotropic Gaussian distribution with zero mean and variance a² 1/d for all experts. These parameters remain unchanged for all experiments.
Training procedure. For each sample size n, spanning from 10 to 10a, we perform 20 experiments. In every experiment, the parameters initialization for the router's and experts' parameters are adjusted to be near the true parameters, minimizing potential instabilities from the optimization process. Subsequently, we execute SGD across 10 epochs, employing a learning rate of 0.1 to fit a model to the synthetic data. All the numerical experiments are conducted on MacBook Air equipped with an MI chip CPU.
E.2 EXPERIMENTAL DETAILS FOR LANGUAGE MODELING TASK
Datasets. We use the Enwik8 and Text8 datasets (Mahoney 2011 for our character-level language modeling task. The Enwik8 dataset comprises 100 million bytes of unprocessed Wikipedia text, while the Text8 dataset contains 100 million processed Wikipedia characters. We further evaluate the word-level language modeling task on the Wikitext- 103 dataset (Merity et al, (016),, which the largest available word-level language modeling benchmark with long-term dependency. It contains
30

---

Published as a conference paper at ICLR 2025
103M training tokens from 28K articles, with average length of 3.6K tokens per article, which allows us to test the ability of long-term dependency modeling.
Metrics. In the main paper, we employ the Bit per character (BPC) (Graves 2013 metric to assess the performance of character-level language modeling tasks. This metric measures the average number of bits needed to encode each character in the dataset. It is caleulated as follows:
T
1 T
BPC(X)
log2 Pa (aa)
where T' is the length of the input string X, P; is the approximate distribution and zy is the character in the input string at location
Essentially, BPC quantifies the average number of bits required to encode each character in the text using the probability distribution predicted by the model. This concept works as follows: characters with high probability get a short bit sequence, while characters with low probability get longer sequence. Then, the next character read from the input and encoded using the bit sequence performance determined from the probability distribution. If the language model is good, the target character will have been predicted with high probability, so the bit sequence will be short. This means that a lower BPC indicates better compression, which in turn demonstrates the model's 's superior ability to predict the next character accurately.
For the word- hevdllevellee language modeling task on the Wikitext- 103 dataset, utilize Perplexity (PPL) (Telinek et al. 1977 as our evaluation metric. It represents the exponentiated average negative liglliiel of sequence and demonstrates how well the model predicts the next word in sequence. More specifically, if we have a tokenized sequence X = (Io, Iz), the perplexity of X is:
1 log po
PPL(X) exp
where po (1rg (x<))) is the log-likelihood of the 7th token conditioned on the preceding tokens Tea according to our model. A lower perplexity score indicates better generalization performance.
We follow the implementation of (Pham al. 2024) to use experts consisting of two linear layers: the first with weights of shape input_dim x hidden_ liidee_dim_eer dim__xpprts and the second with weights of shape hidden_ding slii_ddi__eeee _en__eeerts x input_dim, followed by ReLU activations and Dropout layer with drop rate p 0.1. This architecture ensures the output has shape of input dim input dim while can flexibly reduce parameters compared to a single lincar layer with weights of shape input_ x input dim when the number of hidden dimensions of experts chosen is smaller than dimensions of the input. We consistently apply this architecture with our perturbed and vanilla cosine router across all datasets.
Training setup and hyperparameters. We consider two model configurations: the small and medium setups. The small setup has a total of 15 million parameters with 6 SMoE layers (Dryden & Hoefler 2022), and each layer can learn spatial structure in the input domain and routing experts at fine- iieeggriied level to utilize it. Similarly, the medium setup consists of 36 million parameters with 8 SMoE layers (Dryden& Hoefler 2022). For SMoE layers in both configurations, we employ 16 experts with top-2 gating in both cosine and perturbed cosine routers. To mitigate the representation collapse issue, we adopt the method proposed by Chi al. (022)) Specifically, this approach involves parameterizing the experts using lewer-dimensione embeddings e, E &e and applying linear projection to map token vectors into this compact space, rather than the original high-dimensional hidden space d. Furthermore, we follow Chi et al. (2022) to set de 8, which is half the number of experts in our implementation for the language modeling task.
During training. we use Adam optimizer (Kingma & Ba 2017) with default parameters. We set the number of training steps to 60000 and 80000 for small and medium configurations, respectively. The results averaged over three runs for fair comparisons.
All language modeling experiments are conducted on NVIDIA A100 GPUs. Training the small configuration of the Text8 and Enwik8 datasets takes 11 hours, whereas Wikitext- 103 requires hours. Training Text8 and Enwik8 for medium configurations takes 17 hours, while Wikitext-103 training takes hours.
31

---

Published as a conference paper at ICLR 2025
E.3 EXPERIMENTAL DETAILS FOR DOMAIN GENERALIZATION
Expert types. Similarly, following (Li et al. 2023). the experts utilized in the domain generalization experiments are two-layer feedforward networks, featuring GELU activation and a Dropout layer with a drop rate of 0.1 positioned between the two linear layers. The dimensions of the two linear layers are imput dim hidden dim experts and hidden dim experts input dim, respectively. where hidden dim experts imput paput_dim_> x 4.
Training procedure. All DG experiments are run on NVIDIA A GPUs with 15,000 iterations. The training time on PACS, VLCS, OfficeHome, and Terralncognita is 2 hours, while the training time for DomainNet 7 hours.
F ADDITIONAL EXPERIMENTS
In this appendix, we carry additional experiments on the synthetic data to compare the sample efficiency of the perturbed cosine router to that of the cosine router and the linear router (Nguyen al. 2024c) under the settings where the data are generated from the same regression framework.
F.1 COSINE ROUTER VS. PERTURBED COSINE ROUTER
Experimental setup. We generate the data by first sampling X; ~ Uniform([-1, 11 for i = n....,n. n. Then, we generate Y; according to the following model:
Y; = 8g. (X₂) + 6
where the regression function sc, (-) takes the form of linear router MoE (Nguyen al. 2024c
Softmax((11.) Tx (b)) ReLU ((a)) x +b).
All other experimental details remain the same as specified in Section 5.1 and Appendix
Results. We calculate the Voronoi losses and report the mean values for each sample size in Figure 2a Error bars representing two standard deviations are also shown. In Figure [a the Voronoi losses associated with the cosine router vanish at the rate of O(n Onn00.1)). which is very slow and roughly matches our theoretical results. Meanwhile, those associated with the perturbed cosine router converge to zero at significantly faster rate of On (0.45).. This empirically shows that the perturbed cosine router is more sample efficient than the vanilla cosine router.
F.2 LINEAR ROUTER VS. PERTURBED COSINE ROUTER
Experimental setup. From Table I linear router when the experts are of polynomial forms [aTx x + b)2. Thus, we will incorporate polynomial experts in both fitted models. Additionally,
that the perturbed cosine router outperforms the
generate the data by first sampling X; ~ Uniform([-1,1]4 for i = I....,n. n. Then, we generate Y, according to the following model:
Y, = sc. (X+) + tii
where the regression function sc, (-) takes the form of a linear router MoE with polynomial experts (Nguyen et al. 2024c):
Softmax((87)) + + ((a)) x b;)²
All other experimental details remain the same specified in Section and Appendix E.1
Results. We calculate the Voronoi losses and report the mean values for each sample size in Figure 2b Error bars representing two standard deviations are also shown. In Figure 2b the Voronoi losses associated with the linear router vanish at a very slow rate of O 0.. By contrast, the convergence rate of those associated with the perturbed cosine router of order O(n -0.4).. which are substantially faster. Thus, we can claim that the perturbed cosine router is more sample efficient than the linear router both theoretically and empirically.
32

---

Published as a conference paper at ICLR 2025
Vanilla (r 0.0)
Liosa Rouie: ard
39-
Cal.... G.)
Ei G.)
Perturbed (T-
Pestitel Chdine Borter 0.)
0.9 m00
10-
10
Sample size
Sample size
(a) Cosine router vs. Perturbed cosine router
(b) Linear router vs. Perturbed cosine router
Figure 2: Log-log scaled plots displaying the empirical convergence rates. Figure 2a depicts the empirical averages of the Voronoi losses when using the cosine router (green line) versus when using the perturbed cosine router (blue line). The red dash-dotted lines illustrate the fitted lines for determining the empirical convergence rates. Similarly, Figure 2b depicts the empirical averages of the Voronoi losses when using the linear router (green line) versus when using the perturbed cosine router (blue line). We use the same data samples for those experiments.
33


Reviews:
[{"review_id": "K5Y8zk4VUY", "reviewer": "Reviewer_idQ2", "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", "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?\n\n\n---------------------------\nAfter the response.\nThank 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):\n\n1.\tThe 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])\n2.\tThe 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.\n\n[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.\n\n[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.\n\nAfter the author response again, I decided to change the score back to 6. The reason is in the comment.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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. \n\nQuality: 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.\n\nClarity: The paper is well-written and the implication of the results is easily understandable\n\nSignificance: 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.\n\n2. 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?\n\n3. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZHOsgwISx1", "reviewer": "Reviewer_hEMT", "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.", "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?\n- 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.\n- 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)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "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?\n\nA 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.\n\nIt's possible I am missing something, so I will be happy to increase my score if these concerns are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "edYlZLlUtw", "reviewer": "Reviewer_VrcC", "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.", "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.\n2. 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?\n3. 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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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]?\nAdditionally, as noted by [2], the sparse Mixture of Experts (MoE) approach appears to offer improved generalization capabilities.\n\n[1] Nguyen, Huy, et al. \"Statistical Advantages of Perturbing Cosine Router in Sparse Mixture of Experts.\" arXiv preprint arXiv:2405.14131 (2024).\n\n[2] Li, Bo, et al. \"Sparse Mixture-of-Experts are Domain Generalizable Learners.\" The Eleventh International Conference on Learning Representations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)