Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Authors: Ryotaro Kawata
Kohsei Matsutani
Yuri Kinoshita
Naoki Nishikawa
Taiji Suzuki
Venue: ICML 2025 poster (2025)
Keywords: Mixture of Experts
feature learning
single-index models
Primary area: theory->learning_theory

Abstract: Mixture of Experts (MoE), an ensemble of specialized models equipped with a router that dynamically distributes each input to appropriate experts, has achieved successful results in the field of machine learning. However, theoretical understanding of this architecture is falling behind due to its inherent complexity. In this paper, we theoretically study the sample and runtime complexity of MoE following the stochastic gradient descent when learning a regression task with an underlying cluster structure of single index models. On the one hand, we show that a vanilla neural network fails in detecting such a latent organization as it can only process the problem as a whole. This is intrinsically related to the concept of *information exponent* which is low for each cluster, but increases when we consider the entire task. On the other hand, with a MoE, we show that it succeeds in dividing the problem into easier subproblems by leveraging the ability of each expert to weakly recover the simpler function corresponding to an individual cluster. To the best of our knowledge, this work is among the first to explore the benefits of the MoE framework by examining its SGD dynamics in the context of nonlinear regression.

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Ryotaro Kawata 12
Kohsei Matsutani
* 12 Yuri Kinoshita Naoki Nishikawa
12 Taiji Suzuki
222
Abstract
gated MoE (Shazeer et al., 2017), which was later integrated into transformer- based large language models (LLMs) and further refined in subsequent works (Fedus et al., 2022; Achiam et al., 2023; Georgiev et al., 2024; Jiang et al., 2024; Liu et al., 2024). This kind of MoE enables the ac- tivation of only limited number of trained experts in one forward pass, drastically reducing the inference cost while maintaining performance competitive with other successful architectures of the same order of parameters.
Mixture of Experts (MoE), an ensemble of spe- cialized models equipped with router that dy- namically distributes each input to appropriate ex- perts, has achieved successful results in the field of machine learning. However, theoretical under- standing of this architecture falling behind due to its inherent complexity. In this paper, we theo- retically study the sample and runtime complexity of MoE following the stochastic gradient descent (SGD) when learning regression task with an un- derlying cluster structure of single index models. On the one hand, we prove that a vanilla neural network fails in detecting such a latent organiza- tion as it can only process the problem as a whole. This is intrinsically related the concept of in- formation exponent which is low for each cluster, but increases when we consider the entire task. On the other hand, we show that a MoE succeeds in dividing this problem into easier subproblems by leveraging the ability of each expert to weakly recover the simpler function corresponding to an individual cluster. To the best of our knowledge, this work is among the first to explore the bene- fits of the MoE framework by examining its SGD dynamics in the context of nonlinear regression.
However, theoretical understanding of this architecture is falling behind due to its inherent complexity. Especially, while the mechanism of the initialization, the optimization procedure and the behavior of the router are essentially the same for each expert, it has been repeatedly reported that each expert ultimately specializes in its own way, each contributing to different aspects of the learned task. It is still unclear why such phenomenon happens and why the router can learn to fairly distribute an input to appropriate experts without collapsing to single expert.
To address these fundamental questions, prior work math- ematically studied the mechanism of MoE from the per- spectives of approximation theory on MoE for multi-level data (Fung & Tseung, 2022). statistical learning in Gaussian MoE models (Ho et al., 2022: Nguyen et al., 2023; 2024a) and nonlinear regression (Nguyen et al., 2024b;c), as well as optimization in both classification (Chen et al., 2022; Chowdhury al., 2023) and linear regression, especially for continual learning (Li et al., 2024). However, clear explanation of the success of the MoE is lacking in the con- text of optimization in nonlinear regression which is more general problem than optimization in classification.
1. Introduction
Mixture of Experts (MoE) (Jacobs al., 1991; Jordan & Jacobs, 1993), an ensemble of specialized models equipped with router that dynamically distributes each input ap- propriate experts, has been extensively studied and success- fully deployed in a wide range of scenarios over the past few years. A key milestone was the development of sparsely-
Therefore, in this paper, we focus on such a broader problem setting of optimization in nonlinear regression. We will theoretically study the sample and runtime complexity of MoB optimized with the stochastic gradient descent (SGD) when learning a regression task with underlying cluster structure.
Eqqual contribution 'The University of Tokyo, Japan "Center for Advanced Intelligence Project, RIKEN, Japan. Correspondence to: Ryotaro Kawata "kawatt-ryotaro7255 g.ccc.u acc.u-tokyo.ac. jc. Kohsei Matsutani mhatsutani-kohsell 40@ g.ecc.u-toke woc.u-tokoo.ac.... ap>>>
Contributions Our contributions are summarized as fol- lows. On the one hand, we prove that a vanilla neural network fails in detecting such a latent organization as can only process the problem as whole. This is intrinsically
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
related to the concept of information exponent which is low for each cluster, but increases when consider the entire task. On the other band, we show that a MoE succeeds in dividing this problem into easier subproblems by leverag- ing the ability of each expert to weakly recover the simpler function corresponding to an individual cluster. To the best of our knowledge, this work is among the first to explore the benefits of the MoE framework by examining its SGD dynamics in the context of nonlinear regression.
complexity of n (1d (Arous et al., 2021), where d is the input dimension. Damian et al. (2023) improved this rate to @(d'=) ) by smoothing the landscape. Subse- quently, techniques such as reusing batches (Dandi et al., 2024b; Lee et al., 2024; Amaboldi et al., 2024) or alter- ing loss function (Joshi et al., 2024) enabled to surpass the CSQ lower bound (Damian et al., 2022; Abbe et al., 2023). These approaches improve the sample complexity near information-theoreti limit n x d, which is associated with the generative exponent (Damian et al., 2024). This approach based on the information exponent contributes to deepening our general understanding about the complexity of task and has been applied to specific architectures or techniques, such pruning (Vural & Erdogdu, 2024), pre- trained transformer (Oko et al., 2024b), adversarially robust learning MlousvviHHsseein et al., 2024) and LoRA (Dayi & Chen, 2024). However, the application of this framework to MoE has not been explored yet, and it may hold promise for elucidating the intricate mechanism of MoE. Recently, Oko et al. (2024a) has conducted an extensive theoretical study additive models, where several single-index models form ridge combination. Our setting is analogous to this work, where the data exhibit an additive structure derived from diverse clusters.
Notation P[|] and [x]] denote the probability of an event and the expectation over the randomness of a random vari- able x. O() and o-)) stand for the big-O and little-o nota- tions with respect to d. ((()) and (+) represent the lower and tight bounds. Ö(-,, S(-)) and (()) denotes the upper, lower, and tight bound ignoring any poly-logarithmic con- stant. We call a probabilistic event A happens with high probability (or w.h.p.) if [AA] N 1- d=a with a sufficiently large constant a > 0; the high probability events are closed under union bounds over sets of size polyd.
2. Related Works
Theory of Mixture of Experts. Various aspects of MoE have been theoretically studied in the context of deep learn- ing so far. Ho et al. (2022) and Nguyen et al. (2023; 2024a) studied the convergence rate of expert estimation in Gaussian MoE models for classification, and Nguyen al. (2024b;e) led similar investigation for MoE with softmax gating for regression problems. Chen et al. (2022) pio- neered studies of feature learning with MoE and analyzed the training of nonlinear MoE under a mixture of classifica- tion problems. Building on this, Chowdhury et al. (2023) extended the analysis to patch-level routing, addressing bi- nary classification problems within nonlinear MoE settings. Li et al. (2024) focused continual learning scenarios, but the analysis was limited to linear regression problems and linear MoE. In this work, we consider the broader and more practical problem setting of nonlinear regression problem with a nonlinear MoE model following a gradient-bart optimization.
3. Problem Setting and Preliminaries
In this section, we clarify the problem setting, including the data generation procedure, the formulation of the MoE, and the mathematical description of the training algorithm.
3.1. Data Generation
Let us first formally introduce the notion of information exponent.
Definition 3.1 (Information Exponent). Let {He,} be the normalized Hermite polynomials. The Hermite expan- sion of a gquare-ittegrable function f is given as f(z)= Σg vJ. Hey(z)- The information exponent is defined as IE(f) = k = inf fi aj +0.
The information exponent is defined as the index of the first non-zero coefficient in the Hermite expansion of the non- linear target function(Arouu et al., 2021). The complexity of learning a nonlinear function via two-layer neural net- work optimized by SGD is closely associated with this value (Arous et al., 2021; Ge et al., 2018; Dudeja & Hsu, 2018; Bietti et al., 2022: Damian et al., 2022: Oko et al., 2024a).
Gradient-based Feature Learning Gradient-based fea- ture learning of -ww-dimensional functions using neural networks has garnered significant attention. Subjects of re- search encompasses functions such as single-index models (Dudeja & Hsu, 2018; Ba et al., 2022: Bietti et al., 2022:
Abbe& & Boix-Ad 2022; Mousavi-Hosseim et al., 2023b;
Ba et al., 2023) and multi-inde models (Damian et al., 2022;
The generation process is defined as follows.
Ben Arous et al., 2022; Mousavi-Hosseiie et al., 2023a; Bi- etti et al., 2023; Collins-Woodfff et al., 2023; Dandi et al., 2024a). The information exponent k.,, or leap complex- ity (Abbe et al., 2023), of the target is known to govern its difficulty of learning it, generally requiring a sample
Assumption 3.2 (Teacher Models). Let C' = Oa(1) be the number of clusters. Let f (c = 1,... C) represent the local task specific to each cluster, and let g" denote the global task shared among clusters. A data pair (Fe, ye) in
2

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
the cluster e, where e ~ Unif[I,.... CJ, is generated as Tee ~ N Now Ta).
ve " O(d-1/2) should only introduce a negligible per- turbation d-1/2 to the Hermite coefficients fa, and t of the teacher models.
Ve = [w IO) + seg (w² zz) + +v,
where v ~ NO (2) denotes an additive Gaussian noise that accounts for observation uncertainty. It is assumed to be sampled independently of the input Tic- The scalar p E IR represents scaling factor that modulates the magnitude of the cluster mean vectors U/. E gd-1, and is assumed to satisfy p ~ Ap, where Ap is a sufficiently large constant upper bounded by poly log d. The coefficient Se E R en- codes the influence of the global task within cluster e, and constrained such that Ee = 0 and &e = (11) for all c. fe and g" are univariate polynomials with informa- tion exponent h- > 2 and degree p", and feature indices we; " E Ed-1. We write the Hermite expansion of l and g" as f2 = =E fed He, and g" = Ei= V/ He,, re- spectively. The Hermite coefficients satisfy |5c,t-| = - for all c. The link functions and the index features are
3.2. Structure of the MoE
A MoE consists of its M experts fM. gating function h(x; O) = θJx where θ (01...., 0xx) e Rdx M and a routing strategy that uses the output of the gating function h = (II....,M)) to distribute the in- put to the appropriate experts. For example, for a top- routing, the index of the assigned expert is chosen as m(x) == argmax (m)).. We also define Ty (z; exp(Im(<))) En. exp(hmm(x)) (x)) the softmax gating functions.
In this paper, we will focus on two routing strategies. On the one hand, we define Fi (2; W, Q), the output of MoE follow- ing a top-1 routing weighted by the corresponding softmax gating value, i.e., (x;; W, θ) = (x) Sm(z) (1; Wm... This weighting enables track the gradient of the gating function which is technically impossible with the simple top-1 routing. On the other hand, we introduce the (adap- tive) top-k routing, We adaptively choose k experts for each input æ based on the value of each htm and a thresh- old. Here, we set the threshold to 0 and define the output as FM (a;; W,O) = Σm 1 (mmm (ar) 2 0),5m Wm) in Phase II; formally defined in Section 3.3. This choice of router is related to recurrent problem that the router may fail to determine the appropriate expert when there are sev- eral models playing similar roles. After Phase II with top-1 I routing, a data point x, is no longer routed to experts outside the set of professional experts (formally defined in Defini- tion 4.7) for cluster c; however, competition among the professional experts may still occur. general, this can re- sult in issues related to load imbalance the emergence of redundant experts, which may result from top-k routing with fixed k E N (Zhou et al., 2022). To address this, expert choice routing (Zhou et al., 2022), soft MoE (Puigcerver et al., 2024), and several auxiliary losses, such as load bal- ancing loss, importance loss (Shazeer et al., 2017), and z-loss (Zoph et al., 2022), were heuristically introduced to promote the even distribution of data and encourage diver- sity among experts. There have indeed been prior attempts to vary the number of activated experts depending on each token (Huang et al., 2024; Zeng et al., 2024). With our adap- tive top-k routing in Phase III and IV, we can also avoid this phenomenon without changing the loss as we will show that it prevents the data from being routed to non-corresponding experts and ensures that the experts that could be activated during inference are trained evenly.
Vi.
normalized as E, [f2 (w= 2)² = 1. Eg (wg 2)3] 1, ||=:||| = 1, and |w0 = 1, where 2 ~ Ia).
This model is designed to introduce task interference by Ee = 0, making the learning process more challenging as gradients from different clusters conflict, hindering ef- fective learning. This scenario is closely related line of work on gradient interference in multi-task learning (Yu et al., 2020; Liu et al., 2021; Guangyuan et al., 2023; Zhang et al., 2024b) and recent work on MoE has also addressed this issue (Liu et al., 2024; Yang et al., 2025). Note that all I and g" have the same information exponent ke > 2 and their k"th coefficients have the same absolute values, which implies that all f and g have the same difficulty, making it even more difficult to distinguish each component from the others. We will show that vanilla neural network is incapable of handling such tasks, whereas the MoB can. We also suppose that je* is even, instead of assuming that the products of the Hermite coefficients of the teacher and the student models are positive in Simsek et al. (2024). We also impose condition for each feature vectors w/ and vi as follows.
Assumption 3.3 (Task Correlation). For i, (g) such that i * i', w T "w, = 3). Moreover, for all E [C] U {9} and E [C], ue w = O.
This condition indicates that the tasks for each cluster are diverse. The correlation between two vectors w can be satisfied, for instance, when the vectors are randomly drawn from Unif (sd-1). The cluster signal ve is assumed to be orthogonal to all feature indices for analytic tractability, and we believe this is not necessary condition of our sub- sequent result. While Chen et al. (2022) assume mutual orthogonality among the feature indices, we relax this as- sumption in our analysis. Indeed, a randomized correlation
Importantly, our adoption of adaptive top-k routing is com- pelled by theoretical and technical considerations, and stands in contrast to the classification setting studied in Chen
3

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
et al. (2022); Chowdhury et al. (2023). This phenomenon arises specifically from the challenge of estimating contin- uous function in the regression setting.
Algorithm 1 Gradient-based training of MoE
Require: Learning rates n7, regularization parameter X, sample sizes 7i, T2, T3. T4, initialization scale C2.
For each expert, we consider a two-layer neural network (m; Wm) = Ej=1 Om,jmm on m (wm, + bm. The Her- mite expansions given as 50 (2 (2+ + wm, Vc bmm.)) = Σ== (mii (z).
Initialize um..
Unif((( and am.j
Unif (=1}.
Phase I: Normalized SGD on first-layer of experts
for t = 0 to Ti -1 do
va
Draw new sample (5 Vel.
Moreover, we assume that the activation function om satis- fies the following assumption used in Oko et al. (2024a) to ensure that a fraction of neurons can align with the feature index even though the target functions are unknown.
++1 + w'm. + 17 Vv V
Fi F((x)).
m.j
"m.j wm.// /l// m
[+1
(+1
for j 1,
end for
Phase II: SGD on gating network of router
for L 1= Ti to Ti + 72 1-=1111 Ti) do
Assumption& 3.4 (Student Activation Functions). The activa- tion functions om of the student model and the link function fe and g" of the teacher models satisfy one of the following conditions: (A) am a randomized polynomial activation of degree most (11) as defined in Appendix A.3 and f² satisfies Assumption 3.2 for all = 1..... C, or (B) om is the ReLU activation function, with the additional require- ment that for each fü, the absolute values of the all non-zero Hermite coefficients Jam.iia are (1). In addition to the conditions (A) and (B), we technically assume that the sign the Hermite coefficient Om.jj.. of m (-+ De + +bm.)) is invariant during the optimization since we cannot evaluate the contribution of higher-order Hermite coefficients.
Draw new samples (I': Ve)
0 <- 0% + n' a
E V. F1 (x.).
end for
Phase III: Normalized SGD on first-layer of experts Reinitialize w Wm. Unif
Unif(S" (()) and anl
for = = + Ti + T2 + T3- - do
Draw new sample (I'e v.))
Wm. + w'm + 17 ye V
Wm. FM
w 1+1 + wm // m.
++1 for = I.....J.
m.j
end for
Phase IV: Convex optimization for second-layer of experts
Initialize bm.f Unif([-Co, Cbl) and set új + 8m. w +T+ where dm ~ (n.
3.3. Training Algorithm
To precisely track the model's evolution, we divide the training algorithm into layer-by-layer, similar to previous studies that have researched feature learning in neural net- works (Damian et al., 2022; Ba et al., 2022; Bietti et al., 2023; Abbe et al., 2023; Mousavi-Hosseeni et al., 2023b; Oko et al., 2024a; Lee et al., 2024). Specifically, we consider an algorithm separated into four phases. See Algorithm for the outline.
Draw new samples (1/e vt) T Solve:
{ambm +- argmin
1
Tu 1
Σ
> a 112
Am Ta t=
m
where TB = Ei Tj and TO Σt T2.
Concretely, we start by initializing the weights of the ex- perts as wh.; ~ -d and amj ~ Unif(-1,+1) +1} following Oko et al. (2024a) and Lee et al. (2024) and the weights of the router's gating network On to zero. In Phase I, the first layer of the expert is optimized using a correlation loss, a technique supported by prior studies (Bietti et al., 2022; Damian et al.. 2022: Abbe et al., 2023; Oko et al., 2024a; Lee et al., 2024). For optimization, the spherical gradient, defined as V.S L = (Ia Wm..; TT, VV V.....f.. employed, as explored in Arous al. (2021); Damian et al. (2023); Oko et al. (2024a); Lee al. (2024). Phase II is devoted to the router's gating network which is trained via gradient descent. In this stage, we train the gating network in a single step with a batch size of n é(d). Since the gradient of the gating network h(x; θ) contains differences in the alignment Wm.j" T " between experts after the explo- ration stage, becomes aligned with the cluster mean vector
Ve. Before entering Phase III, the expert weights Wm, and Om.j are reinitialized. While not strictly necessary, this reinitialization helps ensure that the early learning of w does not interfere with or disrupt the effective learning of wg. particularly when f and g are similar functions. Note that if the activation function is ReLU, a random sign flip of Wm,j. as described in Oko et al. (2024a), becomes nec- essary. See Appendix A.3 for details. Finally, we conclude the training with Phase IV, where convex optimization with I-zregularizztion is executed the second layer using noise terms bm_jj to facilitate the estimation of the polynomial of the ReL.U activation function. Note that we use different routing strategy for Phases and II (Fi) and Phases III and IV (FM). ). Refer to Remark 4.12 for details.
We will show that, at the end of each phase, a specific repre-
4

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
sentation of the complex task under consideration is learned, which is possible thanks the idiosyncratic architecture of a MoE. More precisely, at the end of Phase I, some neurons within each expert weakly recover certain vectors and spe- cialize to the corresponding cluster. In the next stage, the router learns to successfully dispatch data to the appropriate expert based on the weak recovery of clusters assigned to each expert. As for the second half of the algorithm, we prove that each expert successfully recovers both the local task and the global task associated with its assigned cluster.
The next theorem shows that whn, never catches the signal " for all t because the gradient for wn is erased by those for wt
Theorem 4.3 (Difficulty of finding the "hidden" signal wy). During the population spherical gradient flow of wimjs for all j = 1, 1,..,,.,, J, we have
sup I Immm 120
T 5 5sd
with high probability.
4. Main Results
Theorem 4.3 indicates that there is an insufficient number of neurons that can align the feature vector wg- As result, it becomes difficult to estimate the function 8..9" s..g"(w5 (wg T. :). This phenomenon is due to the condition Σ. Sc = 0 and that the naive neural network (with polynomial width) cannot utilize the vectors Us. to detect the cluster structure. Interestingly, such difficulty of SGD for the naive neural network has not been shown in prior works studying the optimization dynam- ics of the MoE (Chen et al., 2022; Chowdhury et al., 2023; Li et al., 2024). This was possible thanks to our theoretical analysis based on the information exponent, which appears in the context of nonlinear regression.
In this section, we provide our main results. We first prove that a vanilla neural network fails in detecting the latent structure of our task as it can only process the problem as whole. Next, show that the MoE, on the contrary, succeeds in dividing this problem into easier subproblems by leveraging the ability of each expert to weakly recover the simpler function corresponding to an individual cluster.
4.1. Limitations of the Vanilla Neural Network
4.1.1. MAIN THEOREM
Here, we consider the vanilla neural network
4.1.2. PROOF SKETCH
Jm (2;; Wm) = = Σ am mm.j
(m at + bm
We provide sketch of the proof for the theorem. We demon- strate that a vanilla neural network predominantly aligns with simple tasks wa, which prevents it from aligning with the more subtle one wy. For a comprehensive explanation, please check Appendix B.
as the student model. This also corresponds to the special case of a MoE with only one expert. The size of the vanilla neural network is at most J = (polyd).
The spherical gradient flows of lwiin T "wil are approxi- mately evaluated as
We will demonstrate that there are some scenarios of our teacher model that a single expert cannot solve. For example, consider the following. See Appendix B for further details. Example 4.1. We construct specific problem of our teach model 3.2 as follows. Assume all feature vectors w and w are completely orthogonal for simplicity. Moreover, func- tions are defined as f (w: xxe) = Be, Hea--He--( (w: x.e), for all c, and Seg" (x.g"(x)) = (-1)ot1Hex: (wg xx) for e = 1, and otherwise 0. We assume that there exists at least one pair e') such that sgn Bekk- # sgn Ber.e-. We additionally assume k* 2 5 to prove the result in Lemma B.15. (se)e defined 81 1, 82 -1 and s, 0 (otherwise),
d d de Wing
T "wll ~ 66ns- IJumj T "1
where 7) is a learning rate, under the condition that ac..-mm... > By integration, it takes @ -2)/2) time for the weak recovery.
Now, the most important observation is that the signals of wy are canceled out:
Lemma 4.4. Recall that the inputs are generated as xe = Mie + 2, where z ~ NO, Ia). The Hermite coefficients am,jic. am sm(++tin Ve +bag) are close up to ((d at the initialization. Please refer to Lemma B.4 for the proof.
to satisfy E. Se O. This means the signal w is hard to recover. In short, the data Be are generated as
if 1, Tx,) xx) + Hex- (wz Ie) v H.. (wi
Therefore, when initialized as w T "uyl ~ d 1/2
A2 Her- (wi xe) - Hea- (m; xe) +v c= 2, Be,x Heg- (w: Iz) +
if > 2.
T wel
Remark 4.2. It may be possible that sgn jm :e # sgn am, ket when j # j' because we randomly initialize amj and bm..j.
di
SnJ Σ Bellm wiis T "wgg" "ugik--
5

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
SnJ [aj-'d1//0 1/2 // T "gght SnJ="/umm. T "ug((0+1--1,
4.2. Learning Dynamics of MoE
4. 4.2.1. 1. MAIN THEOREM
On the contrary, the MoE successfully learns the teacher model defined in Assumption 3.2 by enabling the router to appropriately partition the data among the teacher models for each cluster. This is stated formally in the following theorem. We further characterize the sample complexity of this learning process under Algorithm 1.
which implies that the information exponent of g" increases. It takes at least 12n "J((== 21) time for the weak recovery. Therefore, for all j, detecting one of the signals {w becomes easier than wg and each Wm.j tends to align with we rather than "
Next, we identify the subset of tasks C from the entire set [C] that neuron Wmj can align with based on the necessary condition: for all j. there exists e such that &e, Be, am..kk". . > with high probability. Let C, be such subset of C, then we can show that j can only detect te where E Cj-
Theorem 4.6. Under Assumptions 3.2, 3.3, and 3.4, set J = ) as the number of neurons, T1 = 6(d -1) as the number of training steps for Phase 1, T2 = 1 as the number of training steps and n = O(d) as the batch size for Phase I1, T's ((d -1 v de V ()) as the number of training steps for Phase III, and T4 @(c=2) as the number of training steps for Phase IV. Then, under the suitable choices of of and X, with probability at least 0.99 over the randomness of the dataset and initialization.
Based on the above argument, we obtain that for all j. Wm.jj aligns some feature vectors among fw2fes not hidden wg:
Lemma 4.5. For all there exists La S Ö(Jn -1d -2)/2) such that for all t 2 to.
Exc [FM {am I m==1) -((e)) scg (xx))l Sc
1. Iw/m " 2 ((d-1//-21/66)) ) some c E Cj.
We considered the case where each cluster possesses its single-index model while collectively sharing a global single-index model across all clusters. This global task in- duces interference, which attenuates the signal of the shared model. This setting is potentially difficult for a vanilla neu- ral network to learn as shown in Subsection 4.1. The total sample complexity is ( and the time complexity polynomial in d. This complexity the same as learning ingle--idex model by vanilla neural network (Arous et al., 2021) while kernel ridge regression requires (dd') ) (Ghor- bani et al., 2021; ; Donhauser et al., 2021) with respect to d. After the weak recovery of Phase I, the router successfully divides the clusters and enables the expert to learn their target functions.
2. |w/n "wol S O for all c Cj.
3. |wmm "wg] S (dd-1/2)
hold with high probability.
See Lemma 1.15 for formal proof. Intuitively, when the inequalities in Lemma 4.5 hold, then the alignment in the following inequality does not grow because the derivative T "y and
continues to be negative: Let Sm..g = w's
nt m.c = =w'm
T wwe For all l' [1a,1.
(ISmisool + iomeell
Experiments To illustrate the dynamics of the MoE fol- lowing Algorithm I, we focus on a synthetic problem where C in the problem setting of Assumptioo 3.2. We define fi Hez + He; and fī Hez + He, for the local tasks, and g" Hey for the global task. The vectors wi, w² and w were of dimension 200, generated randomly and applied Gram-Schmidt orthogonalization to satisfy Assumption 3.3. As for the student model, the number of experts was set to 8, and the hidden dimension of each expert to 500. The learning rate was set to for all optimization schemes, and Ti 3.5 10 To 300, T3- 107.
dt
7
011ct -k--1
CJ
Σ Cmm.
( ) Ist (1
egC,
so(1)
Σ Be.--amsihte ( cEC,
20(1)
Rd-a" (1/4-11111))))
CJ el' 20,
Smig (at (a-(k" (i(d k2/ //111//))
The alignments of the experts and router at the end of Phase I, Il and III are shown in Figures I and 2. As we can observe, in Phase I, differences among experts arise due to initialization, resulting in variations in the degree of weak recovery for local tasks. In Phase II, the router leverages these differences in recovery, which are reflected in the
where we used the additional assumption A* 2 5, and the definition of Cy. Therefore, the alignment |win.. wgl + T Eoge, |w m_.j "w'l is bounded by 0(d-1/2) for all t. See Theorem B.17 1.17 in Appendix for more rigorous discussions
6

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
w midi " has a larger value by a constant factor with probability least 0.999, which divides the experts into the exclusive subsets that are specialized each cluster
Lemma 4.8 (Following Chen al. (2022); Oko et al. (2024a)). If M 2 C'log C, it holds that
Phase I
(h) Phase II
PIlMe 2 1, Vcl 0.999.
For all m, if Vlogd 2 J 2 C-1 log M, there one neuron Wmjj.
Figure 1. The alignment of the experts after Phase (a) and router after Phase II (b) with the respective feature vectors of each task. In Figure (a), the alignment of Wm. and w (c ....., C) or ny (vertical axis) is computed, and each expert, the distribution of the number of Wo., with larger alignment than max3,e e wn., w reported. In Figure (b), we visualize for each router fa the task with the best alignment in yellow.
Wmj T Wer
max Wmy wel 1/2),
dt o f'
with probability least 0.999.
At the initialization, the inner products only differ by constant. However, when two sequences have initial values that differ by a constant factor, this can cause differences in their growth rates, ultimately placing them in different asymptotic orders. Such a technique has been employed in various contexts (Ben Arous et al., 2022; Chen et al., 2022; Oko al., 2024a). See Appendix C. for details.
Expert
Expur14
Loca
Local
333
0bbu
055
2-0 Laytr
4.5 4.53
3000
2003
2003
4003
(indd Layprr pits
Phase (Exploration Stage). From this phase, we will take for granted that the conditions of Lemma 4.8 satis- fied and will use the term with high probability withing this scenario (i.e., conditional probability). In the exploration stage, one of the neurons in each cluster achieves faster weak recovery for its assigned cluster compared to other neurons, due to the alignment diffcrences introduced during initialization. Now, for m E MMe. each win w follows a gradient flow as
Figure The alignment of experts after Phase III. The align- ment of Wm., (horizontal axis) and w² (c = l,..., C') (vertical axis) is computed. The last row the alignment between Wo.jj and wy.
gradients, as a signal to learn to dispatch the data from each cluster to the corresponding expert. In Phase III, once the router has learned to appropriately allocate the data, each expert can effectively learn both its assigned local task and the global task without signal interference across clusters.
T
d di Iumm.. w nl w w2i
Then we have the following result: Lemma 4.9 (Informal). For all m Mo, there exists some time S Ti ((ad"-)) such that
4.2.2. PROOF SKETCH
In this section, we will provide an overview of how the MoE can detect and learn the latent cluster structure using popu- lation gradient flow, and how our intuition can be extended to the SGD. We proceed in five steps: initialization, Phase Phase II, Phase III, and Phase IV. Please refer to Appendix C for further details in empirical and discretized dynamics.
T "w'l 5((1).
2. |mm.i' w Ö(d-1/2) ( for all (d, ) # (Cn)
3. |w
T wel = O(d for all
Initialization. At the initial state, experts are divided based on the task of the cluster with which they exhibit the highest alignment. We define (for each task) an expert that will eventually specialize in that task follows:
Lemma 4.9 shows that the expert E M, weakly special- ize the cluster c, enabling the router identify experts via weak recovery. This result highlights that, in order for the router to effectively distinguish among experts, a weak re- covery of the feature index is required. This, in turn, implies that a sample complexity of @(d*" may be required dur- ing the exploration phase, implying that a sufficiently long exploration stage warranted before the router can engage
Definition 4.7 (The set of the professional experts for class e).
(Jis ( = jrgmax,,,./ wns T
w
M = {m c
7

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
in meaningful learning. This contrasts with the linear expert setting of Li al. (2024) and the classification framework in Chen et al. (2022); Chowdhury et al. (2023), this finding is rooted in non-convex optimization in linear regression. To prove Lemma 4.9, we leverage the information exponent of the teacher models instead of using the cubic activation in Chen et al. (2022). Compared to the results for additive models in Oko et al. (2024a), we evaluated the growth of wmn
much more subtle argument than theirs. Specifically, we carefully bounded || [[[]] and 0T. = ( (0(1mmllldd//)) /2). Using Stein's 's lemma, v [[[[]l is bounded as O().
Remark 4.12. We use different router algorithms in Phases and II compared to Phases IIII and IV because the size of the set M, is not fixed. Since there is a variation in [Mel for each cluster c, employing a fixed-k top-k algorithm may fail in routing the data to the corresponding experts. On the one hand, if we set the & of top-k as k > |Mel for some e, there might be some o c such that the corresponding input ze is routed to m E _Ma. On the other hand, if we have the is of top-k as k < |.Mel for some c, then there may be expert in the corresponding set Me that is always selected (routed) when Fe. arrives.
T W for all j.
Phase II (Router Learning Stage). Here, we discuss how the router extracts the cluster mean vector Ue corresponding the cluster e from the weak recovery of the experts. We show that the parameters 0m. for some m Mc. become positively correlated with U/- while, the other hand, 6m', for all m'' & MM. become negatively correlated with it. This is enabled by the fact that the gradients of the gating network encode informative signals elicited by the weak recovery of the experts.
Phase III (Expert Learning Stage). In this phase, as the router has learned to dispatch data appropriately, each expert receives and trains only on its designated cluster. Each expert first weakly recovers and then strongly recovers both the local and global tasks of its corresponding cluster. At this point, there exists at least one m E MM such that hm 2 0 and for all m & MM. (xee) < 0 with high probability. Therefore, the teacher polynomials SeD" (w ))
Lemma 4.10. For all m = argmax,, (Ve) E Me and m² & MM.
gT+ T ml
De S -5(1) 5o 211) I
gT T m
Co.
where Σ. Sr = 0, are successfully decomposed into 2(1)) functions and it enables the experts to learn w and g". As for the MoE model, when the input xe is from the cluster c, the MoE model
Proof. (Sketch). Take m & MM. The population gradient for the gating network of the router is evaluated as
M
- 1/2 Vo,, B[4] ~ á Σ
(Zei fāom M =
Σ m=1
1
[1m 2 0] fm(())
W "mm" J "2lk""
m Mo.
is equivalent to
-n(1), due to weak recovery
+
Σ
I "ell""
FM. (am (mmlneM.) =
Σ 1 (xem 2 0] fm( mEM.
m' m'g.M..ji.
22
with high probability. Thus, the MoE model was de- composed into (Fme to which do not share the param- eters because Me n Me = 0, Ve # c' Addition- ally, using hm)((t))) 2 0 with high probability where m = argmax,, Iim (v=) E Mos it holds that
(1).
Therefore, V. 0 < -(1 and lastly we use Lm fm = 0 to bound hing (xx) = -T. To.
This lemma implies that, for Ic = pvc + 2,
Ummii PM (am.))...) = V Wang Smn (Fee am)
with high probability. Hence, Phase III can be completely decomposed into the subproblem of the weak (to strong) recovery of w w " and ",ng,"" " given the inputs {at h in each cluster c. We show the strong recovery of neurons, in parallel with Oko et al. (2024a).
2 0 if m- m <0 if m E Me
Amm (Ix)
with high probability based on the assumption that p = poly log d is sufficiently large. Interestingly, the concept of the information exponent and the weak recovery had essential roles in the router learning.
Phase IV (Second Layer Optimization Stage). In Phase IV, the experts with aligned vectors estimate the link func- tions ft and g" through second-layer optimization.
Remark 4.11. In Chen et al. (2022) and Li et al. (2024), the norm of the cluster signal pvp is as large as the norm of the noise independent of ve. However, in our setting, ||puell2 = poly log d and ||æ||| ~ d'/2 > l|peell2 with high probability. Due to this setup, we had to employ a
First, with some expert m E Me and Wmjj ~ w and Wmny' 2 wy. we construct äm such that
Im = S6 (w, T "xc) + seg" (wi) xe)
8

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
as feasible solution.
cient learning signals for all experts, whereas the adaptive top-k routing employed in Phases III and IV is designed to mitigate competition among professional experts. These observations point to the potential of stage-specific rout- ing strategies tailored to the evolving dynamics of MoE training.
Next, we decompose the whole convex optimization prob- lem into individual subproblems that do not share the experts to employ the results in the standard analysis for additive models in prior work (Oko et al., 2024a).
5. Conclusion
Impact Statement
In this paper, we theoretically showed that MoE can learn the latent cluster structure of problem with sample com- plexity that depends not on the information exponent of the whole task but on the local information exponent of each cluster. In addition, we have demonstrated that the vanilla neural network with polynomial time complexity fails to detect such a structure. While this work contributes to the further understanding of the underlying mechanism of MoE and its success, is still unknown whether the MoE architec- ture is indeed effective to pursue the information-the theoreetc limit. We believe this constitutes promising direction for future work.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
A Acknowledgment
RK and NN were supported by the FY 2024 Self-directed Research Activity Grant of the University of Tokyo's In- ternational Graduate Program "Innovation for Intelligent World" (II)). KM was partially supported by JST CREST (JPMJCR2015). NN was partially supported by JST ACT- X (JPMIAX24CK, and IST BOOST (JPMJBS2418) TS was partially supported by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2115) YK was supported by JST BOOST, Japan Grant Number JPMJBS2418. This re- search is supported by the National Research Foundation, Singapore and the Ministry of Digital Development and Information under the Al Visiting Professorship Programme (award number AIVP-2024-004). Any opinions, findings and conclusions or recommendations expressed in this mate- rial are those of the author(s) and do not reflect the views of National Research Foundation, Singapore and the Ministry of Digital Development and Information.
Implications and Future Directions Our findings offer several insights for designing more effective MoE architec- tures. First, while our analysis demonstrates that MoEs mit- igate gradient interference through explicit partitioning, the number of experts is typically chosen heuristically in prac- tice. This raises the possibility that incorporating gradient- aware routing mechanisms could lead to more principled and efficient expert allocation strategies, as recently ex- plored in Liu et al. (2024); Yang et al. (2025). Second, to prevent competition among professional experts, we em- ployed top-k routing to reduce potential load imbalance. This motivates the design of adaptive routing schemes that dynamically adjust k during training perspective sup- ported by our theoretical analysis in nonlinear regression and recent findings in NLP that adapt k per token (Huang et al., 2024; Zeng et al., 2024). Third, freezing or pruning redundant experts may further alleviate competition and reduce deployment cost, aligning with recent proposals on expert merging (Zhang et al., 2024a).
References
Abbe, E. and Boix-Adserå, E. On the non-universality of deep learning: Quantifying the cost of symmetry. In Ad- vances in Neural Information Processing Systems, 2022.
Abbe, E., Adserå, E. B., and Misiakiewicz. T. SGD learning on neural networks: leap complexity and maddle-to-saddle dynamics. In Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pp. 2552-2623. PMLR, 2023.
Beyond architectural design, our analysis also informs the training process of MoE systems. In particular, we showed that learning a meaningful router relies on observable differ- ences in the experts' weak recovery, which in turn requires a sufficiently long exploration stage due to the non-convex nature of the objective. This suggests that upcycling dense checkpoints pretrained on diverse domains may offer a prac- approach that
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I.,
Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
tical means of accelerating convergence has gained traction in recent large language models (Komat- suzaki et al., 2023; Wei et al., 2024). Finally, our analysis highlights that different phases of training pose distinct chal- lenges. Specifically, the noise introduced during Phase II serves to ensure uniform gradient flow and provide suffi-
Arnaboldi, L., Dandi, Y., Krzakala, F., Pesce, L., and Stephan, L. Repetita iuvant: Data repetition allows sgd to learn digh-dimmnsional multi-index functions. arXiy preprint arXiv:2405. 15459, 2024.

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Arous, G. B., Gheissari, R., and Jagannath, A. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learn- ing Research, 22(106):1-511 2021.
Damian, A., Nichani, E., Ge, R., and Lee, J. D. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., Wu, D., and Yang, G. High-dimensional asymptotics of feature leamn ing: How one gradient step improves the representation. In Advances in Neural Information Processing Systems, 2022.
Damian, A., Pillaud- Vivien, L., Lee, J., and Bruna, J. Computational -tttttiiionnaststtttiialll statistical gaps in gaussian single-index models (extended abstract). In Proceedings of Thirty Seventh Conference on Learning Theory, volume 247 of Proceedings of Machine Learning Research, pp. 1262- 1262. PMLR, 2024.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., and Wu, D. Learning in the presence of low-dimensional structure: A spiked random matrix perspective. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Sys- tems, volume 36, pp. 17420-17449. Curran Associates, Inc., 2023.
Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L.,, and Stephan, L. How two-layer neural networks learn, one (giant) step at a time. Journal of Machine Learning Research, 25 (349):1-65, 2024a.
Dandi, Y., Troiani, E.,, Arnaboldi, L.,, Pesce, L.., Zdeborova, L., and Krzakala, F. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. In Forty-first International Conference on Machine Learning, 2024b.
Ben Arous, G., Gheissari, R., and Jagannath, A. High- dimensional limit theorems for SGD: Effective dynamics and critical scaling. In Advances in Neural Information Processing Systems, volume 35, pp. 25349-255555 25362. Cur- ran Associates, Inc., 2022.
Dayi, A. K. and Chen, S. Gradient dynamics for low-rank fine-tuning beyond kernels. arXiv preprint arXiv:2411.15385 2024.
Bietti, A., Bruna, J., Sanford, C., and Song. M. J. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems, vol- ume 35, pp. 9768-9783, Curran Associates, Inc., 2022.
Donhauser, K., Wu, M., and Yang. How rotational invari- ance of common kernels prevents generalization in high dimensions. In Proceedings of the 38th International Conference Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pp. PMLR, 2021.
Bietti, A., Bruna, J., and Pillaud- Vivien, L. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310. 19793, 2023.
Chen, Z., Deng, Y., Wu, Y., Gu, Q., and Li, Y. Towards derstanding the mixture-of- ex-eersss layer in deep learning. In Advances in Neural Information Processing Systems, volume 35, pp. 2022.
Dudeja, R. and Hsu, D. Learning single-index models in gaussian space. In Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 1887-1930. PMLR, 06-09 Jul 2018.
Curran Associates, Inc.,
Chowdhury. M. N. R., Zhang. S., Wang, M., Liu, S., and Chen, P.-Y. Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural net- works. In Proceedings of the 40th International Confer- ence on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 6074 61 14. PMLR, 23-29 Jul 2023.
Fedus, W., Zoph. B., and Shazeer, N. Switch Transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23 23(120):1-33 2022.
Fung, T. C. and Tseung, S. C. Mixture of experts models for multilevel data: modelling framework and approximation theory. arXiv preprint arXiv:2209.15207 2022.
Collins-Woodfin, E., Paquette, C., Paquette, E., and Seroussi, I. Hitting the high-dimensional notes: An ODE for SGD learning dynamics on GLMs and multi-index models. arXiv preprint arXiv:2308.088) 2023.
Ge, R., Lee, J. D., and Ma, T. Learning one- one-hidden-layer neural networks with landscape design. In International Conference on Learning Representationss 2018.
Damian, A., Lee, J., and Soltanolkotabi, M. Neural net- works can learn representations with gradient descent. In Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pp. 5413-5452. PMLR, 2022.
Georgiev, P., Lei, V. I., Burnell, R., Bai, L., Gulati, A.,
Tanzer, G., Vincent, D., Pan, Z., Wang, S., Mariooryad, S., Ding, Y., et al. Gemini 1 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arxiv:2403.05590 2024.
10

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029- - 1054, 2021.
Liu, A., Feng, B., Wang. B., Wang, B., Liu, B., Zhao, C., Dengr, C.,, Ruan, C., Dai, D., Guo, D., et al. Deepseek-v2: A strong, economical, and efficient mixture-of-experrs language model. arXiv preprint arXiv:2405.0443: 2024.
Guangyuan, S., Qimai, L., Wenlong, Z., Jiaxin, C., and Xiao-Ming, W. Recon: Reducing conflicting gradients from the root for multi-task learning. In The Eleventh International Conference on Learning Representations, 2023.
Liu, B., Liu, X., Jin, X., Stone, P., and Liu, Q. Conflict- averse gradient descent for multi-task learning. In Ran- zato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 18878-18890. Cur- ran Associates, Inc., 2021.
Ho, N., Yang, C.-Y., and Jordan, M. I. Convergence rates for Gaussian mixtures experts. Journal of Machine Learning Research, 23(323):1-81, 2022.
Mousavi-Hosseini, A., Park, S., Girotti, M., Mitliagkas, I., and Erdogdu, M. A. Neural networks efficiently learn low- dimensional representations with SGD. In The Eleventh International Conference on Learning Representations, 2023a.
Huang, Q. An, Z.., Zhuang, N., Tao, M., Zhang, C.,, Jin, Y., Xu, K., Xu, K., Chen, L., Huang. S., and Feng, Y. Harder task needs more experts: Dynamic routing in MoE mod- cls. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 12883-12895. Association for Compu- tational Linguistics, 2024.
Mousavi-Hossein A., Wu, D., Suzuki, T., and Erdogdu, M. A. Gradient- feature learning under structured data. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b.
Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. Adaptive mixtures of local experts. Neural Computation, 3(1):79-87, 1991.
Mousavi-Hosseini, A., Javanmard, A., and Erdogdu, M. A. Robust feature learning for multi-index models in high dimensions. arXiv preprint arXiv:2410. 2024.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. I., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv preprint arXiv:2401.04 X000:2201..0008,,220 2024.
Nguyen, H., Nguyen, T., and Ho, N. Demystifying softmax gating function in Gaussian mixture of experts. In Thirly- seventh Conference Neural Information Processing Systems, 2023.
Jordan, M. and Jacobs, R. Hierarchical mixtures of ex- perts and the EM algorithm. In Proceedings 1993 In- ternational Conference on Neural Networks (IJCNN-93- Nagoya, Japan), volume 2, pp. 1339-1344 vol.2, 1993.
Nguyen, H., Akbarian, P., Yan, E., and Ho, N. Statisti- cal perspective of top-k sparse softmax gating mixture of experts. In The Twelfth International Conference on Learning Representations, 2024a.
Joshi, N., Misiakiewicz, T., and Srebro, N. On the com- plexity of learning sparse functions with statistical and gradient queries. In The Thirty-eigtth Anmual Conference on Neural Information Processing Systems, 2024.
Nguyen, H., Ho, N., and Rinaldo, A. On least square estima- tion in softmax gating mixture of experts. In Forty-first International Conference on Machine Learning, 2024b.
Nguyen, H., Ho, N., and Rinaldo, A. Sigmoid gating is more sample efficient than softmax gating in mixture of experts. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024c.
Komatsuzaki, A., Puigcerver, J., Lee-Thorp, J., Ruiz, C. R., Mustafa, B., Ainslie, J., Tay, Y., Dehghani, M., and Houlsby, N. Sparse upcycling: Training mixture-of- experts from dense checkpoints. In The Eleventh Interna- tional Conference on Learning Representations, 2023.
O' D'onnelll R. Analysis of boolean functions. arXiv preprint arXiv:2105. 10386, 2021.
Lee, J. D., Oko, K., Suzuki, T., and Wu, D. Neural network learns low- wwwddimensional polynomials with SGD near the information thhooretic limit. In The Thirty-eighti Annual Conference on Neural Information Processing Systems, 2024.
Oko, K., Song, Y., Suzuki, T., and Wu, D. Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations. In Pro- ceedings Thirty Seventh Conference on Learning The- ory, volume 247 of Proceedings of Machine Learning Research, pp. 4009-4084 PMLR, 2024a.
Li, H., Lin, S., Duan, L.,, Liang. Y., and Shroff, N. B. The- ory on mixture-of-er -xeeerss in continual learning. arXiv preprint urXiv:2406. 16437, 2024.
Oko, K., Song, Y., Suzuki, T., and Wu, D. Pretrained trans- former efficiently learns don-diiell target functions

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
in-context. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., al. Mixture-of- experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-71144 2022.
Puigcerver, J., Ruiz, C. R., Mustafa, B., and Houlsby, N. From sparse to soft mixtures of experts. In The Twelfth International Conference Learning Representations, 2024.
Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J.,
Shazeer, N., and Fedus, W. ST-MoF: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08900 2022.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, Outrageously large neural networks: The sparsely-gate mixture-of-experte layer. In International Conference on Learning Representations, 2017.
Simsek, B., Bendjeddou, A., and Hsu, D. Learning gaus- sian multi-index models with gradient flow: Time com- plexity and directional convergence. arXiv preprint arXiv:2310. 197993, 2024.
Vural, N. M. and Erdogdu, M. A. Pruning is optimal for learning sparse features in diggh-imessioss. In Proceed- ings of Thirty Seventh Conference on Learning Theory, volume 247 of Proceedings of Machine Learning Re- search, pp. 4787-4861, PMLR, 2024.
Wei, T., Zhu, B., Zhao, L., Cheng, C., Li, B., Lü, W., Cheng, P., Zhang, J.. Zhang, X., Zeng, L., Wang, X., Ma, Y., Hu, R., Yan, S., Fang, H., and Zhou, Y. Skywork-mooe mow: A deep dive into training techniques for mixture-of-expprt lan- guage models. arXiv preprint arXiv:2406.0653 06563, 2024.
Yang, L.., Shen, D., Cai, C., Yang, F., Gao, T., ZHANG, D., and Li, X. Solving token gradient conflict in mixture- of-experts for large vision- languaea model. In The Thir- teenth International Conference Learning Representa- tions, 2025.
Yu, T. Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-task learning. In Advances in Neural Information Processing Systems, 2020.
volume 33, pp.
Zeng, Z., Miao, Y., Gao, H., Zhang, H., and Deng, Z. AdaMoE: Token-adaptive routing with null experts for mixture-of-experts language models. In Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 6223-6235, Association for Computational Lin- guistics, 2024.
Zhang, Z., Liu, X., Cheng, H., Xu, C., and Gao, J. Diversifying the expert knowledge for task-agnostic pruning sparse mixture-offexpprts. arXiv preprint arXiv:2407.09590. 2024a.
Zhang, Z., Shen, J., Cao, C., Dai, G., Zhou, S., Zhang, Q. Zhang, S., and Shutova, E. Proactive gradient con- flict mitigation in multi-task learning: A sparse training perspective. arXiv preprint 1.Xiv:2411.18h6 2024b.
12

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
A. Preliminaries
A.1. Hermite Polynomials
In this subsection, we present key properties of the probabilists' Hermite polynomials that are essential for analyzing functions under the Gaussian measure. For more detailed treatment, we refer the reader Section 11.2 of O' O'Donnell (2021).
Let be the standard Gaussian measure and (u) the corresponding square-integrable function space with respect to For f.g (u), the inner product is defined as (f.g), == Ez- [f())((()
Definition A.1. The ith Hermite polynomial He; I., i E N is defined
x² d'
x²
He,(z) = (-1) exp
exp
dz
2
Lemma A.2. The normalized Hermite polynomials (He./ Vilj. form a complete orthonormal basis for L2
Lemma A.3. The Hermite polynomials satisfy the following properties:
1. Derivatives:
d
He)(:) i1
d=
2. Integration by Parts: For f E (p) and 2, ~ (O, Ia) such that Cov(z, [-1,1].
Be. = p (Ho_
3. Orthogonality: For z, NO La) such that Cov(z, p 1-1,11
(il)p' if i=j
E (2'))
0
otherwise,
4. Hermite expansion: For (H),,,
f(z) L2
(2,
a; = (f, He:)
1=0
A.2. Bihari-LaSalle Inequality and Gronwall Inequality.
In this subsection, we present the discrete version of Lihari-LaSalle Inequality and Gronwall Inequality, which serve as tools for analyzing the growth behavior of nonlinear recurrence relations. These inequalities will be used repeatedly throughout our analysis. The derivation is adapted from Ben Arous et al. (2022).
Let us consider the sequence { defined as
A++1 = A₁ + B(A)
where & > 3 and B > 0. Then we have the following evaluations:
Lemma A.4. We have
Ao
A N
(1 ((k - 2()
Moreover, if A, 1 Vt < T, we have
Ao
A, <
(1 B(1 + 1 2)
13

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Please note that if two sequences start off differing by constant factor, their subsequent growth rates can diverge, leading them to differ in order of magnitude: Let take two sequences
A++1 =A, B(A))*
Ag Bu = =a(())
1
A+11 =A, + +(A,)*
Ao = A 0 0<><
1 B
Then, it takes at most (k - - (Bo) = 1 time to obtain A, 2 Pa((). Let to S b be the first time s.t. A, 2 Ia((). On the other hand,
An
Ao
Alo - An S (1-B(1 ( 2)(AB)) 22)
= od(
4-2
1 (1 1
x
A.3. Activation Functions
In this study, we consider the misspecified setting, where the target function and the activation function are different. However, in order ensure the alignment between neuron and corresponding target task, we expect that the sign of the Hermite coefficients of the target function and the activation function to be the same. Remember that the Hermite expansion of the neuron is expressed as dm.jom (wm.)" z bm.j) >E Om. He, (w 22 and the Hermite expansion of the Vi!
Pe, He, Ho, T2). 2). We assume that least a (11) fraction of neurons (j E [J) satisfy am 2 0 for i = k" and Cm BB,i 0 for k < SiSpp. Similarly, for the global task g, we have J.77% 2 0 for i = k* and S.0mjj > O for k* < <isp p". This condition is satisfied under certain activation functions.
target function of cluster of the local task is expressed as (w 2) D
ViI
For ReL.U activations, the following lemma shows that Om.jj is positive for all i with probability least and the desired condition holds with probability at least over the randomness of the initialization of am.j-
Lemma A.5 (Lemma 15 of Ba et al. (2023) and Lemma 17 of Oko et al. (2024a)). Given degree p E N and ~ I-Co, Co.), the i-th Hermite coefficient of ReLU(2 + bm..j) is positive with probability 1 for all k* sisp p², if C3 is larger than some constant that only depends on p".
For polynomial functions, we randomize the activation functions as (2) EE He,(), where 4., are independent Rademacher variables. The following lemma shows that the randomization of the activation functions ensure this condition.
vi.
Lemma A.6 (Lemma 18 of Oko al. (2024a)). Given degree p' N and b [-Ob, Ob), for each kimin S k* S kauss the i-th Hermite coefficient of am.jjmm (2 + bm..)) is non-zero with probability 2(C, for all ksd S S p². Here, Q hides constants only depending p".
B. Proof of Limitations of the Vanilla Neural Network
In this chapter, we prove how spherical gradient descent using a standard neural network fails to learn some of the signals, introduced in Section 4.1. Here we have only one expert as
1 fm (z; Wm.)
= J Σ
am. (w + bm,)).
jml
From here, we fix m = 1 and J.. Let nmm..c w T wh and Stn.jin = w =o= "hn..- They represent alignment of
the neuron whn, with the signals w and w respectively. Definition B.1 (Restate). We consider the following teacher model:
We have C O(1) classes and the strength of the cluster vector = poly log
14

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Bascl Learning
All feature vectors completely orthogonal.
Additionally assume k* 5.
Teacher models defined as
(we T xc)) Be,k Ho (w Ie), for all c E [C].
seg" (xte) = (-1) Hog- (wg xe) for e 1,2 and otherwise 0, k- even.
We assume that there exists at least one pair (c, c') such that sgn + sgnBerk.
Sc set as
+1 if -=
5e =
if 2= otherwise.
In other words,
f1 Heg- (wi xx) + Hc (wg xe) +v if c=
Ve
32,8 Heg (wy xe) Hey (wg Ie) +v if 2.
Bax Heg- (wi re) +
if 2
lam.j,A-1 O(1), where Om.j.= is the ke th Hermite coefficient of am, a RcLU(. + bm.j).
Remark B.2. We denote the 1"th Hermite coefficients of 0m + " Ve bm.)) as am.jkk .e- It may be possible that sgnOm..j.k.e ja # sgnommjkke if # j' because randomly initialize am..j and bm..j
Assumption B.3. We assume that the size of one student to be at most = Oppoydd).
Outline of the proof. The outline the proof is as follows:
1. We first show that
The Hermite coefficients corresponding to =Hok (wg xx)) cancel out (Lemmas B.4 and C.8),
For all neurons jm there are some tasks [C] such that the signals of w grow (Lemma B.7), the set of such Wm.j is defined Cj.
2. For each jth neuron, the above points imply that there are three types of signals, as shown Lemmas B.9 B.11:
(a) w E Cy: Learnable (1) Ieim.,eel is positive),
(b) w c E [C] \ Cy: Not learnable (g) Inf. is negative),
(c) w Not learnable (the growth rate of the product wr wy is too small compared to (a) because the Hermite coefficients cancel out (Lemma B.4)).
3. We show that all neurons tend to learn the features (a) w c Cj (Lemma B.12).
4. In Lemma B31,, we repeat the argument in Lemma B.12 while keeping the condition of Hermite coefficients in Lemmas B.4 and B.5 until the products (a) become sufficienlly large.
5. We finally show the growth of other products (b)(C) will be blocked (Lemma B.15) once the products corresponding to (a) become too large, additionally assuming k-
15

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
B.1. Characterization of Hermite coefficients
Here we will show that the Hermite coefficients corresponding to HHee (wg xe) cancel out. That is why wy is leamable. Lemma C.8 informally implies that
dt w
Way
=n
-- a-m.jikh from Ho (wi and
+ (-1) amiik 2 )wo wi)h--1
from Hog (wi as)
In (ammsiih',1 amich 2) wight
2nd 1/2 12(w0) T wg)h--1
We see that the growth rate of (mmj.gl lus. wyll is small compared to smgic by a factor of d-1/2, This results in the hardness of learning wy compared to wi c E Cj.
Lemma B.A. At the initialization, the Hermite coefficients of o. (- pw ve bmid) and om (: pw bm.)) are evaluated as
(mmmigsh Omyjie S vd (S O(d-1/3))
by continuity, and J. .el = (1) and sgn(am...c) = sgn(a for all e, E [1...., C] with high probability over the randomness of the random initialization of Wm..
Proof. Remember that the inputs are generated as z|c= pce 2, (0, 1). w "n ee S Vlog d/ V with high probability over the randomness of the random initialization of Wm. since wn,, Unif(S and ve gd-1, Then get
(m..ghh* Om.iihe el S So/log//V
for all c E [C] since
[0mijjk Om.jikt,o -R22- am (w pm Hez) (=)] S Iow 5 del-
Note that km.jk= BE-A((o.)) (m) Hez (=)] and use Lipschitz continuity for ReL.U and use binomial expansion for polynomial activations. By the triangle inequality, we obtain
Om..j.tt. Om.jjk. oe Am..j Om..j.- j0m.j.htee O'mj.* S PV logd// Vd,
Om.jk Z Om.ii - Opp/log d/v and Omjii S Am.j. + /g
Finally, use the assumption that |0km.j.A- 0
will show that the inequality in Lemma B.4 at the initialization continues to be satisfied:
Lemma B.5. Let i' be the K- th Hermite coefficient of am + win. Ve + bm..). Note a am.ight -2)/2 Then we have
Assume bs* > 2, supee(o] K.jje S 11/2+111(22)) for all s S t, and t S 7 n -
u amjkk 2 (dd-1/2) forall all c,j
at arbitrary time u. S
Proof. As discussed in Lemma C.8, we have Jwj val S ((d-1/3) + Ō(1J=1 fo Ining,el 1K at) S Ó(d-1/2). Next, we repeat the same argument in Lemma B.4.
Remark B.6. We also assume
6(1) and the sign does not change for all time (assumed in Assumption 3.4) Lemma B.5 and Remark B.6 imply that we can temporarily ignore the dynamics of af aiikt e So, we omit and denote the coefficient as Omggkt for now.
16

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Lemma B.7. For all j, there exists such that
Be am..k e
ifC 2 2 with high probability.
Proof. Fix j. Use s.. sgnimmjj.e for all c' with high probability and there exists at least one pair (c,c') such that sg # sgn by assumption. These events imply that sgn Am,j.k", # sgn Om.jkktee
Based on the above lemma, we define the set of we which is "learnable": Definition B.8. The set C C [C] consists of the class such that Pex amj.k > O.
We roughly observe that
d
fe, De. elmmggal -10-11 nJ Imnggee c if Cj. j.ae mJ
dt
-nm - M [k--1 if $ Cj,
which implies that C reflects the learnability of the tasks. We more formally have the following result:
B.2. Evaluation of Spherical Gradient Flows Lemma B.9. Assume the conditions posed in Lemma B.5. We have
Σ de
J S -Irmsscl -1 + (ammghtt - Olm.j.A+ Imm.al
1 CJJ
eHC;
o(d-1/2)
In addition, if
2 (ddd111)) 1/2) ) for all e C₁.
Inm.iiel
2 ( (o(d=1/241/(2x))) for all c Cj.
Ism.ana S (((d-1/241(((2) for = 1,2
hold, then we have
--11
(g Kmgjo 2
Σ
CJ
ecc
Proof. By the standard argument of spherical gradient flow (please refer to Lemma C.7 for the parallel discussions in the discretized dynamics), we have
qk* dt K.3.a CJ
(a. (e (Kmmj - (Kmm.j.e))
()) ((w)) (w/) -(w))) KmngieSmiaa)
Back amikee Bes (fm, ((w)) e'cc/e
((w2) (w.) Kemjel (mmje)
Σ
Berk Omgjke el (Bm.je)) ((w3)) (w:) of
ni²
(Bek-am.s.. (Bat-Omsjk.clmmmmi (BBms,,)) k
(Kmm,5.e) (amzk'1 amplo 2) (Em.s.a) k
Kmng.
CJ
17

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Cmssie e
Berk ef )r kimjie =- R.m.j.
=+0(1)
e'ecy M(c)
where we used Definition B.8 that is rewritten as
[ if CEC, otherwise
sgn am.jkk,. ac
and [0mjktel (1) in Assumption 3.4 the last inequality. Now we have shown the first inequality in the statement.
Next, we will show the second inequality. Consider the sum for & Cy. We have the assumption that [km...cl S d- d-1/2-1//2h) x oa(()) for all & Cj and k* even. Then,
d - k.c. Σ sgn (e e4c,
dt
yik* ~ CJ
- -in ( am se egC,
E
(Kim.j.e) -1
amjk-1 - Amjkerl
(Cm.j.g) )
egC,
So(d 1/2),
Σ
Bock Cfmjkte Kmmjee 1h Wimmgl
Ber Olmjkite
W 1 Inim
egc, d'EC, Mc)
ogC, ((CC)M
+1(11)
9(1)
Smaxuoee Iim 18++1
S 7 CJ
Σ
Inmmsiel -- 1 + C²
max exC,
|n +1* + 0(d 12)55mm
eRC,
7 CJ
(1 C² -o.o((1)) Σ
Inm.gial 1k' -1 1/ )
(nmgge 5 oa(1), ve)
oRC,
1--1
π CJ
-c-K++2
Σ
O(d ((d-1/2)](mmsnn
CRC,
where used 2 Ei lai)) < 1 E. |a, for & Σزι and ap...,, an E ER by Jensen's inequality in the last inequality, As for the sum for Cj. we similarly have
d
Σ cEC,
sgn(ssm.ss di timi.
--
x- +2
C² max eEC,
Kmjjj
|++1 16'11_0j -0ddd0/2)
Emij
CJ
oEC,
4-11
> n CJ
Σ Iimm
- 0 max I cEC,
(/2 |Nim.sol
nEC,
(12/((++22/2))66 13/2 She
oa(1)
18

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Next, control the dynamics of Emi.g corresponding wig The growth rate is small because the signals cancel out: Lemma B.10. Assume the conditions posed Lemma B.5. It holds that
d ct dt
Sm.j.gl S CJ
(amgghht1 - J.m.j 22) ck--1 lsm.jgl ofC,
1k" KmJ.cl
Σ cEC,
2/-1/2
for all j.
Proof. We have the population GF as
sgn(Emss) d
Smm. dt nyk" (lam,)tt.1 am.jib (5m -1 (Emji, )²)
CJ
+
Box o 1k' (0-(((())) (wx") (wn) K.m.jaamminn)
=sgn(Emial 1it" CJ
(Omisht.1 - O.m.jk- 2) (Sm...g) jk" -1 (Emj1, (3)
Σ
Amik Emiis
Ed 1/2, Lemma
o=1
(mm (Em/1,)
Σ amsk nEC,
e) (Rmp,e)) + Σ
(- Be ammjle))
lsmssol
CJ
e4C,
where we used
1 if CEC, otherwise
sgn a am
at the last inequality.
Even we ignore [Km.j.el, c E Cj. which has the effect reducing the gradient, the growth rate jtm. Eesc, Inmgge small with the "information exponent" equals to k*:
Lemma B.11. Assume the conditions posed in Lemma B.5. Then we have
k'
sgn(Cm..i)) det dt Smisa + Σ sgn(rim.ii)) d cRC,
dt mmii S oo
|km.jll +
Σ [Nmiieel CFC,
Therefore, it S -
Sm.j.gl + Σ In ol S 0(d egC,
Proof. First, g. S (d211/2) for $ C₁ and 1enggg + Ensc, |Kmm (d-11/2 / 0.
Next, we assume that there exists the time
- such that KKn...gg Eesc, |n sup,egta.)) 15m. Engc, I v 1-2/2 for some 1/(2k*) 8 0 Then the assumptions in Lemma B.9 'We require 1/(2k") to satisfy [o] vel S 1/2
19

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
satisfied. Therefore, have
d sgu(Sm).s) Smmja di
Sgg(kmmp d - Kim... < f |miiill + Σ ofC,
Σ ofC,
IKm.iel
di
fort E [0.
Then we will show the contradiction. Let Sm.jgll
+ Eese, Inm.j It. the dynamics of x', / E [0, is evaluated as
d
2t S (x)e))))
dt
where S poly log d is a constant. By the Gronwall inequality,
30
S (1 (k+ (200)=) ())-17))
1/(6--1)
Therefore, obtain
cT
Sm...gl
+
CRC,
((2-1/2)
1/0-1)
1- Ö ((Em))o) + Eofe, In (nsawjic)) 1- J J-117/ ))
((d-1/2)
δ 111/2 -21/2) 11(2--1)
which contradicts that (5m. + Eosc, j.el 2 d/1// This implies that In S d-1/2 for all & Cj holds if t S Jn- (k'k' -2)/2 and solving the ODE again leads to the desired result.
B.3. Balancing the Race: Learning Before the Hermite Coefficients Deviate
We will show that the alignment "n Cj becomes sufficiently large before the Hermite coefficients om deviate too much using a recursive argument.
The following lemma shows that Wy tends to align with w C Cy::
Lemma B.12. Assume the conditions posed in Lemma B.5. There exists t1 S (((k" such that
1. Σeec, INm. ~ d-1/241/(2x)
2. |n
< O(d-11/)) for all e & C₁.
3. Isth,gol S 122 for 1,2.
Proof. Combine the results in Lemma B.9 and Lemma B.11. Lemma B.9 implies the first condition by Gronwall inequality. Lemma B.11 leads to the second and the third conditions because max(maxxeeg [w'm (ct 1Shnj.g | + Eogc, Jnm.
The intuition the final part in the proof of the above lemma is as follows: The differential equations of af Exsc, Intngie and y² := -oge, INm. + (mmii are
d x' ~' n. pd>nj (xt) x ~ 1/2
dt
20

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
and
d
dt 1/ 2 nd (1)"", 50 === =d
takes at most 77- J(=)) = 11 =n Jd -2) time for x' to grow up to d' 2+ at- and on the other hand, it takes at least 77 Ju')) = 1 -1)/2 (> 17= Ju (k" -2)22 time for y to become larger than (dd-1/2).
Repeating the same argument Lemmas B.5 and B.9 to 12, we have the following recurrence formula:
Lemma B.13. Assume K* 2, A 1/2 and there exists t S 1 dA)( such that
1. Evec, Immit " = d A=A;;
2. Enec, / A// + |mmjjll | S d-1/22
Ei 4 Omiihtl (dd-1112)
3. mm
Then, there exists 4+1 S n Jd such that
1. Eeec, 20 1 ~ d Art1 where A4t1 = - 25
Di+1 |mmid
2. Eoje, Iimasi
Ev-1 4v/
+ Immjal S 1/2
Eit 4r 0 m. S Öw, vo) i ((d-1/2).
3. [0m.jik''e
Proof. Let Inm.jeel S d- A [/, [.1... Following Lemmas B.5 and C.8, we have lw T v.l 20(d-1/2) + Ö(11J~
114
Imm. 1k dt)
50(d-112) + 1/1)
(d
-- # A2(b--2)) -2)
which implies the third inequality. Based on this, the first two inequalities follow from Gronwall's inequality. The differential equations of x' Ecec, " trtt and y = Enec, |N tett
(mm.i.g are
=x' = n (x
x ~ dd
dt
and
d
dt 3º S nn (y 8a - d
takes at most 1) "xx00=--22 17 - (k 2) time for xt grow up d' - - zi- and on the other hand, it takes -Ij 1)/2 (> Ad 2) time for y become larger than (dd-1/2). least 1--1 (1 -
As shown in Lemma B.13, A₁ is shrinking as
-ー k*
A
A +
Ao Ao=
2k*
asymptotically approaching 1/4. Repeating the above argument, we have the following result: Lemma B.14. Assume k* > 2 and take arbitrary A > 0. There exists ta s.t.
Evec, | d- 1/4-2
Eoec, |N + 15mjial S 1/2
T la'a lamisjAtt. 00 onngee | S &Ōwws" ve) S ((d-1/2)
21

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
B.4. Blocking the alignment win T W
We will show that the conditions in Lemma B. 14 hold for to. Note that I Omm.j.k- not assumed in to: Lemma B.15. Take A 1/(8))) defined in Lemma 3.14 and take arbitrary >> to. Additionally assume >> Then we have
Σ eEC,
Im 2 d 1/4-A
and
d
Ismmiss gl
E c&C,
dt
d Proof. Let fa is the first time that at Ecee, I = 0 (note that we can take to such that d satisfying the conditions in Lemma B.12 and this gradient flow stops la > (a). First, because di Eeec, lo 2 and the continuity of the derivative (except for Ecec, Intr.s,a 0), we have a Ecse, Inmm t=s 0 for all & E [ta ta.)
dt EcEe, In 2 0 while
This implies that Eeec, INm.j.el Z d-1/4-A for > la.
Next, we bound the derivative of lsm.joo + Enec, Iim.g.ll- We assume that, there exists E [ta. 0o) such that, supre(taitt (Eim.ssll + Degc, Immssle O(d 1/2+0) with 1/(6k*) 8 O. However, for all [ta.t].
sgn(Emiss) do Smiss Σ s. d m.m..
d
ds
oEC,
Ō (0m.j.tr,1 am 2)
Sm ())m + Σ
(- --xx (mss-- (Km...cc Emmjill
CJ
cEC,
(ō(1), NOT assuming emmma B.5
si(1)
4-11
- Σ
Be, aena-tngttte (f)m.s.)* Emmiill
+ CJ
-52 Σ
m impgell
+ (am a'n J. 2)
OEC,
20(1)
Zd 1/4-
CEC,
2O(1),NOT assuming Lemma B.5
CJ 0 <0
10:1001100 (a-(1 (2-11-22)11/20) - 52(a- ((/(t))))))))
where used -(k- 2)(1/2 - 5) = -(k- 2)(1/2 - 1/(6k))) < (L/4 1/(8k*)) k1/4 A) under the assumption that K²' > 4 (i.e. :-- 2 5). This contradicts with the assumption that (Ie Eefe, In Inmmsle) ~
(poly log 1/24- Therefore, we have (Sim.iisl Enge, Intm,,,e) & (poly log d)d-1/2 for all t E 10, and this leads sgn (
dt fm....9 + Eoge, sgu(sim.,.c) A kmm. S repeating the calculation.
By Lemma B.15, have the following lemma:
Lemma B.16. Assume
The following conditions hold true for all ts:
1. Eeec, [mm_gpl 2 d- 1-14441/(SK))
2. |kim.j.cl 5 O(dd ( for all c C,,
3. | Odd-1/2) 1,2,
Finally we have the following theorem by Lemma 1616:
22

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Theorem B.17. For all 1, 1.... J, c 1, we have
c S ((d-1/2)
sup 100
Sm.j.g
with high probability.
C. Proof of MoE Training
In this section, we present formal proof that the MoE. can learn teacher models (Theorem 4.6), under Assumption 3.2, Assumption 3.3, and Assumption 3.4. We execute the gradient- optimization process outlined in Algorithm 1. Our proof builds upon and extends the reasoning presented by Oko et al. (2024a). We introduce polylogarithmic constants A, and a, A, is of the order polylog(d). while a, is of the order pdl with the following order of strength:
A 02
- A2 A A4 54 As Ō(1).
0
- Ap As, A6 S
A1 is derived from the high-probability bounds the gradient of the experts. d2 is derived from the threshold of the weak recovery of the neuron for the tasks. A2 and As are high-probability uniform bounds the noise and small terms in the gradients of the experts and that of the gating network, respectively. A3 is derived from the upper bounds of the specific components the gradients of the experts. A4 is derived from the upper bounds of the task correlation. 44 and as are derived from the upper bounds of the learning rates of the experts (72) and the router (7,), respectively. Ag originates from the lower bound the Hermite coefficients in the router learning stage. Ap reflects the order of the mean vector scaling in terms of p. given by p Ap
An outline of the proof of MoE training is as follows:
1. Initialization (Appendix C.1)
We first show that, after initialization, there exists a neuron Ith within the set of professional experts M, (Definition C. corresponding to the task c that aligns with a constant factor stronger than other neurons (Lemma C.2 and Corollary C.3).
2. Exploration Stage (Appendix C.2)
After the exploration stage, the neuron Jm which was strongly aligned during initialization, undergoes weak recovery for the task ot it specializes in, whereas other neurons fail to achieve weak recovery and remains at a saddle point (Lemma C.6).
3. Router Learning Stage (Appendix C.3)
After the router learning stage, the router directs the data ne from cluster corresponding to task the experts in M. (Lemma C.14).
4. Expert Learning Stage (Appendix C.4)
After the router completes its learning and experts are reinitialized, the experts m belonging to Me. which now receives the data Fe. achieve weak recovery (Lemma C.21) without being affected by inter-cluster interference and subsequently attain strong recovery (Lemma C.26).
5. Second Layer Optimization Stage (Appendix C.5)
We finally show that by performing convex optimization the second layer of the expert, the MoE achieves c-error with respect to the teacher function (Lemma C.29).
To establish Theorem 4.6, we apply a union bound over multiple events. Given that M = 0(1) and J S Vlog d, the set of events that hold with high probability remains closed under certain union bounds. By combining the aforementioned events, we conclude that the event Theorem 4.6 holds with probability at least 0.99.
C.1. Initialization
To start off, we consider the initial alignment between the neurons and the index features. The following lemma shows that a constant fraction of the neurons aligns with the task of their corresponding cluster by a constant factor more strongly than the remaining neurons. We provide definition of the set of professional experts that depends on initialization.
23

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Definition C.1 (The set of the professional experts for class c).
Fire Cm) = Jeemaxx,.umm... w T
Me = {m Cmd.
The following Lemma holds with initialization.
Lemma C.2. Assume C is (1). Take arbitrary constants > 0 and Aa = + ((logd)) If
M = Clog(C/0)
and
1
M
Vlogd 2 J 2 C
log
8
then [Ml 2 for all e and
W jm W 2
Vd'
2
m.ji.
T
w 2
ofem jfim
max
Aalw T
w2l m
Vdlogd d
1
with probability at least 1 5 with sufficiently large d.
Proof. Fix m. By the symmetry, we have
Pl 2 21121-(1-C7 (1-1-11--1) ---1)M
By union bound,
Pl|Me| 2 1 1 M(1 M 2 1 M exp( - MM/C) 2 1- 6-/3
where M Colo(C//)).
Wm,j (1)) is obtained by
@m. where Wm.j ~ N(0, 1). We have | ~ with high probability (the same argument as Oko et al. (2024a)). Consider the value of Rm.j.c = 10 T m. (T (T~ - w T)ur, and Ee w "n,"" (w"") T Twe- Lid. (((,d=") such that Rm.j.c
[
wD 1/2 because Rmjj. are independent and Eefc "h w (w²) Twe"" S d=1. Therefore, we evaluate the values of Rm.j.c instead of w. m w
Then, for each Rm.jjes there exists Rm.j..
Km.j.
We show that there is some j s.t. Km.j.. is large enough: First,
PiRm.j. < 2d211/2 < 0.9.
Then,
Plmax Rm.j.
2/2 1/2 for some j| 2 1 (0.9)
Taking
M
J ス log 5
8
P Pmmax Rm > 2d for some
2
3M
Following Chen et al. (2022), we have
6
W.
T
W 2
max
1. - O(d-1)
3MJ2C
of cm or 1f
and therefore
(1-
W mi "
max cfc
GM
wm + J2 Wn wi - Odd
with probability at least -5/(3M) and the desired result follows, using 6M Jac 2 (log d) and w WT. w 2 1/2
24

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Since w u O(Vlogd/Vd with high probability and ming [a mm.2 Pe Pe by taking 6 as sufficiently small, we have the following inequality:
max, (a,, i. ,, L
1 + ((d-1/2) with high probability,
Corollary C.3 (Following (Chen et al., 2022; Oko et al., 2024a)). When J C-1 log M and M 2 C' C'log C. for all m, we have at least one neuron Wm.j such that
1
mm W'n Vd
and
(mmms in de ,e I w ) 82- Be k-
2 max[mac m
max Σ
Sel O
n-
max
k2 -2 [mmm u 100
+ alwmsn ) )k--
Joc
ofe or +#
c'EC
with probability at least 0.999, where is small constant S (log d)
Remark C.4. The term d .. defined in Lemma C.7, varies with time t as it is a Hermite coefficient influenced by the mean vector pve the data Te. Nevertheless, Corollary C.3 holds for all t throughout the exploration stage. This is ensured by the bounds on the Hermite coefficients provided in Lemma C.8 and Lemma C.9. Remark C.5. From this section, we will discuss Phase to IV on the event that the initialization was successful.
C.2. Exploration Stage
We train the first layer of the experts. We employ the correlation loss to eliminate the interactions between neurons. The alignment at time denoted as Ke.m..-> defined as the inner product of the feature index w. and the weight of j-th neuron "mm time expressed Ke.m. == w. wm.g. Similarly, kg.m. wy w/m. is defined in the same manner. The purpose of this subsection is to establish Lemma C.6, which serves as the formal statement of Lemma 4.9. Within the expert set Me. there exists neuron j such that the alignment magnitude Ke,m..j with the feature index w satisfies Ke,,.) 2 a2 for some constant 02 > 0. In contrast, for all other expert sets Me' with c' * c, no neuron achieves such alignment; that is, Bee.m. S (dd 2) <, 02 for all j E .Me. Moreover, the remaining neurons in M, also do not reach this level of alignment. To prove Lemma C.6, we first decompose the stochastic gradient update into its population and noise components. Then, by introducing auxiliary sequences, we establish lower bound for eet co.. (i.e., weak recovery) and upper bounds for In'.m.p for all (c,j)= # (Cm and Iyy.ml for all j E [J].
Lemma C.6 (Formal). Consider the expert m MM Let w w S Add Ö(d-1/2) for all e + d' w wyg S A4d-2 \/2 = ( all e and n¹ 1e a,d- Then, with high probability, there exists some time t S Ti
Ô(1c -Id= such that the following conditions hold:
not mih 2 a2
komm < 5A3d 1/2 for all (c.j) # (in.
In/.m.g S 5Agd = Ö for all [J].
Gradient update decomposition. First, we assess the evolution of the alignment by analyzing its population and stochastic contributions. Note that the exploration stage, the router distributes the data the experts with an equal probability of i because the weights 6. of the gating network are initialized to zero.
In Lemma C.7, we will evaluate the update of the spherical gradient descent
De + bm.jj]] K w. T wtt1 TT [wt [wr.m. (Ja wm. w V wi... (m(xx) = m)) Tm (Jre) Jm (w T
w.
[w'n. n/ (Ie w'n w'm. V wt (m((x) = m) Tm (Ie) Fm (w'm.) T De + bm.
Note that La winj win T is a projection matrix used to project the gradient w/e. 1 (m((x)) = m) mm (xe))oam. om (wmn Do + bm.)) onto the tangent space of the sphere Sd-1, Addition- ally, normalization performed to return the vector t the unit sphere.
25

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Lemma C.7. Suppose that o' T/e S a4d & and ne min 2 2 dd= With high probability, the update of n fomm and
" "g satisfies the following bounds.
nf
Re mi
++1
2 Kee mih +
Σ
k.c misin
(win. "w - Koo moin in Hie! mm)))
CM2
C'E(C) 1-
+ is Amm af V6 (H g.mm iin
( T
)x-1 (wie
Wig
Kig mi 1i mg.
Kipe mn A1')
+ w (Ia wmin smis
' wm.
Let + Gn or j # Fim we have
Komma
nf' CM2
Σ
Σ
j Be (el 1-1 (wo wy K.m + is is ammiee (1
C'E(C] i=k"
Int (w: T w/g - -K'e ne.m (gmms))
m.m. " 2 Ag (1)") 2
+ n/'w"" + (La w'm ' wm.
2
I 1+1 S
Kom..
n1¹
Σ
im Bd (He... ( T Nomm mid.mm)) + is amiii Vo (i (fig.m.)) 1-1
"
CM2
d'E(Ci i=k"
In ((7') (n)²) (n')²
(wa T w
- ( +
Toa 2 + wnpwa (Ia w'n w unm T) Wm.
Let j E [T], we have
Kg +
n1
Σ
iā Be mmm) (m
(wn w (11 the m.s) + is is (H 1
CM2
C'E(C] i=k
Int g.m.j An
(1 (nimmm))))
+n n/u (Io w'm w ' Wmn
2
S S1tt1 S g.m.j
Ky.mas
n¹¹
Σ
im Bet a m..)
w. ( s...m.ttte..mmm am (i (
CM2
C'E(Climk-
Inst g...m
A2221 2d
(174) An + n/ u
(1 (K/...))
T (Ia w'n womnywhn w'is T) Wimss
2
E' E'um., represents a mean-zero random variable satisfying Unif (8d-1), with high probability. We can also obtain ( N Ö(1,) and w k/y.m.s Ö(ne) with high probability.
|E E''mmll = Ö(d') and lu TE' s'.w..p = Õ(1). where u
Proof. The population gradient for the first layer of the expert can be represented as a decomposition in the following manner.
E.,. I (m(xe)= m) xm (e))
om (Wm. o (m((x) = m) Fm (Ir) F (wmis T To + bm.s) 2
xee bmns)
ExErr
Σ Ey. [(2 (
Besi Hor(wz T xcc) + 80 >
Vi Ho:(w T xz)) x2
i0 He; ( T xxc) )re)
CM2
Va!
cE(C)
k=k=
26

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Basel Learning
CM
1
Σ
E
Bai
He (wz
+)
se
Yi
(wg T0)
10 Vi!
Hei-i-1(mm (0
1
eE(C]
+ 12
Best He;(10 Vi!
2) +
8c 2
7i
Hes(wng TT 2)
i-1
('7)
i
1
He;- -1 (w (pw ) poie)
He;
Vi!
1
k=
==
1=0
Eg
p
iB vi!
He
(wz (w
T 2))
sc
j.
(
w
CM
Vi!
CE(Ci
k=k"
1=0
Bg
88
87
He; (uig =)
im
He:-r- Hee-r-i(wnn (wim (pw peel w
1
k=
1=0
(
+
Ez
B (w; T =)+86) E
D¹
Va
He, (we g
2)
x
i0i J
2
(i - 1- 2 )= Mpw i-1)
j we)" Wm.))
Vil
Vi!
1
==
1=0
(m)
Eg
T Bod He, (w 2) +
Ta
He (wg 2))
im
1 He; -e S ve))) pri)
Vi!
Vi!
1
==k
1=0
(IV)
where the second equality is due to P(m((1)) m) 37 and Tm(a.) 37 The third equality follows from the definition
Te 2 + pv, where (N, Ia), along with the condition Ve " and ve T " for all (c, e'), as well as the binomial expansion. The fourth equality is due to Stein's 's Lemma and integration by parts.
With the spherical gradient, (III) is negligible since Wmj is unit-norm and w (La Wm. Wm, Wmm. = 0. When
considering Ue V W E[1 (m(xx)= = m) Tm (nn( (ee) om (Wms T xe + bm. (IV) can be ignored because ve T " = 0
for all (c,c"). Thus, we expand (I) and (II).
iBes (i 1)!(w. wm,s) i-1 w (pem lam.jj. (1 (1-1) (pw we)
(1) =
(1-i) VI!
Vi!
k=
1=
p"
8
jil
( (11
Te,
"m))) 1-1 wg.
(II)
(i - 1)
Vi!
1 1
Vil
1=1
Here, we introduce the discrepancy = wmns
between the population and the empirical gradient.
Et
(m(xx) = m) T... (x on (w'm Ie + bm..)
E[1 (m(xe) = m) n. Tm am. mm (wim.s Ic + bmaj)l
T
upf m
We evaluate the empirical update of the alignment.
Ke.m. n/' (La w'n wins T J) win,, 1 (m(xx) m) Fm Jm (win.s T ae + bm.s) |woinss n7' (Id - w'n, wine Wiih. (m(xe) m) mm (x o (w'ms De + bm. \|
me.m.u nn'w"" (Is- wm, Wiss "Wu 1 (m(xx) m) nm (F am jo (w'm.s Be + bm.))
- nr.m,,ll1 1V m 1 (m(()) = m) Tm ( Ve omm (w'm.s T xo + bm.s)ll
2
27

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
(n7)
11 1V Winn (m((x)) = m) am (Ire ( am. Jm (mmm T Tee + bm..))l
2
(1)
Kemm + +n'we T (Ia (Ta
lammiid
1 11- (pwmis T "oe) 'vo))=1) (1-1)
CM2
vi!
<E|C] i=
1=1
1B (i 1)!(w)) T "mm)) ) w + +nww (Ld - wim wims )
Juss
Vi!
lammdi (1 (1-1)) (ow'ns "vo)) 13007
(i 1) ( "mns)" )w
CM2
Vi!
Vi!
<'E(C]] i=
1=1
We e.m 1(7') 2d
(n')² A dd
+ nn (Ia - wmm w
2
Wm.
= -amm + tn (Id w'ns w )
Σ
(aamijd (1 (1-1)
(1 (1-)
(pw'ms T "o6)
+
CM2
Vi!
c'E|C] i= i=k
1mi+1
iB (i 1)l(wcd T wm)s)
(la-wmm ) w + n w.""" (Ja - wmm wims ))
Vi!
Σ
i.mm..j.
Vil (1- (pw'm. T Ue (e) 1=++1
larm.ji (h-1)
i8
(i 1)((ng Toina wmmms+ ug
CM2
Vi!
e'ECc
k= 1(n')²A22 A 2d
w/ (11')
(1))² A1 3d7
+ w (Ia - w'm Wmj T mm.
2
2
p'
K.m + CM2
Σ
2ā'm Bet, (k's (we²
ne n Re mmm) + is To g.m))) (n( i-1
O'E(C)
i=k-
In p (17)) 2d 2 + n/w (Ja wmm wmss
(wt wg me (
w T ' "mn.""
where (i) is due to Taylor expansion, Cauchy-Schward incquality, and the orthogonality property of the Hermite polynomials. In (ii), we used the expansion of (I) and (II) and |/ wit, 1 (m(zz) m) mm ye mommm..... (w'm. De + bm. I Aid which holds with high probability. Here, we introduced m. defined by
idt id
m.j.e
d.km...d Vil
Jammj (1 11-1 :) (1: 11- i (pw'no Twc)
Vil
Vt!
1=1+1
Furthermore, it can be equivalently rewritten as follows:
& m.j....
E2 amm.j'm (wim. + pumm Un + He, ())]
Note that the definition of ām.j..c given here differs from that of Om.jáic in Appendix B. In Appendix C, we define the Hermite coefficients by incorporating the perturbation induced by prom.j Ue. and formulate them in a manner involving a first-order derivative. To distinguish this modified definition, we introduced the tilde notation.
In the same way, we obtain an upper bound as follows:
2+1 S Nemm.
Σ
Σ
Be Abe m.s) m (w: " ki m.s) + is Vol ( ji-1
CM2
i'E
Inc
m A
(17t) A dd
(w. T w/y ni (omm))
2 + (Ia w'n, Wns T "wo"" E' wmss
28

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
When (c,5) Jm), (0' ( A agd= 7
S (emms("") 2 since we have 1ee mi 2 d and 1/ I
We obtain an upper bound on the difference of Kam. over a single step.
c.m. Ki Kamm
" (Ia -mm T ") Wewn. 1 (m((z)) = m) Tm ( Fm (wims T Ie + bm.s)l
+ In' (emm))(('))
2 V/Ve,,, (m((z) m) Tm (Te) ( 5o (w'm. T Ic + bm.iil
+ (178)3
V whins (m(xx) = m) Tm om (w T 24 + bm.)) win 1 (m(Ix)) = m) am (I)) (2o))))c0mmmmlm (m T Too De + bmJ)ll
Wu (Ia - wnn w T "
2
= Ö(ne)
All terms the RHS are bounded by Ö(1e) since
w (Ia wmn u VV Wiin, 1 (m(xx)) = m) am (ne)x jom T Ze + bm.))l = =
T xe + bmg)ll = O(d), and I (Ia wins T 'V wit,, 1 (m(xx) = m) Tm (m (m. (win.) T
VVe... "h... 1 (m((xe) = m) nm
Tee + bma)l
V wett.. 1 (ml(a)) m) Tmm (za)yoam.j (wmm. Ie + bm.))l = (dd with high probability. We establish similar statements for Kg.m.j. We obtain lower bound follows:
12t k. uz, (o,,TTT,, (Ia w'i V wi. 1 (m(xe) m) mm (Fe))emm mm (wims T Iee Dmi) | 17 (Id - wtt w/ns Veu,,
wt, 1 (m(xe) m) Tm (m(e))) ml (w'ms + bm.iil
Ni yym - -n'wg" T (la wiis Vwin, 1 (m((x) m)) n. (I Well), (w'ms Iee + bmi)) (Ja-
|nt "(
/ Wantn 1 (m(x))) = m) mm (m om (umm T dee + bmsill
2
(17²3
| whin, 1 (m((xe)= = m) nea Mee)yymmmm j mn (w's Dee + bmn.s)ll
2
(60
N/t Mj + n/ n/'u"" I (Ia w'n w'm 1) CM² C'E(Cl i=
1 (1-1) (ou'ma T (ea)-1) Vi!
(1=1)
l=1
iBe (i 1)!( T "wmmi) ) w + n nrung" (Ia w
s T)
Vi!
D0
lam
(1- (1-1) (pwtiis Tucy (wo)t-) i50 (i 1)l(wng T (awTT "mim)) 2
CM2
Vil
Vi!
<'E(C) i=
1=1
Int M.g.m,1(77 ( 2
(n1)²AA
A 3
+ r/ n'uu" T (Ia w'm umm
E Wm.t
2
kym. n1'w9' T(da unmwumm 1)
iammija
[l0m.jd (1 (2-1)
(ownn. Te
CM2
V (-
c'|CC ink"
Imitl
1. (i 1)!(w( T "m.)))
6-1
")00e wx + + T (Ia
wh,, w T)
Vi!
lam.ja (1-
(pwing T "vo)) "oe))-o)
18
1)!(wng (winn) w
CM2
VI!
1 (1-
Vi!
c'E(C)
k=k*
1=++1
29

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Ind g.m.
((77))
A
(17')² A
d + n (1a - s ")='
2
2
Wm.ii
g..m
ry
im Be
(wg wa - g m.) + is am Mi )-1
CM
d'ElC]i=A i=
1 (17) A1dd 2d
(1)))\11 a + 17 nz (Ia " um 1mm.3"
(1- m)) )))
2
where (i) is due to Taylor expansion, Cauchy-Schwand inequality, and the orthogonality property of the Hermite polynomials. In (ii), we used the expansion of (I) and (II) and ||V "nt.. 1 (m(x m) mm (m (fe)yeam,, (w (wmm) De + bm.s)ll I A which holds with high probability.
T
In the same way, we obtain an upper bound as follows:
nj CM2
1+1 "gmm. I m.
Σ
iā Pe mm) 1 wa kt s mm) ) + is a Ti (i ggm.
d'E(C] i=
In
(17)) A1 + 17 wy (Ia wm wmm T) Wm.""
-1- (kg. )))m
2
Similar to Ke.m.j. we obtain an upper bound on the difference of kg.m.j over single step.
1ggm Ni ml o/ lug" (Ia wm. V) Wee, (m((x)) = m) Tim (m (w'ms Ie + bm..))
V wm.: (m((ie) m) Fmm (T (oee))))emmmmmm.. (um De bm.))ll
(17))3
[wg (Ja - w wmis T Vww,, (m((1e = m) Tam (Ie ym (wim Ic + Um..))
+
2
V whin 1 (m(xx)) m) am (m (wim De bl,))l 2
Ö(ee).
Note that E Wm.. are mean-zero sub-Weibull random variables, and their partial sums exhibit strong concentration behavior.
Weak recovery for the corresponding cluster. Building on Lemma C.7, we establish Lemma C.6.
Now we show that the mean vector pv, does not significantly alter the Hermite coefficients of the activation function when ne m ((d
Lemma C.8. Suppose that |ve T win..! = (d b), In)mm ò and Ingg. = Ö(d b) for all 0,1....,
-1) Then, by setting n¹ ne I agd 4 obtain that [ātt1 m.j.hi.c am.jii = O(d b) with high probability.
Ö 2), INC,m.ي Ō(d-b) and [ng. (d-+t for all s = 0, U,1...., Then, by leveraging the evaluation of the gradient update in Lemma C.7, we obtain that
Proof. Consider the case where amji 0. Suppose that l0s "wwh
v'p CM2
|ec " [t 5 Wee T "hin | +
Vi+ Tā Bec (H m)) (we Twing)
+| n'P
Σ
Vi Is mnsam (H(g.mm) (ve "mmm))?)
CM2
1=k-
Jee T Tw w A 2d
(
3d
+ n'ue (Ia win, wins })=" wm.
2
2
30

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
rpp CM22 p"
I 1vo T "umn j
+
Vp' + [(max fa ((
Bespeeeesssssselllaaam + max (s f.
==0
+ 1e ToTw
w
A
Ar + +/ V Wi, Wmimd ')=" win.
2
2
I |vo T "n +
NeP
Vp* +
[(max ( j. Bo,il max In 1k" + max (s Mi max In mm 167)) !l
CM2
6,8
max, lvc T Wm. ((1e) )A2 (1e
dd
+t
7e (Ia ; T) mm. 8 0
2
2
b)
where the last inequality is due to m m Yal max, !n)mml ) Ö(d-k*) max, 1ve. w ((n,)* A1dd ad
CM2P P" vp" + (mmxx, j...j...c max, Imām / +
Ö k- + +2), and (,, 3d
O(d ((d=4++3) +8), ). since have nh sad $, loo wmyl Ö(d- and ke.ma. Ö(d $). Note that le ww wm = O(d with high probability.
Additionally, when S O(dk
Σ s=0
7 T (Ia - wm "imm T) "mms Ō
with high probability. Recall that
(lamdid (1-1)
(pwi.. Too))-))
i imm.j.e
(1 (
VII
Vi!
1=+11
Since pie Tu "mmn = (dd-))d))))))))))))))))))) $), the series Ei=i+1 (lomasd (1-)) T Nee) " decays exponentially with respect to d.
Together with w = Ö(d-à 1) and k&.,, O(d 1) with high probability over initialization randomness, obtain j amgil = (dd-) with high probability by induction. The proof can be similarly established for the case where Omgi. < 0 since h²t is even.
Furthermore, we demonstrate that when vi.m.. grows asymptotically larger than those of other neurons whose K Kümmj have not increased significantly, Ve Tw",, "mm.. also becomes larger, leading to increase in the Hermite coefficient am
Lemma C.9. Consider neuron which satisfies Omm.j,k" Be.k. 0 and J. > 0 for k* i S p". Suppose that Ive T "mm.s = (dd=+), Ke,m, = =((d=t), and [n@.m.) = ((d+t) for all 8 = 0, 0.1....,'S5
. = Ō -1). Then, by
setting Tf = [/ S. agd- 4, we obtain lā mmjiicll |om...sl = (dd= 2) with high probability.
Proof. Consider the case where am.j. s = 0, 1....................... .,, it, we have
0. Suppose that nem.. = (d(d-)) 2 O and Ing.m.sl = Ö(d for all
T W 2 ve
T "h
+
NeP
Vi + lā ti (e.. '' (ve T (wmm))?)
CM2
==0
b=
Th (k m.))) ' - (ve "imm." )3}
+
Vi + T 180 miie
k=
maxx
Ive ( 2d 2d
(16) (a))AAA" A
+
Σ
NoVe T (Ja wm. T) um.
2
2
a=0
31

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Basel Learning
2 Ue
Two
" m.j
+
neP CM
vk + lā j.kk oo m.) (wo T "wmm))
==0
MeP
Σ p" pp
max SS m. Pilm 16' (vee T win
CM2 8=0
maxs Me w
umm ((Te) A (n6) (
Ve T (Ia w 1) wmd
2
2
s=0
= [d ())
where we used
CM2 p"vvv""v vp² + max. Ise Milki mm = =(ddk')
(n +++ and max, 10 " (m
(n A, Pad 2 when
since we have Te S and "r, 1ee w = Ōd- and Ka.mjj = Ö(d b). In addition, -1), we have that o NeVe T (La - wmn wmm T) um.51 - Ö( ((d +) with high probability.
Note that Iue "w (dd with high probability, Thus, combined with the assumption that sgn (mm is the same for all and p = (()), it holds that
iō m.jic
3.1m...d Vi!
00
lammj (1 (1-1) (pw'n. Tua)
i0m
(dd #).
Vi!
Vi!
=fiid 2)
The same holds even when am,.i. < 0 by considering the upper bound in the same manner.
We show that even w/..) Tu, increases, the coefficient at e.m. remains bounded by polylogarithmic function in d.
Lemma C.10. a'jjii pwr Tuo = Ō(1).
L Eg (om.iio'n (wis + pwm.. Uo He) T =)) = Ö(1), where ~ _N(0, Ia) and
Vil
Proof. Define Z == w/..j Tz ~ NNO0 1), since |uum = Also note that pu'n. T de Ō(1), given ||e/|| 1 and p ((). By the Cauchy Schwarrz inequality, have
Ez am (Z+ + pwmm. T "ve + 5m He: s V Ez am. 2 pw Ve + bm.s) I'In Ez He, (Z)²
12lumm (2 + pumn. Vo bms) I-
We now provide casewise bounds according to the activation function dm. amm() (.) is the ReLU function (.()) = max (0.-).
Ez Jam.j (em.)220 20 (Z + pw T Ue + bm.j) lam.l-
If (' is degree-p polynomial d' (: bm)) = EP- Cg)- with p = 0(1).
p-1
Ez o'm (Z pwm Tv. bm.
Jamaly
Ez
CalZ pw T Too
q=0
p p-1p-1
lamJly
G,,C,IEz [12 + pw'n. T To)
32

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
2(p-1)
lamoly
Bz
((z + pwm Tuc)"
2(p-1)
11 R8 Z65 powin Tvol" Toal--s
=lammly
Σ
C ==
=O())amml
where the constants Cg. Cr, and Cu arise from the binomial expansion. Combining both cases, obtain
_ Ez (amoo (2 + ppmmm T + Hee = O(1)
0 demma,,
Vi!
as desired.
Remark C.11. For the sake of conciseness in the exposition of the proof, we omit the superscript in am Based on Lemma C.8, Lemma C.9, and Lemma 1010, the bounds the subsequent lemmas are properly justified, regardless the variations in the coefficients a..m..j.""
To prove Lemma C.6, We introduce auxiliary sequences that provide the following bounds.
wowwy Lemma C.12. Consider the expert E Me. Let w we S A4d = ( for all # d, w wy S A4d ) for all and of Tle S and $ For all 0, 1.....,, suppose that
Ke mth az
|:Gmml I Ni for all (c,5) # (cm Jin)
In S A2Azd 1 for all (c, j)+ + Fmn.
In leg.m.il
S Ni m.i mii for all j.
leg.m.l S A2Asd for all j.
Then, by introducing auxiliary sequences (Pr) and (Qi)tt) characterized as follows:
ppt Pf + CM3
kt am ke er Bor x- (1 with PP = (1 (2) mi
and
1=+1= = Qi + (1+ 02)
17. k* max max ee
3 3e
Σ
se am |(00))
CM2
c'E(C]
1 fa $ k- am.jk on Bcr, k (. 10- -6-1 d with Qi (1 + 0z) maxfmax Iicm Ing.m
Aa CM2
cer mi lower bounded by Pf for all s = 0, 1, .,, + with high probability. For all (c,5) # (Cmsim | m and In gg are upper bounded by Qi for all = 0, 1......tt with high probability.
Proof. Suppose that e mm.... S a2 | S ke in for all e + ct or j & Jü, and I Ingm S A2 for all j E [J] for all s = 0, l...,
We first deduce the auxiliary sequence (P
oo mi 2 Ke
ry" CM2 tE|C) k=
Σ
Σ
Simm (Ve ) (+ T Tw: xir mi (.
33

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
+ the )²
(we:.
wn - rie monim Rg m.mmiin
A1 2d
ce minn
Wen T (Ia w min Wm ')=² w
2
+
CM2
17
am ve on Pon (1 mommmmm
mi (M Me miji ) )k--1
p²p²²
max Jim MoPoua max I epem
-1
max l T w ofcn
M
"pp²²
1pp²²
max in max In momm c,
M max [sel max jmm 2,0 Milli m. 8,6
-=-1 max mm m0
M2
open
11'p²2
max [scl max de Will m. b,c
1
man (17)) 2d + 17 u T [ wmmn Wmin e Wm,j""
M
Ke mn + CM2
17t
am do oo Bogni (1
ci
mmm (M
m.mm )k-
ppp²²
max max Jim &,,e, max |scl max jāmjin s..7 mi ) (re,
M
1,c
mmxx w T "w'l + max In of
max open
ww wÿl + Ing.msi
ofem
T T(a w w wmns im T)=² uin"" mm 11" "aAi" 2d - + 17" u/
We used conditions w w = Old 1) for all + d' and w w = Ō(d for all (n )² $ az² S jag
x'p'² M2
max (maxe, jmm.t A, maxe [sel max,, lām.ston.. (nin ) (maxefct Iw2 w
+
maxefet, I + maxeyen I ugll (ng.mst.)))
ann" hit Cmmiink K or Ben (Neg mi JOMIN
yk
-1
and
ee m. 1 A1 2d === I 4CMsh k dm ez Bet x (1i2e m,.j
-²-
Hence,
A 1+11 So Mi.
2 1 minn
+
3 -02)
11¹
ki amiiinh fe Be it (K
(Nom mi )
k--1
CM2
17 "
(la
wmnjz Wmnii Ji
T)
Wm.
3
2 Kic m
+
(1
U2
ki Orm BBee (wrn CM2
4
=0=0
n W
(Id
w
m.nn woy
Wm..p
We bound the noise term. Note that = has sub- Weibull tail.
If I Ag
cm m. m.
2-22""
Σ
11 wo T (La wmas wminn )= wois S TTAIVaS S m.: S a2 min
a4Aino
=0=0
with high probability.
If s > A et mini
22-28
11" w T (La wh um wo' Wm.jo E Wm.3 I A1s
A - noSA (ne (m k--1
=0=0
a2
k2 a Jt fot Ber (Ned
(1 mm
4CM2
s'm0
34

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
with high probability. Therefore, for all O, .....,& k m.
can be lower bounded as
02)
--11
R mi 2 (1 Q2)
02) 1 m.i
+
k" S'm en Be, (1 cmmm
CM2
s' ===0
With the aid of an auxiliary sequence (PF))=0) where Pp (1 (2)) miin. and
pet1 Pf
11º
k2 Jm.iink k Peg h (PP)
CM2
Ki m. is lower bounded by Pf for all 0,,,,.... 1. Next, we deduce the auxiliary sequence (Qi)
For the upper bound of maxefe, jE m.cmm.h we have I - kmmm S Am with high probability. Thus, the sign
of s++1 Ne.m. is the same as that of n2 na. or hemm AiTe. Similarly, the sign of g.m.. is the same as that of nc g.m.j
s.mm S A17]e-
Fix m [M].
We show that the following bounds hold for all 0,1,...,,t 1.
maNefen, m= is upper bounded by Qi, for all j E J
1+1 where the sequence (Qi,e) a defined recursively as follows:
Ui. (1 + 02) max(ks zdt 31.
and for s 2 0,
0 = =Qie- (1 +1+ a2)
CM2 k" max jām el x (Qiiec) k--11
77th
g*- + As TI.
k" am.ik or Ber x- (w mish )k- d
CM2
with high probability.
In " is upper bounded by Qi,, for all j Ji.
1+1 is defined recursively follows: where the sequence (Qi,,) ==
Pi = (1 + 02) max(les mm 2
and for 0,
11 Qi, + (1 02)
7 1"1
Sermjik wh (Q 1--11
(Qi.))* + A3 r7"
k" ōm.jk cr Beg vs (K (o d
-11
CM2
CM2
C'E(C]
with high probability.
max is upper bounded by R2 for all j $ J =
where the sequence (Ri,)) (+1 is defined recursively as follows:
Rhe (1 a2) max max p - of
35

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
and for 0.
R = Ri,c + (1 + a2)
CM2 k" max Jām. de Se S (Rf, ) -1
77t
with high probability.
Ing m is upper bounded by Rí. for all j $ Tm 1+1 where the sequence (Ri s=o is defined recursively as follows:
Rf.g = (1 + 02) max mmm b
and for 2 0.
R Ri (1 + 02)
1-1 Σ
Sm of )
85
CM2
d'E[c]
with high probability.
Furthermore, for all 0.1....,t 1, Qi,e Vig Ri and Ri. are upper bounded by Qi.
t
We sequentially present each bound.
For c + cm and j E Im
|sctmag max { A1Mes Kemms
17a CM2
Σ
Σ
i Be, (Fe.
CE(C) t=
k 11m. A
Ar
+ Sem )
W
w
+
2
2
+ w. T (Ia um Wmis
T)= wmin
S max { ALeer | rem..
rya
11
p
i. Ben, s (m)) (w= wa:
(w
CM2
i=k-
io tmm.j.toeess 11= (n², )
+1
imm a m.) (wo)" wiy kam the kimmm)))
d'fc,ot, t=k*
im (w= wg]] +1
is a. Te
d'ElOTi=k"
C'E(CI
I ((7))) A1
A1
+
+ +w= (La win, W T) wmis
TTT(la
< max { AA mamm CM2 (p-² maxx maxl@mmsiiian Ben i (1i (Kie ) -1 -wit" wen
11a
+ k" max Beck || II 1k--1 + 2² max a jui Bel mi
Cp' max max(@m Ber d'stccct
max I
|k" d'
(max I w
max cipeee
m))))
di
Cp22
max S Vil | -1 d.
w. "Ll (lw,"
In ln..m.))))
36

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
1N Kamml (
2d
(178)* Ai
dd
+
+ 11 w/ (Ia - w "mmmmmg T) wine
2
2
S Smax { (AmTos Fommi
+ (1 +1 (2)
11
he max jūmmjikt Bel d Be |1
1k
CM2
c¹
11¹
)k- - ( (0 A1 3d
(1"")" dd
+ Ag CM2 ".M2h"mm... Am s Pog b- (Ki ( m.jin
2
2
+ :²= (Ia s T) '''umm Wmm)
S mmxx { AAmes max !nom.l
opem
2 + (1 +
6--1 1 .(2 ()
a2)
CM2 k" max @mmj.. e Bet max
max oftn
me
3
s'=0
+ A3
Σ
k" Öm.jk com Bep k- (152 min
1e
d
max
Σ
np We (Ia w mi wiis "="" wml) B
CM2
8'=0
Since, w uh Od for all + d' wy T W Ō(d-2) for all e, k.mm.j "q. = (d- for all j E [J], the term ks max;|ām.j.k ,ae,k mmmam
(dd for all (c, * (c Im) and subsumes the remaining terms within the are subsumed with weight Joz (,
expression cMa2 (.) with wight jaz in the fourth inequality. I (r)) 2d
2
2
by 11" < and & and mm A2Agd & in the fifth inequality. For the noise term, if s < Azdh
Σ
np we T (la W w E Wm.
1 S SnAAIVs V S jazd
max ofom e'=0
22
with high probability.
If s > Agd -1
02 k² 3CM2 d jmm.j Be w n2 at= -
n w T (Ia w Wmis T ) wms| S AcAiVSs
max cafice
with high probability. Thus,
max INomml exte
S (1 + d2)
max
max Ino offer
m } }}
1 k-
+ (1 (1+)a)
CM3 hi max (amm et Ber xe max max estem
Inee
Qi. ]
s'=0
n
+ Ag
k" ām.j.k .c,, Beg, Ben (1 mijon ) 3
CM2
x'm0
In contrast, the following inequality holds for Qi
1p¹
@ < (1 + (2) max max I ofce
5 d 3) + (1 a2)
Σ
k² max jāmj,k Be ((Qi"') k--1 *
CM2
e
8'=0
A3 Σ =0=0
k2 (m..jjkk on Beg be (Net or mina
k2
d 1.
CM2
Therefore, by induction, we establish that max IN5 | is upper bounded by Qi, for all j E Jn and s = 0, .....,t+ with high probability.
37

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
We apply a similar procedure for g and j E Tm-
Ing.ma s max { AAiTes
11¹
Pe ( w (om
(wz
CM2
C'E(C] i=k*
INc h
1(10) )
(n))) dd
+ Be (1
(1
(ni (m..m.))))
2
2
+ wzz T (Ia wm. Wmnj T)=²
Wmi
wmis
5 max {
AAiTes gg. +
11.
ii om Beg, if (K m... ) (wng Wenn - g m.)
CM2
im Pe (1(.m m (wg) w - ki m. sti m.))
d df i=k*
p
is o. (1 (Kig ((mmm))11
e'Elcli=k-
+ N.gmm Ani2) +
(172)
A1
3d
W/g (Ia - Wm. w ') wmms }
[3'wm//
2
< max { (Aiees ket mam CM2
npe
(R* atmmmss oz Bezik h (Ki mn) ) -1 (110g woz + mm)
+ p² 2 ma.iim Ben ol( m))) (lug woe + (ng.m.))
+ up2 max am Be efc,
max 1K miml" /k c#tom
(max cfem
lni u + max In mll) e'ecr
+ k-
Σ
2s am o - 1x--lkng.m A--1 -1 + Cpp
max
(se V.ll. ei
e'Elo
1ggmml(7") A(7))) + (17²) A1
(Lu wmy w T)E² wmm) }
11 nuug
S max { (AiTes gmg. (1+ 02)
2
k*
Se Tk* Imax & (Ing 1--1 1 , 2
CM2
c'ECc
+ As Σ
112 CM2h k2 am.ikk eom Bor, k (Nig m.jj
d +
Σ s'=0
nuu (La wmn Dimmms
s²=0
Since, w T w = Ö(d=t) for all + d', w w W = Ö(d for all c, kam ( for all (c,j) * (Gno and
-1 subsumes the remaining terms within the expression X (+) with weight ja in the fourth inequality. Int.m.s 1(n" 3d
"g.m.. = ( for all j E [J], the term k2 max, Maxx/Eeecy Sermm 26- mm
"A, a are also subsumed
2
with weight jāzz by 17 S a & and | & Agd 1 in the fourth inequality.
By providing an upper bound for the noise term in the same way,
2azd- 4, if S Aza daylas 3CM
np wny (Ta - w T)=² Wm.j
<
k*) Eeeeco 8 Ser ,e ((d 1
if > Agd
with high probability. Thus,
IN mg. S (1 + az) max (Ing.m. fd }}
2
38

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
n
max { [ng.msl vig) k-
(1 + 02)
E
CM2
CM2
k*1
Σ
Sel
s'=0
e'ElC]
A3
k" a.m. com Ben, n (Ni ma
) d
K' -
CM2
s'=0
In contrast, the following inequality holds for Qi.gi
11g11 & (1 + 02) max mmax
( g.mm 1 + (1 02)
11
T Mk (Q 6-11
CM2
opem
s'=0
e'E(c]
A3
k Gmig on Bog k (15 midin
w-
d
CM2 s'=0
Therefore, by induction, we establish that | ااا. is upper bounded by Qi,g for all j E Jú, and s = 0,1,.... + 1 with high probability.
We also consider similar argument for j & Jm
Instrag max { (Ames Mams + ry"
ia (e, 1 (w:T w (.m.mkm..mm..
CM2
cE(C) i=k
Inc 1e ((n")) (1²²) (11)" 3
(w: Wig - ne m...
2 +
1f wh T (Ta- ums wis T)=² Wmil }
S max { [Ames | Remms + CM2
1f"
1 im dmm.stheBesss (1 - (Kr
if ( mm)) (ww wys kir Kid.
i=k" c=
is J. (1i (Hig..m)) (w= (gg) +
is am ( (ngg.m))))
C'ECC) i=
e' e'E(C] i=
(171) A1 K2 Mammm(rr ( 21
+ wx (la u T)=² Wm.
2
mmx
A
17"
(k' (k"āmmji M -oBa I1 - p²² max (mmsstaeesllllmmm.
CM2
Op2
max d'pe,t
am Ber c'yt
max (Ne mmm 1 (max luw T wayl + (nG dtt
+pp
max j ji Bor allki.mmm) 1/" -lw Twull + lmg. c'Fc,i
IN (()m) 2
dd
+ w (la w T)E" wm.d
2
2
S maxx { AAlle
max o ofet
+ (1+ (2))
17
k" max l Ber
CM2
max
max ofce
1 (d
1y w T (Ld - w " T) wmal + max
of
w w Ō(d-2) for all c, N.mm Since, w w. Ō(d for all + d', w
Ō(d-2) for all (c,5) # (ch Im) and
net k.m = Ō(d=t) for all j E [J], the term 1s max,|ām.jit j.k' B.Ba 115.mmm)
16- subsumes the remaining terms within the
39

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
expression Cara (-) with wight a in the fourth inequality. |K2,,, (n7)) 5az by rf I aqd क and kk.mm A2A3d in the fourth inequality. By upper bounding the noise term in the same way, we have
( +
are also subsumed with weight
2
2
222 8, if S A2d
Σ w (Ia i wins ") wm. s'm0
02nec k maxe jūm.j. of Be N 1-1( () 3C
if > A2d -1
with high probability. Thus,
mm _a })
max Inmmm (1 + (2) max (max cE|C]
ofem
Immm Ri,e] x--1
(1
U2)
CM2
ks max jāmigh oe Bet |max
max epcn
In contrast, the following inequality holds for Rije.
Rie S (1 + (2) max max
(1+ me d +1+ 02)
CM2 s'=0
Σ
17
k" max jām.j. e Be Be (
K-11
Therefore, by induction, we establish that in is upper bounded by Ri, for all j $ Jm and & = 0,1, + 1 with high probability.
Insti S max
(AiTes
mggm +
iāmdi a (m)) ) (wg we - kt
CM2
e'ElC1 k==
).) 11 - (n (3) + In (0m)))
An
+
2 +
2
+ w/y T (la - win Wiis T) wmil }
I maax/ { (Aiees mg.ms +
n1a CM2
i Bet, (1)), ( w ki -
C'ECc i=
is ame.ee Ti(k mm. 1 - (K (KG.m.))3))
C'ECC) imm-
1 1(7") And (n7))
+ w (Ja wm,s Wmis T)E² mmss
2
2
< maax (Aile
Ni Kigmm
+
CM2
17a
(Cp"2
max jamisshee Ber, max In22 mms" (max lwo " + Ing.m.s)) c,
k*
Sel am..i te Tk- Tk [|ng. -1 + Op22
max (s Wilgg.mm sl
d'ECc
Ke g.m.
(
(17²)s
Al
3
w (La
w un. T) "mmm
2
2
1jst kil
se amm e vk- max (nrmmsl 1'1 1 (d
< mmx
ALTe: max ofem
+ +11+
CM2
Σ
3
s'=0
d'e[c]
+ max
11
W
T (Ia
w
m.
E wm
40

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Since, w w ((d-)) for all c', wa w (d-bb) for all e, mamm ( b) for all (c,j) + (m and ng.m.. = 2) for all j [J], the term k max: Eeel Serttmm se To |1 jk' -1 subsumes the remaining terms
within the expression cara (.) with wight az in the fourth inequality. 16 ( + (o") A. are also subsumed
with weight 502 by 1/² I aad & and m I A2A3d in the fourth inequality. By upper bounding the noise term in the same way, have
agd 3, if S A2d
1f wng (Ia - w Wis )=" mmms
S
if Azd agtlee k. "E 8 Sr0m (em.itteeeo--l =) 3CM2
with high probability. Thus,
1 g. $ (1 + a2) max {
[ng.m.ll
} d x
2
8 Gm of Tk- max { [ Rig k1
+ (1 02)
Σ
CM2
='=0
<E(c]]
In contrast, the following inequality holds for Rij
1.g < (1 + 02) max max Ing mgg.m 2d dd + (1+ 02)
11
k21
Se , No Me- (Bi
k'
CM2
of
s'=0
c'ECC
Therefore, by induction, we establish that I=s.m..] | is upper bounded by Rí for all j Thh and 0, 1, + with high probability, Finally, we consolidate the auxiliary sequences Qi, Vi.g 12, and Ri into unified auxiliary sequence Qi, which serves as an upper bound for |n and (ngth..) for all (c, * (mm))n)) where m E MM.
ve k jm.jjk" .ct, Ber e Clearly due to Ag E C. by the upper bound provided by Qi, and the upper bound provided by Rf., is subsumed by the upper bound provided by Qi.g. Consequently, for all (c,j) # (cn Jm), nem.. and Kg.m is upper bounded by Qi for all s = 0, 1....,, + 1 with high probability.
(N (No......... )
--11 d > O, the upper bound provided by Ri, is subsumed
(Qi.g) is expressed as (+1
oc+1 = Qi + (1+ a2)
CMak k" max jām.ght ee Ber E.
Seāmj 7h 111)))))**
e'E(C]
112
k" amijkt er, Ber, A- (natt min - d
+ As
CM2
with oh = -11 + a2) max (max, IN IN g.m. ld
Based on Lemma C.12, we prove that (ne.m.g for e + ct or & Tm and lgg.m.p remains upper bounded throughout the trajectory by induction.
Lemma C.13. Consider the expert m E Mc. Let wt w/ S Add O(d-1/2) (/2)fe for all & c'. w T " < And-
1(2) for all e, and 17" = ne S and 22 For all s = 1,1.... t, suppose that
ni mim S Saz
max(Ing Iigg 1 S Ni m.s for all (c,j) F (cns jm))
Ing ml S 4 for all (c, # (Cis )

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Then, if we have p-=1
cm m.. (22
max[Ins f+1 ( S Ki m.n. for all (c, # (m),
max (Inetm (n S 4Azd 2, for all (c, + Jm).
hold with high probability.
Proof. To begin, consider the the case when
(1 + 02)
ks max j. ee de &e
Σ
Sm (11))
CM2
d'E(C]
A3 CM2 k ji ,e, Beg, N= (4 dr
holds for all = 0.,,
Due to Lemma C.12,
lf+1 s Q + (1 + 2a²)
7
k² max Gm... .et Bel Be
Σ
Seam (000m......m11100000000 ((0))*
CM2
e'E(C]
holds for all 0, t. By applying Lemma A.4 to Qi
Of
(1- (eh 2) + 3ag) max/max.: jām.j, Be x= EEvejo od -|0 holds for all 0, 1, where we used (1 + (1 + 20z) 2h ke max(maxe jām.jj e Be
Eeetcr 8 Se oe of 111 -1 14202
By applying Lemma A.4 to P..
Pp
Pe N
(1 noki (k+ 22 a2), j.k sn Beg (--M- -
holds for all 1.,,...tt.
From Corollary C.3, Lemma C.8, and Lemma C.9, it hold that Q < Pp and (1 +
3az) max (max/ jām.j.k' (max maxx(mmxx/// m I m. }d=b)) < <1 - (1-a2)am a2) se,
Edeco Seāmjk v .
so, Beg, d (N mi
mi b" which establish that Of < Pf.
Thus, Q++1 < pt+ which indicates that max (max, s In b Ni miin for all (c,j) A (cmo
Since P < K mm I a2
t < Ne (PP)((PP) -k- -(22) (a2))=k* -
ks (k- - 2] a2) a2) j.k Bem
ne CCMM2(1 + 50z) max(max, l e BB se I.I ie +2 am enn ke
k* 21 + 3az Jmax(mmxxa. jāmjj, Bel Eneeo Beamiik of
In contrast, Q++1 Agd- holds only if
- CMM ((Q9) +2 (Azd 1) -2422
-x++2
12 b:s (1=" 21 + 30z max mmxx jām,s Bor 1. 8 am "e Te-
42

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
7 (C (Q9) +*+2
kz k* 2)(1 + 342) ma [Ckm,jA-" Be Ar 12 2 Se. of ne
where we used (Aad- )) - I a2(Ki mmji ) ) (2() -k' (P( S a2((1)) -2
From Corollary C.3, Lemma C.8, and Lemma C.9, we have
1 je of II maxxmax lammml Ing.mll 2 5d-$3) 2)) -'-2
(1 + Sag) maax)axx/mxxx ke Be, k.ll
Σ
e'E[C]
< (Om.j.k i Ber, dk (1og (R min
However, this leads to the contradiction that
max j se Be Re Eld Seamstr (PP) +2
be
Ne (1 -22) +2 (1 502)
Om.ji en Ben
Thus, Of > Azd does not hold, which implies that Qi < Agd 2.
Next, consider the case where
(1 + a2)
"1"
k max (max Jā e Set Ber k-bb
se Smm..j ((Q9)
CM2
e'E(C]
A3 CM2
17*
k j. / or Beg vo (11) d
holds for all
0, ,,,,,,,
71 but
(1 + (2)
17,
k" max (ām.j.k" et (e
Σ
Se am (()))) (
CM2
c'E(C]
A3 17º
CM2 k" Cm.j,k oo Bo, Mo (PP) d
holds for $ t. Here, suppose that
O++1 S Qi + 202
Te ke "m..jk st Beg (1 CM2
holds for all T I,..., T25
Then,
T2
0 -TT
2Ag
Te CM2
k' amj.k* et Ber A- (Po -ld
T1
2Ag
(prxt (patt1_"T) PT1)
(1 a2 d'
ām.ijk cg. Ber
Ag (1 + az)d?
pm
max jāmjkk ot 30 Evelg 8 oe
2Ag (pr2+1-1) 1 -P
(1 - 1=)2
ām,ih ot Bog A.
Aa (1 + az)d
p
m mmxxel jōm.j B Eeeeg 8 amajjieeeeee
43

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Thus,
(1 + 02)
The CM2
k" max(max ee Ret Σ
Semmmi N (QQ2++1 f c'((c]
?le
k- jm.jjkk Sa Beg d (P 1 d
2 A3 CM2
We have O++1 S Pt since
(1 + 02) maxfmax, jāmgi .o EeE Se am e dd) db) मंं 2
A sot Beg
Therefore, obtain max(max, |nemm lng.m |t I In the same way,
Ke min for all (c, + (m)
s+1
=--1="" + Σ
24g
ne CM2
hi dm oo Bog b- (PS) -Id
S=
2A3
$ 0 +
(pt+1 - Pr.))
(1 a2 a2)d3
2Ag
$ Q7 +
(1 - 02 )d P(11
$ Q7 + +33dd 1
s 4Agd 3
where the last inequality is by Q7 s Asd from the first case.
Finally, we establish Lemma C.6.
Proof of Lemma C.6. Suppose that Ti L(nck 2)(1 5am) ,et, -M-
and
Ke, J. S 02 max m. lngmm.) < k² ".g.
m m.jt and max ( l I Asd 3, where (c,j) * (cin
for all 8 = 0, 0,,,,., Tj. Then the bounds given by Lemma C.12 and Lemma C.13 hold for all = 0, 1,.,, Ti with high probability.
Thus, by Lemma A.4 and Lemma C.12.
P
(1 (k* 2)(1 u2 0m on Bo k C-1M-2
However, obviously when t- Ti.
pd
2 Pf >
(1oc (k** 2)(1 a2 jm ot Beg A (P)
2 (nch (k+ - 2)(1 a2 vr)) ott Ber &. C-IM-2) MM 1.
This leads to contradiction as
T3 oom min < 1. Since max Ing kg.mm. S 4Ayd 1 from Lemma C.13,
Indt oe mi
11-1 Kep Mi
S Allle S Ala S Ag
Thus,
+ b
h-1
I S 5Agd 1= = Q(d
min
44

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
C.3. Router Learning Stage
This subsection dedicated to proving that, after the exploration stage, the router successfully learns dispatch the data x to the appropriate experts m M, with high probability. At this stage, employ single-step training for the gating network, as in Damian al. (2022); Ba et al. (2022); Dandi et al. (2024a).
Lemma C.14. Take T7 Tr S (55, n (d)) and T2 1. Then, any m Me satisfies hin ( +7++ maXm Be (Ze; +2)+ < O with high probability.
To prove Lemma C.14, show that, for experts not belonging to the set of professional experts, the alignment between the cluster signal Ue and the weights of the gating network 6m. represented as i,,m == vc 0 is upper bounded.
Lemma C.15. Take nTT Thr I az, n ((d), and T2 1. Then, for all m & Ma. we have
S
c.m
Ti -a² k" a5.Appl (c)rs) < max [[M] 'c.m" a2 k' az Aspll (c.MT))
Here, we describe an important property of the softmax router: by initializing the weights of the gating network fm zero, we ensure that their sum over all experts is also zero.
Lemma C.16. Eme[M) 02h,+22 Eme[M] 0 where T2 =
Proof. We define the gradient with respect Om for an input pair (xta v²)
as
Von Le = Vo... I [1 (m(x²) = m) nm (1.)))fmmmm))
= (1 (m(x²) = m)) - nm (= Di
The sum of gradients over the entire batch is
1 Σ Von Ce,s Y1
(1 (m(x²) = m)) Fm (mt) (a) 0. mE[M]
Σ
7 ME[M]
This result, combined with the initialization T 0 and the fact that Eme(an T (m())) 1, completes the proof.
Based on the above properties of the gating network, we establish Lemma C.15.
Proof of Lemma C.15. The population gradient for the gating network of the router can be expressed as
-Vo [C = Vo E, . [1 (m(xx) = m.) n. (fa))efff
= E, Bey (1 (m(a.)) = m)) um ( (Fe (Fr))t))
We decompose the key components of the population gradient.
E.E,, (eeffm(z.)
Σ Σ
B..
Σ
D"
Bes
Hew
(w.TTa xx) 8c
Σ
Y
He:(w xn))
Σ
am (em xxc) xrc
CJ
Vi!
Vi!
Vil
JE[C]
k=
E.
Best
Hos 2) +
Ho₂
(wo *
T 2))
Om i He; (Wiis (2 + pue))
CJ
Vi!
Vi!
CEEC
JEM
==
+E.
p²
Bes Her TT) 2) + 8c Σ
2i
He:( T2)
am.ji He, (Wm.) T(z + pro)) puo)
Vi!
Vi!
E2
1B He; (w:² T =))
am.i.
Hei-i( (Wm.s T=) Toa )w?]
CJ
CE|C j=1
i=
1=U
45

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Eg
8c
He 1 (wg T
2)
Σ
He; (ee=i(Wm. Tz) (p
Tuo)) Two)))u ww²
(I)
Bus lew((w TTTT20 22)+ + 8c
E
He, (w) 2) 2))
Om.ii
Σ
(i - ()HHe) (Wmm.) T (ow Tue) (orommsTTee?? Wm.j
+
k=
km
1=0
1=0
Eg
Bes
He 2) + se Σ Vi!
Her(
2) () amiii
Ho; Heg-i(ums xpwrm. Tue))) (oc)).msTe)))))))))) pole
vi!
(IV)
where the second equality follows from I.. = = + pMle, where = ~ NO Ia), and de Twy and no wy for all (c, c²). Third equality follows from the binomial expansion, Stein's 's Lemma, and integration by parts.
From now, we look the alignment Ve T0 Thus, (I) and (II) are negligible since u w 0 and ve Twğ for all c') We expand (III) and (IV).
1- 1-i-1 1 (owms Toc) 1 Be,i (i + 1)!(wms ww))mmm Wm. Vi!
am.ild
(II) =
VI!
t=
1=++1
Okm.ji!
(proms Tuc)
1-i-1)
867 (i + 1)((mm T "wg) wms
Vil
1---11 1 1
Vi!
1=
1=++1
n¹
p
Σ
Vi + j Reees "mis +
vi + ji f. (F((ggmm))) Wmir Σ i=k
k=
Omm ( 1
VT! (1 1-i) (pwms T Tvc) 1) Bad ill (wmms T "w") pve
(IV) =
Vi!
1=
p²
0
am..j ( 1
(-1)(au)))) (pwims Tuo) 11)
1
) 8c
il(wmms "wi't pvc
VII (-
Vi!
F
(e pee ΣΣ
Σ
( ( press
i=li=k
where we used the orthogonality property of Hermite polynomials both (III) and (IV). We introduce the discrepancy Ee... between the population and the empirical gradient.
Eo... = V8 Canhh + Von Eyjel Ceaal
Next, we evaluate the empirical update of the alignment. Note that c.m vo T and nm = Em exp(e
exp(ern Tane)
Ti+Ta
.T
- T
Ve
Ler
c,m
Ti l.m
+ 71 Ven E.E. [1 (m(xe) F. (m (xx)) + + Fiv T Te...
Tive
Ti em
B [1 (m(xe) m) T. (1 W S F ( )adl
e'lc
46

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
v
17
Bx (m((e) E Ma) Am(az) (Zee Ji (F We Smc of (Sarr )atee)
C
d'E(C)
Σ
Be, [1 (m(xe) & Mc) mm(ay) (Be (m We f )ee) = 'B..
C
e'E|C)
We derive upper bound for (A)] and [C)), and a lower bound for (B) for & Me. We first derive upper bound for (AA and (C)]-
T CJM
((A) S
Σ
Vi+ la ma a 2mm))) (wiis T ve)
C'ECC j=1 M.
+ Vi d. sel (ggmm))) (wmm ve)
Σ
Σ
am.ss Ts )' @m.jice (img)) pl
+ CJM 1=1 m. km
5 1/ p Vp* +1
max md M. di
jmm Set
max mEMenn's mn
max Io T mgMori
Vel
+ Vp" +1
max mgM. 5n
max In gigmml meMoii
max m.M...
T lwmis vel)
n
0
T' mn p ptp pd m&More,
max
max
max
! p)
C
m&MuuJ...c
m M.M.iiie
mgMoc.
= Ö( +),
TTT
(C) &
Σ 3 Σ
Vi la Be (wm ve)
CJM d'E(C] j=1 mgM, i=k"
Vi + lâm (1 (wi.
T vo)
F Σ Σ
Σ
Lm Ban (w mm )''p + ampssse mmm p
CJM
1=1 me
m k=
1/ (p" Vp* +
max m.M J.
lmm.jsi.e
max M
In m.mjj
max meMoit
lwiti val
+ p" p" +1
max mgMeat' Js
a (mm.
max mEMMcw
Ti gimm
max mgMo..
I mis vel)
n
p
max meMierhiiee
a Jm..
max mgMcced
IN Ti cf p+p meMedhie
max
max Jmm.
In mgMed
In P)
C
Ö((nd
where we used (2(m) m) = and Tm (xee) Tm(xy) (Ie) M for all m E [M]. Furthermore, by Lemma C.6, all terms in the RHS are upper bounded by Ö(n, since < Ö(d for all (c, # (Cr and n.g.m.." = (dd-))
for all j [J]. By Lemma C.10, Jām.j,t.cl are by at most ((1)
Next we derive lower bound for (B). Therefore,
(B) CJM3
ΣΣ eE|c j=1
Σ
Σ
Vi Tâm s.e Bet Bet, (nd (wune T (ee)
E.M k=
47

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
+ vi + lam' J.s Set (o gomt a) (w T va]
1
Σ
amm Sii aimm 5) p am' Si 9.mm g)
CJM3
j=1 m E.M. l=
(am'liiiik Back (Fem K p)
CJM
m'EM.
-JlMelp"
max Jtim
Jám' J.h.Beed
max
(n sll*
j7 #Jm..'MMM,
(mamm) k P
-JIMelp" Jie
max
amm S.
max j.m'
T max wo vel m'j
-J|Melp" vp* +1 max d.m di
jmm' J. max mn Il)"
d'mml
-JMMelp vp* +1 max jâm & d.m Jid
Ti T ( ml max mem s vel)
23
m'EMoi
DrP
a2 AgΩ
CJM3
0,
where used that k..m.. 02 for all Me and j Tm âm' Jin ,' ,BB,, 2 As, K2.ms (dd-1/2 for all (c,5) * Jm) and ky " = ((d-1/2) for all j E [J]. as shown in Lemma C.6. Among the terms on the RHS of cJis ), all terms except for the first one are smaller than the first term by at least an order of Ö(d- 2)-
Thus, the case where m $ MMo we have
c,
T1+T2
T1 c.m
02
Adsi(
CJM3
PPP
+
77,00 T
<
Ti
02 Ano
TTP
+ +
/ve"" To...
c,m
CJM3
< -2
k' A
CJM3)
MrP
+ + =00ml
which holds with high probability.
Since |ve To.., fecia Uee T [[1 (m(x²)) m)) ) Tie ( Ö(1), by applying Hoeffding's
inequality with n = Ōdd) sample and C = 0(1), we have
C
Go -mm
=
Ue
pm. 1
Lach + vo
Va..
Σ
E Cool
Ö(d ))
with high probability. Therefore, by noting from Lemma C.16 that
max T2
c,mm 2
Σ
T2 c.mm O,
m' [[[M]]
mn E'E(M)
we finally obtain that
Ti+Tzs c,m
-as 12 Ass
TrPP CJM3
< max
camm [[M
TI+T: :- Agl TrP
CJM33
m M[M]
max
c.m' 02 x-
TitT
05 Ag
CJM3
We establish Lemma C. 14
48

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Proof of Lemma C.14. In Lemma C.15, we demonstrated that the cluster signals of clusters not assigned the router's gating network aligned with high probability be negative. Here, aim show that, upon observing new data Te pte where ~ (0, Ia), the data is not dispatched to the experts not assigned to the corresponding cluster with high probability. With high probability, we have
110 10 +T -
110
0% = 7r
Vem n
Lociee
=m
Vo
1
Esjel Sa, + = )|
C
c=1
511
1
E, ByjckCal1
+ T/ |
C
$ ag
where we used the previously expanded (I), (II), (III), and (IV) along with ||w2|| ||gll Te pcc and 0T: NO, | (|?), it follows that, with high probability.
|m ||
||ve|| 1. Since
8T. 2) (sA5s(pp)
For Mas we have, with high probability,
(=
021+12)= 0211:11
4T2+ Dc p1
T1422 +8t 001++22 asp (
a² Apps2
1 CJM3
+As0( +AgO(Vlogd) <
<0
where we used p a2*' A&ll oonrs)
A.O(
By combining this result with
EmE[MM +Ti+ 0. have
hm (x 22++) z
1
Am (x (zit
max
m
Me
m'
Σ
Am (Iri Ti+Ta
Mel
m MM
S-azt as Ag52( G
M - MMM
p²
02 K' Ag
MMll
CJM3
0
(2ci (T++22)
for m & Ma as desired.
C.4. Expert Learning Stage
In this subsection, we discuss the individual learning of experts in the context of receiving data from their assigned clusters with high probability. As Appendix C.2, we introduce A; and a; with the following order strength, but they do not necessarily have be the same.
A S 2--22-1 A2 x A4 04 < Ag Õ(1)
The proof in this subsection follows the same structure as that in Appendix C.2 and is based on the proof by Oko et al. (2024a). We define as f Ti T2 + 1.
49

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Re-initialization. Before entering the expert learning stage, the expert weights Wmj are reinitialized. Although this initialization is not strictly necessary, it is performed ensure decent path for the alignment that the product of g" and the activation Hermite coefficients is positive. If the Hermite coefficients of f and g are identical, aligning with wy may become challenging, as w is already aligned.
The re-initialization satisfies the following condition:
Lemma C.17 (Following Lemma2 of Oko et al. (2024a)). When J 2 & log 5 for each m in Mo. we have at least Juin neurons Wmj such that
M.
w "m.
T We 2
1 Vd
18 (wmni wx -² 2 (s | T wgl 2 + ar(wms T wx)
k-
k2-2
with probability at least 0.999 with sufficiently large d, where ae is a small constant, where de (log Likewise, when J Z Jmin polylog(d) for each m in Ma we have at least Jin neurons Wmj such that
T Wm..
9
2
|se (wmm.
T 10" wgy
2 1e Ilms wel k--2 + aa(umm wg)
Vd
with probability at least 0.999 with sufficiently large d, where ag a small constant, where 0g S (log (logd)h* d)k--2
Based on re-initialization, we define the set of neurons that comparatively align with the indexed features w and wg.
Definition C.18. We define the set To as the set of indices j that satisfy the given conditions:
Jn = T=={jell jE[J] jj we 2 Vd
[Bax-1(um
T w
2-
2 (een--llw% wel k²
+ adlwin T 2--2)
w())
Similarly, define the set Ig as the set of indices that satisfy the corresponding conditions:
Jg = { J j un
T wg) k'-2 2 118c,--ll
w k'
T k' + tg wgg) (
|scTe=1
Vd'
Remark C.19. In the rest of this section, will discuss Phase III and IV on the event that the re-initialization was successful.
Adaptive top-k routing. Importantly, in this subsection, we conduct an analysis similar that in Appendix C.2, but employ different routing strategy. As demonstrated in the previous subsection, the router does not route data to experts where m & MM. however, it cannot definitively determine which expert among those where E MM. should receive the data. Therefore, resolving conflicts within m E MM without knowing M. or ve requires an alternative approach. One solution is to choose experts for each F, by the following strategy:
Expert m is in top-k if and only if (x) 2 O.
The complete MoE model, incorporating the adaptive top-k routing, can be expressed as (Zei {am) ) == EM 1 [hm (x.) 0] fmmmmaa (xx). This routing strategy mitigates load imbalance, which would otherwise disrupt data routing among experts m E MM under top-1 routing.
Lemma C.20. There is least m Me such that fm is correctly routed with high probability over the randomness of Te. In other words, for some fixed m Mo on the randomness of Ic,
P[ze is routed to the set of experts including m] d
In addition, fnn is never chosen for all m & Me when given De with high probability.
Proof. By Lemma C.14, with high probability, we have
M MM-Me M '² MMl CJM3)
)
max hm (xei 0Tx) 2 as 1" asA
>0
50

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
and
0² <0
hm (Fei 022) < -a2" a5
CJM3
for $ Ma.
The proof is complete, since employ adaptive top-k routing, where expert m is activated if and only if hm (2t)) 2 0.
By Lemma C.20, the overall MoE model is, with high probability, equivalent to EMe fam ImEme) == EmEm, [hm 2 0] (xa).
Following Oko et al. (2024a), we sequentially demonstrate the following:
For 13,1 S T3,1, Je neurons achieve alignment of ((1) with wy and 1g neurons achieve alignment of (2(1) with " i.e., weak recovery.
For f3.2 I Ts,2, Je neurons achieve an alignment of - O(1) with wh and J9 neurons achieve an alignment of 1 1 (()) with w.y.
In a total time of (Ts,1 - t3,1) (T,,2 f3.2) T3,3, Je neurons achieve an alignment of 1 with w and Jg neurons achieve an alignment of - with wy i.e., strong recovery.
Weak recovery. In the same manner as the exploration stage, we begin by evaluating the stochastic updates for the alignments, ni ke.m. and nit g. Subsequently, we derive the lower bound of keem.. for j Je and ky.m., for j Tg as demonstrate that there exists
well as the upper bound of mm Il for j Te and In'.m..l for j E Jg. Furthermore, point in time when k.m for E Tr grows to constant level, while k.gm.. for j Te remains at the saddle point.
Lemma C.21. Consider the expert in E M, and j E Jo. Let w T " S $ = O(d-1/2) and of = "e and Then, with high probability, there exists some time t3,1 $ T3.1 = θ(n -Id" such that the following conditions hold:
R Remmi 2 022 and
fs.1
S 5A3d =
=0(-111) (dd-122)
The same argument applies symmetrically when exchanging e and g.
Similar to Lemma C.8 and Lemma C.9 Appendix C.2, we provide a bound on the Hermite coefficients influenced by the mean vector.
Lemma C.22. Under Adaptive top-k routing and re-initialization in Lemma C.17, suppose that Iva wmm.. = Ô(d In | = ((d=t) and (ng.m.) = (d for all = T ((tk" Then, by setting vj' 1]c S and 4
we obtain that m...ic am.jal Ö(d à) with high probability.
Proof. Since the gradient in Lemma C.8 changes only by an order of polylog the proof follows in the same manner.
Lemma C.23. Consider a neuron which satisfies Omm.j.k- Be,k- > O and am.3. > for k' <i p² Under Adaptive top-k routing and re-initialization Lemma C. C717, suppose that lvo wm.ss = ÿd }), k.,m, = (dd and Ing gg = ((d=1) 1) for all = 0, 11...,,. T T = Ö(akk-1) Then, by setting 17 = S and "2, we obtain that
jattt. Jam.aal = +)(d+t) with high probubility.
Proof.
Since the gradient in Lemma C.9 changes only by an order of polylog the proof follows in the same manner.
51

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Remark C.24. For the sake of conciseness in the exposition of the proof, omit the superscript l in ã mm.j... Based Lemma C.22, Lemma C.23, and Lemma C.10, the bounds in the subsequent lemmas properly justified, regardless of the variations in the coefficients ā a.m..j.."
Proof of Lemma C.2/ We begin by evaluating the stochastic updates of the experts E MM. Since the router learns to dispatch the data Te to the experts m MM. with high probability by Lemma C.20, we have F(m((x)) E MMe) - d for some A > 0. For clarity, we will henceforth assume and write P(m((.)) = m) = 1 throughout the remainder of this proof. Thus, we consider the probability conditioned on the event that the data is routed with high probability.
By analyzing the gradient update, as in Lemma C.7, we obtain that
Kemm 2 kemm + +1/w²² w (Ia w "inwmm T ") wm... ([I(((()))) m) m) m (w'ms Dee bm.)l
Ik2
| Vmm,, Wm (m((xx) = m am. (w'ms T De + bm )l
2
(n)
| Wm. ((m(xx) = m modamm I' Im (w'ms Do + bmi))
2
(La - w m. w'mj T
+ 17 "
2 N m. + 1/ w
T
p
20
lammid (1- (1-11)
(pw'm T '0o) ()
1)!(w T wmm) 5-11
VI!
(1-i)
Vi!
1=
IN' "(mm)((r"""
2d A
(17²)² ad3
+ n w (Jd win wins )
2
2
D¹
2 mij + n//
Σ
in immer (1 (khe
+ i.s y-1 (w] wiz tie f ggmms)
|k 1(7') 2d
(171)2 dd3
Ai
2 + vi w T (Ja win wm.d ")E"ma"
a (m)(a)) = m)yelm om (w Te + bmg)] + Wm., (mm(x.) m.yeaam. mm (w'n.) xe bm..) in the first inequality. We used similar decomposition to Lemma C.7 in the second inequality and adopt the notation iosm Vi
where we introduced mean-zero random variable 3 - Va
Dc + L1mi11
lammal (1-) (ow,, Tuo)
kat amde in the third
inequality.
In the same way, obtain upper bound of k²cm...
1+1 S k.oms + n/ww?" (Ja whn w " "m. E (((((()))= = Wem (m T De + bmng) Re het
(600)) 2d (n1² (nn)) 2d
+ nn" T (La wnn wimms T)
2
I m.m + 1/
Σ
(w:nn D.em (e. 21 (nC A.s... ( (w - oet mem.s
i=k-
1
(1 ( A 3d + w/ T (Ia - w/n, winss )==
+
2
Similarly, we carry out the corresponding calculations for rot "g.
1+1 Kig.m.j 2 Ke.m. + n¹ $=k
R..m.ii,.BB.. (He. a (w= "ng - nit oe. mmm))) (1 (n) |?
isa +i (e( ) i-1 (1
52

---

is defined recursively as Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Basel Learning
Inc gmm..
An
A12dd
An
d
T (1a - w u ")umm. + 17 w/'
2
and
p¹ Σ
1+1 Ky g.m. S Kams + r/
(tām.Sit. (Kam)) (wo
wi (w] 1e s
+ iisiimii g))
1-1
(K (m
+ 9.m. ( A2
nnt
(17") A
+ tn (Id w Wm.s D)
2
Next, we introduce auxiliary sequences to establish the following bounds for ko.m. and ni m... The derivation follows the same approach as the proof of Lemma C.12.
Consider one neuron j E Te and suppose that Ke,,. S 02 | S Kem.. and Ki g.m.j. I A2A3d à for all s = 0,1,...,t Then, k²m..j is lower bounded by Piil, for all s = O, ......,,t + 1, where the sequence (Pii,c) s= is defined recursively follows:
1+1
p6 (1 - a2) a2) and
ps+1 Pitt. + 11"a2"" Ōm.sk* for,e (Pi,ee) th- for s 2 0,
with high probability.
While, Kg.m., is upper bounded by Qin.g for all 0, 0, ..... , + 1, where the sequence (Qin.)) ++1 is defined recursively as follows:
Qmg = (1 a1+a2) a2) max I Ing.m.s 2 -d $ and
Qui Qin.o + (1 + 02)1) k (8 7 (Q )
+ Aar k "m..j. aBe (H 1 -1 -Id for 2 0,
with high probability.
Consider one neuron j E J, and suppose that Ny ggmm. I a2 Ikamol S Ny mig. and K I A2Asd à for all = 10..,. t. Then, x g.m. is lower bounded by Pii.. for all 0, 1,..,... ...... t + 1, where the (Pi..)) is defined recursively as follows:
1+1
Pill, (1 02) and
ps+1 Pii.g + 17 "2kh scāmijk Lee
eeee (Pil.))* "11, for s 2 0,
with high probability.
[+1 ++] While, k.m.. upper bounded by Qin.. for all O, 1,...,, where the (Qin,e) ==0 follows:
Q (1 + 02) max Inamml 2
a } and
Qut! Qin,c + (1 a2)n" ks" Jmm.s & BeB, ((Qin.)) K-1
Agn k** ".āmmjk c (n). d for 2 0,
e
with high probability.
53

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Using these auxiliary sequences, we can deduce the following. The derivation is the as that of Lemma C.133
Consider one neuron j E Te and take nf' Te I a,d= Suppose that k.z.m.j S ag, lng.m.pl I Ke. and l S 4.Add hold for 0, ..1,...,.. l.. Then, if memm S 02, lngtm.s S "eam.." and Ing.m.s! 4agd 2 with high probability.
Consider one neuron j E Tn and take of = S aad Suppose that ng.mm S 0z k S kg.m.. and |+يmm I 4A3d hold for 0, 1, 1.....t.. t.. Then, if g.m. I 0z, Nemmm S =+t1 g. and | I 4Azd with high probability.
( Ka.m. I d2 with Te, and 1 S Ke. and Ing kkg.m.g S A3d with j & Te for all 1,,, T3,1. Then, the above bounds holds for all s = 0, ..... T3,1 with high probability.
Suppose that T32 (ncck (k* 2)(1 5k" (amm.j.k ,.B,, ) (Pi).e)
Thus, by Lemma A.4,
po M1.
Kem. 2 Pin.e (1 - (k* - 2)(1 02 )2 Bo, (PD)s))-->
However, when t Tı,
Pmc
Kemss Pi 2 (1 Tekk 2) a2) am h Be (Pi) (Hi )
1
2 (nch" -2)(1 (2) ( ,eBkk))) ) > 1.
Ti This leads contradiction as k. S 1. Since | lkg.mll $ 4Agd
S Allle S Alasd 2 S Azd
Thus,
S |n [x
- + Me.... - K Ko,m
S. 5Agd
f3.1 This concludes that there exists some time 13,1 S TT,1 = O11 -Id -1d" such that memm. > 02 and gmg. I 3Ayd à for j E J. with high probability.
k S 02 with j E Jg. and | S m. and [ss,mms S A3d 5/ with & J, for all s 0, 1.., T3,1 and then, there exists some g. > 02 and fee 5A3d for j E Jg with high probability.
In the same way, suppose that T3.1 = [(Tck*(k* - 2)(1 (s (P
time 1a,1 S T3.1 (1%- -1d _d= such that k 13.1
Transition from weak strong recovery. Next, we show that the neuron for E To aligns with w up to large constant 1 C2. We denote t- t 13,1-
Lemma C.25. Consider the expert m E M, jnd E To. where j satisfies Lemma C.21. Let w T wg S Aad o(d
and r/ = 1i. I a $ Then, with high probability, there exists some time t,22 5 Ta, = Ô(n -1) such that 1 ta. 1 02. The same argument applies symmetrically when exchanging e and g.
ke Z
Proof. Once the alignment reaches a constant level, the projection onto the spherical constraint via (Ig - Wm.j wT. weakens the signal, requiring the reconstruction of the auxiliary sequences discussed Lemma C.21.
54

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
Consider experts Me-
Suppose that, for j E Je. Ki.m.j S 1 0z. Ing.m.p S k2 j' and Kig.n.j S AgA3d 3 for all 0, ,,...,t.. by using a similar inequality evaluation as in Lemma C.7 and introducing a mean-zero random variable E
t. Then,
-V (m((m)))) m) mollmmm jmm (wmm + bm..ill + V Wim. (m((xe) m)yaam.. Jm (w'ms + bm.)).
( Hom. + ri"
Σ
(mmmsicPesstemm.. 11 - (Kom))))
+ am (.
(w T wy - m (mm))))))
I ((7)) A1dd (17") + w2 T (La wms w )=
2
Homms 11 k jmjjk fem.j. (Kom 1h- (1 1 komm)) )2 pp2 max se0miii (k - -p22 12 m (m)m)))) (n=)) Ke dl w/ T (Ia wms wms )
mwo wgl
Since uz wiz S Asd In S A2 and mom 2 bd p 2p2 maxi se0 (m)) ) 1p max, sammmniciill S 2027 " aBax (N (n ) j A 2idd 1221 vc am j. . ( mm))
1 "ww wÿl S
"az))""mmskk aassk-aaast (K (1r))?ng 5
and
For the noise term,
a ifs S A2d xork++]
T )= wmss
rf W (Ia w W
Haon k , x (K)) - if Azd
='=0
with high probability. Therefore, by noting that (1 - 11--in kr ke )3) S Ful. kem can be lower bounded as
Kem. 2 (1 - (22)) (1-22)mm0
+ u2
Σ
r1 k2 Sem.jidi,,BBeaeee (Kr. W--1
6'=0
By introducing auxiliary sequence (Pit,t) ++1 where
pro (1 a2)x and
P's + 02n1" k² km.jkk" Be (P ( for s 2 0,
k--1
MlI
then kemm. lower bounded by Pit. for all 0, 1, .,.. with high probability. In addition, with the same proof as Lemma C.21, by introducing auxiliary sequence (Qmig) (Q
)+1 where
2m. 6Azd 4, and
**--1 Qin Glü.g + (1 a2)m" (ol))"""""ssaa k=mmmm,,... To (Qmi.))**
+ A37 k2 "mm.jkk B ()) -1d d-3 for s 2 0,
then ki "gmy is upper bounded by Vii for all s = 0, 1,..,...44 ++ with high probability.
Similarly, for a neuron j E Jg. re mj.. is lower bounded by Pilig for all s = 0, .t, and kam. is upper bounded by t+1
1+1 Qü,e for all = 0,1...., + with high probability. (PMng) = and (Qii.e) == are defined as follows:
M1..g (1 a2)) m. and
55

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Basel Learning
P
Pil.g + a27j k" s.amm. e (Pil.g) xx
yk-1 for 2 0.
ill.g
Yi. = 6Azd 2, and
I1I.c - Viu + (1 a2)11 k" lāmms .B,, 1Q 6
+ Azri hi* s.0m oo (K ( ) - -1d for 8 2 0.
Note that the periods the notation of the auxiliary sequences (Pii.e) intentionally used distinguish them from the auxiliary sequences in Lemma C.21.
++1
==0" (m...g
++1 ==0² (Rii.g)
) a ==0"
H+
==0 are and (Qii,,) ) H+
We prove the following arguments by induction using this auxiliary sequence, in the same manner Lemma C.6.
Consider one neuron j E Te and take nf = Te I and 4 Suppose that ka.mj I 02, |n I S me and A hold for all 0, 0,1....,f Then, if kem... I 1 a2, 1kg.m.s Ke,,..j and Ing.mal S AgA3d with
lkg.m.g I I high probability.
In the same way, consider one neuron j J, and take of ?e S agd 5. Suppose that &.gmm.f S 1--22, Inc.m.s S Ngg.n.j and m S A2 Azd-= hold for all 0, 1, Then, if n I 1 0z In I k I+1 with high probability.
g.m.j" and In S A2A3d
Consider one neuron j E Te. Suppose that ke.m.j $ 1 02 hold for all s = 0, ..... T3,2, where
T3.2 L(nck" (k* k. BBe, k- (Pil) y kk- 33
However, at / = T3,2,
pmo
K.mjj 2 Pif..c 2
(1 - Maki (h* 2))22mm j.k Bo k (PM,e) ( (3-))19) Fi
>
(Nak h2t (k* 2)02ūm j.k Be x (Pillo) yku -2 -yyoo
This leads contradiction. Thus, there exists some time 13,2 T3,2 such that l2. amm.j
a2 The same proof applies to
ta. N..m.. for j E Jg by taking T 1lr.k" (k* 2)028,.0m.j.. 7 (PrD k'
illee) J.
Strong recovery. Finally, we show that the neuron for j E J. amplifies the alignment with w and the neuron for j E Jy amplifies the alignment with wy and we establish strong recovery (n 2 1 --c > 02 for j E Je, and ne "ym.. 2 1 1 a2 for j E Jz).
Lemma C.26. Consider the expert Ó(d-1/2), n¹¹ ne S agd- 4/2 for t S (T3,1 1,,1) (Ts,2 13.2) and n¹ me S minf a1 odd as for (20
MM and j E To where satisfies Lemma C.25. Let ". w S Asd
(I,,1 t3.1) + (T3.2 t3.2) St (Ta,1 ta.1) + (Ta.2 t3.2) + T333 1. Then, Ke mi
where T3,3 Ô(E ne'
- 11) holds with high probability.
The same argument applies symmetrically when exchanging e and
Proof. Consider experts m E MM. Suppose that, for j E Te, we have 1 - 202 & k.ms S for all & =
Then, we have
Kem. Ko.m.j + 11² kt ām.i.kk Be (1 (No.m.s) ( (ni
+ n
is is. -
(w: wg
Ke Re 2AIdd we (Ia w wnnjums T) Sum.
T
56

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
2 Kemm + nf k Lo (1 (Nam, ( mm) 111
-nepp" 2 max (m V Vlwe wÿl( (1- (N,,m.))
-Tepp" 2 max s (500 (ggmm.mmm 1k' "gg T (18 -u²
w boeeeoonmn umm.hhemm
i
- - AA d + ne W. T (Ia w wins )=uum
Σ
6
n k* āmjjk "amski - kem ) + Σ
ne T (Id - winwwin Diems
2 ke.m..j +
5
s'=0
6'=0
Since K.m.. 2 1 - 3az, we have wig (Ia w )uin. S v Vbaz and kg.ms = (wg T (wo T
Wmm))
+
" (Ia - w. T "wmm I vooz + -(dd))) $). Hence, we obtain n7" a (1 (F (K c.m.
12.
C x. (1 ke.ms)) by kams 2 3az Mep'² max; Isaam Willu wÿl(1
(no (mom kk
"mmii BB.
,e (1
kzmm.)) by w "u
Ö(d nepp max; Is di Allkg
w (La
wwwe ptamm S 1rr am
amjik BB (1 me.m.)) by max, (s. Val/ am .Be, S vaz
when
ammii pwm.s T
O, and nem 2Ard S 3on* amatk o (1 k....) by ne S ged In addition, we have Es'= W0 (Ta- wins )=
>
age + f amaid eBo m)) with high probability. Therefore, if 1 - 3a2 S k S 1 - 5. for all 8 = 0, I....
et1 c.mm.
02
c+
Σ
merk" Ōm,jkk BBB (1 m
3
s'=0
0
a2
c
+ snek Cem.keeeeeB.kk. .eBe,
1
3
3
and 1 302 5 Kemm hold.
For 0 < s (Ta,1 - f3,1) + (Ta ta,2) - 1, it holds that ke., 2 1 - 2az by taking n/¹ = Tie S and &
Take 1/" = Nee min ( 44 ed g9 2 and suppose that ni 1 and for all (Ta,1 13,1) + (T3.2 13,2) S 1 S
(Iz,1 d3,1) (Ta,2 t3,2) + T3,3 1, where
T3.3 = laz(nerck' āmij. "e,,k ) 1 + 1.
Then, it holds that k. 2 30² + 5/2 ten kn Üm.j,k B for all (T3,1 - (3,1) + (T3.2 - 11,2) S15 (Tx,1 - 13,1) +
(Ta,2 13.2) + T3,3 - 1. However, at t= = (T f3,1) + (T3,2 - 13,2) + T3.3, 1 - 30² + 3 tone "m.j.k'. ,cBc, 2 1, which
leads to contradiction. Thus, there exists some 4g3 < (Ta,1 la,1) + (Tg,2 - (3.2) + T3,3 such that Kc.m.j 2 1 - 3. When
there exists some time k..m.j 1 for t 13,3, Rm.m. 1- - 20g since |x/,m., N. Aine' Thus, Nomm. > 1
holds for all f3.3 S t < (T3,1 t3,1) + (Ts.2 t3.2) + T3,3 until ni.m > á holds. By recursively applying this step. obtain the desired result.
C.5. Second Layer Optimization Stage
We have i.i.d. test-time inputs XT = (x', ..... TT) and we extract Te inputs X (x/ (x.....,,,?) .TTe)
x.. in the cluster e from X",
Note that Σe T. T. We know that each X is successfully routed to the expert m M. with high probability over the randomness of T
C. APPROXIMATION OF SINGLE INDEX POLYNOMIALS
We suppose om are ReLU functions.
Lemma C.27 (Following Damian et al. (2022); Oko et al. (2024a)). Fix c and the corresponding expert m E Me such that for all tep hn ( 2 0 with high probability. Suppose that by ~ Unif([-Ca, Cb]) with C. = Ö(1). Let ha(z) be
57

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
polynomial with degree 4 O(1), E Unif (S4-1(1)), w w. Then there exists d1:... "1:...,,20 02N R such that
1
ajom m0mmm (w (Txx + +b))
1
ajom (w- "=a + bb)) - he(w" x/()) =((N-).
sup I.
2N
2N
Moreover, we have Σ/ o (NN and Σド layl = Ō(N).
We obtain similar approximation results for polynomial activations (see (Oko et al., 2024a) for details). Using Lemma C.27, we show that .for all c, there exists some m E MM and om such that fm can approximate ft + seg".
Lemma C.28 (Following Oko et al. (2024a)). Let am, m = m1=1,-.. 1,...,M. M be ReLU activations or polynomial activations. Fix c and the corresponding expert m E Me such that for all to. hmm (x 2 with high probability. Assume 2 Jmmippoly log d. There exists some parameters am (0mg)) such that
2 Σ om (1 m + bb) fo (w see rog"("" a/cc) 2 50([mmial 12
1
1l [ (a(ee) 2
Σ
Te
J
m
where |a 12 = Ö(Ja)Jminl))
10 11 Ö(JVC) and Enez = 0 for m' E Mo {m}.
The main difference between the proof in (Oko et al., 2024a) is that we may have superfluous experts m' E MM [m}. However, we only need to put " = 0 for all m' E MM {m] and j = J.....J.
Proof.
C.5.2. OPTIMIZING THE SECOND LAYER
We present the result of optimization of the second layer: Lemma C.29, Suppose that J = (JJminool log d). There exists >> 0 such that the ridge estimator åm satisfies
Σ (x 2 0] Jm (166 - (w", xe)) - 8:0 (no" s.)) 50 (1((ninn (|Jmin|]+++)) +c) T72
poly log
Exg.a
T
with probability at least 1 - od(.). Therefore, by taking |Jminl Ô(E and T Ō(E=)), we have Ö(E) loss.
Proof. Let Al {fām flmmmmem EmEm |lāmlrr I Ememillate ]I We know that the router h exclusively route X to the subset of experts Mo with high probability. Therefore, the minimization problem of the empirical 1,2 loss is decomposed as
M
2
1
Σ
Σ
(am))') 2 0] Sm, (x²) y²²
min (anc) EA Mr
T
m==1
27
Te
min (am M)....))
Σ
Mmme (2 2 0) s a (x/2) y/:
Te
to
Am'EMe
27
Te 1
< min
Σ
Σ
1 [ (2/2e) 2 0| fi F (x y.--
Oc aim hom where (am M. EAiMe
To
E.Mo
Σ
To
min
1
Σ to
Σ
1 [m (1. 2 0] Jm -yt
T
(a Jm EAAio
T
m'EM.
=.f.
BBa [fo-vo)
58

---

Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based Learning
where used Bmm (x 0 m' E Me with high probability, MgnMee for all c', and Ve, {am Im A'm {ām) Im A M) Therefore, the optimization is performed in parallel for each subset of parameters {am, Mc)
and we can bound the population loss as
Ec Ex llfe Velll
Ex, l1fe Vell Ex,. uell} v
E., [1fe yc)3
sup
A.. EAMo
Applying Lemma 14 of (Oko et al., 2024a), we have
l(fe = =
Jetal'+ee
Jmial
Apprxximation
Concentration
and the generalization error is bounded as
- B., llfe vell] so =0
( 1 + c
poly log Ve
Ex, [lfe Vel]
sup Gnd EA-: (omgs ni E,M,
with probability at least o1(1) for the ridge estimator.
59


Reviews:
[{"review_id": "QKHPTL3yDj", "reviewer": "Reviewer_YoCg", "summary": "In this paper, the authors investigate the sample and runtime complexity of Mixture-of-Experts (MoE) optimized with the stochastic gradient descent when learning a regression task with an underlying cluster structure of single index models. In particular, they show that a single neural network cannot detect a latent cluster structure. On the other hand, MoE is capable of doing that as it aggregates the ability of each expert to weakly recover the simpler function associated with an individual cluster. All the results are proved rigorously. Finally, they conduct some numerical experiments to justify the theory.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Reviewer YoCg,\nWe sincerely appreciate the reviewer's thoughtful and constructive feedback. Below, we carefully address the raised concerns and questions.\nDue to character limits, please see our response to Reviewer **wqTg** for references.\n## Model parameters\nFollowing the reviewer's precise feedback, we will revise the manuscript to provide clear explanation of $w^*_c$, $w^*_g$, $s_c$, and $\\nu$.\n## Router\nPlease refer to the **\"Router\"** part in our response to Reviewer **bEDR**.\n## Implications\nChen et al. (2022) present a failure of a single expert in binary classification where noisy features confound the true label. In contrast, our work focuses on a **regression** setting and reveals a different failure mode where **gradient interference** across clusters attenuates the signal, increasing the information exponent and hindering optimization. Our findings highlight how MoE mitigates interference by explicitly partitioning the network. Such interference has been recognized as a major challenge in multi-task learning (see lines 125-127 in the left column).\n\nFurthermore, we demonstrate that, for the router to distinguish between experts, a **weak recovery** of the feature index is necessary. Our analysis reveals that this requires a sample complexity of $\\tilde{\\Theta}{(d^{k^*-1})}$ during the exploration stage, implying that a sufficiently **long exploration phase** is essential before the router can begin meaningful learning. This contrasts with the linear expert setting in Li et al. (2025), as our findings arise from non-convex optimization.\n\nFinally, unlike in classification (Chen et al., 2022; Chowdhury et al., 2023), in the regression setting, **competition among experts** within the set of professional experts $\\mathcal{M}_c$ may lead to load imbalance after the router has learned to assign inputs. This observation suggests a practical motivation for several heuristic approaches (see lines 122–131 in the right column).\n## References\nAccording to the reviewer’s suggestion, we will incorporate references Chen et al. (2022) and Allen-Zhu & Li (2023) as inspiration for our data structure, and Liu et al. (2024) and Tang et al. (2025) as practical works on gradient interference in MoE in Assumption 3.2. We will also include Arous et al. (2021), Ge et al. (2018), Gudeja et al. (2018), Bietti et al. (2022), Damian et al. (2022), and Oko et al. (2024a) in lines 096–098 in the right column.\n## $v_c$ and $g_j$\nWe appreciate the reviewer pointing that out. The notation $v_c$ is correct, but $g_j$ is a typo. We will correct the typo in the final version.\n## Q1.\nIn our setting, feature vectors are sampled from $\\mathbb{S}^{d-1}$ rather than from $\\mathbb{R}^d$, in order to normalize signal strength for the analysis of simultaneous learning of multiple signals. This setting is frequently employed in prior works (e.g., Arous et al., 2021; Dudeja et al., 2018; Bietti et al., 2022; Damian et al., 2022; Oko et al., 2024ab). We would also like to note, in high-dimension, a uniform vector on $\\mathbb{S}^{d-1}$ converges in distribution to $\\mathcal{N}(0,d^{-1}I_d)$, in each coordinate, by concentration of measure.\n## Q2.\nAs the reviewer correctly pointed out, in $\\hat{F}_m$, the number of activated experts is not fixed and $k$ is a stochastic variable. There have indeed been prior attempts to vary the number of activated experts depending on each token (Huang et al., 2024; Zeng et al., 2024).\n## Q3.\nWhen analyzing the gradient, the correlation between the **nonlinear** target and the **nonlinear** activation arises. To decompose this correlation, we employ the Hermite expansion, which forms an **orthogonal** basis in the $L^2$ space under the Gaussian measure, as first used in Ge et al. (2018) and Dudeja et al. (2018).\nSpecifically, we have\n$$\n\\mathbb{E}\\_{z \\sim \\mathcal{N}(0, I_d)} [ f^*\\_c({w^*\\_c }^\\top z) a_{m,j}  \\sigma_m({w_{m,j}}^\\top z + b_{m,j}) ] = \\sum_{i = k^*}^{p^*} \\alpha_{m,j,i} \\beta_{c,i} {\\langle w_{m,j}, w^*_c \\rangle}^i.\n$$\nThis expresses the interaction as an expansion in powers of **alignment**. Crucially, the first non-zero order $k^*$ in the expansion (i.e., information exponent) **governs** the signal strength.\n## Q4.\nIn our study, the top-1 routing used in Phase II constitutes a sparse MoE. Although the adaptive top-$k$ routing employed in Phases III and IV allows $k$ to vary, it remains a sparse MoE in the sense that not all experts are activated. While adaptive top-$k$ is motivated by technical challenges in the theoretical analysis of nonlinear regression, we believe that end-to-end training with top-1 routing is feasible with the incorporation of auxiliary losses to control load balancing, which presents an interesting direction for future work.\n\nWe will incorporate the revisions identified through this review into the final version. We would be happy to address any further questions during the discussion.", "claims_and_evidence": "Yes, the results are supported by both rigorous proofs and empirical evidences.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I did not check the theoretical proofs carefully. However, I read the proof sketches and the arguments made sense to me.", "experimental_designs_or_analyses": "Yes, I checked the experimental setup on page 6. The experiments are for justifying the theory only.", "relation_to_broader_scientific_literature": "The theory in the paper is to understand the underlying mechanism of Mixture-of-Experts and how it works.", "essential_references_not_discussed": "Most of the relevant references have been cited. However, I encourage the authors to cite at least one reference for the teacher models in Assumption 3.2. If this model is first introduced by this paper, then the authors should elaborate on the formulation of the data generation process. \n\nAdditionally, in lines 096-098 in the right column, the authors should add references for the claim \"The complexity of\nlearning a nonlinear function via two-layer neural network optimized by SGD is closely associated with this value\".", "other_strengths_and_weaknesses": "**Strengths:**\n\n1. The problem considered in this paper is of interest due to the recent success of Mixture-of-Experts in large-scale models. \n\n2. This theoretical results are corroborated with rigorous proofs. \n\n**Weaknesses:**\n\n1. The writing of Section 3.1 and Section 3.2 is not really good. \n\nIn Section 3.1, the writing in Assumptions 3.2 should be improved. In particular, the authors should provide an intuitive explanation why they have to involve the Hermite polynomials and Hermite expansions when before or after introducing them. The assumptions on model parameters, e.g., $\\omega^*_c$, $\\omega^*_g$ are not explained thoroughly. Additionally, the roles of parameters $s_c$, $\\nu$ are not introduced, making it inaccessible for people who do not have background on the considered problem.\n\nIn Section 3.2, the introduction to the formulation of MoE is quite complicated as the authors discuss different gating functions and routing strategies alternately without clear separation, making it difficult to digest.\n\n2. Although the paper considers a different problem than that in (Chen et al., 2022), the practical implications are quite similar, that is, an MoE is capable of performing something that a single expert cannot do. It would be more interesting if the theory can provide new insight into improving the MoE performance. \n\n3. **Major issue:** The mixture-of-experts formulation considered in this paper is far from practice. In particular, in the formulation of $\\hat{F}_M$ in line 119 in the right column, whether an expert $f_m$ is activated or not depends entirely on the positivity its respective gating value $h_m$ rather than the magnitude rank (TopK) of the gating value as in the literature (Shazeer et al., 2017). This point makes me confused about the relevance of the theory in this paper.\n\n**References**\n\nChen, Z., Deng, Y., Wu, Y., Gu, Q., and Li, Y. Towards understanding the mixture-of-experts layer in deep learning. In Advances in Neural Information Processing Systems, volume 35, pp. 23049–23062. Curran Associates, Inc., 2022.\n\nShazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017.", "other_comments_or_suggestions": "1. In the Assumption 3.2, I think $v_c$ should be $v^*_c$. Please feel free to correct me if I am wrong.\n\n2. In line 111 in the left column, should $g^*_j$ be $g^*$?", "questions_for_authors": "1. Why do the parameters $\\omega^*_c,\\omega^*_g,v^*c$ belong to the set $\\mathbb{S}^{d-1}$ rather than $\\mathbb{R}^d$?\n\n2. How do you guarantee that $k$ experts will be activated given the router $1[h_m(x)\\geq 0]$ in the formulation of $\\hat{F}_m$? Based on the router formulation, there could be more or less than $k$ experts activated per input.\n\n3. In high-level, what are the reasons for involving the Hermite expansion in this case?\n\n4. Can the current results be generalized to the sparse MoE setting in (Shazeer at al., 2017)? If not, what are the main challenges?\n\n**References**\n\nShazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017.", "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": "WBzUk68hb5", "reviewer": "Reviewer_VrRr", "summary": "This paper presents a comprehensive analysis of the sample and computational complexity of Mixture of Experts (MoEs) when optimized using stochastic gradient descent (SGD) for a regression problem.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Reviewer VrRr,\nWe sincerely thank the reviewer's insightful and constructive comments. In the final version, we will separate the presentation of theoretical results and experiments, and further elaborate on the writing in the supplementary material. For the theoretical results, we will include a high-level summary, detail implications, and highlight key differences from prior works. We clarify each point and question below.\nOwing to the character limits, we respectfully refer the reviewer to our author response to Reviewer **wqTg** for the references.\n\n## Implication\nPlease refer to the **\"Implications\"** part in our response to Reviewer **YoCg**.\n\n## Key differences from existing works\nPrior MoE optimization studies (Chen et al., 2022; Chowdhury et al., 2023) focus on binary classification with noisy features. In contrast, we analyze regression and highlight how **gradient interference** (see lines 125-127 in the left column) across clusters weakens the signal and effectively increases the **information exponent**, revealing a clear separation between vanilla neural networks and MoE in Theorem 4.3. Although Li et al. (2025) also tackle regression, their linear setting does not account for the **non-convex** exploration phase we uncover.\nRegarding information exponents in feature learning, Oko et al. (2024a) assume additive task structure within single data points, whereas we study **additive structure across clusters**. Theorem 4.7 is the first to apply nonlinear sample complexity analysis to MoE, showing how the router’s gradients reflect expert specialization via **weak recovery** differences.\n\n## Real-data experiments\nWe certainly appreciate the importance of experiments on real-world data, as the reviewer has rightly pointed out. However, our study is theoretical in nature. Real-world datasets inevitably contain noise, which makes it challenging to isolate specific phenomena in a controlled and interpretable manner. Therefore, we conducted controlled experiments using synthetic data. Within the line of work on the complexity of learning low-dimensional nonlinear functions in specific architectures, Dayi & Chen (2024) and Oko et al. (2024b) also rely on synthetic experiments based on Gaussian data, while Vural & Erdogdu (2024) do not provide empirical results. In the final revision, we will include the more detailed learning process of MoE based on synthetic experiments.\n\n## Q.\n### Structure of MoE\n* **Gradient-aware routing:**\nWhile our analysis shows that the MoE mitigates gradient interference by explicitly partitioning the network, the number of experts is typically selected heuristically in practice. This raises the possibility that integrating gradient-aware routing mechanisms could lead to more efficient and adaptive MoE architectures. Recent works also explore the role of gradient interference in MoE (Liu et al., 2024; Tang et al., 2025).\n\n* **Adaptive routing:**\nTo prevent competition among professional experts ($\\in \\mathcal{M}_c$), we adopted top-$k$ routing to avoid potential load imbalance. Thus, dynamically adjusting the number of activated experts $k$ is expected to stabilize training. This insight stems from our theoretical studies in nonlinear regression and aligns with recent NLP research (Huang et al., 2024; Zeng et al., 2024), where $k$ is adaptively varied based on tokens.\n\n* **Freezing (pruning) redundant experts:**\nTo mitigate competition among professional experts, reducing redundant experts is also worth exploring, potentially lowering deployment cost. Recent work has investigated expert merging in this context (Zhang et al., 2024).\n\n### Training process of MoE\n* **Upcycling from dense checkpoints:**\nWe argued that meaningful router learning requires **differences** in weak recovery among experts, which in turn necessitates a **long exploration stage** due to the non-convex optimization. In this light, upcycling dense MLP checkpoints trained on distinct domains may offer a practical approach to accelerate convergence. This idea has seen several empirical successes in recent LLMs (Komatsuzaki et al., 2023; Wei et al., 2024).\n\n* **Stage-specific routing:**\nThe noise $r_m$ introduced during Phase II ensures uniform gradient flow and sufficient signal for all experts. In contrast, the adaptive top-$k$ routing in Phase III and IV addresses competition among professional experts. This distinction—different challenges occur at each learning stage—suggests that stage-specific routing strategies may be helpful for effective MoE training.\n\nWe welcome any further questions or concerns that may arise and would be pleased to provide clarification.", "claims_and_evidence": "Based on the data model proposed in Assumption 3.2, all the claims appear highly likely to be correct.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "The Appendix lacks a clear structure and primarily consists of a dense collection of theoretical results without sufficient explanatory text. As a result, I find it difficult to follow most of the proofs.", "experimental_designs_or_analyses": "Good.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "1. A numerical justification using real data is needed to support Assumption 3.2.\n\n2. The theoretical results and experiments should be presented separately, or at least, a high-level summary of the theoretical findings should be provided. Currently, it is difficult to discern the major implications of your theoretical results. Additionally, it would be beneficial to highlight key differences from existing works, such as Theorem 4.3.\n\n3. Real-data numerical experiments are necessary to further validate the theoretical findings.\n\n----------------------------\nPost-rebuttal comments:\n\nI will keep my original score (3). I appreciate the new theoretical insights introduced in this paper; however, these insights need to be further validated through experiments on real data. Without such empirical justification, the impact remains limited.", "questions_for_authors": "Any thoughts on improving the structure or training process of MoE from your theoretical results?", "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": "ocXnQNvcpH", "reviewer": "Reviewer_bEDR", "summary": "This paper theoretically studies the learning dynamics of Mixture-of-Experts (MoE) models in nonlinear regression tasks with an underlying cluster structure. The main contribution of the paper is listed below\n\n- Proves that a standard neural network fails to detect and exploit the latent cluster structure, while MoE successfully separates and learns simpler subproblems.\n- Analyzes the SGD training dynamics of MoE, showing that experts weakly recover cluster-specific functions, and the router learns to assign data correctly.\n- Establishes polynomial time and sample complexity, demonstrating MoE’s efficiency in learning clustered tasks.\n- Proposes and theoretically validates a multi-phase training algorithm for MoE, alternating between expert learning and router learning to optimize specialization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Reviewer bEDR,\nWe sincerely thank the reviewer for the thoughtful and constructive feedback. We will incorporate the suggestions into the final version. Below, we address the reviewer’s comments. For references, please see our response to Reviewer **wqTg** due to character limits.\n## Multi-phase training\nWe agree that the multi-phase training procedure deviates from practical applications. However, multi-stage optimization is often employed to reveal the intrinsic nature of complex optimization problems and is commonly used in studies of feature learning (Arous et al., 2021; Ge et al., 2018; Gudeja et al., 2018; Bietti et al., 2022; Damian et al., 2022; Mousavi-Hosseini et al., 2024;Oko et al., 2024ab). Joint training of the router and experts is promising, but in nonlinear regression, it poses intractable non-convex challenges and is left for future work.\n## Assumptions and practical relevance\nSince obtaining negative results for vanilla neural networks is mathematically challenging, we followed the approach of using gradient flow, as in Mousavi-Hosseini et al. (2024) and Simsek et al. (2025). In addition, while Chen et al. (2022) assume all signals are orthogonal, our analysis only requires orthogonality to the mean vector. Our theoretical results remain rich in implications. Please refer to the **\"Implications\"** section of our response to Reviewer **YoCg**.\n## Router\nWe would like to clarify that our theoretical routing mechanism is not an oversimplification, but rather a **principled approach** to addressing fundamental technical challenges in MoE without practical heuristics.\nIn Phase II, training is performed using top-1 routing with added random noise $r_m$. After this phase, **a data $x_c$ is no longer routed to experts $m' \\notin \\mathcal{M}_c$** (where $\\mathcal{M}\\_c$ is the set of professional experts) due to the negative correlation between its router weight $\\theta_{m'}$ and the cluster center vector $v_c$.(Lemma 4.11 and Lemma C.16). However, it remains possible that **competition among experts $m \\in \\mathcal{M}_c$** leads to load imbalance if we employ top-1 routing. This phenomenon differs from the classification setting in Chen et al. (2022) and Chowdhury et al. (2023), and stems from the need to estimate a continuous function in the regression setting. \nWe analytically determine the cause and employ an **adaptive top-$k$** routing in Phases III and IV, where an expert $m \\in \\mathcal{M}_c$ is activated if it satisfies a sufficient condition $h_m(x_c) \\geq 0$, ensuring that it covers the data distribution of cluster $c$. The value of $k$ is **adaptively** determined based on the data, a strategy that is also adopted in practice (Huang et al., 2024; Zeng et al., 2024). Importantly, while the adaptive top-$k$ arises from a theoretical and technical challenge, practical implementations often incorporate auxiliary losses or dense training to address load imbalance (see lines 126–131 in the right column), ensuring consistency between theory and practice.\n## Q.1\n* **Expert choice routing (ECR)**:\nSince ECR (Zhou et al., 2022) uses the same gating network as our token choice routing (TCR) (Shazeer et al., 2017), both methods identify clusters through differences in expert recovery reflected in the gradients. The key difference is that TCR selects top-k experts per token, while ECR selects top-k tokens per expert, potentially dropping tokens unselected by any expert. We believe similar theoretical results are attainable if token drop is mitigated, e.g., via expert capacity constraints.\n* **Dense MoE**:\nDense training (Pan et al., 2024) activates all experts without applying top-k selection to the softmax outputs. Since all experts are activated, the gradient for expert $m$ depends on a comparison with a weighted average over all experts, potentially causing gradient competition. Specifically, when weak recovery occurs across experts $m \\in \\mathcal{M}_c$, differences in gradient magnitudes can distort the update direction and complicate theoretical analysis. Therefore, similar guarantees in the dense setting demand extra assumptions or modified algorithms.\n* **Hash-based Routing**:\nSince hash-based routing (Roller et al., 2021) is non-learning, misaligned partitioning of the cluster structure may induce gradient interference and hinder theoretical guarantees.\n## Q.2\nIf all signals ($w^*_i$, $v_j$) have overlaps of order $\\tilde{O}(d^{-1/2})$, then by carefully bounding the resulting overlap terms, we believe the MoE can still reliably learn the underlying cluster structure. Specifically, the gating network's gradient includes additional terms of magnitude $\\tilde{O}(d^{-1/2})$ relative to the main signal, arising across task–neuron combinations.\nIn contrast, if the overlap is of order $\\tilde{\\Omega}(1)$, the dynamics become profoundly more complex. Thus, we are unable to provide a definitive conclusion in this case.\n\nWe are happy to clarify any remaining questions during the discussion.", "claims_and_evidence": "The paper provides strong theoretical proofs under specific assumptions like population gradient flow and orthogonal features. However, these conditions may be too restrictive for real-world scenarios, and the claimed failure of vanilla networks is shown under specific setups, so it may not be generalized.", "methods_and_evaluation_criteria": "The paper’s chosen regression setup (with latent clusters and a shared global component) generally aligns with its theoretical goals, making the methods and evaluation criteria reasonable for investigating MoE’s advantages in that specific setting. However, the multi-phase training procedure is not particularly relevant to state-of-the-art deep learning applications, though it provides valuable theoretical insights.\n\nMoreover, the authors considered two routing mechanisms: top-1 routing and threshold-based routing, the latter of which is not a standard approach. These choices were likely made to simplify the theory, but they resulted in an oversimplification.", "theoretical_claims": "I briefly reviewed the proofs, and the theoretical results appear overly complex and dense, making them harder to follow than necessary.", "experimental_designs_or_analyses": "The experiments are limited in scope but provide valuable insights.", "relation_to_broader_scientific_literature": "This paper takes a step toward understanding MoE training. While the theoretical scope is overly simplified, it serves as a solid starting point.", "essential_references_not_discussed": "The authors could add additional", "other_strengths_and_weaknesses": "- The application of theoretical ideas from feature learning to MoE is novel and interesting to me.\n\n- The presentation of the theoretical results and problem setup is unnecessarily complex. I strongly encourage the authors to refine the notation for better clarity and readability.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How do the theoretical results change with different routing mechanisms?\n\n2. How might your analysis extend to scenarios where the cluster signals and global signal are only partially orthogonal or have overlapping directions, and do you expect the router to still reliably separate clusters under weaker assumptions?", "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": "NIaaCLOR6I", "reviewer": "Reviewer_wqTg", "summary": "The authors consider the problem of learning a mixture of $C$ tasks, plus one global task, with a mixture of two-layer network experts. They first establish that a single network is unable to recover the global signal, by constructing a special instance of this class of targets. They then turn to analyze the mixture of experts, trained with a stagewise algorithms in which different set of parameters are sequentially trained, and prove that this model is in contrast able to . A fine analysis of features learned at each stage is reached, with notably the first layer neurons specializing to different subtasks, and the routing weights then learning to point towards relevant experts, with the global task finally being learned.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Reviewer wqTg,\nWe sincerely thank the reviewer for their thoughtful evaluation and insightful questions.\n\n## Q.1\nAs the reviewer correctly noted, the conditions for the theorem rely on the inequality $M \\gtrsim C \\log (C / 0.001)$ in Lemma 4.9. In connection with this, while we initially described $M$ and $C$ as polylogarithmic in the input dimension $d$, we have since realized that ensuring reliable initialization requires both to be $O(1)$. We will revise the final version accordingly. This refinement is made for theoretical clarity and does not affect the main conclusions or practical relevance of our work. Aside from this adjustment, no other major changes to assumptions or results are necessary.\nAs discussed in Oko et al. (2024a) in the context of feature learning, the Hermite coefficients of the student model can be perturbed by random initialization, a factor not considered in prior work on MoE classification (Chen et al., 2022). Our analysis builds on this observation to study robust specialization under such randomness. However, if the number of clusters $C$ is too large, the number of expert-to-router combinations grows rapidly, making the specialization problem significantly more challenging.\nWe stress that this assumption is mainly to ensure **theoretical reliability**. In practice, MoE models often incorporate techniques such as auxiliary diversity losses in routing to promote **expert diversity** and avoid **premature collapse** (Cai et al., 2024). Our assumption reflects these practical strategies, and the insights into learning dynamics and specialization in MoE remain valid and relevant.\n## Implications from our theoretical results\nWhile our study is theoretical in nature, it also carries practical implications. These implications are discussed in the **\"Implications\"** part of our response to Reviewer **YoCg**, which we kindly invite the reviewer to refer to.\n\nDue to character limits, we respectfully make use of this space to provide the references.\n\n### References\n* Allen-Zhu & Li. (2023), Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning, ICLR.\n* Arous et al. (2021), Online stochastic gradient descent on non-convex losses from high-dimensional inference, JMLR.\n* Bietti et al. (2022), Learning single-index models with shallow neural networks, NIPS.\n* Cai et al. (2024). A survey on mixture of experts. arXiv preprint arXiv:2407.06204.\nChen et al. (2022), Towards understanding the mixture-of-experts layer in deep learning, NIPS.\n* Chowdhury et al. (2023), Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks, ICML.\n* Damian et al. (2022), Neural networks can learn representations with gradient descent, COLT.\n* Dayi & Chen. (2024), Gradient dynamics for low-rank fine-tuning beyond kernels, arXiv preprint arXiv:2411.15385.\n* Dudeja et al. (2018), Learning Single-Index Models in Gaussian Space, PMLR.\n* Ge et al. (2018), Learning One-hidden-layer Neural Networks with Landscape Design, ICLR.\n* Huang et al. (2024), Harder Tasks Need More Experts: Dynamic Routing in MoE Models, ACL.\n* Komatsuzaki et al. (2023), Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints, ICLR.\n* Li et al. (2025), Theory on Mixture-of-Experts in Continual Learning, ICLR.\n* Liu et al. (2024), GRIN: GRadient-INformed MoE, arXiv preprint arXiv:2409.12136, 2024.\n* Mousavi-Hosseini et al. (2023), Gradient-Based Feature Learning under Structured Data, NIPS.\n* Oko et al. (2024a), Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations, COLT.\n* Oko et al. (2024b), Pretrained transformer efficiently learns low-dimensional target functions in-context, NIPS.\n* Pan et al. (2024), Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models, arXiv preprint arXiv:2404.05567.\n* Roller et al. (2021), Hash Layers For Large Sparse Models, NIPS.\n* Shazeer et al. (2017), Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR.\n* Simsek et al. (2025), Learning Gaussian Multi-Index Models with Gradient Flow: Time Complexity and Directional Convergence, CPAL. \n* Tang et al. (2025), Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model, ICLR.\n* Vural & Erdogdu. (2024), Pruning is Optimal for Learning Sparse Features in High-Dimensions, COLT.\n* Wei et al. (2024), Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models, arXiv preprint arXiv:2406.06563.\n* Zeng et al. (2024), AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models, ACL.\n* Zhang et al. (2024), Diversifying the Expert Knowledge for Task-Agnostic Pruning in Sparse Mixture-of-Experts, arXiv preprint arXiv:2407.09590.\n* Zhou et al. (2022), Mixture-of-Experts with Expert Choice Routing, NIPS.", "claims_and_evidence": "The paper is primarily theoretical in nature. The authors bolster and illustrate their theoretical insights -- namely the specialization of experts and routing weights at different stages-- in convincing numerical experiments in Figs. 1,2 .", "methods_and_evaluation_criteria": "The paper is primarily theoretical in nature.", "theoretical_claims": "I have not checked the proofs in detail. From my reading, the claims seem sound.", "experimental_designs_or_analyses": "The experiments presented in Fig. 1, 2 seem sound.", "relation_to_broader_scientific_literature": "The work is similar in spirit to the prior key work of (Chen et al., 2022) in the investigation of feature learning by mixtures of experts, although the task and setting differ. To the best of my understanding, the novelty of the present work lies in the fact that the complexity of the target function can be understood through its effective information exponent, thus connecting to a rich literature on this topic in the simpler setting of learning single or multi-index models with single networks (e.g. Arous et al., 2021). This allows for insightful explanations of the inability of single experts to learn the function, as discussed in 4.2.1.", "essential_references_not_discussed": "To the best of my knowledge of this line of works, I am not aware of an essential reference the authors fail to cite.", "other_strengths_and_weaknesses": "I am overall in favor of acceptance. The paper establishes a clear theoretical result for an interesting model, and connects to the body of works on information exponents, which I believe paves interesting bridges for future works. The theoretical results are clearly exposed, and sufficient discussion is provided whenever needed to build intuition (e.g. in subsection 3.3). I did not however carefully check the proofs, nor do I have the expertise to provide an educated reading of the latter.", "other_comments_or_suggestions": "I do not have major comments or suggestions.", "questions_for_authors": "1. (Clarification question) Are there any conditions on the respective scaling of $M$ with $C$ for Theorem 4.7 to be applicable ? Is it the condition of Lemma 4.9?", "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:

Meta-review: The paper investigates the sample and runtime complexity of Mixture-of-Experts (MoE) optimized using SGD when learning a regression problem with an underlying cluster structure of single index models.

All reviewers thought that the paper made some contributions to the theory of understanding MoE. The reviewers commented that the theory was somewhat limited and may not have too much impact practical since it was simplified, but they all found it a nice starting point to provable and theoretically understanding MoEs. As such, I recommend it for acceptance. 

I will advise the authors to consider the comments made by the reviewers especially regarding cleaning up some of the proofs in the appendix.
Decision: Accept (poster)