Statistical Perspective of Top-K Sparse Softmax Gating Mixture of Experts
Authors: Huy Nguyen
Pedram Akbarian
Fanqi Yan
Nhat Ho
Venue: ICLR 2024 poster (2024)
Keywords: Mixture of Experts
Maximum Likelihood Estimation
Voronoi Loss Function
Primary area: learning theory

Abstract: Top-K sparse softmax gating mixture of experts has been widely used for scaling up massive deep-learning architectures without increasing the computational cost. Despite its popularity in real-world applications, the theoretical understanding of that gating function has remained an open problem. The main challenge comes from the structure of the top-K sparse softmax gating function, which partitions the input space into multiple regions with distinct behaviors. By focusing on a Gaussian mixture of experts, we establish theoretical results on the effects of the top-K sparse softmax gating function on both density and parameter estimations. Our results hinge upon defining novel loss functions among parameters to capture different behaviors of the input regions. When the true number of experts $k_{\ast}$ is known, we demonstrate that the convergence rates of density and parameter estimations are both parametric on the sample size. However, when $k_{\ast}$ becomes unknown and the true model is over-specified by a Gaussian mixture of $k$ experts where $k > k_{\ast}$, our findings suggest that the number of experts selected from the top-K sparse softmax gating function must exceed the total cardinality of a certain number of Voronoi cells associated with the true parameters to guarantee the convergence of the density estimation. Moreover, while the density estimation rate remains parametric under this setting, the parameter estimation rates become substantially slow due to an intrinsic interaction between the softmax gating and expert functions.

Published as a conference paper at ICLR 2024
STATISTICAL PERSPECTIVE OF TOP-K SPARSE SOFT- MAX GATING MIXTURE OF EXPERTS
Huy Nguyen
Pedram Akbarian
Department of Statistics and Data Sciences
Department of Electrical and Computer Engineering
The University of Texas Austin
The University of Texas Austin
Austin, TX 78712
Austin, TX 78712
huynm@utexee Shyyn@utexas.com edu
akbarian@utexas. .dd
Fanqi Yan
Nhat Ho
Department of Computer Science The University of Texas at Austin Austin, TX 78712
Department of Statisties and Data Sciences The University of Texas at Austin
Austin, TX 78712
fanqi. yan@utexas.eeeee..mmeeee
minhnhat @tt
ABSTRACT
Top-K sparse softmax gating mixture of experts has been widely used for scaling up massive deep-learning architectures without increasing the computational cost. Despite its popularity in real-world applications, the thcoretical understanding of that gating function has remained an open problem. The main challenge comes from the structure of the top-K sparse softmax gating function, which partitions the input space into multiple regions with distinct behaviors. By focusing on a Gaussian mixture of experts, establish theoretical results the effects of the top-K sparse softmax gating function on both density and parameter estimations. Our results hinge upon defining novel loss functions among parameters capture different behaviors of the input regions. When the true number of experts ki is known, we demonstrate that the convergence rates of density and parameter estimations are both parametric on the sample size. However, when k; becomes unknown and the true model is over-specified by Gaussian mixture of k experts where k > ki., our findings suggest that the number of experts selected from the top-K sparse softmax gating function must exceed the total cardinality of a certain number of Voronoi cells associated with the true parameters to guarantee the convergence of the density estimation. Moreover, while the density estimation remains parametric under this setting, the parameter estimation rates become substantially slow due to an intrinsic interaction between the softmax gating and expert functions.
INTRODUCTION
Introduced by Jacobs al. 1991 and Jordan Jacobs (1994), the mixture of experts (MoE) has been known as a statistical machine learning design that leverages softmax gating functions to blend different expert networks together in order to establish a more intricate model. Recently, there has been huge interest in a variant of this model called top-K sparse softmax gating MoE, which activates only the best K expert networks for each input based on sparse gating functions (Shazeer et al. 2017 Fedus et al., 2022a; Chen et al. 2023). Thus, this surrogate can be seen as form of conditional computation (Bengio 2013) Cho & Bengio 2014) since significantly scales up the model capacity while avoiding proportional increase in the computational cost. These benefits have been empirically demonstrated in several deep learning applications, including natural language
processing (Lepikhin al. 2021 Du et al. 2022
Fedus et al. 2022b Zhou et al., 2023 Pham et al.

---

Published as a conference paper at ICLR 2024
2024), speech recognition Peng al.) 1996; Gulati al.) 2020; You al. 2022) computer vision Dosovitskiy et al. 2021 Riquelme et al., 2021 Liang et al. 2022 Bao et al. 20222, multi-task
(Rives et al. 2021 Chow et al. 2023; Li et al., 2023] Han et al. (024). Nevertheless, to the best of our knowledge, the full theoretical analysis of the top-K sparse softmax gating function has remained missing in the literature.
learning (Hazimeh al. [202]] Gupta al., 2022) and other applications
In this paper, we aim to shed new light on the theoretical understanding of that gating function from a statistical perspective via the convergence analysis of maximum likelihood estimation (MLE) under the top-K sparse softmax gating Gaussian MoE.
There have been previous efforts study the parameter estimation problem in Gaussian MoE models. Firstly, Ho et al. (2022) utilized the generalized Wasserstein loss functions Villani (2003 20088 to derive the rates for estimating parameters in the input-free gating Gaussian MoF. They pointed out
that due to an interaction among expert parameters, these rates became increasingly slow when the number of extra experts rose. Subsequently, Nguyen al. (2023) and Ngguee al.] (2024) took
into account the Gaussian MoE with softmax gating and Gaussian gating functions, respectively. Since these gating functions depended on the input, another interaction between gating parameters and expert parameters arose. Therefore, they designed Voronoi loss functions capture these interactions and observe that the convergence rates for parameter estimation under these settings can be characterized by the solvability of systems of polynomial equations (Sturmfels 2002).
Turning to the top-K sparse softmax gating Gaussian MoF, the convergence analysis of the MLE, however, becomes substantially challenging due to the sophisticated structure of the top-K sparse softmax gating function compared to those of softmax gating and Gaussian gating functions. To comprehend these obstacles better, let introduce the formal formulation of that model.
Problem setup. Suppose that (X1, Y)), .,,,...,(4 (Xn, Y,) E R x IR are i.i.d. samples of size n from the top-K K sparse softmax gating Gaussian MoE of order ks, with the conditional density function
h. 9G. (.. = Σ Softmax X, K; K; (a)( x b7)
(I)
where G, = Eie exp(F%s (o (z)) is true but unknown mixing measure of order k2 (i.e., linear combination of k. Dirac measures 8) associated with true parameters (9ha- oz bf, of) for i (1, ). Here, hn (X, b) == a X + b is referred to as expert function, while we denote f((1/4, r) as an univariate Gaussian density function with mean j and variance a (The results for other settings of f(-1u, o) are in Appendix Ep. Additionally, we define for any vectors v = (V1...., )k.) and u = (u1...., us.) in Rd that Softmax(1:) == exp(x:) //E; exp(e;) and
TopK(u,, K;u,) = U/ + uis
oz is in the top K elements of v;
-00
otherwise.
More specifically, before applying the softmax function in equation (T) we keep only the top K values of (8i1) X, X...., ( X and their corresponding bias vectors among fi..-... - Po.. while set the rest to oo to make their gating values vanish. An instance of the top-K sparse softmax gating function is given in equation (3). Furthermore, linear expert functions considered in equation (I are only for the simplicity of presentation. With similar proof techniques, the results in this work can be extended to general settings of the expert functions, including deep neural networks. In order to obtain an estimate of G., we use the following maximum likelihood estimation (MLE):
Go E arg max =
Σ (Y/())) F
Under the exact-specified settings, i.c., when the true number of expert is known, the maximum in equation (2) is subject to the set of all mixing measures of order k., denoted by εx, (9t = {G kh
(Bus, S1i, a, b, 0,) E 22. On the other hand, under the over-specified settings, i.e., when hi is unknown and the true model is over-specified by a Gaussian mixture l experts where k > hi,, the maximum is subject to the set of all mixing measures of order at most k
EL= exp(Sex (
i.e., O₂ (2)= = (G (G=>X Ek exp( (o b..o.) < k' < k, (Bni, Bio, ai, bs ov) E 2}.
Universal assumptions. For the sake of theory, impose four main assumptions on the parameters:
2

---

Published as a conference paper at ICLR 2024
(U.1) Convergence of MLE: To ensure the convergence of parameter estimation, assume that the parameter space 52 is compact subset of R x R R x R x R., while the input space X' is bounded.
(U.2) Identifiability: Next, assume that 31,, O₂2 and fos. 0 that the top-K sparse softmax gating Gaussian mixture of experts is identifiable. Under that assumption. we show in Appendix C that if G and G' are two mixing measures such that (Y ga (Y|]) for almost surely (X, Y), then it follows that that G = G". Without that assumption, the result that G = G" does not hold, which leads to unncessarily complicated loss functions (see [Proposition Nguyen et al. (2023)). D)).
(U.3) Distinct Experts: To guarantee that experts in the mixture ① are different from each other, we assume that parameters (ai, 0i, oi, (ak D/, .x,) are pairwise distinct.
(U.4) Input-dependent Gating Functions: To make sure that the gating functions depend on the input X, we assume that at least one among parameters 9i2...., ,i is different from zero. Otherwise, the gating functions would be independent of the input X, which simplifies the problem significantly. particular, the model (I) would reduce an input-free gating Gaussian mixture of experts, which was already studied in Ho et al. (2022).
Challenge discussion. In our convergence analysis, there are two main challenges attributed the structure of the top-K sparse softmax gating function. (1) First, since this gating function divides the input space into multiple regions and each of which has different convergence behavior of density estimation, there could be mismatch between the values of the top-K K sparse softmax gating function in the density estimation 9a. and in the true density gG. (see Figure D. (2) Second, under the over-
specified settings, each component of G. could be fitted by least two components of G2. Therefore, choosing only the best K experts in the formulation of 9a. (Y is insufficient to estimate the true density 9G. (X) As result, it is of great importance to figure out the minimum number of experts selected in the top-K sparse softmax gating function necessary for ensuring the convergence of density estimation.
Contributions. In this work, we provide rigorous statistical guarantees for density estimation and parameter estimation in the top-K sparse softmax gating Gaussian MoE under both the exect-speciffed and over-specified settings. Our contributions are two-fold and can be summarized as follows (see also Table I):
1. Exact-specified settings: When the true number of experts k. is known, we point out that the density estimation 92., converges the true density gG. under the Hellinger distance the parametric rate, that is, Ex [h(ga (+Ix),
(...). SG. (+IX)) ((n3//3) Then, we propose novel Voronoi metric D, defined in equation (5) to characterize the parameter estimation rates. By establishing the Hellinger lower bound Ex [h( (-IX), 9c. (-1X)) 2 Dj (G G.) for any mixing measure G E Ex. (9), we arrive at another bound D, (G,,, G.) ) = On 11/2), which indicates that the rates for bj, estimating true parameters (xp( Pic, bb, of are of the optimal order ((n-1/2).
2. Over-specified settings: When h is unknown and the true model is over-specified by a Gaussian of it experts where k > k., we demonstrate that the density estimation 9G converges to the true density gG. only if the number of experts chosen from 98., denoted by K, is greater than the total cardinality of a certain number of Voronoi cells generated by the support of G.. Given these results, the density estimation rate is shown to remain parametric the sample size. Additionally, by designing a novel Voronoi loss function D2 in equation (8) to capture an interaction between parameters of the softmax gating and expert functions, we prove that the MLE G,, converges to the true mixing measure G. at a rate of -1/2) Then, it follows from the formulation of D2 that the estimation rates for true parameters 01, a,. bj.03 depend on the solvability of a system of polynomial equations arising from the previous interaction, and turn out to be significantly slow.
High-level proof ideas. Following from the challenge discussion, to ensure the convergence of density estimation under the exact-specified settings (resp. over-specified settings), we show that the input regions divided by the true gating functions match those divided by the estimated gating functions in Lemma (resp. Lemma 2 Then, we leverage fundamental results on density estimation for M-estimators in (van de Geer, 2000) to derive the parametric density estimation rate under the Hellinger distance in Theorem []] (resp. Theorem 3). Regarding the parameter estimation problem, a key step is to decompose the density discrepancy 98. XY - 9g. (YlX) into a combination of
3

---

Published as a conference paper at ICLR 2024
Table 1: Summary of density estimation and parameter estimation rates under the top-K sparse softmax gating Gaussian MoE model. Here, the function F(-) stands for the solvability of the system of polynomial equations ( and C is a Voronoi cell defined in equation 4 If [cy] 2, then we have F((C,)) 4. Meanwhile, if |C}] 3, then we get F([C,)) 6.
Setting
9c. (Y
exp(96,)
$i, b;
03.0z
Exact-specified
((n-1/2)
((n-1/2)
Onn 1/23 O (n 11/22/2((C)D))
Onn-1/2
Over-specified
(
O(n-1/4
O(n O(n-1/F(IC))) 1)
2
Figure I: Illustration of two partitions of the input space with respect to the TopK function in the density estimation 9G, (left) and the true density ga. (right) under the exact-specified settings when
k. = 3 and K = 1. Here, the regions labelled as 1, and 1 contain X E X' such that (311 T X and TX X are the top-1 elements of ( ((3)) x)²_ and (( X)²_ respectively. Other regions are defined similarly. Assume that 31: + S as n - do for any i E (1,2.3), then the regions 1n: 2n 3n should respectively match their counterparts 1 2,2",3 3 to guarantee the convergence of 98. to IG.- Lemma reads that this property holds when the sample size n is sufficiently large.
(
linearly independent terms. Thus, when the density estimation ga., (YlX) converges to the true density gG. (Y) the coefficients in that combination also tend to zero, which leads to our desired parameter estimation rates in Theorem (resp. Theorem 4
Organization. The rest of our paper is organized as follows. In Section 2 we establish the convergence rates of density estimation and parameter estimation for the top-K sparse softmax gating Gaussian MoE under the exact-specified settings, while the results for the over-specified settings are presented in Section 3 Subsequently, we provide practical implications of those results in Section before concluding the paper and discussing some future directions in Section 5 Finally, full proofs, numerical experiments and additional results are deferred the supplementary material.
ml Notations. For any natural numbers m 2 n,, we denote [n]] == {1, 21, n) and (*)) micm
Next, for any vector u, v E Rd, we let |u| Σí= u; u! :- ui!... uu!...ual, u" = of ud and denote ||u||| as its 2-norm value. Meanwhile, we define |A| as the cardinality of some set A. Then, for any two probability densities fi and fz dominated by the Lebesgue measure v, we denote V(fi. f2) == 1 l faldv as their Total Variation distance, whereas h22 fo) := 2. [(Vfi "dv represents the squared Hellinger distance. Finally, for any two sequences of positive real numbers (a,,) and (b,), the notations a = Obb,)) and a₂ S 6n, both stand for an I Cb, for all n e N for some constant C > O, while the notation an = (b,,) indicates that the previous inequality may depend on some logarithmic term.
2 EXACT-SPECIFIED SETTINGS
In this section, characterize respectively the convergence rates of density estimation and parameter estimation in the top-K sparse softmax gating Gaussian MoE under the exact-specified settings, namely when the true number of experts k. is known.
To begin with, let us introduce some essential notations and key results to deal with the top-K sparse softmax gating function. It can be seen from equation (T) that whether (at) X +b belongs to the top K experts in the true density ga. (YIX) not is determined based on the ranking of (8i.) T X + %
4

---

Published as a conference paper at ICLR 2024
which depends on the input X. Additionally, it is also worth noting that there are total of q := (x) different potential choices of top K experts. Thus, we first partition the imput space X into regions in order to specify the top K experts and obtain an according representation of 9G. (Y|X) in each region. In particular, for each & E [q). we denote {&i, f2...., .x
as an K-element subset of [k.] and (fK+1,-... fx.,) [k:.] (61, l2. ........ (x). Then, the f-th region of X is defined
X2 :=:=4 { E X:: (31.)
(giv) TT, 2, Vi (e,,...,kk) lx.,"' [fK.I.....,.,,
for any & E [g]- From this definition, we observe that X E X3 for some & E [ul such that {4,,.... {x} = [K]. then it follows that TopK((())) (X ; (%x) = (f%,) x Sb. for any i E [K]. As a result, (at) X + bi..... (ax)
) X b& become the top K experts, and the true conditional
density 9G. (Y is reduced to:
K
exp((3T))
9c. (YIX)
Σ
[(YI(a)) X + 0,
Pj=1 exp(( TX X + Po.)
Subsequently, we assume that the MLE Gn takes the form Gr = Σ exp (36)) (āp.ay ar 88 app)")
where the MLE's component WT"' ( 3r, oi b o.") converges to the true component w := (B Sig, at bf. ) for any i E [k.]. We figure out in the following lemma a relation between the values of the Topk function in ga. (YIX) and 9a. (YIX):
Lemma 1. For any E [k.]. let 3i6, Bis He such that - S my for some sufficiently small > O. Then, for any 8 E [9), unless the set X2 has measure zero, we obtain that X2 = Xe where
Ne : fx E X (B1s)Ta x 2 (BIx) a. v,Vie (xI,....kk),7 (lx+1....,, L..
Proof of Lemma is in Appendix A.3 This lemma indicates that for almost surely X, Top)( X, K; (2,) = ($1) X Fox equivalent to TopK X, K; (K) = (3 Por for any i E [k=]] and sufficiently large n. Based on this property, we provide in Theorem the rate for estimating the true conditional density function gc.:
Theorem (Density estimation rate). Given the MLE G. defined in equation Q the convergence
rate of the conditional density estimation 92, to the true conditional density 9a. under the exact- specified settings is illustrated by the following inequality:
P (Ex (+X)), 9a. 9..x)) > CVlog(n)/n) S n
where C' > 0 and 0 are some universal constants that depend on Ω and K.
It follows from the result of Theorem that the conditional density estimation 9a.. converges to its true counterpart gG. under the Hellinger
Proof of Theorem D can be found in Appendix A.I
distance at the parametric rate of order ( This rate plays critical role in establishing the convergence rates of parameter estimation. In particular, if we are able to derive the following lower bound: Ex [A (-))), gG. (-IX))] Z (G, G.) for any mixing measure G E Ex, (I), where the metric D, will be defined in equation ( we will achieve our desired parameter estimation rates. Before going into further details, let us introduce the formulation of Voronoi metric D2 that we use for our convergence analysis under the exact-specified settings.
Voronoi metric. Given an arbitrary mixing measure G with k' components, we distribute those components to the following Voronoi cells generated by the components of G. Manole & Ho (2022):
C, = CG = {i E [k'] ||0, 07 S ||0; 6 Vj' $j}.
(4)
where 0: == -B1i; a, a;) and 0z = ( ag, b7, 0z) for any j E [k..]. Recall that under the exact-specified settings, the MIE G₂ belongs to the set &s. (9). Therefore, we consider k'' = k: in this case. Then, the Voronoi metric D, is defined as follows:
K
Di (G, G=)
max
Σ LEC
exp(30s) (14 | || + ||Abse, + || 1l)
((..)
1=1
Σ Ece,
exp(30x) exp(( (5)
5

---

Published as a conference paper at ICLR 2024
Here, we denote A,12, B1c -ie, Aaie; := a, -i, Abue := b, - by and Aou; o, oi Additionally, the above maximum operator helps capture the behaviors of all input regions separated by the top-K sparse softmax gating function. Now, we are ready to characterize the convergence rates of parameter estimation in the top-K sparse softmax gating Gaussian MoE.
Theorem 2 (Parameter estimation rate). Under the exact-specipee settings, the Hellinger lower bound Ex ( ( 9G_ (+XX)) 2 Di (G, G.) holds for any mixing measure G E &k, (9t) Consequently, can find universal constant C1 > 0 depending only G.,, Q and K such that
(Di G.) > C1 Givlog(n)/m) & Snaa
where e1 > 0 is universal constant that depends only Ω
Proof of Theorem is in Appendix A.2 This theorem reveals that the Voronoi metric D1 between the MLE G, and the true mixing measure G., i.e. Di (G... G.), vanishes at the parametric rate of order O(n-1/2). As result, the rates for estimating ground-truth parameters exp(86)) $1, i aj 0y of optimal at Ö(n for any E [
OVER-SPECIFIED SETTINGS
In this section, we continue carry out the same convergence analysis for the top-K sparse softmax gating Gaussian MoE as in Section 2 but under the over-specificed settings, that is, when the true number of experts k, becomes unknown.
Recall that under the over-specified settings, we look the MLE G, within Og (S2), i.e. the set of all mixing measures with at most & components, where k: >> k., Thus, there could be some components (8% af by of) of the true mixing measure G. approximated by at least two fitted components (B a. by o of the MLE G... Moreover, since the true density gG, (Y|X) is associated with top K experts and each of which corresponds to one component of G., we need to select more than K experts in the formulation of density estimation 98, (YIX) to guarantee its convergence to
ga. (Y In particular, for any mixing measure G - exp( (111))81 us consider a new formulation of conditional density used for estimating the true density under the over-specified settings follows:
O&()), let
XY = Σ Softmax(TopK((I))) X. - X o.).
Here, (Y) is equipped with top-K sparse softmax gating, where 1 K < is fixed and might be different from K. Additionally, the definition of MLE in equation (2) also changes accordingly to this new density function. Then, we demonstrate in Proposition an interesting phenomenon that the density estimation Iö., converges ga. under the Hellinger distance only if
K Ix Σø ICe, where C₁, is Voronoi cell defined in equation
Proposition 1. K < ( E/=1 [Ce, then the following inequality holds true:
inf Ex (+ (.)) (-xx) 0. (E),
Proof of Proposition[ deferred to Appendix B.3 Following from the result of Proposition I we will consider only the regime when max C(k-] Σ [Ce, < K < k in the rest of this section ensure the convergence of density estimation. As the number of experts chosen in the density estimation changes from K K. necessary partition the input space X' into q == (*) regions. More specifically, for any E [4]. we denote (21, (,,),. as an K-element subset of [k] and (7 [ := [k] \ [21 .e...................... }x). Then, we define the 7-th region of X' as follows: x= fa E X: (B1x)
(Biz) x, Vi (l,, exc, i' {lKtD..., (x)))
Inspired by the result of Lemma I we continue to present in Lemma 2 relation between the values of the TopK functions in the density estimation 96 and the true density go..

---

Published as a conference paper at ICLR 2024
Lemma 2. For any j E [R.]] and i C,. let Pic, E R that satisfy || 3i, S nj for some sufficiently small ny > O. Additionally, for (e, Ej= Ce, S K S k, we assume that there exist & E lqi and ê E [$] such that {Z1.....
.xkt = Cia U... U Uix. Then, if the set X2 does not have measure zero, we achieve that X²- = =x,
Proof of Lemma is in Appendix B.4 Different from Lemma we need to impose on Lemma 2 an assumption that there exist indices E [g] and E [7] that satisfy (fI...., -57 = Ce, U... UCx- This assumption means that each component (Bi, Q. by a..) corresponding to the top K experts in 9a.. must converge to some true component which corresponds to the top K experts in gG.- Consequently, for X X7, if ((p)((()))
(81,) + f6 holds true, then achieve that TopK T X, K; (%)) (9%) T + 3ü and vice versa, for any j E [k.] and i E C2-
) T K; Poy)
Given the result of Lemma 2 we are now able derive the convergence rate of the density estimation GG. to its true counterpart gG. under the over-specified settings in Theorem 3
Theorem 3 (Density estimation rate). Under the over-specified settings, the conditional density estimation 9G, converges to the true density 9G. under the Hellinger distance at the following rate:
PP Ex(b(a (1) (-1x), (..))....110. (4x))) > C' (Voo((n)/)) S n
where C" > O and d' > 0 are some universal constants that depend on S2 and K.
Proof of Theorem is in Appendix B. Although there is a modification in the number of experts chosen from Ta. Theorem verifies that the convergence rate of this density estimation gc. under the over-specified settings remains the same that under the exact-specified settings, which is of order Onn 1/2). Subsequently, we will leverage this result for our convergence analysis of parameter estimation under the over-specified settings, which requires us to design new Voronoi metric.
Voronoi metric. Regarding the top-K sparse softmax gating function, challenges comes not only from the TopK function but also from the Softmax function. In particular, there is an intrinsic interaction between the numerators of softmax weights and the expert functions in the Gaussian density. Moreover, Gaussian density parameters also interacts with each other. These two interactions are respectively illustrated by the following partial differential equations (PDEs):
)x((XY) = Ys(X,Y), 0,,0b
0s(X,Y) = 2
Oa
do
062
where we denote s(X, exp X). [(Yla X ++b a). These PDEs arise when we decompose the density difference 98, (Y|X) - gG. gG X) into a linear combination of linearly independent terms using Taylor expansions. However, the above PDEs indicates that derivative terms which admit the forms in equation (6) are linearly dependent. Therefore, we have to group these terms by taking the summation of their coefficients, and then arrive at our desired linear combination of linearly independent elements. Consequently, when 9a_ (YIX) ga_ (Y)X) converges to zero, all the coefficients in this combination also tend to zero, which leads to the following system of polynomial equations with unknown variables [215+ 22j: Æ3j1 24js 25, m
m
Σ
2 Zgi 24 O.
a a ag! ag!
(an (0 (.).
for all (m,72) E Nd x N such that 0 < mml < 0
< 12 < -- lm and |m + 1/2 2 1, where an 02 ?, n,,00l + 0g + 04 = 72)
Jmmm == ((01, 02, 03 (14) E Nd x No x N N
Here, a solution to the above system is called non-trivial if all of variables 2s: are different from zero whereas at least one among variables 23; is non-zero. For any m 2 2, let F(m) be the smallest natural number such that the above system does not have any non-trivial solution. In general case when 2 1 and m 2 2, it is non-trivial to determine the exact value of F(m) Sturmfels (2002). However, when is small, Nguyen et al.) (2023) show that F(2) 4 and F(3) 6. Additionally, since F(m) is a monotonically increasing function of m, they also conjecture that r(m) = 2m for any m 2 2 and
d > 1.
7

---

Published as a conference paper at ICLR 2024
Given the above results, design Voronoi metric D2 to capture the convergence rates of parameter estimation in the top- sparse softmax gating Gaussian MoË under the over-specified settings as
P2( =
max ((R.
exp(S6i) [IA II ||Aau, || +| |abie, + |
jE[K], sECi, 12,1=1
+ Σ Σ exp( (oooi/(11880 |Issie, ( 12 + || || jE(K], iECr, Ice, 1e,1
+ |biie p(|Ce, + |Bore,
P(iCe
K
exp(3o:) - -exp( |
(8)
iEC,
for any mixing measure G E On (12). The above maximum operator allows us to capture the behaviors of all input regions partitioned by the top-K sparse softmax gating function in JG.. Then, we show in the following theorem that parameter estimation rates vary with the values of the function F(.)
Theorem 4 (Parameter estimation rate). Under the over-specified settings, the Hellinger lower bound Ex [h(aa(-)X (x), 9G. (-Ix))] 2 D2(G, G.) holds for any mixing measure G E (x()). As consequence, we can find a universal constant C2 depending only on G.,, Q and K such that
P P(Ds(Gn,..) > C2. C2V/log(n)/m) 2 02
where 02 > 0 is a universal constant that depends only on Ω
Proof of Theorem 4 is in Appendix B.2 A few remarks on this theorem in order.
(i) Under the over-specified settings, the MLE Gn converges to the true mixing measure G, at the rate of order (n2222 under the Voronoi metric D2. Let us denote Cу C2 (G,), then this result indicates that the estimation rates for ground -trutt parameters exp(3i;), $1,, aj. bj. fitted by only one component, i.c. |CJ) remain the same at Ö(n 1/2) as those in the exact-specified settings.
(ii) However, for ground-truth parameters which are approximated by at least two components, i.c. |CT] > the rates for estimating them depend on their corresponding Voronoi cells, and become significantly low when the cardinality of those Voronoi cells increase. In particular, both St and bj admit the estimation rate order ((n-1/2((()))) Meanwhile, the convergence rates of estimating af and of are of the same order Onn On ((n(((C(C])) C) For instance, assume that (gi aj. bj. oj) has three fitted components, which means that (C}) = 3 and therefore, r(IC)) ) = 6. Then, the estimation rates for :iz;00 bg and 0z are Onn and ( -1/6), respectively.
PRACTICAL IMPLICATIONS
In this section, we present three main practical implications of our theoretical results for the use of top-K sparse softmax gating function in mixture of experts as follows:
No trade-off between model capacity and model performance: In the top-K sparse softmax gating Gaussian mixture of experts, since the gating function is discontinuous and only a portion of experts are activated for each input scale up the model capacity, the parameter estimation rates under that model are expected to be slower than those under the dense softmax gating Gaussian mixture of experts (Nguyen et al. 2023). However, from our theories turns out that the parameter estimation rates under these two models are the same under both the exact-specified and over-specified settings, As result, we point out that using the top-K sparse softmax gating function allows us to scale up the model capacity without sacrificing the computational cost as well as the convergence rates of parameter and density estimation.
2. Favourable gating function: As mentioned in Section due to intrinsic interaction between gating parameters and expert parameters via the first PDE in equation (6). the rates for estimating those parameters under the over-specified settings are determined by the solvability of the system of
8

---

Published as a conference paper at ICLR 2024
polynomial equations (7). which are significantly slow. However, if we use the top-1 I sparse softmax gating function, i.e. activating only a single expert for each input, then the gating value is either one or zero. As result, the previous interaction no longer occurs, which helps improve the parameter estimation rates. This partially accounts for the efficacy of top-1 sparse softmax gating mixture of experts in scaling up deep learning architectures (see (Fedus et al. (022b))
3. True/ Minimal number of experts: of experts, which can be partially addressed using theories developed from the paper. In particular, suppose that the MIE Gn consists of k components. When the sample size n goes to infinity, every Voronoi cell among CT. ......, CE. contains at least one element. Since the total number of elements of those Voronoi cells is kn the maximum cardinality of a Voronoi cell turns out to be Rn - k, + 1. This maximum value is attained when there is exactly one ground-truth component (gi aj 03)
Another challenge is to choose the true/ minimal number
fitted by more than one component. An instance for this scenario is when CT = kn, h- + 1 and |CTI = ... |eg_ = 1. Under this setting, the first true parameters Pi1. bj and aj, of achieve their worst possible estimation rates of order On -1 +)) and (1/11/ -1)1))) respectively, which become significantly slow when the difference k, k. increases. As a consequence, the estimated number of experts k,, should not be very large compared to the true number of experts k..
5 CONCLUSION AND FUTURE DIRECTIONS
In this paper, we provide a convergence analysis for density estimation and parameter estimation in the top-K sparse softmax gating Gaussian MoE. To overcome the complex structure of top-K sparse softmax gating function, we first establish connection between the outputs of TopK function associated with the density estimation and the true density in each input region partitioned by the gating function, and then construct novel Voronoi loss functions among parameters to capture different behaviors of these input regions. Under the exact-specified settings, we show that the rates for estimating the true density and true parameters are both parametric on the sample size. On the other hand, although the density estimation rate remains parametric under the over-specified settings, the parameter estimation rates witness sharp drop because of interaction between the softmax gating and expert functions.
Based the results of this paper, there are few potential future directions. Firstly, mentioned in Section 4 a question of how to estimate the true number of experts k. and the number of experts selected in the top-K sparse softmax gating function K naturally arises from this work. Since the parameter estimation rates under the over-specified settings decrease proportionately to the number of fitted experts k. one possible approach to estimating k. to reduce k until these rates reach the optimal order Ö(n -1/2). That reduction can be done by merging close estimated parameters within their convergence rates to the true parameters (Guha et al., 2021 or by penalizing the log- likelihood function of the top-K K sparse softmax gating Gaussian MoE using the differences among the parameters (Manole & Khalili 2021). As a result, the number of experts chosen in the density estimation K also decreases accordingly and approximates the value of K. Secondly, the theoretical results established in the paper are under the assumption that the data are assumed to be generated from a top-K sparse softmax gating Gaussian MoF, which can be violated in real-world settings when the data are not necessarily generated from that model. Under those misspecified settings, the MLE converges to mixing measure G E arg mingeo, (2) KL(P(Y [X) ga gY (X) where XX) is the true conditional distribution of Y given X and KL stands for the Kullback-Leibler divergence. As the space O2(2)) is non-convex, the existence of G is not unique. Furthermore, the current analysis of the MLE under the misspecified settings of statistical models is mostly conducted when the function space is convex (van de Geer 2000). which is inapplicable to the current non-convex misspecified settings. Thus, it is necessary to develop a new analysis and a new metric to establish the convergence rate of the MLE to the set of G. Finally, since the log- ligleiill function of the top-K sparse softmax gating Gaussian MoE is highly non-concave, the EM algorithm is used to approximate the MLE. While the statistical guarantee of the EM has been established for vanilla Gaussian mixture models (Balakrishnan et al.) 2017; Dwivedi et al., 2020b a). to the best of our knowledge such guarantee has not been studied in the setting of top- K sparse softmax gating Gaussian MoE. A potential approach to this problem is to utilize the been widely used in previous works to study the EM algorithm [Balakrishaan et al. (2017); Ho et al. (2023). We leave that direction for the future work.
potilatiiiottt--mpp-- analysis that has

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of Machine Learning.
REFERENCES
S. Balakrishnan, M. J. Wainwright, and B. Yu. Statistical guarantees for the EM algorithm: From population sample-base based analysis, Annals of Statistics, 45:77-120, 2017.
H. Bao, W. Wang, L. Dong, Q. Liu, O-K. Mohammed, K. Aggarwal, S. Som, S. Piao, and F. Wei. VLMo: Unified vision- languae pre-training with mixture-of-men wwla-mmmddlllppppeeeee In Advances in Neural Information Processing Systems, 2022.
Y. Bengio. Deep learning representations: Looking forward. In International Conference on Statistical Language and Speech Processing, 2013.
Faicel Chamroukhi, Allou Samé, Gérard Govaert, and Patrice Aknin. Time series modeling by regression approach based on a latent process. Neural Networks, 22(5-6):593-602, 2009. Publisher: Elsevier.
T. Chen, Z. Zhang. A. Jaiswal, S. Liu, and Z. Wang. Sparse moe as the new dropout: Scaling dense and self-slimmable transformers. In The Eleventh International Conference on Learning Representations, 2023.
K. Cho and Y. Bengio. Exponentially increasing the compilyyttooommmtttttiint ratio for conditional computation in deep learning. arxiv preprint arxiv 1406.7362. 2014.
Y. Chow, A. Tulepbergenov, O. Nachum, D. Gupta, M. Ryu, M. Ghavamzadeh, and C. Boutilier. A mixture-of-expert approach to rl-based dialogue management. In The Eleventh International Conference on Learning Representations, 2023.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum Likelihood from Incomplete Data Via the EM Algorithm. Journal of the Royal Statistical Society: Series (Methodological), 39(1):1-22, September 1977. ISSN 0035-9246.
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
N. Du, Y. Huang, A. M. Dai, S. Tong. D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. Yu, O. Firat, B. Zoph, L. Fedus, M. Bosma, Zhou, Wang, E. Wang, K. Webster, M. Pellat, K. Robinson, K. Meier-Hellstern, T. Duke, L. Dixon, K. Zhang, Q Le, Y. Wu, Z. Chen, and C. Cui. Glam: ICML, 2022.
Efficient scaling of language models with mixture-of-experts
R. Dwivedi, N. Ho, K. Khamaru, M. J. Wainwright, M. I. Jordan, and B. Yu. Sharp analysis of expectation-maxiiizatii for weakly identifiable models. International Conference on Artificial Intelligence and Statistics (AISTATS), 2020a.
R. Dwivedi, N. Ho, K. Khamaru, M. J. Wainwright, M. Jordan, and Yu. Singularity, misspecifi- cation, and the convergence rate of EM. Annals of Statistics, 44:2726-2755 2020b.
W. Fedus, J. Dean, and B. Zoph. A review of sparse expert models in deep learning. arxiv preprint arxiv 2209.01667, 2022a.
W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23:1-39, 2022b.
A. Guha, N. Ho, and XL. Nguyen. On posterior contraction of parameters and interpretability in Bayesian mixture modeling. Bernoulli, 27(4):2159-2188 2021.
10

---

Published as a conference paper at ICLR 2024
A. Gulati, J. Qin, C. Chiu, N. Parmar, Y. Zhang, J. Yu, W. Han, S. Wang. Z. Zhang, Y. Wu, and R. Pang. Conformer: anvvolutio-auummettee Transformer for Speech Recognition. Proc. Interspeech 2020, pp. 5036-5040, 2020. doi: 10.2 htt1:///meerppeccc22000------
S. Gupta, S. Mukherjee, K. Subudhi, E. Gonzalez, D. Jose, A. Awadallah, and J. Gao. Sparsely activated mixture-of-experts are robust multi-task learners. arXiv preprint arxiv 2204.0768, 2022.
X. Han, H. Nguyen, C. Harris, N. Ho, and S. Saria. Fusemoe: Mixture-of-experts transformers for fleximodal fusion. arXiv preprint arXiv:2402.032266 2024.
H. Hazimeh, Z. Zhao, A. Chowdhery, M. Sathiamoorthy, Y. Chen, R. Mazumder, L. Hong, and Ed H. Chi. Dselect-k: Differentiable selection in the mixture experts with applications to multi-task learning. In NeurIPS, 2021.
N. Ho and X. Nguyen. On strong identifiability and convergence rates of parameter estimation in finite mixtures. Electronic Journal Statistics, 10:27 00:271-307, 2016.
N. Ho, C. Y. Yang. and M. I. Jordan. Convergence rates for Gaussian mixtures of experts. Journal of Machine Learning Research, 23(323):1-81, 2022.
N. Ho, K. Khamaru, R. Dwivedi, J. Wainwright, M. I. Jordan, and B. Yu. Instability, computational efficiency and statistical accuracy. Journal of Machine Learning Research, 2023.
R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton. Adaptive mixtures of local experts. Neural Computation, 3, 1991.
M. I. Jordan and R. A. Jacobs. Hierarchical mixtures experts and the EM algorithm. Neural Computation, 6:181-214, 1994.
Kwon and C. Caramanis. EM converges for mixture of many linear regressions. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 1727-1736. PMLR, 26-28 Aug 2020.
J. Kwon, W. Qian, C. Caramanis, Y. Chen, and D. Davis. Global convergence of the algorithm for mixtures of two component linear regression. In Alina Beygelzimer and Daniel Hsu (eds.), Proceedings of the Thirly-Seoond Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 2055-2110. PMLR, 25-28 Jun 2019.
J. Kwon, N. Ho, and C. Caramanis. On the minimax optimality of the em algorithm for learning two- component mixed linear regression. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 1413- PMLR, 13-15 Apr 2021.
D. Lepikhin, Lee. Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. In International Conference on Learning Representations 2021.
B. Li. Y. Shen, J. Yang, Y. Wang, J. Ren, T. Che, J. Zhang, and Z. Liu. Sparse mfixture-of-expeer domain generalizable learners. In The Eleventh International Conference on Learning Representa- fions, 2023.
H. Liang, Z. Fan, R. Sarkar, Z. Jiang, T. Chen, Zou, Cheng, C. Hao, and Z. Wang. M Mixture-of-Exxe
Vision Transformer for Efficient Multi-task Learning with Mccel-Accelerator Co-design. In NeurlPS, 2022.
T. Manole and N. 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.
T. Manole and A. Khalili. Estimating the number of components in finite mixture models via the group-sort-fuse procedure. The Annals of Statistics, 39))3344333 2021.
11

---

Published as a conference paper at ICLR 2024
H. Nguyen, TT. Nguyen, and N. Ho. Demystifying softmax gating function in Gaussian mixture of experts. In Advances in Neural Information Processing Systems, 2023.
H. Nguyen, TT. Nguyen, K. Nguyen, and N. Ho. Towards convergence rates for parameter estimation in Gaussian-gated mixture of experts. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, 2024.
F. Peng. R. Jacobs, and M. Tanner. Bayesian Inference in Mixtures-of- and Hierarchical Mixtures-of ofxtures-of-Expprts Models With an Application Speech Recognition. Journal of the American Statistical Association, 91(455)953-966 1996. ISSN 01621459.
Q. Pham, G. Do, H. Nguyen, TT. Nguyen, C. Liu, M. Sartipi, B. T. Nguyen, S. Ramasamy, X. Li, S. Hoi, and N. Ho. Competesmoe effective training of sparse mixture of experts via competition. arXiv preprint arXiv:2402.025266 2024.
C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, Jenatton, Susano Pint, Keysers, and N. Houlsby. Scaling vision with sparse mixture of experts. In Advances in Neural Information Processing Systems, volume 34, pp. 9883-8995. Curran Associates, Inc., 2021.
A. Rives, J. Meier, Sercu, Goyal, 7. Lin, J. Liu, D. Guo, M. Out, Zitnick, Ma, and R. Fergus. Biological structure and function emerge from scaling unsupervised learning 250 million protein sequences. Proceedings of the National Academy of Sciences, 118, 2021.
N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expprts layer. In International Conference Learning Representations, 2017.
B. Sturmfels. Solving Systems of Polynomial Equations. Providence, R.I, 2002.
H. Teicher. On the mixture of distributions. Annals of Statistics, 31:55-73, 1960.
H. Teicher. Identifiability of mixtures. Annals of Statistics, 32:244-243 1961.
H. Teicher. Identifiability of finite mixtures. Ann. Math. Statist., 32: 32:1265-1269 1963.
van de Geer. Empirical Processes M-estimation. Cambridge University Press, 2000.
C. Villani. Topics Optimal Transportation. American Mathematical Society, 2003.
C. Villani. Optimal transport: Old and New. Springer, 2008
Z. You, S. Feng, D. Su, and D. Yu. Speechmoe2: Mixture-of-experts model with improved routing. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7217-7221, 2022.
Y. Zhou, N. Du, Y. Huang, D. Peng, C. Lan, D. Huang, S. Shakeri, D. So, A. Dai, Y. Lu, Z. Chen, Q. Le, C. Cui, Laudon, and J. Dean. Brainformers: Trading simplicity for efficiency. In International Conference Machine Learning, pp. 42531-42542, PMLR, 2023.
12

---

Published as a conference paper at ICLR 2024
Supplement to "Statistical Perspective of Top-K Sparse Softmax Gating Mixture of Experts"
In this supplementary material, we first provide rigorous proofs for all results under the exact-specified settings in Appendix A while those for the over-specified settings are then presented in Appendix B Next, study the identifiability of the top-K sparse softmax gating Gaussian mixture of experts (MoE) in Appendix[[] We then carry out several numerical experiments in Appendix Dito empirically justify theoretical results. Finally, establish the theories for parameter and density estimation beyond the settings of top-K sparse softmax Gaussian MoE in Appendix
A PROOF FOR RESULTS UNDER THE EXACT-SPECIFIED SETTINGS
In this appendix, we present the proofs for Theorem | in Appendix A. while that for Theorem then given in Appendix A.2 Lastly, the proof of Lemma is provided in Appendix A.3
A.1 PROOF OF THEOREM
In this appendix, we conduct a convergence analysis for density estimation in the top-K sparse softmax gating Gaussian MoE using proof techniques in (van de Geer 2000). For that purpose, is necessary to introduce some essential notations and key results first.
Let Pk. (+) (9)(Y[X) G E Ex. } be the set of all conditional density functions of mixing measures Ex. (Ω). Next, we denote by (e, Pk. (S2), |l II1) the covering number of metric space (Pa. (I2),I || IIi). Meanwhile, HB(e, Ps. (52), h) represents for the bracketing entropy of Pk. (S2) under the Hellinger distance. Then, we provide in the following lemma the upper bounds of those terms.
Lemma 3. If is bounded set, then the following inequalities hold for any 0 < < 1/2:
(i) log (n, Pk. (9)), -1 11)) & log(1/n))
(ii) Hn Pk. (R), h) x log(1 (1))).
Proof of Lemma
is in Appendix A.1.2 Subsequently, we denote
Px. (12) (_(=+G.)/(((x)) G &x. (9))
2/.2 (1)) (X( G e &x. (92)
In addition, for each 8 > 0, we define Hellinger ball centered around the conditional density function 9a. (Y) and met with the set Pk. ( as
///² (0, 8) = 12 E PM/ (S2) hgg.G.)) 53.
To capture the size of the above Hellinger ball, van de Geer (2000) suggest using the following quantity:
(8,6, (,.2((,6))))
(2) = S 2 (1, 11/2 (0,,t),1 IIDdd 8,
where V max L, 83. Given those notations, let us recall a standard result for density estimation in van de Geer (0000)
Lemma 4 (Theorem 7.4, van de Geer 2000 Take (8) 2 JB(6, (1, (2,5)) such that (82 is non-increasing function of Then, for some sequence (ō.) and universal constant which satisfy V 2 cvl (5), we obtain that
IP Ex Moaa (-X). (3.)).a..123 (1x)).9. >6) cexp(-no8 (
for any 8 2 8n
Proof of Lemma can be found in van de Geer (200).. Now, we are ready to provide the proof for convergence rate of density estimation in Theorem Appendix A.I.I
13

---

Published as a conference paper at ICLR 2024
A.1.1
MAIN PROOF
It is worth noting that for any 0, we have
(t Pl,2(0, t), Il -I) & Hp(t, Pk. (2, t),
Then, the integral in equation
(9) is upper bounded as follows:
/ /2/11
JB(8, (2,)(2)))) (2,
<
(t, (2,
log(I/t)dt Vō, (10)
/213
where the second inequality follows from part (ii) of Lemma 3
As result, by choosing (6) 5 /log(1/6), we can verify that (6)/6² is non-increasing function of ò. Furthermore, the inequality in equation (10 indicates that (6))) 2 JB(ó, Pk_ (22, 6)).
Next, let us consider sequence (6..) defined as 6n == Vlo(n)/nn. This sequence can be validated
to satisfy the condition / vńo² 2 (8) for some universal constant c. Therefore, by Lemma 4 we reach the conclusion of Theorem 1
P Ex P(Ex (N (-I)), JG_ (-1x)) CVlog(n)/n) Sn
for some universal constant C depending only on f2
A.1.2 PROOF OF LEMMA 3
Part (i). In this part, we will derive the following upper bound for the covering number of metric space (Px. (9), Ili)) for any 0 < < 1/2 given the bounded set Q:
log N (1,, Pk. (I)) - I11) x log(1/n).
To begin with, we define + == {(a, b, E Rd R x R+ (Bo, a, a, b, σ)) E D}. Note that $ is a bounded set, then θ also admits this property. Thus. there exists an 17-cover of θ denoted by {(8p. B,)
Rd (Bo. 81 a, b, σ) E 9. and A,, be an 7-cover A. Then, it can be validated that 10, S On (o(n-(dt]kk and [A, s On (d+3)k.).
Op Additionally, we also define
E Ek. (2)) we take into account two other mixing measures. The first measure G" = Σi= exp(3ur) (B.,,,,,,, where (a,, J,) G, is the closest points to (a;, (a,,bi,0 o,) in this set for all i E [kx]] The second one = Eh=
Subsequently, for each G = -
in which (Boir Bie) E A, for any i E [R:.]. Next, let us define : (9a PA. (1)): (For B1)) E,, (a,
J.) E G,, Vi E [k:]],
then it is obvious that 9G T. Now, we will show that T is an r-cover of metric space (Pr, (2), I1-11r) with a note that it is not necessarily the smallest cover. Indeed, according to the triangle inequality,
119G galli < |96 11 + llse sall1-
(11)
The first term in the right hand side can be upper bounded as follows:
k.
/ AxY
l69. 11 S
Σ
F(Y x x b,, f(Yia, x b, o,) ((x,)))) d(x,Y)
h-
(10:- - + 6b 6,1l + -|o, 10.1---1))e 7.l1)
xxy
(|a: ail + 11:; - -b + 10.
o..1)
(12)
Next, we will also demonstrate that ||9a - galla Z n7. For that purpose, let us consider q = k (x) (,,
K -eelement subsets of (1,.... k). which assumed to take the form (2,,2,,- .,.--.,,
for any e E [4]
14

---

Published as a conference paper at ICLR 2024
Additionally, also denote {lx+I....,,f) [x] {l....,k) (li...., .xx] for any & e [g]- Then, we define
X = E X: 31 8 : (x,,...,,k)).
:= (z E x: B1 x 2 Bipz: : (l,,....x)) E (lx+h...,, (x.))
By using the same arguments as in the proof of Lemma[]in Appendix A..3 we achieve that either
X₁ = X₂ or Xe has measure zero for any E [4]. As the Sofumax function is differentiable, it is Lipschitz function with some Lipschitz constant L 2 O. Since X' is a bounded set, we may assume that ||X|| < B for any X E X. Next, we denote
K
K
(X) (B1e.0 x + Poe. 1=1"
(x(X) = BIe, + Boe, s=1"
for any K-element subset {f1,... (k) of (1...., k.....,... Then, we get
[Softmax(T((X))) - Softmax(* (x))] S L
(x( - #x(Xi]l
K
S L
Σ
(1,31e. - Bie, 11 || + 130e. - Booe.)
L
nB +n)
2n.
-97ll1 S it follows from the above results that Back to the proof for |190
97l11
190 (YIX) (Y x)) ()) a(x,))
xxy
I0 (Y
(V|x) d(x,Y)
f=1
Xexy
Softmax(me(X)) - Softimx ||V|T, x + be, Fe.) d(x, V)
(13)
From the results in equations 12 and is an c-cover of the metric space (Px. (S2). (11). Consequently, we achieve that
deduce that ||9c 9alli S n. This implies that
Nn,, Px. (S2), I11) - A n 6 (133)))+)))) which induces the conclusion of this part
log N(, Pk. (9), |l- II1) S log(1/n).
Part (ii). Moving to this part, we will provide an upper bound for the bracketing entropy of Pk. (2)
under the Hellinger distance:
Hn(n, Pk. (9). S log(1/n).
Recall that ④ and are bounded sets, we can find positive constants -7 S a TX S 7 and
W1 5 5 u2. Let us define
Q(
Virai
exp
Sug
y2
for Y 2 27
for YY| < 2y
Then, it can be validated that f(Y(a b, S Q(X,Y) for any (X,Y) E x
Next, let S 7) which will be chosen later and (gI,..., .9N] an 7-cover of metric space (Px. (2), || [I1) with the covering number N = N(T, Pk. (2), [la). Additionally, we also consider brackets of the form (Y V (X] where
(Y) == max (Y T, 0)
/ (Y max(g. (Y + Q(I[)))
15

---

Published as a conference paper at ICLR 2024
Then, we can check that Pk. (22 S U (V VY and A. (Y|X) - (Y(() 1
min Q(Y x)), Let = max V8ug) } log (L/T), we have for any i E [N] that
V w = 1 J|1
I (YIx)
(X dxdy+ + 1 [ - (XX]xx)ddxxx
-
52
S ST +xp
S'T
2u2
where S' is some positive constant. This inequality indicates that
(S(S" Pk. (2)), I11) - log N Ps. (9), |l 11)) - log(1/7).
n/ = By setting =
we obtain that (n, Pk. (9), || IIl)) S log(L/n). Finally, since the norm || || is upper bounded by the Hellinger distance, we reach the conclusion of this part:
(n Ps. (O)). & log(1/n).
Hence, the proof is completed.
A.2 PROOF OF THEOREM 2
Since the Hellinger distance is lower bounded by the Total Variation distance, that is 2 V. we will prove the following Total Variation lower bound:
Ex [V (1)), (+X)) 2 GG,
which is then respectively broken into local part and global part as follows:
Ex V (ga(-1X), (-1X))]
inf (I)
> 0,
(14)
Di G.)
Ex (V gG. (+X))
inf (E),,
> 0,
(15)
(G,
for some constant e >> 0.
Proof of claim 14 It is sufficient to show that
Ex (- (G,
lim e-0 GEE. (G.G)E
inf
> 0.
Assume that this inequality does not hold, then since the number of experts k. is known in this case, there exists sequence of mixing measure G. = E exp(X8 (rop Ag.o) &x. (S) such (
that both Dj (G,, G.) and Ex[V (-IX), gG. ((x)))/Di(C.. G.) approach zero as n tends to infinity. Now, we define
CỤ CG
[k:.] : ||w - will S |w will Vs +}},
for any j E [R:.] Ass Voronoi cells with respect the mixing measure Gn. where denote wi = (3T, of b" a and 3 == (8r aj. by d;)) As we use asymptotic arguments in this proof, we can assume without loss of generality (WLOG) that these Voronoi cells does not depend on n, that is, C = CT. Next, it follows from the hypothesis (G.. G.) + 0 as n + oo that each Voronoi cell contains only one element. Thus, we continue to assume WLOG that C₁ = {j} for any j E [kw]], which implies that (ST), a.. 65, 63.05) (Big. ag. (),)) and exp exp(28)) + exp(3i, ) as - Dc.
Subsequently, to specify the top selection in the formulations of go.. (X) and ga. (YlX), we
divide the covariate space X' into some subsets two ways. particular, we first consider q == (8.)
different K-element subsets of [K:.], which are assumed take the form (e.....,& [.),, for l E (q). Additionally, we denote [Lx+1....,.e } == [ka] \ two following subsets of X:
Then, we define for each E (g)
X = I E X: (87)) Tx 2 (8t,) Tx Vj E (l....,,.kk Lkt.J' E (lx*t..., (ea
X2 = fDE X: (Bi,) Tz 2 (Bi Jx x Vj (lr....,kkk (lxti...f Br. 1.
16

---

Published as a conference paper at ICLR 2024
+ oo for any j E [k=]], we have for any arbitrarily small n 0 that || }i || - ? and |bū P,yl πy for sufficiently large n. By applying Lemma | we
Since (88)- , + (80),,9,))
obtain that X2² X2 for any & E [q] for sufficiently large WLOG, we assume that
Di Di G.)
Σ exp (i) (1 + +I I| + || + [1 + exp((()) exp
where we denote A := P1 - fi. day := a. 0j Ab = : bg and Ao; == o² - -,
Let l [g] such that , { {I....,K? }}. Then, for almost surely (X, E X2 Y, can rewrite the conditional densities gG., (Y|x) and gc. (Y|X) as
K
exp((9)))
T (hx)
= 9G,, (x|x)
X + bb, o;)) (())
1=1 exp((3t) X + (()))
K
exp(((11) X + ((ss)
9c. (YIX) =
f(Y((a)) x o
Tx
exp ) TX X + Biy))
Now, we break the rest of our arguments into three steps:
Step 1 - Taylor expansion:
In this step, we take into account H, == E exp((3t) + [b:)] (XIx)
- 9G. (Y|x)
Then, Hn can be represented as follows:
K
Hn = exp (3m)
+b) exp(((()) X))((((())) -xxP(()))))xxx(((())) + 12.
exp()
exp(((F)) X)ga,, -xp(x)))(((((x)) XX)G,,
K
+
Σ
exp( - exp(3
exxx((()) X)9,, (Y [x)]
By applying the first-order Taylor expansion to the first term in the above representation, which is denoted by An, we get that
K
cxp(
An
a1
(A
a!
glan taatasf ff xx((c))+++(+.00)) + (X,,
X t0 exp(() X)
Th ahloal losltts doas
where Rí (X, is Taylor remainder that satisfies Rí (X, Y)/D' 0 as co. Recall that f is the univariate Gaussian density, then we have
f aga
1 S f
(Y((a)) x o( 2
(Y)(a))
T x ob
oh
which leads to
exp (3))
An
(AB7) (Aa))aa (Ab)) (Aa))) (Ao7))a
2as
|a|=1
alaz|tan+20
ta
exp(()7) Tx)
(a;) T X + 0), + RX(X
ah 02l
exp(
(A jai (Ab))))
2a
|mi 1m1---m =1 aEJ
002 f
X exp(((1)))
Tx))
(Y x 03 RX,,,),,
(16)
th
17

---

Published as a conference paper at ICLR 2024
where we denote 7/1 a1 02 E N 72 l0g + a3 20x N and index set
In = {(0)) E No x N N x N 0n + 02 1, 03 + 204 = - |o2/]-
(17)
By arguing in similar fashion for the second term in the representation of 1l,, also get that
exp(P6r
B Σ
(APT) X exp(((1)) (X)G, R2(,,)),
1=1 In|=1
where Y) is a Taylor remainder such that R2(X, Y)/D G.) + 0 - co. Putting the above results together, rewrite the quantity Hn, as follows:
2m
Hn =
Σ
Σ
U m2
= exp((31, x)
OhT
1=
K
+ Σ
Σ
W X exp((((T)) X)ga. + +Xxx,) (18)
0shisl
in which we respectively define for each E [K]
that
exp(
U Gm
(A
2
b.y W = exp(19t) ?
(A
for any (71 72) (Os, 0) and hy * Oa. Otherwise, UÑ -W exp()))) exp(
Step 2 Non-vanishing coefficients:
Moving to the second step, we will show that not all the ratios U..yy.n. as n goes to infinity. Assume by contrary that all of them approach zero when n - oc, then for (m, m2) (0z, 0), it follows that
na /Di (Gn, G.) tend zero
K
K
Un
(x)) - exp((())
Σ
mm
0.
(19)
D D G.)
DT (Gn: G.)
==1
==1
Additionally, for tuples (m, 72) where E fei, e2..., Caf with = (O....,0) 0.
0,
and 1/2 0, we get
K U
Σ exp(
Σ
Di (Gn G.) - 0.
Di (G G.)
By using similar arguments, we end up having
exp($T)
|| + || + + + + 0.
DG G.)
Due to the topological equivalence between norm-l and norm-2, the above limit implies that
Di G.)
Σ exp((3)) || + ||
| + |Ao"
+
(20)
(G,, G.) - 0. which is Combine equation (19) with equation 20). deduce that D, (G., G.)/D, (Gn,
a contradiction. Consequently, at least one among the ratios U on DD (Gn G.) does not vanish as tends to infinity.
Step 3 Fatou's contradiction:
Let denote by M. the maximum of the absolute values of Un Un. /D, (Gn, G.) and W /D, (Gn, G.). It follows from the result achieved in Step 2 that 1 1/m, Co
18

---

Published as a conference paper at ICLR 2024
Recall from the hypothesis that Ex (V (1x). (() G.) 0 Thus, by the Fatou's 's lemma, we have
0 lim Ex (V (ga,, (- (11)). (11x)]
liminf 99c, ( -G. (Y X)
AxdY. This result indicates that 19a, (Y|X) - 9a. (Y[X))/D/ G.) tends to zero as n goes to infinity for almost surely (X, As result, it follows that
Di G.)
(i G.)
H.
9G. (YIX) 9G. (Y
lim 7-000 Myy Di (Gn, G.)
lim n-000
O.
mn Di (Gns G.)
Next, let us denote U wy [m, D1 (G,, G.)l a note that at least one among them is deduce that
Ts.m and W "/s.//mm /mm, D1 (Gn, G.) + Ki with From the formulation of Hn in equation (18). we
8Th (Y X + b,,
Teman X exp(((i,) X)
oh!
m Va
Kan X exp(((11)) X)9G. (YIx) O,
(21)
0ghii31
for almost surely (x, Y). This equation equivalent to
Σ
f
(ai o;) + Ki, KG. (.. (YIX)
Tim th oh
Eiu 02
x exp(( xx) 00
for almost surely (X,Y). Note that Si....., Six admits pair-wist wise different values, then (exp((3i,) X) Tx) E [K]] is a linearly independent set, which leads to
2T (Yl(a)) X + b.,o)) + Ks, (Y X 0,
Σ oSlm g
Σ
Timin t2 The
for any i E [K] for almost surely (X, Y). is clear that the left hand side of the above equation is a polynomial of X belonging to the compact set X. As result, get that
8 Temina dh?
Σ
(a T X + of, + Kim.g.. (. =0. O,
02
for any i E [K]. S |ml < 1 and almost surely (X, Y). Since (aī, bj, (i.aj)..... (axx s Uis ok) have pair-wise distinct values, those of ((a1) + by, o), ((ax) + ok are also wair-wise different. Thus, the set
(a) T X + by, af), gc. (Y :1 [K] is linearly independent. Consequently, we obtain that Tim.. = Kj.7 0 for any i [K]. |m|++) + S 2 and S |r| S 1, which contradicts the fact that at least one among these terms is different from zero.
Hence, we can find some constant é' such that
Ex (( (-X))
inf GEE. (G..GE
> 0.
Di(G,G,
Proof of claim 15 Assume by contrary that this claim is not true, then we can seek sequence G", E (2) such that (G,, G.)
ee and
lim Ex IV (goz, (-1X), (-XX)) 0,
(G'," G.)
which directly implies that Ex [V (goy (+X), (-IX))] + 0 as n + co. Recall that Q is a compact set, therefore, we can replace the sequence G¹, by one of its subsequences that converges to a mixing measure G' &x. (2). Since D, (G',, G.) e'. this result induces that D, G.) > é'.
19

---

Published as a conference paper at ICLR 2024
Subsequently, by means of the Fatou's 's lemma, we achieve that
lim Ex Ex[2V(ga; (1x), 2 limmnnf gc,, (V - gG. (V|X) (u)())(())
It follows that 9or (Y 9G. (Y])) for almost surely (X, Y). From Proposition 2 we know that top-K sparse softmax gating Gaussian mixture of experts identifiable, thus, we obtain that
G" = G.. As a consequence, D (G", G.) = 0, contradicting the fact that Di (G', G.) 0.
Hence, the proof is completed.
A.3 PROOF OF LEMMA T
Let 7І M,e, where is some fixed positive constant and M, will be chosen later. For an arbitrary E [q]. since and f2 are bounded sets, there exists some constant oz 2 O such that
min (oi) x (8),)) cjE,
(22)
where the minimum is subject to Li E {li..... .....,,k))))) and 1' {lx+1,...... ,..
{. }. We will
point out that cj > O. Assume by contrary that oz O. For E X2, we may assume for any 1 i j ki that
(Bie.) x 2 (Bie,)
Since of 0, it follows from equation
22 that (fiex) JTx (fient) TT: or equivalently
(iie Biex+)) 1 Tx x==1 O.
In other words, X is subset of
2 = fr E:: (3iex Biex xx 0).
Since 3iex BIek=1 * Oa and the distribution of X is continuous, it follows that the set Z has measure zero. Since X C Z, we can conclude that X2 also has measure zero, which contradicts the hypothesis of Lemma T Therefore, we must have ež > 0.
As X' a bounded set, we assume that ||ж||| S B for any x E X. Let XT, then we have for any i fli.... .x} and i' E (4x+1..... lx. that
B (B1-- (Bi))Bi) x (BR) J
MMeB + (Bie)² JT che
-MeBB + oge (Big - Bie) x + BT
2 B + +opE + BT.
By setting M, S & we get that æ E X₁, which means that X77 C Xr. Similarly, assume that there 2B'
exists some constant e₁ that satisfies
(8))) x - = che.
min 3,3,2'
Here, the above minimum subject to Xr i E (6..... tx} and il (lK+1,...,,e, l. If
M, < Ct. then we also receive that X₂ C Xi.
2B'
Hence, we set M, 2B min(c;, ce). we reach the conclusion that X; = Xt.
B PROOF FOR RESULTS UNDER OVER-SPECIRE SETTINGS
In this appendix, first provide the proofs of Theorem B.I and Theorem in Appendix B. and A.p2 respectively. Subsequently, we present the proof for Proposition in Appendix B.3 while that for Lemma 2 is put in Appendix B.4
20

---

Published as a conference paper at ICLR 2024
B. PROOF OF THEOREM! 3
to demonstrate the result of Theorem 3 Recall that under the over-specified settings, the MLE G. belongs to the set of all mixing measures with at most > k. components, i.e. (2 Interestingly, if we can adapt the result of part (i) of Lemma to the over-specified settings, then other results presented in Appendix [A.] will also hold true. Therefore, our main goal is to derive following bound for any < < 1/2 under the over-specified settings:
In this appendix, we follow proof techniques presented in Appendix A.
log. (n, Pa( I11) 2 log(1/n),
where PI( = (gglY (Y X) G (S2)). For ease of presentation, we will reuse the notations defined in Appendix A. with Ex_ (22 being replaced by (2)). Now, let us recall necessary notations for this proof.
Firstly, we define + = ((a,b, E IRd x X R+ (Bo: 31, a, b) E Q), and 6, is an n-cover of
O. Additionally, we also denote A == {BB, 31) E R x R : (So. 31 a, b, o) E S2} and D, be an 17-oovrr of A. Next, for each mixing measure G Ei=1 exp(8o: )6 0z.... b.oo.)) E (2)), denote G' Ex exp($o:)0 (B. in which (a,, b, o,) O is the closest point to (a,, b,, o,) in this set for any E [k). We also consider another mixing measure G = E exp( (3, O₂ (22) where (Bois Pin) T, is the closest point to (Soi, 31x) in this set for any E [k.]. Subsequently, we define
L: tJā E P&(S)) PI) E A,, (ā,, De, F.) E E
We demonstrate that C is an n-cover of the metric space (Ps(S), || [11, that is, for any Ta E Px there exists density Ig E L such that |1øG Talli n. By the triangle inequality, have
1īa Talla 1 S 11āG Io + 1o Tall..
(23)
From the formulation of G², we get that
1o Io 111 S
Σ
[(Y1a, x x b,, - f(y X + b, F. d(),)
X
[m
(la: T.ll -z,l
1b, -b, + lo; σ: d(x,)
xxxy
$n
(24)
Based on inequalities in equations
23 and
24
it is sufficient to show that I1 1zyl1 S 1.. For
any E [ let define
X = {z E X (Bnx) Ta x 2 (Bwr) Ix vie .i.....,,,khe E 1dt)
X = x: (11)) 2 (Z, vie Lx),i (xx41..., 43).
Since the Softmax function is differentiable,
a Lipschitz function with some Lipschitz constant L 2 0. Assume that ||||| < B for any X E X' and denote
K
Ty(X) = (BTz,0 Bozi) K s=1"
(x) = TT x + 8T,
1=1"
for any K-element subset (i.... .. [x] ) of (I....,,) k}. Then, we have
||Sottmax<xx(X)) - Softmax TT (x)) $ L. |xt(X) - xx(1Il
S L L
(11917,, B -1 -| +1307, -Bor.
1B n)
S 77
21

---

Published as a conference paper at ICLR 2024
By arguing similarly to the proof of Lemma in Appendix B.4 receive that either X₂ X2 or Ne has measure zero for any i E [ As result, we deduce that
7
1o - Ic. 111 S
Io -
TxXy
7=1
K
<
(07lmax(77(X)1)
(xX(m J(Y a.. x oz))d(XX
XxxY
2=1
iml
5.
Thus, is an n-cover of the metric space (Px(22) II [11), which implies that
N(n. PA(),, II)) & A, x 1 - On (0n-1(++))) (4+1)k) x Onn (1)134)3))) O(1 O -2d++4))).
(25)
Hence, log N (2)), II 111) S (1)))
B.2 PROOF OF THEOREM 4
Similar to the proof of Theorem in Appendix A our objective here is also to derive the Total Variation lower bound adapted to the over-fitted settings:
Ex [V (-x()))) gt. (+x))] D2(G,G.).
Since the global part of the above inequality he argued in the same fashion as in Appendix A will focus only on demonstrating the following local part via the proof by contradiction method:
Ex (- 9G. (-1X))] > 0.
lim
inf CEO.L))
(26)
(GG,,G))
Assume that the above claim does not hold true, then we can find sequence of mixing measures Gn : Di=i exp( aa oo)) E Ox(S)) such that both D2 G.) and Ex[V (-IX).9G. 9G. (-I G.) vanish when goes to infinity. Additionally, by abuse of notation, we reuse the set Voronoi cells Cj. for j E [ks.], defined in Appendix(A) AA] Due to the limit D2 (Gn:; G.) + 0 - co, it follows that for any j E [k.], we have Liec, exp(?%) exp( and (St a,, be,,b,,,00 ) + (gi aj UJ. 03)
k.,
for all i C,. WLOG, we may assume that
D2(.nn G.) = Σ Σ
cxp( |1AS%,11 P + | 1 + A ((IC)) + |A
jEIK EEC 10,1>1
Σ JE[K] SEC
exp(
(Bt)
||8%
||
A69
|
exp (2h) exp(
EC,
1c1=1
Regarding the top-K selection in the conditional density ga., we partition the covariate space X' in a similar fashion to Appendix A More specifically, we consider q = (tx) subsets {li,..., LK} of {1...., k. for any E E [q]. and denote {Ex+ fx_ = [k=]] {li...., Ex}. Then, we define
Lx).j' X = f x x ( 2 (Big)) Vj E (fI....,ex).) E
(..)). ................
partition method of the covariate space for the weight selection in the conditional density 9G.. In particular, let K E N such that max (1,1%,, Σ [c,, < K < ie and q : (*). Then, for any E [g), we denote (Z1,...,Za 44)
for any L E [g)- On the other hand, need to introduce
as subset [k]] and (R.t1.... [x] := [k] \ (l...... .......e ...)).. Additionally, we define
x= = D X: (0%)) Tx 2 ( x, Vi (Z
i' (Tx+)...,,44] [.].
22

---

Published as a conference paper at ICLR 2024
Let X E X2 for some & [q] such that (4,...., (x} (1...., K3. If f21.... CC... UCK (1x)))/222 G.) A 0 as tends infinity. This contradicts the fact that this term must approach zero. Therefore, we only need to consider the scenario when there exists E [g] such that {?,.... (7) = C, G.U...U UK. Recall that we have (00, (1)) - (B) 31,) n - oo for any j [k.] and E Cj. Thus, for any arbitrarily small 1/ > 0, we have that 11 Fi || f and 100 Po) ny for sufficiently large n. Then, it follows from Lemma 2 that X = A for sufficiently large n. This result indicates that X E X."
for any 8 E [7]. then (x,,
Then, we can represent the conditional densities 9c. (Y and 9G.. (Y for any sufficiently large n as follows:
exp((31,) (x
X + (8,)) (xp(((1 TX
x x
X + ()3z)
K
exp(( T X + (0t)
Do. (Y]X) =
[(Yl(a))) T X + 69 o.)).
2 Evecy exp( Tx + Pbr)
iEC
Now, we reuse the three-step framework in Appendix A
Step 1 Taylor expansion:
Firstly, by abuse of notations, let us consider the quantity
K
Hn exp(( x+,,)) (YIx) ga. (Y(x).
Similar to Step in Appendix
we can express this term as
H,
exp(Xit)
(X)((((()) of exp((3t)) -xx((() + b])
j=1 ieC;
exp( (3)) exp(((])) X)9G. (V exp(((()) "x (xxx)
j=1 iEC;
Σ
exp -
(T)((tt)))((((((()))) x + b;, exp((Bt,)) T X)9an (x,x)
== An + B, + En-
Next, we proceed to decompose A, based on the cardinality of the Voronoi cells as follows:
An Σ Σ
exp(3%) iiC,
exp(((7))T (T x))((((a)) X + b7 az) -xxp((01)) (x)(((())) Tx x + tbb
ΣΣ exp( exp(((9)) (X)(( x + D9 o.") ( Tx
x - 0j.
ACC
By applying the Taylor expansions of order 1 and F(CC)) to the first and second terms of An respectively, and following the derivation in equation (6)). we arrive at
An Σ EEE
Σ #CC, 15 + 52
exp(
(A n
2
f
x exp (
(Y
x + Rx(X
ThT
Σ Σ
exp(3())
)ai (An (Aoo))
(Doj)aa
20
J 1 tEC, 18mm It na ( (2)) mEJni
0
t exp(((t) Tx):
X)
(Y((a)) x bj,
x + RX(,)))
ah
23

---

Published as a conference paper at ICLR 2024
where the set Jm..z is defined in equation (77) while (X, is Taylor remainder such that (X ()n G.) - 0 as n - do for i E [3,4]. 4). Similarly, we also decompose B, according to the Voronoi cells as A, but then invoke the Taylor expansions of order and 2 to the first term and the second term, respectively. In particular, we get
exp(((.) ?!
Bn ΣEΣ
(APR) X7 exp(((())) X)sa, (VIX) + Rs(X,
1:1/1=1 iEC, h-|=1
(xp()%)
Σ EED Σ
(AP X exp((())) (X)_)))))GG))...... (Y Ro(X, Y),
7
3-C,, >1 ieC, 15his2
where Rx(X, Y) and Re Re(X, Y) are Taylor remainders such that their ratios over D2 D2(Gn, G.) ap- proach zero as n + oc. Subsequently, let us define
(xp(
S Smm
Σ
Σ
(Ao,)
2
iEC, aEJni.72
exp((%t)
3n Σ
(A
y
iEC
for any (m, 12) + (0a, 0) and h1 + Og, while for (M1 72) (0,, 0) we set
S $0.00 TT,o. Σ exp(P(i) exp()
As consequence, it follows that
D)(
8m on ta (Yl(a)) T X o;
Hn Σ
Σ
Smm
X exp((3(t)
1=1
1
K
141
Σ
Σ
Ty X exp(((3)) X)go, (YIX) R&(X, Y) Ro(X,Y).
(27)
1-/=0
Step 2 Non-vanishing coefficients:
In this step. will prove by contradiction that at least among the ratios S Smm (D G.) does not converge to zero as n + cC. Assume that all these terms go to zero, then by employing arguments for deriving equations [9) and (00), we get that
exp (p0) exp()
D2(Gnn,o)
xE
Σ exp( (1A81 || (
|| + +
+ [ooi)]] 0.
Combine this limit with the representation of D2 G=), we have that
Σ
Eex((B1
[1 ((C))) + Aa
" re0
A ((C))
+ Ao rde, 0.
D2 G.)
This result implies that find some index jr [K]] :Cprl that satisfies
Σ exp (3t)) (IA ( || + || <(
+ | ((C,))) +I reg0l
A 0.
D2(Gnn, G.)
IEC,
For simplicity, we may assume that j = 1. Since ST... my / (2(( G.) vanishes as n + oo for any (2)) E Nd x N such that 1 < |m + 72 < F((Cs)), we divide this term by the left hand side of the above equation and achieve that
(xp 20
Eiec, EaEdni
(A jan (Aafi (i (AbR) (Aofi) )ai
- 0, (28)
Σiec, exp(3%)) (1AAR
+ |ADg (((()) + A RICT)
24

---

Published as a conference paper at ICLR 2024
for any (71. 72) No x N such that < 1m 722 F((C1)).
Subsequently, we define M, == max ( ||, || 1/2, A 1/2 )1} and Pa == max(exp({8t) (6;:) E C1}. As a result, the sequence exp( (Bü:)/p Pn is bounded, which indicates that can substitute it with its subsequence that admits positive limit 2 == 1im exp()%:)/Pr Therefore, at least one among the limits 25. equals to one. Furthermore, we also denote
(A M/M + 216 (AM, - 221 (A6 + 234 (2 246-
From the above definition, it follows that at least one among the limits 21s: #2s: R3s and ²i: equals to either -1. By dividing both the numerator and the denominator of the term in equation (28 by Pn Pn H1t we arrive at the following system of polynomial equations:
25 211 22 231
02 0.4
0,
an 02! ag. 04!
iEC o oEJu1.02
for all (71. 72) N : Mil 7 F((C1)). Nevertheless, from the definition of F([C1 )), we know that the above system does not admit any non-trivial solutions, which a contradiction. Consequently, not all the ratios S)mm.. / (Dn: G.) tend to zero as n goes to infinity.
Step 3 Fatou's contradiction:
It follows from the hypothesis that Ex (V (1.))).cc())))/GCGcGcc - 0 Co.
Then, by applying the Fatou's ' lemma, we get
(-1X), (-1x))] D2(G,n G.)
IG.. (Y - -G. (X)x) dxaY,
0 lim
lim inf n-000
2
D2 G.)
which implies that Ia. (YIX) 96. (YIX))/D2((,,, G.) + 0as
for almost surely (X, Y).
Next, we define mn as the maximum of the absolute values S j.mm/// //apa//DD (D,,,,,G.) follows from Step 2 that 1/m,, A co. Moreover, by arguing in the same way as in Step 3 in Appendix A we receive that
Ha/ Hm/(m, D2 G.)l +
(29)
as n + 0c. By abuse of notations, let us denote
na/ / D2(Gn, G.)] > Tjm.m
Ty (m D2 G.)] - Kjn
Here, least one among Tym K5n non-zero. Then, by putting the results in equations 27 and 29 together, we get
K 2F(
f (Y((a)) T X + ++,
Σ
X exp((3), X)
Tam
Oh
10
+2 0
K H+gge
Kig X2 cxp((3))) X)9G. (Y|X) = 0.
Arguing in a similar fashion as in Step 3 of Appendix A we obtain that Jjm = Kjjy = for any j E [K]. 0 < Imll + 72 2F((C, and 0 < hi < + 121, This contradicts the fact that at least one among them non-zero. Hence, the proof is completed.
B.3 PROOF OF PROPOSITION 1
Since the Hellinger distance is lower bounded by the Total Variation distance, i.c. 2 V, it sufficient to show that
inf Ex (xIIV 0. GEOX(N)
25

---

Published as a conference paper at ICLR 2024
For that purpose, we first demonstrate that
lim inf 2 (E) Ex (x (..x)) > 0.
(30)
Assume by contrary that the above claim is not true, then can find sequence Gn Σ exp(X% (1 02 o0)) E (x())) that satisfies D2(Gn, G.) - 0 and
Ex Ex[V (-1x), ga. (-1x)) - 0 when n tends to infinity. By applying the Fatou's lemma, have (1 (11),, (-
lim Ex (x n-to0
1 lim inf Ig., [X] - 9c. (Xxx))(XXX))
(31)
2
xxy
The above results indicates that Tg. (Y - 9c. (X - 0 as n co for almost surely (X, Y). WLOG, may assume that
K
Σ
C1,
|cil + |C2l + ....+ Ccdl
(L
Let us consider X XT. where E [q] such that {li,.... .x] (I...., K}. Since D2 D2 G.) converges to zero, it follows that (S% a² o.!) + (5i ,. aj. b5,b),,)) and Eiec, exp(Bb) exp(36,) for any i E C and E [k.]. Thus, we must have that X E N2 for some [q] such that (Z,...., (5 Ci U UCK. Otherwise, Ig. (Y 9G. (Y x) A O, which is contradiction. However, as K < Σソ أー [cyl. the fact that (E,...., { = C1 U.... U Cg cannot occur. Therefore, we reach the claim in equation (30). Consequently, there exists some positive constant e' such that
inf GEO, (G)
Ex [ (+x))) >
Given the above result, it suffices to point out that inf (G) Ex Ex(V(Gc( (-1x)) >
(32)
We continue to use the proof by contradiction method here. In particular, assume that the inequal- ity (32) does not hold, then there exists a sequence of mixing measures G², E Ox(S2) such that D2 Gx) >
and
Ex (V (1IX), (-xx)) + 0.
By invoking the Fatou's lemma in equation
we get that TG: (YIX) 9G. (YIX) + 0 -> for almost surely (X, Y). Since Ω is compact set,
substitute (G..) with its subsequence which converges to some mixing measure G" E Ox Then, the previous limit implies that Te (Y] = gG. (X) for almost surely (X,Y), From the result of Proposition 2 in Appendix C we know that the top-K sparse softmax gating Gaussian MoE is identifiable. Therefore, we obtain that G' = G., or equivalently, Dz G.) =0
On the other hand, due to the hypothesis D2(Gr,, G+) > e' for any E N, we also get that D2 (G', G.) > g > 0, which contradicts the previous result. Hence we reach the claim equa- tion 32 and totally completes the proof.
B.4 PROOF OF LEMMA 2
Let 7, MYc, where is some fixed positive constant and M4, will be chosen later. As X and are bounded sets, find some constant ct such that
min FJJ
( (Bi,)) dje,
where the above minimum subject to x E Xr E {lr..., ext and j' E (lK+1.... 4. By arguing similarly the proof of Lemma T Appendix A.3 deduce that ot 0.
26

---

Published as a conference paper at ICLR 2024
Since X' is a bounded set, may assume that ||a||| I B for any X. Let X2 and ? E l such that (f1,-.., { Ct2 U Uii-- Then, for any i {l,,..., (x) and i' (RRDD..., (x). we have that
97.20 = (31, - --īj) x
xx+ (Biz)'
2 M,eB + (Biy) ''x cge
= M,eB + oge + (giz - - + Bl1
2 -2 + +.. + Pipe
where j E (.).....k) and j' E ((K+1...., (k.) such that i E C, and i' E Cgr. If Mj S %
then
2B
we get that x E Xx, which leads X; -..
Analogously, assume that there exists some constant ce 2 0 such that
(Bt)) (8 = ciE,
min JJ.
where the minimum is subject to Zz. i E (li,..., (x) and i' (1R-11--.;44 .,).. Then, if M, S Ce then we receive that X₂ Xt. 2B'
As consequence, by setting M, 2B
min(er, ce), we achieve the conclusion that X Xi.
C IDENTIFIABILITY OF THE TOP-K SPARSE SOFTMAX GATING GAUSSIAN MIXTURE OF EXPERTS
In this appendix, we study the identifiability of the top-K sparse softmax gating Gaussian MoE, which plays an essential role in ensuring the convergence of the MLE G, to the true mixing measure G, under Voronoi loss functions.
Proposition 2 (Identifiabiliity. Let G and G' be two arbitrary mixing measures in O₂ Suppose that the equation (c(Y|X)= = (Y holds for almost surely (X,Y) E X Y, then it follows that G = G¹.
Proof of Proposition First, we assume that two mixing measures G and G' take the follow- ing forms: G Ex= exp(3ur)0 (pnaa, b..o.) and G' Σk= (x Recall that for almost surely (X, (X,Y), then we have
ga (Y x)) = g0 (Y
(oft T K; Boo)) s(Y(a X + bz)
Σ Softmax(TopK(({)) xx - + 6,, of)
K; (())
(33)
Due to the identifiability of the location-scale Gaussian mixtures Teicher (1960 1961 1963). we get
that & = k and
(Softmax(To Xx; { x, K; (.)): :16 [A]]
[M]]
(oottmax(Top(((()))
for almost surely X. WLOG, we may assume that
Sottmax(Topp((355)) T X, K; Box) = Softmax(TopK((B)))
8c))),
(34)
for almost surely X for any E [k]] Since the Softmax function invariant to translations, it follows from equation (34) that B1 B1 + U1 and for 86 Uo for
i1 Rd and up R. Notably, from the assumption of the model, we have Bi = 31x = O₂ and 3ox = 3y = O, which implies that Vi 0g and to 0. As result, obtain that B1s 8/. and Bo: 8o. for any E [k].
27

---

Published as a conference paper at ICLR 2024
Let consider X Xe where & [g] such that fei,..., {l...., K}. Then, equation 33 can be rewritten as
K
exp(Ba)exx((
bz) = exp(Soi) exp(81 x4,,00))
(35)
for almost surely (X,Y), Next, we denote J1, ........JJ Jm as a partition of the index set [lt]. where m I k, such that exp()0:: exp(Snc) for any i, E Jj and j (m]. On the other hand, when i and i' do not belong to the same set Jj. we let exp(3oi) o exp(Bor). Thus, we can reformulate equation (35) as
ΣΣ Σ j=
exp(3x))xxp((11)
xxp()) bz)
m exp (Bs)) exp exp(()x))((((((c)) x b/,0())
iEJ,
for almost surely (X, Y). This results leads to {((a.) X + bb, o,) :EE J,} (((a)) T X + +b,0)) i Jz), for almost surely X for any j E [m]. Therefore, we have
{(a;, oz) : i J,} = ((4, (1 : J/}.
for any j E [m]. As consequence,
G G= -=ΣY Σ
exp(Box)((11.
exp(P6.)0 = G¹.
3=1 SEJ,
j=1
Hence, we reach the conclusion of this proposition.
D NUMERICAL EXPERIMENTS
In this appendix, we conduct few numerical experiments illustrate the theoretical convergence rates of the MLE G,, to the true mixing measure G., under both the exact-specified and the over- specified settings.
D.1 EXPERIMENTAL SETUP
Synthetic Data. First, we assume that the true mixing measure G, = Σ exp( (0 of order k. = and associated with the following ground-truth parameters:
661 -8. 6G O,
81 25, 3 = 0,
oi -20,
bj 15,
σ = 0.3,
0 20,
6 = -5,
0 = 0.4.
Then, we generate i.i.d samples ((X,, Yo) by first sampling X,'s from the uniform distribution Uniform(0. 1] and then sampling Y ss's from the true conditional density 9G, (Y)X) of top-K K sparse softmax gating Gaussian mixture of experts (MoE) given equation (I). In Figure 2 we visualize the relationship between X and Y when the numbers of experts chosen from gG. (YTX) are K = 1 (Figure 2a and K = (Figure 2b). respectively. However, throughout the following experiments, will consider only the scenario when K 1, that is, we choose the best expert from the true conditional density 9G. (X)...
Maximum Likelihood Estimation (MLE). A popular approach to determining the MLE G, for each set of samples is to use the EM algorithm Dempster et all (1977). However, since there are not any closed-fomm expressions for updating the gating parameters 3o., B1 in the maximization steps, we have to leverage an EM- based numerical scheme, which was previously used in Chamroukhi et al. (2009). In particular, utilize simple coordinate gradient descent algorithm in the maximization steps. Additionally, we select the convergence criterion = 10 and run maximum of 2000 EM iterations.
28

---

Published as a conference paper at ICLR 2024
15
10.0
10
5
2.5
0.0
0.
-2.
-5.0
-7.5
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
(a)
K=1
(b) ==
Figure 2: A visual representation showcasing the relationship between X and Y. along with their respective marginal distributions when K and K
Initialization. For each [k., k. + 1}, randomly distribute elements of the set (1,2, k} into k, different Voronoi cells C,, each contains at least one element. Moreover, we repeat this process for each replication. Subsequently, for each j E [k.]. we initialize parameters 31 by sampling from a Gaussian distribution centered around its true counterpart 31, with small variance, where Cy. Other parameters Box, a,, b,, oi are also initialized in similar fashion.
D.2 EXACT-SPECIFEED SETTINGS
Under the exact-specified settings, we conduct 40 sample generations for each configuration, across spectrum of 200 different sample sizes n ranging from 10 to 10Sh It can be seen from Figure 3a that the MLE G, empirically converges the true mixing measure G. under the Voronoi metric D, the rate order O(n -1/2), which perfectly matches the theoretical parametric convergence rate
established in Theorem 2
D.3 OVER-SPECIFIE SETTINGS
Under the over-specified settings, we continue to generate 40 samples of size for each setting,
given 100 different choices of sample size n [102, 10 As discussed in Section 3 to guarantee the convergence of density estimation to the true density, we need to select know, existing works, namely Kwon et al. 2209 Kwon &
2 experts from
the density estimation. As far Caramanis (2020): Kwon et al. (2021), only focus on the global convergence of the EM algorithm for parameter estimation under the input-free gating MoE, while that under the top-K sparse softmax gating MoE has remained poorly understood. Additionally, it is worth noting that the sample size must be sufficiently large that the empirical convergence of the MLE returned by the EM
algorithm aligns with the theoretical rate of order Onn-1/2) derived in Theorem
E ADDITIONAL RESULTS
In this appendix, we study the convergence rates of parameter estimation under the model I when a probability density function an arbitrary location-scale distribution. For that purpose, we first characterize the family of probability density functions of location-scale distributions
F: (f( a,),), (a,b,o) E e).
(36)
29

---

Published as a conference paper at ICLR 2024
3.6n 0.54588
6.2n -0.3253
D1(G,, G.)
6-
D2(Gn. G.)
log(loss)
log(loss)
-1
2
2
0
-3
-2
6
10
10.0
10.5
11.0
11.5
log(sample size)
log(sample size)
(a) Exact-specified settings with K =
(b) Over-specified settings with K = 1 and K = 2
Figure 3: Log-log scaled plots illustrating simulation results under the exact-specified and the over- specified settings. We analyze the MLE Gn across 40 independent samples, spanning sample sizes from 102 to 10 The blue curves depict the mean discrepancy between the MLE Gn and the true mixing measure G,, accompanied by error bars signifying two empirical standard deviations under the exact-specified settings. Additionally, an orange dash-dotted line represents the least-squares fitted linear regression line for these data points.
where h1(X, a) == a X + b stands for the location, denotes the scale and θ is compact subset of d R x Rt, based on the following notion of strong identifiability, which was previously studied in Manole & Ho 2022) and Ho & Nguyen (2016):
Definition (Strong Identifiability). We say that the family F is strongly identifiable if the probability density function f (X, a, ()) is twice differentiable w.r.t its parameters and the following assumption holds true:
For any k > and k pairwise different tuples (a1, (i),,,,,.. (ax, bu, σn) E O, if there exist real coefficients a 82 for E [Ks]] and 0 S &1 + &g S 2, such that
(8)
ghi+ea f (Y (X, (x), 0,
(
in
an 0
E1 E1=8a==0
for almost surely (X, Y), then we obtain that a 12 0 for any E [k:.] and 0 S S + + 2
Example 1. The families of Student's 's t-distributions and Laplace distributions are strongly identifi- able, while the family of location-scale Gaussian distributions is not.
In high level, we need to establish the Total Variation lower bound Ex [X[V[ 9G. (+XX)) Z
D(G, G.) for any G E (2))). Then, this bound together with the density estimation rate in Theorem I (resp. Theorem 3) leads to the parameter estimation rates in Theorem (resp. Theorem 4
Here, the key step to decompose the difference 92. (YlX) - 9g. (Y|X) into a combination of linearly independent terms using Taylor expansions. Therefore, we have to involve the above notion of strong identifiability, and separate our convergence analysis based on that notion.
Subsequently, since the convergence rates of maximum likelihood estimation when F a family of location-scale Gaussian distributions, which is not strongly identifiable, have already been studied in Section 2 and Section 3 we will focus on the scenario when the family F. is strongly identifiable in the sequel. Under that assumption, the density f( o) is twice differentiable in (f1, o), therefore, is also Lipschitz continuous. As a consequence, the density estimation rates under both the speci-spe and over-specified in Theorem I and Theorem B still hold true. Therefore, we aim to establish the parameter estimation rates under those settings in Appendix E. and Appendix E.2 E.2 respectively.
30

---

Published as a conference paper at ICLR 2024
E. EXACT-SPECIRIED SETTINGS
In this appendix, by using the Voronoi loss function D (G, G.) defined in equation 5 demon- strate in the following theorem that the rates for estimating true parameters exp Pia ot 6 of parametric order (n- which totally match those in Theorem 2
Theorem 5. Under the exact-specified settings, if the family F is strongly identifiable, then the
Hellinger lower bound Ex[h(ga(: (X), 9G. [-[X)]] Z Di G.) holds for any mixing measure G E &k. (It.. Consequently, we can find universal constant C3 >> depending only on G., and K such that
P P(Di G.) > C3 Csyllog(n//)) S n,,
where > 0 is universal constant that depends only on Ω
Proof of Theorem is in Appendix E.3.1 E.3.1
E.2 OVER-SPECIFEED SETTINGS
In this appendix, we capture the convergence rates of parameter estimation under the over-specified settings when the family F is strongly identifiable
Voronoi metric. It is worth noting that when F is strongly identifiable, the interaction among expert parameters the second PDE (6) no longer holds true. As a result, it is not necessary to involve solvability of the system 77 in the formulation of the Voronoi loss as in equation (8) Instead, let us consider the Voronoi loss function D3 G.) defined as
(2 G.) =
max (
Σ
Σ
exp( | Il + |a ll + Abbe, + Aooze I]
EE[K], iEC2 [0el==1
+ Σ
Σ
exp(6ni) ||B 112 11A 12+ 1/bbbe, 12 ++ Ao PP]
jE(K(, kECr, Ci.y |>|
exp(Boi) - exp( ) (37)
for any mixing measure G E ( Equipped with loss function, we are ready to illustrate the convergence behavior of maximum likelihood estimation in the following theorem:
Theorem 6. Under the over-specified settings, if the family F is strongly identifiable, then the Hellinger lower bound Ex (x(g ga. (-1X)] Z D2(G, holds for any mixing measure G E (I)), As consequence, we can find a universal constant C4 > O depending only on G., Ω and K such that
P (DDs(Gn,G.) G.) Ca Coylog(n)/n) S nn
where e₁ > 0 is universal constant that depends only Ω.
Proof of Theorem is in Appendix E.3.2 Theorem 6 indicates that true parameters 3 aj b7 of which are fitted by a single component, share the same estimation rates of order O(n-1/2) as those Theorem By contrast, the estimation rates for true parameters Si at bū of which are fitted by more than one component, are of order O(n-1/4). Notably, these rates are no longer determined by the solvability of the system 77.. Thus, they are significantly faster than those in Theorem 4 The main reason for this improvement when F is strongly identifiable, the interaction among expert parameters via the second PDE in equation (6) does not occur.
31

---

Published as a conference paper at ICLR 2024
E.3 PROOFS OF ADDITIONAL RESULTS
E.3.1 PROOF OF THEOREM
Following from the result of Theorem and since the Hellinger distance is lower bounded by the Total Variation distance, i.e. h 2 V, it is sufficient to demonstrate for any mixing measure G E (x(O) that
Ex [ (g) 2 Di G.),
which is then respectively broken into local part and global part follows:
Ex (va((1X),ga... (-X)) 0, Di(G,G,)
inf GEE, (GG.
(38)
Ex (9a( (- (X), gG. g.. (-XX)]
inf GEE,, (BP...D))))
0,
(39)
Di Di G.)
for some constant e' > O. Subsequently, we will prove only the local part (38 while the proof of the global part (39) can be done similarly to that in Appendix A.2
Proof of claim
38 It is sufficient to show that
Ex [V (1) (1x)) > 0.
lim - GEE. (G.G))
inf
GG,
Assume that this inequality does not hold, then since the number of experts ki is known in this case, there exists sequence of mixing measure G, = Σh. ((p() az ag..go & (12 such that both (G,, G.) and Ex (Vig,, (-IX), ())))) G.) approach zero as n tends to infinity. Since D2 (G Gx) - 0 as - oo, each Voronoi cell contains only element. Thus, we may assume WLOG that C2 = {j} for any j E [k,]. which implics that (3%) aj, 67,0)) +
(Bi az. 0j) and exp(3i)) + exp($6)) as n + co. WLOG, we assume that
K
Di (Gn; G.) = Σ eep((()))
(I4 + || + || + +1A00 + eex((6))) -xx((xx
where we denote A88 = PT Bi. Dag = a. n.. A = by by and Dot = of of
Subsequently, by arguing in the same fashion as in Appendix A.2 we obtain that X = X7 where
X²" {o x x: (Br)) (Bih))
T x Vj E
(lx+1........ fa.1)
X2 = {s E x: (8i,) x 2 (Bi T x Vj E (e.....,, lxt.J' E (lxti.... 1mm1)
for any E E [g] for sufficiently large n.
Let E [4] such that {fi,..., {x} {1,... K). Then, for almost surely (X,Y) E X2 x V. we can rewrite the conditional densities fG., X) and 9G. (YJX) as
exp(( X + (
9G. (Y|X)
f(YY(a))TT)) X+ ob")
i=1 1=1 exp( X (
exp(($1)) X + 3
9g. (YIX)
((Y Tx X x+
(xi((()))) TX + Bbッg
Now, we break the rest of our arguments into three steps:
Step 1 Taylor expansion:
32

---

Published as a conference paper at ICLR 2024
In this step, we take into account H, := Σx exp((5i.) Tx + (bi) (Y|x) 9G. (X]
Note that the quantity Hn, can be represented follows:
Hn
exp()()))
(x)(_))))))(((((xxxxxxx)))>) of) +f exp(((1)) ())(((()))))) + o7)
K
exp()
(X)gG. exp(((9))) XX).,. xx) (x) (X
K
exp(Phi) - exp((%t) exx(((8))) Tx " x x exp(((()) X)9a, (1x))
1m1
By applying the first-order Taylor expansion to the first term in the above representation, which is denoted by An, we get that
K
exp(Pii)
An
(A31)
(Aoo))
a!
aaz tastaif
x Xartas exp((31.) X)
(Y)(a)) T + bb, + Rr(X,
oh poaltts dg
where (X, Taylor remainder that satisfies Rí (X, Y) (X, - 0 as n + xc. Let
71 01 a2 No. 1p2
|azl ag N and 173 04 N, then rewrite An follows:
no
exp(3%.) a!
A, Σ
Σ
(A (Da
===
lmn +1p=1 1/s
th exp((3ti) X).
ah!' 00 (Y) T X + (2, + R1 (40)
where we define
Imimsoos ((a)) N² x No N a1 + 02 m/, |02| + 03 72, a4 13)
(41)
By arguing in similar fashion for the second term in the representation of Hn, we also get that
exp( (APR) X exp((?))] X)9G. x) R2(X,Y),
B, =
where (2(X, is a Taylor remainder such that R2 (S(XYY)DD(G,, G.) - 0 as + co. Putting the above results together, we rewrite the quantity Hn as follows:
K
2- Bl
g f
Hn =
U ma.ns
X exp((87)) X)
The 00 (Y((a)) x x
1=1
Vi =
Im
+
K
W
X exp(((1))) X)aa. + (X,Y) +
(42)
0shisi
which we respectively define for each E [K] that
exp(
U 8..1..2.7
(Ao,)aa,
a!!
aE
exp(
W 1,
(A.R)
?
for any (n, 72, na) # (0,, 0) and |y| * 0g. Additionally, U2 - W... cxp(8%)) exp((6))
33

---

Published as a conference paper at ICLR 2024
Step 2 Non-vanishing coefficients:
Moving to the second step, we will show that not all the ratios U ma /D1 Ge) tend to zero n goes to infinity, Assume by contrary that all of them approach zero when oc, then for (m 122 175) = (Oa. 0, 0), it follows that
K
K
U
exp((6i) - exp(
Σ
M2.97
+ 0.
(43)
D (G.: G.)
Di (Gn: G.)
==1
==1
Additionally, for tuples (71 72, 78) where E e2.,....... edt with ej = (0,... O.
0, .,....0)
j-th
and 7/2 = 178 O, we get
K
K
U pmm.ool
exp(A))AARM
Σ
Σ
- 0.
(i G.)
D DG G.)
m eet
By using similar arguments, end up with
Σ exp(( || + || || + +ل6, | O.
Di G.)
Due to the topological equivalence between norm-l I and norm-2, the above limit implies that
Σ ( || + || + + + AAoI]
- 0.
(44)
Di G.)
Combine equation 43) with equation (44), deduce that Dj (Gns G.)/Di(L (Gn G.) O, which is
a contradiction. Consequently, at least one among the ratios UR Up mamo D (Gn, G.) does not vanish tends infinity.
Step 3 Fatou's 's contradiction:
Let us denote by my the maximum of the absolute values of UP mnm / (Gns G.) and WT,/Dr (G,, G.). It follows from the result achieved in Step 2 that 1 1/m, + xo.
Recall from the hypothesis that Ex (V(G,, (-1X), 9G. (-IX)) Di (Gn, G.) 0 n Thus, by the Fatou's 's lemma, we have
Ex [V (gG,, (-1X), 9G. (+1x))]
lim 19G., (1. (Y XX) n0000
lim 22000
dxaY.
Di G.) This result indicates that 19G, (Y|X) 9a. (Y|X))/D, (G,, G.) tends zero as n goes to infinity for almost surely (X, Y). As result, it follows that
Di G.)
lim n-000 Mi D(G G-)
Hn
Σ
T x (%:) lim 99a, (Y|x) 99. (Y 0. (G,,, Gz)
U .ra.n/ [m,, (G,, G.) + Tem mam and /T7, D (G,,, G.)] Kig with a note that at least one among them is non-zero. From the formulation of Hn in equation (42) deduce that
Next, let us denote Um
K
f dh dons (Y)(a)) x + +,,0,
main t exp((81t) T X)
1=1 Ma=
Ii +1==0
K
Xx)o. Kin X exp((3i) Txxx(( (Y (Y = 0,
1=1 1ihhlli
for almost surely (X,Y). This equation equivalent to 2 mm
Σ
Ta exp((3i.) X)
Oh Th Dars (Y X + b7,00')
mm 1
Up 7/8
(x))) kim exp((97)) X)9a. ()1x)(111)) x x = =0,
34

---

Published as a conference paper at ICLR 2024
for almost surely (), It is clear that the left hand side of the above equation is polynomial of X belonging to the compact set X. As a result, we get that
k
m
92:9 exp((31.) X) gm ff on dons
Σ
Σ
(Y((a)) X + bb,
1=1 1/2 +
+ Ki.m 9G. (Y exp( x) 0,
for any E [K]. 0 < 1 I and almost surely (X,Y) Since (aj, bj. oi)........ (0xx bx: ok) have pair-wise distinct values and the family F is strongly identifiable. the set
a f Thi Th da
(YI(a)) X + 0f) E [K]. 0 S 72 + 2 22mml
linearly independent w.r.t (X, Y). Consequently, we obtain that Tum .2. = K4.m for any i E [K]. 0 |m| S and S 72 + 7/s 2 |/1,, which contradicts the fact that at least one among these terms is different from zero.
Hence, we reach the desired conclusion.
E.3.2 PROOF OF THEOREM
Similar to the proof of Theorem 5 in Appendix|E.. E.3.1 our objective here is also to derive the local part of the following Total Variation lower bound:
Ex[V(Tc sa.
(.....(11))) (+x)) Ds(G,G.),
for any G E (x) In particular, we aim show that
Ex ( (+Xx)) > 0. D2(G.G.)
lim inf - (.....)
(45)
Assume that the above claim does not hold true, then can find sequence of mixing measures Gn == Σ (xp ja (p))) E Ox(S2) such that both D3 G.) and (V (-1X). gG. (1x))) ()(])////GGG//) G.) vanish when n goes infinity. Then, follows that for any j E [k.], we have Liec, exp(5%) + exp(3),) and (3 a,, Dy. of)) - (5 o7 0j. 0z) for all i Cj. WLOG, may assume that
Da G.) = Σ Σ
|IAS% || - | + | || exp(
JELK iEC 1c,l>1
K
Σ
exp(B(r)
||S%; + || || + AbG + |
exp(33) -ex(((6))
JELK iEC, 1c,1=1
IEC
Subsequently, let X E X2 for some E E [g] such that (l.....,(k) [kt = {l,..., Kj. where
X = f= X (Bi,) 2 (Big)) Tx, Vj e (eI....,xx).) E (lK+I......4. 1n.3)
Then, for any 7. E iq), we denote (?1,..,, [2k) as a permutation of (1,...,k) and
X7 = for X: (0))) (8re) xx Vi (Z....., Ta....,,TT,,r' I' E (Ix41....,,xx
If (x. U.C...UUC for any E [7].
then V (IX), (4x)) G.) A 0 as
n tends to infinity. This contradicts the fact that this term must approach zero Therefore, we only need to consider the scenario when there exists E [q] such that {?.... [ = C, U... UC By using the same arguments as in Appendix B.2 we obtain that X E XP.
35

---

Published as a conference paper at ICLR 2024
Then, can represent the conditional densities gc. (Y]X) and gG., (Y]X) for any sufficiently large n as follows:
9G. (G.() = exp((
K
T + (65)
5(Y)((5)) x + [ bbb
j=1 =1=1 exp (31g)) T + Pog')
K
exp( X + 501)
9G. (MX) = Σ Ey'=
[(Y((a,)) xx + o.))
expl (}1p))
X + (pP))
j=l;EC;
Now, we reuse the three-step framework in Appendix E.3.1
Step 1 Taylor expansion:
Firstly, by abuse of notations, let us consider the quantity
K
Hn exp(((1,) + xb))] Ia. (YIX)
IG. (YIx)
Similar to Step in Appendix A we can express this term
Hn =
exp(
-
j=1 3EC
K
exp(
iEC;
K +
- exp( -
-xxp - (T)((())))glg,.. X)ao,. xxxc((())xxx)x
j=1
sEC,
:= An + B, + E,.
Next, we proceed to decompose A, based on the cardinality of the Voronoi cells as follows:
A = Σ Σ
exp(>?))
(xxp((()Tx)))(((X) ())((((())) X + """""""")) exp((P1)) (x)(((())))) x + oi))
Σ
Σ exp(
"X)(((((a)) x + 87 o, exp(( (a)(( + by o;
By applying the Taylor expansions of first and second orders to the first and second terms of Ans respectively, and following the derivation in equation (40) we arrive at
ll Σ
A, = Σ
exp(
(A )0 (Aoo))
(Do
ad!
3:C,|=1 1 kEC, Da 0 mm +
a
to
x X cxp((3i,) X)
00 (Y Tx + (j,03} + Rx(X,Y)
33Σ
1-18 Σ
exp(
(A )m (Aa )m2 (Ab oo (Aoo)) oM
a!
#C//11 iEC, ===0
aEt m2 92:73
are tns f
X exp((3i,) x)
(a( X 13, + (X,,) Th t2 ao
where the set I'm va.ua is defined in equation 41 while (, Y) is Taylor remainder such that
RX, /)/Dx(G, G.) + 0 as n + Do for é E [3.4 4). Similarly, we also decompose B,, as
exp( (ART) X exp((81,)) X)gG. (Y|X) + Rs(X, Y)
Bn Σ
?!
S4C, m1 VEC, 1=|=1
exp(3n) (A X exp(((B)))) ( (1IX) RR((X
Σ ΣΣΣ
?!
HCCsSS1EC, 15h1s2
36

---

Published as a conference paper at ICLR 2024
where Rs (X, Y) and R6(X, Y) are Taylor remainders such that their ratios over D3 P3(Gns G.) ap- proach zero as n - oc. Subsequently. let us define
exp(
S n2 =
(A 01
(A
ad!
sEC, aEE
Ty = exp(
ABt)))
--S j.7,0,01
?!
iEC
for any (m, 72, n7) + (Od, O, 0) and |1| + Od, while for (71, 73)) = (0a, 0) we set
S = -T,o0 = Σ
exp(P%i) cxp()%%)-
As consequence, it follows that
141
2
Hn Σ
Σ
Σ
S 3.m.n2mn
X exp((31,)) ())
(Y]( X x)b
The
j=l
===
mml+
K
Σ -D
Th X eTexp((31)) (99,,((1)) + Rs(X,Y) + Rs Rx(X,
(46)
Step 2 Non-vanishing coefficients:
In this step, will prove by contradiction that at least one among the ratios ST. m2. DDa G.) does not converge to zero as - oc. Assume that all these terms go to zero, then by employing arguments for deriving equations 443 and (44). we get that
K
1 Dg G.)
exp( (x0) exp(
EC,
Σ exp( (1 || + + + +
+ O.
-, IEC
Next, let ej = (O..... 0.
O,....,0 0) for any j E [d]. Then, we have
j-th
K
K Σ
U 0.ol
exp(
Σ
+
Da G.)
Dz Dz(Gnn G.)
i=1
==1 m 2ea)
Similarly, we also get that
K
K
cXp(
+
exp( 2 +
Da G.)
Dz(Grn Gx)
Moreover, note that
K
U
Jmm 2,
Σ
exp(
= Σ Σ
0.
Da G.)
Dg G.)
mE (201..., 20a)
Gathering all the above limits, we obtain that 1 Dz(Gn,, Gz)/Da G-)
0 as n oo, which
is a contradiction. Thus, at least one among the terms Sy j.m....m mp. (Da Gx) does not converge to zero as n >> 00
Step 3 Fatou's contradiction:
follows from the hypothesis that Ex (V (1XX),o, (L.. - as + oc. Then, by applying the Fatou's 's lemma, we get
Ex Ex (- (-1)) (-1x))
Gc. (YIX) 9G. (Y
0 lim
lim inf n-400
Dz G.) dxdY,
Dz(Gnn G.)
37

---

Published as a conference paper at ICLR 2024
which implies that IG.. 9G. - (Y.x)))/))))) G.) - Os for almost surely (X, Y). Next, we define m, as the maximum of the absolute values S J.m.ma/Da(G... follows from Step 2 that 1/m, A co. Moreover, by arguing in the same way as in Step 3 in Appendix E.1 we receive that
Tm/ G.)] +
(47)
as 00. By abuse of notations, let us denote
Smmmm Dm
D3 (n, G+]] + N2.
Here, at least one among mam is non-zero. Then, by putting the results in equations 46 and together, get
111 201 121))
Σ
X exp(((i,)
Tx)
X).
(Y)(a)) X + 03.
ahe 00%
1m
1+ 1+ >1)
3x.00 X exp((31g) X)9c. (Y
In/=0
for almost surely Y), Arguing in similar fashion in Step of Appendix E.1
we obtain that [1e,)> >1)-
T2-9 0 for any j E [K]. 0 S Imil + + S 2(1 [ [1c,
and 0 < 7
This contradicts the fact that at least one among them is non-zero. Hence, the proof completed.
38


Reviews:
[{"review_id": "gKfSIAwBML", "reviewer": "Reviewer_qHcA", "summary": "The paper provides theoretical analysis on the convergence properties of sparse softmax gating mixtures of Gaussian experts. The key contributions are proving convergence rates in two cases - when the number of experts is exactly-specified, and when there is an overspecified number of experts. \n\nFor the exact-specified case, the authors show a convergence rate of $O(n^{-1/2})$ to the true parameters under maximum likelihood estimation. This is an important theoretical result as it quantifies the sample complexity.\n\nFor the overspecified case, the convergence rate depends on the cardinality of the Voronoi cells induced by the gate activations", "questions": "- How do these convergence results compare to prior analysis on softmax vs sparse gating? Does this theory suggest one gating approach over the other?\n- Due to the instability of the EM algorithm, there is no error bar provided in Figure 3b. Can we still consider this experiment to be reliable to to justify the theoretical results?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The theoretical analysis of convergence rates for MoE models is novel and useful. Understanding sample complexity of different MoE architectures is valuable.\n- The paper is technically strong, with detailed proofs of the main results. The analysis for the overspecified case considering Voronoi cells is creative.\n- The assumptions are clearly laid out, making the results easy to interpret and apply.", "weaknesses": "- As noted, the paper is quite dense with heavy notation. More intuition and examples earlier on could make it more accessible.\n- The writing in the universal assumptions section is unclear, and should be revised for readability.\n- The implications of the theory for practitioners could be expanded on more in the discussion. Guidance on model design is lacking.", "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": "vF4PpHqh8g", "reviewer": "Reviewer_X5io", "summary": "This paper analyzes Top-K sparse softmax gating mixture of experts. In particular, they focus on Gaussian mixture of experts and present theoretical results on the effect of MoE on density and parameter estimations.", "questions": "1. What are the practical implications of the work, if any, for the use of Top-K sparse MoE as a means of conditional computation?\n2. How does this work compare to that of [1]?\n\n[1] Towards Understanding Mixture of Experts in Deep Learning https://arxiv.org/pdf/2208.02813.pdf", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Analysis of MoE is a highly relevant and impactful direction of research.\n* The analysis seems thorough and rigorous at a high level and the presented results are interesting.", "weaknesses": "* The paper is quite dense, and it would have been helpful to outline high level ideas before delving into notation-heavy math.\n* Chen et al. [1] also analyze MoE, but in the deep learning setting. How does the analysis presented here differ from this work?\n\n[1] Towards Understanding Mixture of Experts in Deep Learning https://arxiv.org/pdf/2208.02813.pdf", "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": "4H3CnPdmwh", "reviewer": "Reviewer_dwg3", "summary": "Based on the Gaussian mixture expert model, the influence of top-K sparse softmax gating function on density estimation and parameter estimation was analyzed. Novel loss functions and Voronoi metrics were used to characterize the behavior of different regions in the input space.", "questions": "Considering W2, I would like to know the difficulties in generalizing the results of this article to MoE methods without Gaussian assumption.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1. The writing is clear and concise, making it easy to understand theory.\n\nS2. From the author's introduction, this paper is the first to perform convergence analysis for maximum likelihood estimation (MLE) under top-K sparse softmax gated Gaussian MoE. The results of this paper are innovative and helpful in inspiring new expert mixture system designs.", "weaknesses": "W1. The theoretical results of this paper lack experimental verification.\n\nW2. It only considers Gaussian mixture expert models and does not consider other types of models.", "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)