Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Authors: Quentin Fruytier
Aryan Mokhtari
Sujay Sanghavi
Venue: ICML 2025 poster (2025)
Keywords: Expectation Maximization (EM)
Mixtures of Experts (MoE)
Mirror Descent
Primary area: theory->optimization
TL;DR: We show theoretical convergence properties of the Expectation Maximization algorithm for a general class of Mixtures of Experts by viewing the algorithm as projected Mirror Descent with unit step-size and Kullback Leibler divergence regularizer.

Abstract: Classical Mixtures of Experts (MoE) are Machine Learning models that involve partitioning the input space, with a separate "expert" model trained on each partition. Recently, MoE-based model architectures have become popular as a means to reduce training and inference costs. There, the partitioning function and the experts are both learnt jointly via gradient descent-type methods on the log-likelihood. In this paper we study theoretical guarantees of the Expectation Maximization (EM) algorithm for the training of MoE models. We first rigorously analyze EM for MoE where the conditional distribution of the target and latent variable conditioned on the feature variable belongs to an exponential family of distributions and show its equivalence to projected Mirror Descent with unit step size and a Kullback-Leibler Divergence regularizer. This perspective allows us to derive new convergence results and identify conditions for local linear convergence; In the special case of mixture of 2 linear or logistic experts, we additionally provide guarantees for linear convergence based on the signal-to-noise ratio. Experiments on synthetic and (small-scale) real-world data supports that EM outperforms the gradient descent algorithm both in terms of convergence rate and the achieved accuracy.

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Quentin Fruytier! Aryan Mokhtari'
Sujay Sanghavi
Abstract
larly effective when the feature space can be divided and processed by specialized models, known experts. Instead of relying single, large model to handle all input-output mappings, MoE utilize an ensemble of specialized experts. Each expert is responsible for a specific subset of the input space, allowing the system to efficiently route inputs to the most appropriate expert. This partitioning enables each ex- pert to focus on mapping its designated inputs outputs using separate, optimized model. By leveraging multiple specialized experts rather than a monolithic model, MoE can achieve greater scalability and flexibility. This mod- ular approach not only enhances computational efficiency but also allows for improved performance, as each expert can be finely tuned to handle its particular segment of the input space effectively. Real-world applications of MoE such as Sparse MoE span across various domains, including language translation, speech recognition, recommendation systems, and more (Fedus et al., 2022; Ma et al., 2018; Hinton et al., 2012; Liu et al., 2024).
Classical Mixtures of Experts (MoE) are Machine Learning models that involve partitioning the in- put space, with a separate "expert" model trained on each partition. Recently, MoE-based model ar- chitectures have become popular means duce training and inference costs. There, the par- titioning function and the experts are both learnt jointly via gradient descent-type methods on the log- niglliieihoood In this paper we study theoretical guarantees of the Expectation Maximization (EM) algorithm for the training of MoE models. We first rigorously analyze EM for MoE where the conditional distribution of the target and latent variable conditioned on the feature variable be- longs to an exponential family of distributions and show its equivalence to projected Mirror De- scent with unit step size and Kullback-I Leibler Divergence regularizer. This perspective allows to derive new convergence results and iden- tify conditions for local linear convergence; In the special case of mixture of 2 linear or logistic experts, we additionally provide guarantees for linear convergence based on the signal-to-noise ratio. Experiments on synthetic and (small-scalc) real-world data supports that EM outperforms the gradient descent algorithm both in terms of con- vergence rate and the achieved accuracy.
In its most generic form, training an MoE model involves training both (a) the parameters in the individual experts, and (b) the gating function that routes inputs the appro- priate expert. Typically, these are both learnt jointly by first formulating the final loss function as applied to the ensem- bled output, and then minimizing this joint loss function (via SGD or its variants) over the parameters of the gate and the experts.
In this paper we investigate, primarily from a theoretical perspective, the training of classical MoE as defined by Jacobs al. (1991) using classic algorithm: Expectation Maximization (EM). As opposed to SGD-based methods which are agnostic to whether a parameter is in the gate or in an expert, EM first formulates two separate problems one for the router, and another for the experts - in a specific way. then solves each problem in isolation and in parallel, and then collates the outputs to arrive the updated set of gate and expert parameters.
1. Introduction
Classical Mixtures of Experts (MoE) (Jacobs et al., 1991; Jordan & Jacobs, 1994) are a crucial class of parametric latent variable models that have gained significant popular- ity in deep learning for their ability to reduce both training and inference costs (Chen et al., 2022). MoE are particu-
'Electrical and Computer Engineering Department, Univer- sity of Texas at Austin, Austin, Texas, United States of America. Correspondence to: Quentin Fruytier <ddl76@ my.utexas.edu Aryan Mokhtari <mokhtari @ @awstin.uuuttii edu>> Sujay Sanghavi <sanghavi mail.utexas.eda edu...
For our theoretical results, we consider the setting of gen- eral MoE where the joint distribution of the target and latent variable conditioned on the feature variable belongs to an exponential family of distributions. We then narrow in on two simpler instances of MoE models: Mixture of Linear Experts (where each expert a linear regressor) and Mix-
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
ture of Logistic Experts (where each expert is logistic regressor). The router in each case is linear softmax.
particularly regarding the convergence of EM iterates to the true model parameters and the rate of this convergence. An adopted framework in the past decade, introduced by Balakrishnan et al. (2017), interprets EM as a variant of the gradient descent algorithm. For latent variable models with a strongly convex EM objective that satisfies a condition known as "first-order stability," was shown that EM iter- ates converge linearly to the true parameters. Subsequent works utilized this framework to show a local linear rate for Mixtures of Gaussians and Mixtures of Linear Regressions (Balakrishnan et al., 2017; Daskalakis et al., 2017; Dwivedi et al., 2018; Kwon et al., 2019; Kwon & Caramanis, 2020b; Kwon et al., 2021; Kwon & Caramanis, 2020a).
The primary finding of this paper is to unveil the correspondence between EM for general MoE to projected Mirror Descent. A similar correspondence was first discovered by Kunstner et al. (2021), but was limited to generative models for which the joint complete data dis- tribution belongs to an exponential family of distributions; this did not include MoE. As such, our contributions can be seen as a generalization of (Kunstner et al., 2021) to all generative models for which the conditional distribution of the target and latent variables conditioned on the feature variable belongs to an exponential family of distributions. We next state the details of our contributions.
Main Contributions:
In a recent work, Kunstner et al. (2021) proved that EM- where the complete data distribution belongs in an exponen- tial family is equivalent to the mirror descent algorithm with unit step-size and Kullback-Leibler (KI.) divergence regularizer:
1) In Theorem 4.1, we show that when EM is applied to general MoE, the iterates are equivalent to the ones gen- erated by projected Mirror Descent on the conditional likelihood function with unit step-size and KL diver- gence. Here, the projection is an expectation moment projection over the parameter space. By leveraging this correspondence, in Theorem 4.2, obtain sufficient conditions for which EM applied to general MoE con- verges to a stationary point or the true parameters. We further characterize the explicit convergence rate for each of the considered settings.
p(a;; 1) p(<: の)) ) (1)
KL. [p(x; 0)1p((r Φ)1 = Exje log
This led to the first non-asymprotic convergence rates for EM, independent of parameterization. While this charac- terization of EM included Mixtures of Gaussians, it failed to extend to Mixtures of Regression or MoE. Still, the au- thors analyzed the setting where the distribution of the latent variable conditioned on all other variables is exponen- tial family distribution for which they showed EM con- verged sub-linearly to a stationary point. Our work extends these findings to MoE, obtaining sufficient conditions under which EM converges sub-linearly and linearly to the true parameters.
2) Next, in Theorem 5.1, we narrow in on the special cases of mixtures of 2 linear or logistic experts and show EM is equivalent to Mirror Descent with unit step-size and KL divergence, without requiring any extra projec- tion. This perspective allows us to recover classic MD convergence results which we contextualize for our setting in Corollary I. Finally, we characterize the sufficient conditions for convergence in terms of the Missing Information Matrix (MIM) in Theorem B.2 and, subsequently, the signal-to- the generative model in Theorem B.4.
There have also been works attempting explore and un- derstand how to fit MoE. The foundational paper (Jordan & Jacobs, 1994) was the first to empirically use EM and EM-variants to fit MoE, yielding encouraging results. Then, the follow-up paper by Jordan & Xu (1995) showed that for MoE and Hierarchical MoE with strongly convex negative log- ligeelkllhood objective, EM and EM-like iterations con- verge linearly to the true parameters where the rate constant depends on the eigenvalues of the Hessian matrix. How- ever, the objective is generally non convex, raising doubts about whether the necessary assumptions for their result hold, even locally. Other works have remarked that the na- ture of the gating function creates a form of competition between the experts during training that can lead to local minima. Bayesian methods for MoE include variational learning and maximum aposteriori (MAP) estimation. But, Yuksel et al. (2012) noted that these solutions are not triv- ial due to the softmax gate not admitting a conjugate prior and are prone to getting stuck at local minima. Makkuva et al. (2019) analyzed a variant of the EM algorithm for MoE which consists in 1) first recovering the expert param-
ratio (SNR) of
3) Finally, on synthetic and small scale proof of concept real world datasets, we observe that EM outperforms gradient descent both in terms of convergence rate and the achieved performance. As well as supporting our theoretical results, this re-iterates the power of the EM algorithm for fitting MoE previously suggested by (Jordan& & Jacobs, 1994; Jordan & Xu, 1995).
1.1. Related Work
The EM algorithm (Dempster et al., 1977) is a powerful tool for fitting latent variable models. Previous research on EM has demonstrated that, under mild smoothness assumptions, its parameter iterates converge to stationary point of the log-likelihood objective (Dempster et al., 1977; Wu, 1983; Tseng, 2004). Subsequent research on EM introduced new analytical frameworks to provide specialized guarantees,
2

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
eters using a tensor decomposition method, then 2) whilst freezing the experts, utilizing EM to fit the gating function's 's parameters only. They proved that their approach recovers the true parameters at a nearly linear rate. Then, Becker et al. (2020) experimented with an EM variant algorithm for Gaussian Mixture of Experts called Expectation Information Maximization (EIM) which featured an extra information projection step. They obtained promising empirical results on both synthetic and real world datasets. Our work ex- tends upon previous works by making the direct connection between EM for MoF: and projected mirror descent. We also unveil the sufficient conditions for EM to converge sub-lineat sabblinerrly to a stationary point or true parameter, and for special cases, characterize these sufficient conditions with respect to the SNR of the generative model.
In this paper, we will consider three different settings. The first, General Moł, is the most general setting we con- sider. It comprises all MoE where the distribution of y, 2 conditioned on æ belongs to an exponential family of distribution and enjoys a natural re- re-parameterization ply, 2/æ, 0*) = ply, z; 02) with 0: = n(æ 0*) E Ô. That is to say that the conditional distribution can be written as
p(y, = x. 0*) x exp ( (< 2), 0 + A(0)), (4)
where syy 2), 2z. and A(-) are, respectively, referred to as the sufficient statistic, natural parameterization, and log par- tition of the exponential family of distributions. This setting includes the popular cases where p(y) [x, 2) is Gaussian or multivariate Bernoulli. In the former case, the exponential family of distribution corresponds to Gaussian Mixture.
The second setting, Mixture of Linear Experts, is when
2. Mixture of Experts
plyl 2 = i;; exp { (y x B,)²
Next, we formally describe the setting under consideration for the Mixture of k Experts. The notation used throughout is summarized in the Notation Section of the appendix.
2
as the density function of the normal distribution. This is equivalent to assume that the target x 33 + e where e is additive zero mean Gaussian noise with unit variance.
Data Generation Model: First, the input or feature vector variable x E Rd is sampled based on a probability density function p) Second, given the feature vector æ, latent variable z E [k.], responsible for routing x to the appropriate expert is sampled with probability mass function P(2|z). Fi- nally, given the pair (ar, =). the target y is generated accord- ing to the probability distribution p(yl æ., 2). Hence, the com- plete data distribution is p(z, y. 2) = p(y æ 2)p(2(=))p(() and the joint input-output probability distribution can be written
Finally, Mixture of Logistic Experts, is when the density function p(y) x., 2 i (3;) can be written as
exp(x (3))
P(y = 1 1|x, = = i; (3) = 1 + exp(xx (3;))
(6)
and Ply = 1|x, ==;;;;;;;;;))))) =
Pvy = 1/x, ==; Bi).
Maximum Likelihood Loss: Given the assumed data dis- tribution of (ax,)) y). our goal is to find the set of feasible parameters |3 E Rgd and w E Baxk that maximize the log likelihood function. For ease of notation we detine θ = (B, w) E 2dd the concatenation of all parameters. Specifically, from (2), the expected log likelihood is
p(z,y) = ppxx)
Cp 2( Σ =E[k]
(2)
This is the most general form of data generation under mix- ture of k-experts, but here we focus on the case where x is sampled from a unit spherical Gaussian distribution, i.e.
]
x ~ NO, Ia), and P(z|z)) = P(= [z;; w²) is parameter- ized by a set of vectors (wi,..., w and can be cast as
Ex [())p(=)== Ex,
log
2(ylæ,))P (2(()))) =E|K
Σ
P(2 = =,x;;
e00
E [k].
(3)
Given that only the last term of the sum depends on param- eters O = (w, (3) the negative log-likelihood objective function, L()), that we aim to minimize can be written as
ΣjE(K] e0 '7'
In other words, the probability mass function of the latent variable is the softmax function between the inner product of x with the parameter vectors concatenated as w" = (uri..., wi]) E dd
]
Ex.Y Y log Σ p(yllx,2,,))P(1 w) =E[
Regarding p(y|æx,2), the probability distribution of the tar- get variable y conditioned the input variable x and latent routing variable
Note that as we will discuss in detail, for both mixtures of linear logistic experts, the above objective function is known to be non convex with respect to O. In fact, this is generally true for Mixtures of Gaussian, Mixtures of Regressions, and Mixtures of Experts. In the next section we will discuss the use of the EM algorithm for solving this optimization problem.
(i.e., expert i), we also as- sume that it is parameterized by a vector 33 and we have p(ylæ, 2 = i) = p(y|0,, 2 = i; (;))) Thus, given the con- catenated vector 3 = [Bi.... .... 3 E Rexk we can write 2r z) = p(yla,2; 2; )))
3

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
3. EM for Mixtures of Experts
3.1. EM for Symmetric Mixture of 2-Experts
Next, we present the EM algorithm for MoE. EM takes a structured approach to minimizing the objective [()) in (7). Each iteration of EM is decomposed into two steps as fol- lows. The first step is called "expectation": For current parameter estimate 0e, we compute the expectation of the complete-data log- oeeliielihood with respect to the latent vari- ables, using the current parameter estimates 0' and denote it by (0))0",, i...
So far, we discussed EM for the most general form of MoE. Next, derive EM for the special case of Symmetric Mix- ture of Experts (SymMoE), the focus of our analysis in Section 5. SymMoE is a simplified version of MoE where (i) the number of experts is restricted to 2, represented as & E {-1, and (ii) the experts are symmetric around the linear separator, i.c., 3 := 3 = - This sym- metric structure simplifies the probability density functions introduced earlier, making the subsequent analysis easier to follow. We explore these simplifications in detail below.
Q(0(0") = -Ex [Bz [logp(æ y, 2; 0)]
(8)
Then, in the second step called "maximization", we simply minimize the objective Q(0(0') (or maximize -Q()))")
As we are restricted to two experts, the expression for P(z |a;; =²) = P(= = - 11 |1;; w²)) is
with respect to 0 E Ω and obtain our new parameter as
e0 Tw wi + e w
g++1 == argmin @') PEG
(9)
P( = 1|2; w") e
Since loy le; 0) log p(y] x;; 33 + log logp(z|as (= w),
For ease of notation, we define vh == w wy and repa- rameterize the probability mass function of a given a as
it follows that the EM objective (8) is linearly separable in the parameters 33 and u. Thuss we can rewrite ((0)() as the sum of two functions that depend only on G and w, respectively. Subsequently, the EM update (9) is obtained as the concatenation 0 = (w'+1, 1 where
P(=la; exp T
(11)
1 e 5
Thus, under this simplification, the EM update of the gating parameter w is now given as
+1 argmin- Ex,Y 0. [log p(=|=((;)))) 2 w)ll weRd
1 2
3+ argmin
-8 Ex, E 0 [ogg(y|>z B)]]
w (+1
argmin Ex w.y0 log
BERd
wERd
While the above minimization problem is strongly convex, it does not have closed form solution. In our experiments, we use gradient descent to obtain w
EM has two well understood characteristics: (a) its update always minimize an objective that is an upper bound on the likelihood, and (b) fixed points of the EM update are also stationary points of the likelihood. We now show this below. The original objective, (()). can be decomposed into the difference between the EM objective and another function that is bounded below by O. Specifically, for any 0, Φ E Ω
For the special case of a symmetric mixture of linear experts (SymMoLinE)) the expression for p(y) x 2; 3²)) given in (5) can be simplified as
(y za Gm22
((0)= - Ex, y [log(p(s(<;; 0))]
plyl,,2 2; x exp
(12)
2
- Fx y E. El [log (og(p(ylar (1/ 2; 0))
p(= x; ()) ]
p(y, 2 0)
Under this simplification, the EM update of the expert pa- rameter 33 is now obtained more compactly as
--xxx y E u. log
B Ex, [(2p(= 1|x y: 0*) 1) my)
Denoting H())))) == Fxx,,1E212...d [logp(=)p) xr, y: 0)). it
follows that
Similarly, for symmetric mixtures of logistic experts (Sym- MoLogB) with y E (-1, 13. the expression of p(y (xx 2;) given in (6) can also be simplified as
(()) = Q(0))) ())
(10)
where (0 bounded below by Thus, Q(0(<)) acts as an upper bound on the negative likelihodoo
exp (12 x T33
P(ylæ, "u,,,,,;3)) 2;3²)
(13)
It c 8
Next, applying Jensen's 's inequality shows that (0)))) is
minimized at 0 = & where (0)))) = O. Consequently, it
Under this simplification, the EM update of the expert pa- rameter (3 is now given as
follows that the negative lokelikel gradient matches that of the surrogate EM objective at & = 0, i.e. VC(@)) =
VQ(0 0). This suggests that any stationary point of ((0) is also a stationary point of the EM algorithm and vice versa.
3 argmin Ex. v
Ez v...or
log
1 1
e
B
BERd
yo c
2
TT

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
3.2. EM for Deep and Sparse MoE
vergence regularizer: The Bregman divergence induced by differentiable, convex function h R + R. measures the difference between h(O) and its first-order approximation at 0' i.e.
So far, we have derived EM for the foundational formula- tions of MoE, initially proposed in (Jacobs et al., 1991). While EM is straightforward to derive in these cases, the same does not hold for deep MoE-and and especially not for Sparse MoE. A deep MoE consists of / 2 2 MoE blocks, as defined in Section 2, stacked sequentially. In this setup, the input x passes through the first MoE block and then sequentially through all subsequent blocks produce the output y-
DA 6*) = h(0) - h((0)= (Vh(0') 0*). (14)
Thus, the iterations of MD are derived by minimizing the following expression:
- (( + (VL(0'), 0²) + (0' 0).
(15)
For completeness, and to encourage future work large- scale applications of EM for MoE, we propose a formulation of EM for deep and sparse MoF in Appendix D.4. Instead of solving a separate variable problem at each layer-as as done in classical MoE-we posit that the latent variable E [b] should represent the entire sequence of experts selected across the network. This allows us to construct the EM objective in Equation (8).
Finally, in projected mirror descent, the update completed by projecting the solution obtained from minimizing (15) onto a subspace. As mentioned, this scheme is reasonable if the function approximation using the Bregman divergence serves as an upper bound for the function (()). This can be ensured when the step size is sufficiently small, and the condition of relative smoothness is satisfied (Lu et al., 2018).
An EM-like solution can then be derived for Sparse MoE, where the loss is computed solely from the sequences of experts observed through greedy expert selection. This formulation provides principled approach to training deep and sparse MoE models using EM.
In the upcoming theorem, we formally establish that for general MoE, minimizing the EM objective function de- fined in equation (9) is exactly equivalent to minimizing the subproblem associated with a single step of MD, as de- fined in equation (15), followed by a projection step. This result demonstrates that the EM update for General MoE is essentially performing an MD step in a specific natural re-parameterization space, then projecting the resulting so- lution onto Q. The proof of this result is involved and is deferred to Appendix A.I.
4. Main Result
In the previous section, we derived the EM update for both MoE and SymMoE as the solution to minimizing the EM objective in (8). We further demonstrated that this solution can be decomposed into the concatenation of the respective solutions to two minimization sub-problems. In this section, will show that this update is exactly equivalent to per- forming a single step of the projected Mirror Descent (MD) update on C()) with a unit step size and the KL divergence as a regularizer. To illustrate more clearly that minimizing Q(0)0') in (8) corresponds to a projected MD step on the loss (()) at the point 0¹, we first provide a brief overview of the core concept behind the MD update.
Theorem 4.1. For General MoE, there exists natural reparameterizatioo O. E (17(- 0) 0 E Ω} with
((0) = Ex [L()=)]
(16)
and mirror map A(Oz) such that the EM update in (9) simplifies and is equivalent to the expectation moment pro- jection,
argmin Ex KL [pl 2. =zt1) 10 pu =/m(ae, 0)]] (17) BEN
In most gradient-based methods, the next iteration is ob- tained by minimizing an upper bound of the objective func- tion. For example, in Gradient Descent (GD), the next iterate is found by minimizing the first-order Taylor expansion of the objective at 0², with a squared norm regularizer. Specifi- cally, for minimizing [(),, the GD update with step size r) at 0' is equivalent to minimizing the following function:
where for each x, Øc+1 is obtained from the following MD step,
argminnVVL(O)).. ub 0%) + DD(tb 02), (18) EE
with L(0m) being 1-smooth relative to A(Oz). Further, vabi, 1b2 E Ωi the divergence function D 1b2) is equal to the KL divergence on p(y, 2, 21b);
((0')= + (VC()"), 8t)
10
2m
This function indeed serves as an upper bound for f(-) if 7 S I/L, where L is the Lipschitz constant of (())..
Da(Wi. 422) = KL[p(y, 2 2|1b1))
(19)
Mirror Descent solves a similar sub problem where instead of a squared norm regularizer, we employ the Bregman Di-
Proof sketch. We utilize the assumed property that the con- ditional distribution p(y. 2|12, 0) belongs to an exponential
5

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
family of distributions to decompose the EM surrogate as
Next, we briefty introduce key definitions that will be used later. We say 0! is initialized in locally average-convex region of [(0) with respect to the random variable X, if there exists a convex set θ S Ω containing 0', 0 such that for all Φ, 0 E e
Q(010') - Q(0'00")) Ex [(VL(O)), 0, 02) + (0,, 02]] (20)
The above derivation follows similar steps to Kunstner et al. Theorem and is provided in Appendix A for completeness. We note that because ge = ,y( 0) is not necessarily linear in 2, we cannot conclude that minimizing (1) with respect to results in a direct MD step. Instead, recall the point-wise (in æ) MD iterate, 0t+1 := argming, (VL(02), 0x -0") + Da 0z). Differentiating and setting equal to holds that
Ex [L 2 Ex [L( (VL(02), 62 0z) (23)
where Oz == 17( 0). Furthermore, θ is called a-average- convex relative to h if
Ex [L(dw)] 2 Ex [L())) EEx (VL( On Oz) aDD On)] (24)
VA(0=**> s(x; 0¹)
(21)
Now, thanks to the previously shown correspondence be- tween EM for general MoE and projected MD, we are able to present novel convergence properties of the EM Algo- rithm when applied to General MoE. The theorem that fol- lows provides sufficient conditions and explicit rates for the convergence of EM iterations to a stationary point or true parameters. The proof, adapted from (Lu et al., 2018) 8), is a bit more involved due to the nature of the extra projection step. It is included in Appendix A.2.
Using (2) and the decomposing (L((2)) = VA((2)) (<x; )² we see that (1) can be further simplified to
Q(010') 1(0510))=
Ex - AV Oz 0%) + A(0)) A(()))
Because A(0++1) and A(0%)) only depend on 0', minimizing the above with respect to is equivalent to minimizing the following with respect to 0
Theorem 4.2 (Convergence of EM). Assuming A As- For general MoE with re- re-parameterization given by 0, := 7x 0), strictly convex mirror map A(0z). and if for all 0c+1 01+1, .0 {0% 0z)
Ex (aa02 0z 12
(22)
Finally, substituting the Bregman Divergence induced by by the KI. divergence yields the claim (this derivation is also included in Appendix A.I for completeness).
DA (w . 0
Ex
Ex DA (0 10tt1,0+1) + DA (0z, (0ts00zt)))) 02t1)],
(25)
We remark that our proof for this result is not merely limited to the setting of MoE, but is satisfied for any mixture for which the distribution p(y, = ax 0) of the target and latent variable conditioned on the feature variable belongs to an exponential family of distribution.
then, the EM iterates (0' heam satisfy:
1) Stationnarity. For no additional conditions,
[((') -((0)) (26)
min Ex [D 0z11]] :E(T]
4.1. Convergence Guarantees for General MoE
2) Sub-linear Rate to 0-. If e is initialized in θ, a locally average-convex region of ((0) containing 0*, then
Before moving into the convergence results, we specify the necessary assumptions previously outlined in (Kunstner et al., 2021) for the iterations of the EM Algorithm to be well-defined for our class of MoE models. See Appendix A.2 for more in-depth discussion of the implications of these assumptions.
C( -((0) < Ex ( 1100000.. 0.)
(27)
T
3) Linear Rate to 0-, If 0¹ is initialized in O S Ω, locally average-convex region of L()) relative to A()) that contains 0*, then
A.. The conditional distribution p(y, =00x)) is a steep, min- imal exponential family of distribution and m7(x,)) :) is continuously differentiable function.
[(0") L()*) I -a) D 02)] (28)
The above condition on the initialization to belong in lo- cally average-conver region is satisfied trivially if L(0z)) is almost surely convex relative to A(0a). Such an assump- tion is stronger, but more in line with standard sufficient conditions for optimality of MD.
Ag. The optimal objective function value is bounded below, i.e., (()*) -00, on the constraint set Ω
Ag. The following sub-level sets Pog = 1Φb E Ω
@( < ()) are compact.
6

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
As noted in the literature, EM's convergence is sensitive to initialization. If 0² is initialized within a locally average- convex region of C(O), the EM iterates for the MoE problem will converge sub-linearly to the true parameter. However, e is in region where (()) is strongly average-convex relative to A, the iterates will converge linearly. This last assumption is different from that of prior work, which typi- cally require 0' be initialized in locally strongly convex region.
of EM can directly benefit from this paradigm for solving each iteration's convex optimization subproblem (i.e., Equa- tion (9)) where a GD-style method is typically used. Scaling laws for the mini- batch paradigm suggest that reducing the batch size should be accompanied by proportional reduc- tion in the learning rate (Shuai et al., 2024; Malladi et al., 2022; Goyal et al., 2017). However, since EM equivalent to MD with fixed learning rate of 1, this kind of modular tuning is not directly applicable to EM. This does not imply that there is no optimal batch size for EM. Rather, extending theoretical guarantees to stochastic and mini-batch settings can be approached through the framework of stochastic mir- ror descent (SMD) and mini batch MD (MBMD), both of which have been studied in the context of composite opti- mization (Duchi et al., 2010). We highlight this as an open direction for future research on scalable implementations of EM for MoE.
4.2. Discussion of Main Result
The results show that the EM update (9) for general MoE is equivalent projected mirror descent with unit step- size and KL divergence regularizer on the complete data distribution. We offer the following additional remarks.
First, if we have oracle access to the EM updates for w and 3. EM requires no hyper-paramettrs, unlike GD, which sensitive to the step size. This can be especially advan- tageous for cases where the B-update has closed-form solution (as is the case for linear experts), making EM's benefits over GD more evident. Additionally, while GD regularizes progress based on the Euclidean distance be- tween iterates, EM adjusts progress based on the divergence between probability distributions across iterations. This is often more suitable for latent variable models, where small Euclidean changes may cause large shifts in the mixture distribution, and vice versa.
5. Special case of SymMoLinE and SymMoLogE
In this section, we narrow in two special cases: the Symmetric Mixture of Linear Experts (SymMoLinnF) and the Symmetric Mixture of Logistic Experts (SymMoLogE) models. We first show that minimizing the EM objective function defined in (9) is exactly equivalent to minimizing the subproblem associated with single step of MD as defined in equation (15), with a step size of v 1. This time, we fully specify the mirror map and show it is strictly convex. Unlike the previous result, this correspondence between EM and MD does not feature the extra projection step that was present for general MoF. This allows us to more easily characterize the sufficient condition of EM for linear convergence to the true parameters, relating to the MIM and SNR of the generative model (see Appendix B.5 and B.6). The proof is provided in Appendix B.2.
Second, whereas previous analysis of EM for various set- tings hinged on various types of analyses ranging from verifying obscure conditions to - less reproducible at scale direct proofs, the connection to MD that we unveil greatly unifies the process of analysis and provides more intuition as the inner workings of EM. In particular, as we will discuss in Section 5, our framework for analysis allows to easily provide intuitive conditions for lincar convergence to the true parameters that are based on the MIM and subsequently, the SNR of the generative model.
Theorem 5.1. For SymMoLinf: and SygMologE, there is a mirror map A(o) such that the EM update in (9) is equivalent to
Third, while Jordan & Xu (1995) also demonstrated that the EM algorithm for MoE converges linearly to the true parameters, the sufficient conditions they provided are more restrictive. Specifically, their analysis requires the Hessian to be negative definite, and the convergence rate depends explicitly on its cigenvalues. These conditions are similar in nature to those typically required for GD-type methods. In contrast, our sufficient conditions for optimality align with those of MD, which more accurately captures the conver- gence behavior of EM, as established by the equivalence shown in Theorem 4.1.
argmin VL((0)) 0-) 11)++720 +DA(0,0') Onn
(29)
where, Vé, O E Ω the divergence function (0 0') is equal to the LL divergence on the complete data:
Da 0) KLlp(a, || V,;
(30)
In particular, in the case of SymMoLinE,
(x gg)
+ +log (1 (1+
A(O) - Ex
(31)
2
Finally, large-scale applications often favor a mini-batch training paradigm, as it tends to yield better performance for a given computational cost. Large-scale implementation
while in the case of SymMoLogE.
A(0) Ex [Log ((1 ( (1+c007) (1+60T
(32)
7

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Finally, in both cases, the map A A()) is stricily convex in O and [(0) is 1-smooth relative to A(0).
Log- inzrinood Eror
Log: lestihoodd Emrr
101
EM
12"
Oytamaatn Ereer
As shown in the proof of the result, is evident from (41) and (45) that p(ar,,, v,=00) does not belong to an exponen- tial family of distributions for either SymMoLinE or Sym- MoLogE. Therefore, this result does not simply follow as a corollary of (Kunstner et al., 2021, Proposition 1), but stands as an independent finding, introducing another class of latent variable models where EM is equivalent to MD. A couple of follow-up remarks are in order.
25-1
643
120
332
(a)
(b)
Figure I. Convergence of objective errors L(0') L(0") and L( - (() in Fig la and Fig Ib, respectively, averaged over 50 instances when fitting SymMoLinE.
First, as the loss is 1-smooth relative to A, this validates the choice of 1 for the Mirror Descent update, and subse- quent convergence results from the MD literature. Specifi- cally, in Corollary B.1 of Appendix B.4, we contextualize convergence results from (Lu et al., 2018; Kunstner et al., 2021) for SymMoLinE and SygMoLogE that feature 1) a guaranteed sub-linear rate of convergence to stationary point at no additional assumption, 2) sub-linear rate of con- vergence to the true parameter if initialized within convex region of the loss-function that includes the true parameters, and 3) a linear rate of convergence to the true parameter if initialized within a region of the loss function that con- tains the true parameters and is strongly-convex relative to the mirror map. Then, in Theorem B.2 of Appendix B.5, we further characterize the assumptions required for linear convergence by relating the relative strong convexity of the objective to the eigenvalues of the MIM. Lastly, in Theo- rem B.4 of Appendix B.6, we characterize the existence of the local region of convergence as a function of the SNR of the generative model for the cases of SymMoLinE and SygMoLoggE. We then conclude in Appendix B.7 with a discussion of the implications of the results.
legression Paramettrrs
Recovery Paramaters
Gradient
Stenarcal Erice
641
10c
332
Haratan Murr Bet
Isoatton Muntberr
(a)
(b)
Figure 2. This figure shows the progress made towards the truc |3 o---a"ll2 and [| por--w"le in figures 2a and 2b respec- [|w"|
parameters, L tively, averaged over 50 instances when fitting a SymMoLinE
Synthetic Dataset. We created the synthetic dataset so as to simulate a population setting of SymMoLinE. We sampled 10 data points from an SymMoLinE with known additive unit Gaussian noise (i.e. .N(0, 1)) and true param- eters 3 w E 10 that sastisfy ||@*|12 | - 4. Subsequently, we run full-batch EM, Gradient EM, and GD for 50 iterations and report the results on the training set averaged over 50 instances. Each time, re-sampling the true parameters, initial parameters, and whole dataset. The initial parameters, are randomly initialize within a neighbor- hood of the true parameters, and are consistent across all benchmarks.
6. Experiments
In this section, we empirically validate our theoretical re- sults by comparing the performance of EM with Gradient EM (Algorithm 2), and Gradient Descent (GD, Algorithm 3). Recall that EM for MoE obtains its next parameter iter- ate as the concatenation to the solutions of two minimization problems. Instead, Gradient EM obtains its next parameter iterate as the concatenation of a single gradient update on the respective uub-mininization problems of EM. This dif- fers from GD that obtains its next parameter iterate as the gradient update on the negative log-likelihood objective 0. We evaluate these methods on both synthetic dataset and the real-world Fashion MNIST dataset, consistently report- ing significant improvements for EM and Gradient EM over GD. We also provide mini-batch CIFAR-10 10 experiments with more than 2-experts in Appendix C.I. Note that our aim is not to achieve state-of-the-ee accuracy, but to reiter- ate that EM can be more suitable than GD for fitting specific models.
Figure shows the objective function progress. EM requires fewer iterations to fit the mixture compared to both Gradient EM and GD, with Gradient EM also outperforming GD in fitting time. Figure 2 illustrates the progress toward recovering the true SymMoLinE parameters. Once again, EM requires significantly fewer iterations to fit the mixture compared both Gradient EM and GD, with Gradient EM also taking considerably less time than GD.
Overall, we observe that all three algorithms exhibit a linear convergence rate, both in optimizing the objective function and fitting the true parameters. This aligns with our theo- retical results for MoE and is consistent with findings for Mixtures of Gaussians and Mixtures of Linear Regression
8

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
in high SNR scenarios. To validate our results further, we perform a paired t-test (Ross& & Willson, 2018). For EM and Gradient EM compared to GD, we obtain a T-statistic 2 22 indicating that the difference in final accuracy is statistically significant (p-value ~ 0.000).
Table I.
Performance for single Logistic Expert
Accuracy
Random Invert
Single Expert Single Expert
83.2%
No
10.2%
Yes
Validation Experiment on Fashion MNIST. For the small scale proof of concept experiment on Fashion MNIST (Xiao et al., 2017), we alter the dataset so as simulate a mixture of 2 Logistic Experts. To do so, we perform inversion transformation on the images at random with probability 2- Effectively, the transformation inverts the images from a white article of clothing on a black background to a black article of clothing on a white background. As shown in Table 1, the single expert on the original Fashion MNIST dataset reaches an accuracy of 83.2% on the test set. Meanwhile, the single expert cannot achieve better than accuracy of 10.2% on the altered dataset. This suggests 2-component MoLogE is appropriate for fitting the altered dataset, so long as the ground truth partitioning is linear in image space.
Table 2.
Performance for 2-Component MoLogE
Accuracy
Cross Entropy
EM
78.5%
0.827
Gradient EM
66.0%
1.29
Gradient Descent
62.4%
1.30
The 2-component MoLogE to be trained consists of one Linear gating layer of dimension 2 x 28 28, and 2 logistic experts of dimension 10 x 28 x 28 each. We randomly initialize each linear layer to be unit norm and execute the algorithms on the same datasets and with the same initial- izations. For Gradient EM, the only additional code needed over GD is to define the EM Loss function appropriately, and then perform a Gradient Step on the Gating parameters and the Expert parameters separately as describe in Algo- rithm 2. For EM, for each iteration, we perform several gradient steps in inner loop to approximately recover the solutions to the sub-problems described in (9). We report our findings for the full-batch iteration of the respective algorithms in Table 2 and Figure 3.
(a)
(b)
Figure 3. Test accuracy and objective function, 1 En= 1 and L(0') in 3a and 3b, respectively, averaged over 25 instances for 2-component MoLogE train on Random Invert FMNIST.
7. Conclusion
In this paper, we theoretically addressed the problem of Mixtures of Experts (MoE) with the use of the EM algo- rithm. We first showed that the EM update for MoE could be interpreted as a projected Mirror Descent step on the log-likelihood with a unit step size and a KI. divergence regularizer, extending the result of (Kunstner et al., 2021) beyond complete data distribution belonging to an expo- nential family. Building on this, we characterized different convergence rates for EM in this setting under various as- sumptions about the log-likelihood function and specified when these assumptions held. Lastly, we empirically ob- served that EM can outperform gradient descent in both convergence rate and final performance.
In Table 2, we report the respective final test accuracy and cross-entropy loss values after 100 iterations of EM, Gradi- ent EM and GD for fitting a 2-component MoLogE on the altered Fashion MNIST dataset, averaged over 25 instances. We see that EM boasts a much improved final test accuracy that nearly recovers the single expert accuracy on the origi- nal unaltered Fashion MNIST dataset of 79.2%. Meanwhile, Gradient EM also registers an improvement over GD. In Figure 3, we report the progress made on the accuracy and objective function for the test set over the 100 iterations, av- eraged over 25 instances. As was observed in our synthetic experiment, EM takes considerably less iterations to fit the mixture than both Gradient EM and GD, where the former also takes considerably less time to fit the mixture than GD. To validate our results further, we perform a paired t-test. For EM and Gradient EM compared to GD, we obtain a T- statistic 2 17 indicating that the difference in final accuracy is statistically significant (p-value ~ 0.000).
9

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Impact Statement
T. N., al. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. IEEE Signal processing magazine, 29(6):2-97, 2012.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Jacobs, R. A., Jordan, M. I., Nowlan, J., and Hinton, G. Adaptive mixtures of local experts. Neural computation, 3(1):79- 87, 1991.
Acknowledgements
Jordan, M. I. and Jacobs, R. A. Hierarchical mixtures of experts and the EM algorithm. Neural computation, 6(2): 181-214, 1994.
This research was supported by NSF Encore Tripods (2217069), NSF's AT Institute IFMI. (2019844) and NSF Grant 2007668.
Jordan, M. I. and Xu, L. Convergence results for the EM approach to mixtures of experts architectures. Neural networks, 8(9):1409-1433 1995.
References
Balakrishnan, S., Wainwright, M. J., and Yu, Statistical guarantees for the EM algorithm: From population to sample-based analysis. 2017.
Kunstner, F., Kumar, R., and Schmidt, M. Homeomorphic- invariance of em: Non-asymptotic convergence in kl di- vergence for exponential families via mirror descent. In International Conference on Artificial Intelligence and Statistics, pp. 3295-3303, PMLR, 2021.
Becker, P., Arenz, O., and Neumann, Expected infor- mation maximization: Using the i-projection for mixture density estimation. arXiv preprint arXiv:2001.cm 086682, 2020.
K Kwon, and Caramanis, C. EM converges for a mixture many linear regressions. International Conference on Artificial Intelligence and Statistics, pp. 1727-1736. PMLR, 2020a.
Chen, Z., Deng, Y., Wu, Y., Gu, Q., and Li, Y. Towards understanding mixture of experts in deep learning. arXiv preprint ar arvi:22208. 020222000.28/3,.. 2022.
Kwon, J. and Caramanis, C. The EM algorithm gives sample-optimality for learning mixtures of well-separated gaussians. Conference on Learning Theory, pp. 2425- 2487. PMLR, 2020b.
Daskalakis, C., Tramos, C., and Zampetakis, M. Ten steps of EM suffice for mixtures of two Gaussians. Confer- ence on Learning Theory, pp. 704-710. PMLR, 2017,
Dempster, A. P., Laird, N. M., and Rubin, D. B. Maximum likelihood from incomplete data via the EM algorithm. Journal of the royal statistical society: series B (method- ological), 39(1):1-22, 1977.
Kwon, J., Qian, W., Caramanis, C., Chen, Y., and Davis, D. Global convergence of the EM algorithm for mixtures two component linear regression. In Conference on Learning Theory, pp. 2055 2110. PMLR, 2019.
Duchi, J. C., Shaleh-Shwee Shwartz, S., Singer, Y., and Tewari, A. Composite objective mirror descent. In Colt, volume 10, pp. 14-26. Citeseer, 2010.
Kwon, J., Ho, N., and Caramanis, C. On the minimax opti- mality of the EM algorithm for learning two-component mixed linear regression. In International Conference on Artificial Intelligence and Statistics, pp. PMLR, 2021.
Dwivedi, R., Khamaru, K., Wainwright, M. J., Jordan, M. I., et al. Thcoretical guarantees for em under misspecified gaussian mixture models. Advances in Neural Informa- tion Processing Systems, 31, 2018.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseck-v3 technical report. arXiv preprint arXiv:2412. 19437, 2024.
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(120):1-39, 2022.
Lu, H., Freund, R. M., and Nesterov, Y. Relatively smooth convex optimization by first-order methods, and appli- cations. SIAM Journal on Optimization, 28(1):333-355 2018.
Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: a70e:000667077 2017.
Ma, J., Zhao, Z., Yi, X., Chen, J., Hong, L., and Chi, E. H. Modeling task relationships in multi-task learning with multi-gate mixture- mixture-of-expptts -f-eeess In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery data mining, pp. 1930-1939, 2018
Hinton, G., Deng, L., Yu, D., Dahl, G. E., Mohamed, A.-r., Jaitly, N., Senior, A., Vanhoucke, V., Nguyen, P., Sainath,
10

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Makkuva, A., Viswanath, P., Kannan, S., and Oh, S. Break- ing the gridlock in mixture-of-experts: Consistent and efficient algorithms. In International Conference on Ma- chine Learning, pp. 4304 4313. PMLR, 2019.
Malladi, S., Lyu, K., Panigrahi, A., and Arora, S. On the sdes and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
Orchard, T. and Woodbury, M. A. missing information principle: theory and applications. In Proceedings of the Sixth Benkeley Symposium on Mathematical Statistics and Probability, Volume I: Theory of Statistics, volume 6, pp- 697-716. University of California Press, 1972.
Ross, A. and Willson, L. Basic and advanced statistical tests: Writing results sections and creating tables and figures. Springer, 2018
Shuai, X., Wang. Wu, Y.. Jiang. X., and Ren, X. Scaling law for language models training considering batch size. arXiv preprint arXiv:2412.01505. 2024.
Tseng, An analysis the EM algorithm and entropy- like proximal point methods. Mathematics of Operations Research, 29(1):27-44, 2004.
Wu, C. On the convergence properties of the EM algo- rithm. The Annals of statistics, pp. 95-103, 1983.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms, 2017. URL hwtpthtt//// //attps///rxiv.... org/ /bs/ 7708.07747.
Yuksel, S. E., Wilson, J. N., and Gader, P. D. Twenty years of mixture of experts. IEEE transactions neural networks and learning systems, 177 2012.

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Appendix
Appendix Contents
A) EM, Projected Mirror Descent, and General MoE
A.I) EM is Projected Mirror Descent for General MoE.
A.2) Convergence Results for EM applied to General MoE.
B) EM, Mirror Descent, and Log and SymMoLinEE
B.1) EM is Mirror Descent for SygMoLogE and SymMoLinE
B.2) Proof of Theorem 5.1 for SymMoLinE
B.3) Proof of Theorem 5.1 for SymMoLogE
B.4) Convergence Guarantees of EM for SymMoLogE and SymMoLinE
B.5) Satisfiability of Conditions from Corollary B.1
B.6) Correspondence Between the MIM and SNR for SymMoLinE and SymMoLogE
B.7) Existence of Locally Convex Region
C) Additional Experiments
C.I) Experiment on Grayscale CIFAR-10
D) Algorithms
D.I) EM for MoE
D.2) Gradient EM for MoE
D.3) Gradient Descent for MoE
D.4) EM for Deep and Sparse MoE
Notations
We summarize here the notations used throughout the paper. For clarity, we distinguish between different types of mathematical objects (e.g., vectors, random variables, distributions) and follow standard conventions where possible.
The Kullback-Leibler (KL) divergence of a distribution from distribution is denoted by KL(gl1z) == I (x)) log (g(#)/p(z)) dr. We use lowercase letters such as to denote continuous probability density functions and uppercase letters such as P to denote discrete probability mass functions. The Euclidean (or (2) norm vector denoted by 112. use the compact notation [k] := (1.)....,k).
We denote vectors using bold lowercase letters (e.g., x), and random variables using uppercase letters (c.g., X). Bold uppercase letters (e.g., X) are used to represent either vector-valued random variables or matrices; the distinction between the two is clear from context. For matrix M, we denote its tth eigenvalue by X, and the corresponding eigenvector by u, The minimum and maximum eigenvalues of M are denoted by Amin and Amax. respectively. We use J to denote the d x d identity matrix, and e; to denote the 7th standard basis (unit) vector in Rdd
Expectations are written as Ex(f(X) = pf=)f(() dæz, where the distribution of X implicitly defined by the context. When needed, we make the dependence on parameters explicit by writing X; 0, where 0 denotes the parameters of the distribution of X. The notation V(,, E) refers to the multivariate normal distribution with mean vector μ and covariance matrix Σ.
12

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
A. EM, Projected Mirror Descent, and General MoE.
In this Section, we provide all the complete proofs and discussions relating to results from Section 4.
A.1. EM is Projected Mirror Descent for General MoE.
In this section, we provide the full and detailed proof of the main result, Theorem 4.1. For ease of comprehension and in the hope that this will provide useful insights into other types of non-exponential family mixtures for which EM is also connected MD, we prove our result following the same general ideas that of Kunstner al. (2021, Proposition 1).
For ease of reading, re-state the theorem below:
Theorem 4.1: For General MoE, there exists a natural re pa-aameteerizaiion 0₂ E {n(,, A E 92} with
(()) = Ex [L())=)]
and a mirror map A(O₂) such that the EM update in (9) simplifies and is equivalent to the expectation moment projection,
[ y. = 0t+1 | =/n(ar, 0)]]
argmin Ex KI.
AEN
where for each 0++1 is obtained from the following MD step,
(rgmin(VL(02), 2b 0z)) Da(., Oz).
with L(0x)) being 1-smooth relative to A(2)). Further, Vehi, 2/2 E Ω. the divergence function (1 (b2) equal the KL divergence on p(y, le):
DA(M1 1b22 KL/p(y = 2/1/1]].
Proof. The EM is centered around iterative minimization of the surrogate upper-bound ( For conditionally exponential family of distribution, We can decompose it in terms of the sufficient statistic and tog-partiiion:
Q00@0) = Ey.
Σ
x, x; (() 09)
-
Σ =/<; 0)) + +n(p(()))) P(zly, xx;
--y
Σ ((S(y, 2), 0m) ,(0%)) + In(p(=))) (acp))))P((ppp) P(=ly, xx; 04)
Ex
(Ex| 12 (S(y, 2)]. Oz) A(0²) In(p(ar))
s(m:0')
Ex
(s(;; 0'),0z) + A(Oz) In(p(z)))
Therefore, the above also implies
Ex [V) Ex [VQ(0' (0.)]
(33)
= =Ex VA(O=) s(x; 0*)]
(34)
Simple algebra then shows that
Q(010))
Q(00)0000 0²)
Ex | (s((( (x; 0¹), 0z) + A(Oa) (s(x;; 0') A((2))
Ex [ (s(x;; 0¹), ), ON 0z) + A(O0)) - A(O'2)
Ex [-(s(x;;) -A(O0) 02 0²)) + DA (0n, (0,:0')
Ex [VL(( O2 O',) (O,, 0')]
13

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
where i) adds and subtracts (VA(02), 02 0-)) This is especially important EM minimizes Q(0 1(0 in each iteration. Now recall that [()) Q(010') H()00') and H 0*) ( 0*) s O, it follows that
(()) ((0') Q(000") Q(0")0') ((000') - -100000
S Q(000)) Q(0'(0")
Ex (VV 0n 0 DA 01))
Thus far, have shown that the EM iteration under the considered setting is equivalent to minimizing the following upper bound on (()), w.r.t. 0:
(( + Ex [(VI(0)) 0-- -O2) Da(O=, 0a)]
(35)
Now, recall 0 argming, (VI(0)) Oz 0z)) + (A(0,. 0%) which is the outcome of mirror descent step. Differen- tiating and setting equal to O, it holds that
VA(0=+)) = s a:: 0*).
(36)
Using this and decomposing vl) above, we see that (35) is equal to
[()*) + Ex [VA(O*) s(x; 0,)), 0. 0z) + D 0*, 0²2]
= ((0") + Ex (VA(0") - VA(0+1)) 02 0z) Da 0%) -VVA((++) Oz 0²)) + A(0m) A(0%))
= -(0') + Ex
Thus, minimizing (35) with respect to 0₂ equivalent minimizing
Ex D 0z*1
(37)
Substituting the Bregman Divergence induced by A by the KI. divergence yields the claim.
remains to verify that (a 0=+1) = KI. [p(s [pl. =|0++1 I[p(,. = 00m))] and that the function I.(Og). is 1-smooth relative A(Oz). This follows directly from previous work by Kunstner al. (2021) since I.(0g. is the expected negative log- ledlikkl of y|æ where A(Oz)) is the log-partition of the exponential distribution p(y, 2/0=). For completeness, the derivation is as follows. For Y, Z belonging to exponential family of distribution, the KL divergence can be decomposed directly using (36) as follows to obtain the Bregman Divergence:
p(U; 2;;) Oz) )
p(y, z; 0(+1)
KL ply. = 0 ))pp(,, =00m]
Ey.z 00:01-1
log
(S(y,,)) 2),
0++1 - -Ox) + A((0)) A(0++1)
(VA(OL+)) Ø1+1 0m)) + A(Oz) A(0
(A 01+1).
A.2. Convergence Results for EM applied to General MoE.
In this section, we provide the full proof of Theorem 4.2. Before we begin, we recall and discuss the regularity assumptions previously made. Recall assumptions A,, A2, and Ag:
A The conditional distribution p(y, 2 0z) is steep, minimal exponential family of distribution and 17(x,-) is a continu- ously differentiable function,
Aq The optimal objective function values is bounded below, i.c., _((0) -o0, on the constraint set Ω,
Ag The following sub-level sets 2o == (Φ E Ω @(+10) < Q(0(0)) are compact.
14

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Assumption A₁ serves ensure that L()) Ex [L(0=)] is differentiable, that the EM surrogate is also differentiable and has solution. It further serves guarantee the mirror map A is smooth, ensuring that projecting into the dual space is well-defined. We note that if the re-parametrization function is continuously differentiable in 0, it will hold that A is satisfied for the popular case that p(g, 2/z) is Gaussian mixture (this includes MoE with Gaussian experts). Next, A2 and Ag are classical optimization assumptions that serve to guarantee the solution of the M-step is unique and exists within the constraint set S2, thereby ensuring the EM iterations are well defined.
Further, we make the additional remark that the projection, (17), in Theorem 4.1 can be seen to be equivalent to the following projection over the space of functions on æ.
n1 0²+ 0+1
argmin Ex D 01
(38)
We can see that the set {n((: 0) E [2} is not necessarily guaranteed to be convex. Such a result would require the e--parametrization function 17, 0) to be affine. In situations where this set is not convex, we cannot take our weak generalized Pythagorean identity (25) for granted, and thus we include it as an extra assumption be satisfied for these convergence results. Still, convexity is not necessary for our generalized Pythagorean inequality hold. For instance, ensuring 0*+1 is in the relative interior of (m(: 0) f E Ω}, or simply satisfying the inequality in expectation will suffice, but may be difficult to show.
For ease of reading, we re-state the result below: Theorem 4.2: For general MoE with re- parameterrizatin given by 0x := 17( 0). strictly convex mirror map A(02), and for all 0++1 0 Db E 102,0 ).
Ex [D (D ( 0= Fx [DA (00+1 : 0 + DA (0x: 0
then, the EM iterates (0') )ee(T] satisfy: I) Stationnarity, For additional conditions,
min Ex [Da(02 0z+1 ((0)) C(O))
SE[2]
T
2) Sub-linear Rate to 0" If 0 is initialized in θ, a locally average-convex region of (()) containing 0", then
Ex [D 02)
[( [()*) <
T
3) Linear Rate to 0- If 0 is initialized in θ C Ω. locally average-strongly convex region of L()) relative to A(0) that contains 0" then
[()") - (())) < (1 -a) (0 0%)]
Proof. The proof is divided into three parts that correspond to each of the three sub-results of the corollary. To aid the reader's comprehension, we re-state the cosine law for Bregman divergence (also known as 3-point lemma).
Lemma A.1 (cosine law for Bregman divergence). Assume the mapping A is proper and convex. Then, for all a, b, e á holds that
it
Da(a, Da(a,)) Da(c,b) (VAC)) VA(c), c)-
(39)
Part 1): Stationarity.
We begin by utilizing the result from Theorem 4.1 that the conditional log-likelihood L(O,) is 1-smooth relative to the mirror map A(O.). For all O2 E Ω.
1(0'+2) < L(O%)) + (VL(())) 0t+1 0%))) + DD O.).)
Taking the expectation on both sides with respect to the random feature variable æ yields:
1(0 < L((') + Ex [(VI(O))
g++1 0%) (A'-0-1 0%,)
15

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Combining (33) and (36) then plugging into the above, we obtain
(( S (()*) + Ex (VA(0)) VA((0= 0,00+1 6z)) D ,0,)
Decomposing (0++ 0z) above and canceling appropriate terms yields
(( - C(0") + Ex (VA(0t+)) ) 0 + A(O A((*))
can then be checked that i) is equal to Ex 1 ,0-11) - Da(0 05-1)
Therefore, substituting the above and
utilizing the inequality (25), follows that
C( S C(0") - Ex [DA 0z]
Re-arranging the terms and averaging over T'iteraaions yields the claim:
Ex [D = S -(0 -((")
T 1(0 -((0" L(0)) L(OT) D 0 - Σ
T
min KT
T
Part 2): Sub-linear rate to 0* We begin by utilizing the result from Theorem 4.1 that the conditional log- liellikel L(O,) is 1-smooth relative to the mirror map A(Ox). For all Ox E Qt then apply expectation with respect both sides yielding:
(( < ((0") + Ex ((VL(0)) 0+ 0t)) D00 [==
We then add and subtract 1x (VV
0:)] to obtain
1(0++)) < C(0') Ex (VL(O)) 0++ 0n + 0% 0%) + DA ,0')
We then use the average local convexity assumption, (23), and obtain
c(0)+1) < (()") + Ex (VL( 0 - -00) (A 0.0
Combining (33) and (36) then plugging into the above, obtain
((0 < L("") + Ex (VA(O')) V VA(O+))) 0 0%) + (x 000)
(40)
We now decompose (A(+ 0.) using Lemma A.1 with 0C+1, 0º and == 0++1 and obtain
(VA(0))) V.A )) 0++1 0% + (A 0%)))
= =VA(())) VA(0* +++ 0++1 0 + +D (0))) D (0 00+1 0t+1) (VA(0)) VA(0*+1), 0 0
=((++1)) A(O%) + (V 0¹ 0%) + (- e++1 0% A(
A(Ø++1)
We now add and subtract AVA(O)) 0 and (A(0 0 and group terms to obtain
(1 0 + ( t+1 0%) + (VA VA(O"'1)) 0,+1 0
We now apply Lemma A. again to ii) with a 0 b 0º c g++1 and obtain the sub-result:
(VA(0)) VA 0'+1 0%) + (1+0 0²
(0 #1 0% D Ø'+1) (0'+011 0
16

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Plugging the above equality into (40), we obtain
c( S (())) + Ex [Da(0 0²)) Da(O= 0z+ + ) D ,02+1)
Then, using 25, we obtain
1() - C(0") Ex [Da(0 022) - DA 0 0²+11]
Re-arranging, then averaging over 7 iterations yields the claim:
T
T
T(C(0") (()))) -
Σ (c(0") - -(())) -
Ex D 0²)) D 0=+1)]
1=1
==1
Ex (A 0z) DA 0
Ex ( 0=)
-(0") C(0))
T
T
Part 3): Linear rate to 0"
We begin by utilizing the result from Theorem 4.1 that the conditional log-likelihood L(O,) is 1-smooth relative to the mirror map A(Ox). For all 02 E Ω2 then apply expectation with respect to on both sides yielding:
(( < (()") + Ex ((VI())) 0
0z) + (A .0,
We then add and subtract [x ((V)( 01)] to obtain
(( v L()') Ex [VL(())) 0 0 - 00 + D .0"]]
We then the local a-strongly average-convexity assumption, (23), and obtain
c( s [(0")++++ Ex (VV 0 - + D 002)
a 01)]
Then, following the same steps for the sub-linear case, we Combining (33) and (36), utilize the cosine law for Bregman Divergence D twice, then apply (25) obtain:
L( s C()") Ex [D 0%) DA 0 aD 0²) $ L())) Ex [11 - D (% 0²²)]
Unraveling the recurrence over T' iterations, yields the result:
(((T) - L()") s (1 - "a)")"xxx [D 0.)
This completes the proof.
17

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
B. EM, Mirror Descent, and SymMoLoggE and SymMoLinE.
In this section, we provide all results, proofs, and discussion pertaining to EM for symmetric mixtures of logistic or linear experts.
B.1. EM is Mirror Descent for SygMoLogE and SymMoLinE
In this section, we provide the full and detailed proof of Theorem 5.1. For case of comprchension and in the hope that this will provide useful insights into other types of con-enxoeetiall family mixtures for which EM is also connected to MD, we prove our result following the same general ideas as that of (Kunstner et al., 2021, Proposition 1). We split the proof into two parts (SymMoLinE and SymMoLoggE) which be found in Appendix B.2 and B.3.
For ease of reading, we re-state the theorem below:
Theorem 5.1: For SymMoLinE and SymMoLogE, there is mirror map A(0) such that the EM update in (9) simplifies and equivalent to
( 0*) (2,0, 0"), BEN
where Veb Q the divergence function (0 0*) is equal to the KL divergence on the complete data:
Da(Ø. 0) = KLIp(æ, v, 1)) v;;;
In particular, in the case of SymMoLinE,
(x Tx
3)² 2
+ log 1 cx v0/u)
A(O) Ex
while in the case of SygMoLogE,
A(O) = Ex log ((1+ (1
em "")]
Finally, in both cases, the map A(0) is strictly convex in O and C(0) 1-smooth relative to A(0).
B.2. Proof of Theorem 5.1 for SymMoLinE
Proof. Recall that we consider a component SymMoLinE (see Section 3.1) where = E (-1,1) is the latent unobserved variable, and
1) x LO,,
2) ()) = exp( 2 T 1te Tw
(u-em 2)
exp
3) plyla (uyl,,20 2;
V2
We begin by deriving near exponential form of the complete data probability density function p(æ, 2, v: 0):
p(x.y. y, 0)
= plylæ, 2; 0) ()=(2)))pp())
expplogr(yycz z; (3) log Piz|æ; w) + log p(x)
exp
(u zar ]
log(2rr) +
+
x T w - "woog + =Tw w-= + log p(x)
2
2
2
exp
2²
+
yexx
T
B
22
(xxT 33
2 +
w log(1 e= "n)) + p(z)(x))
1
log (2tr)
2
2
2
yex
[
(a:
(3)
log(1 x TTw) logp(x)
y2
log(2r) }
exp
w
2
2 2
18

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Thus have recovered the decomposition,
w
p(a y, 2;
exp
+ a y, 0)
(41)
year
[3]
S(..
where in a(x, y, 0), the feature variable cannot be linearly separated from the parameter 9:
a(x, y, 0)
w
(x:
log (1 (1+ + c ( p(g)
y²
log log((2) (2x).
2
2
At this point, we pause and discuss the implications of the obtained form. First we recall that for random variable U to belong exponential family, it must satisfy
p(u; 0) h(u) exp ((((u), 0) A
for some s()),(-). 0, A()) that are called the normalization function, sufficient statistics, natural parameters, and log- partition function respectively. To clarify, note that 1) A(O) must be function of the parameters only and cannot depend on and 2) (()) must be function of the variable only and cannot depend on the parameters. In other words, it must be that hu)) and A(0) are linearly separated inside the exp. With the above in mind, we see that SymMoLinE is not an exponential family. Further, we remark that the above formulation showing S(æ, y, 2) is linear with (w, does not extend beyond the symmetric setting of the Mixture of Linear Experts; for k 2 3, this relationships becomes non-linear. This turns out to be problematic for showing EM is equivalent to MD for 2 3. Lastly, note that taking the expectation of a(x, y, 0) over (z.y) y) p(x, y: )*) yields,
(a 3))
+ +log (1+ 11+e0 c²0 ] -C
Ex,y
= -Ex
2
= - - -.
where the above follows from Ex x w
0 and C' == -Ex y [log p(x) - 1²² [ogg(2T]] is not a function of the parameter O. With the obtained form (41), we now continue with the proof.
Part a): Show EM is MD, i.e., argmingen Q(010') = argmingen (VaiiinnnnVeVi (VC(0), 0C) + DA 0*).
Taking the appropriate expectation, the EM objective Q can be written
Q(010') Ex,y y
[2 or [log p(g v, 2;
Ex.y y Ezj 0 [(5(x, y. 2), 0) a(x y, 0)]]
-Ex, y [a(x, y, 0) Ex,y [B [(S(x, 0\]
A(0) (s(0*), +
where s(0') := -xxyy go [S(æ. y, 2)]. As consequence, it is also true that
VQ(0"10)) = VA(0') s(0').
(42)
Continuing, we use the above to simplify the expression for Q(0(0') Q(0'10") that will subsequently give us the MD loss:
Q(010') 1') = A(0) - (s(0"),) (s(0"))
(s(0') 0²) (VA(0)) 0¹) AVA(0"). 0²) A(0) A(0*)
(VQ(0"00)) 8¹) + (*)0,0")
(VC(0), 0*) DD
19

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
where first adding and subtracting (VA(0"), 0 0*) then i) follows from (42) and ii) follows from VL(0)= VQ(00) (see Section for the derivation).
Finally, the first part of our result follows trivially
argmin Q(0l0") argmin Q(0 -()'0'))((''))))
den
dE
Part b): Show Da(0, 0) KL 20
This result follows simply from decomposing KI. (I. v, 0 Ip(a, v, 2; 6)] as follows:
Ex x log pla, 2: 2; 0)
2;
p y, y, 2;
logp(0) y²
(41) (s(0), Φ) A(O) + A(0) + Ex A(t) A()) (s(0),
log log(2T) (TT)
A(Ø) - A(0) - (VA(0), Φ 0).
where i) follows from the fact that Φ 0 minimizes i [logp(æ, y, 6)]. To see this, use Jensen's 's inequality:
- 2 [logp(æ, y, 2; 0)
Jensen's llog F. x 1 [p(x, y, 2; 0)] = loog
p(a, 0) dadzdy
2.1.
Jensen's log
(2yy p(., U: 2 f)dadzdy
=- log(1) 0.
Finally, taking the derivative with respect to Φ and setting equal to completes the proof:
θ E x 1 [log p(x, y. (
0
00
0
5(x, v. 2) + =Ex,Zy10
a(x y. Φ)1
d
=00) VV)(O)
Part c): Show C(0) 1-smooth relative to A(O).
The function C()) is said be 1-smooth relative to A()) for all Φ, 0.0, it holds that
(() < (()) + (VL(Ø, 0 +) + (0,
Recall the following from Section 3. The objective function ((0) is related to the EM objective Q( by (10),
(()) Q(0lØ) (0)(),
where H()00) 2 and ((000) = 0 for all o, 0 f. Consequently, it then holds that for all o, 0.0% Ω
((0) Q(010)
(43)
((0) < Q(1()).
(44)
Recall also from part a) that Q(0(Ø) Q((()) (VL()) 0 Φ) (0, Φ). Then, the claim follows naturally from the above as follows:
(44) (()) S ((0())
(((l0) + (VL()), 6) + (x)0,))
-(()) + (VC()) 0 6) (0,
20

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
follows that (()) is 1-smooth relative A()).
Part d): A()) and the MD objective is convex with respect to 0.
Here, we will show that the mirror descent objective is strongly convex in 0. It is important for this to hold so that the iterations of MD are well-defined; the minimizer of a strongly convex objective exists and is unique.
Note that the mirror descent objective, (VC(O), 0*) + D 0*), , is strongly convex in 0 if A(O) is strongly convex in 0. Therefore, since A(0) given (31) is twice continuously differentiable, strongly convex with respect to 0 if and only if V ((0) = rd. for some 0. We begin:
0a
(x T (3)
V 700)
E
log 1 e:0 (
002
2
0²
(x
T B)
Ex x
log
1 cx0
002
oe T
E x
0
(1teoT Fo
0
Ex
ex T to
E
xx
(I+ex 10
0
Ia
where the last line follows from the assumption that æ is sampled from unit spherical Gaussian distribution: Ex xx T] Ia for x ~ (N, Ia).
From the above, we see that A(0) is strictly convex, and it strongly convex if Ex
FF
(I+c= T w)²
e=
t rld for some
r >> This follows from Lemma B.5 where we show its eigenvalues are bounded below by min Thus, it holds that
{ ? || , |w
V A() min Ω
Ω
Izd-
||1|||
Restricting the feasible set Ω to be all 0 E 2d with ||0||2 S N for some N E [0, oo), it holds that A(0) is strongly convex with respect to 9 on f2
With part proven, this concludes the proof of Theorem 5.1 for SymMoLLnnE. We now prove the same for SymMoL.ogE, referring this section where necessary.
B.3. Proof of Theorem 5.1 for SymMoLogE
Proof. Recall that consider 2 component SymMoLogE (see Section 3.1) where z E [-1, 1) is the latent unobserved variable, and
1) a ~ Ia,,
exp + x T 1tce
2) ((zla;))
3) P(y1,,2;3)) expi( xel1 Tx 14e=
21

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
We begin by deriving near exponential form of the complete data probability density function p(æ, 2, 0):
p( 2, yi 0)
Ply 2; ()) p))
exp(log P(yla, 1,,,,,,;,91 2;13) + log P(z|ar; u) + log p(a:))
1-2411
exp
log
exp(x 13} 1 + cm 1
1 c emi
+
2 +
2
x T w - -logg( + t w)
"=""w) + log ((x)
-exp
1/2 +
log
exp(x T [3]
+
yE+ +1
log
+ 2 +
x: T w - -lo((1 1 + c e²²00 w) ) log p(xr)
2
1 1 ex
2
+ 600 3
2
1z +
a 3 log [ + cx
+)+
2 +
x w log(1 + c=ww) w) + log +logp(x))
exp
2
2
yER
x
w
+ 3
log (1 + c:00 0) (1 + c00 v log p(2))
exp
2
Thus have recovered the decomposition,
Ma 2
p(a, y; 2; 0) = exp
+ al y, },(
(45)
w
where in a(x, y, 0), a cannot be linearly separated from the parameter O:
axx, y. 0) x (w + B) log [1+ e² 0) (1+c="ug w)] +lggp(z)) p(x).
2
Similar to SymMoLinE, we can see that p(ær, y, 2; 0) does not belong an exponential family of distribution. Also, note that taking the expectation of a(x, y, 0) over (x, y) p(æ, , y: 0*) yields,
Ex.yla(x, 0)] = Ex log (1+ e 8) (1+cm1=)]]
= A(0)) C,
where the above follows from Ex now continue with the proof.
(w++)
and C == x [(xpp( is not function of the parameter O, We
From here on, the proofs of part a), part b) and part e) follow identically from that of SymMoLinE, will refer to Appendix B.2 for those proofs. We will now show part d).
Part d)) and the MD objective is convex with respect to 0. Here, we will show that the mirror descent objective is strongly convex in 0. is important for this to hold so that the strongly convex objective exists and is unique.
iterations of MD are well-defined; the minimizer
Note that the mirror descent objective, (VL(0), 0 - 0¹) + (0, 0*). is strongly convex in 0 if A(0) is strongly convex in 0. Therefore, since A()) given in (31) is twice continuously differentiable, it is strongly convex with respect to 0 if and only if V A(0) ž r 12a, for some > 0. We begin:
V A²A(0)=
82
Ex
[og ((1+ (1 (1+ ex e0Tw)]
802
82
(lgg (1 (1+c0 + log (1 cx se"ww)
= Ex
002
Tw
E
a.
(1+e= ww)
0
To
Ex
xx
(1te= To²
Ex
DD
(1+e=Tw)
0
c
Ex
wx
(1+e= Tay?
22

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective From the above, we see that A(O) is strictly convex, and it is strongly convex if Ex wx T
rla and
(I+e
To poy
Ex
= rla for some > 0. This follows from Lemma B.5 where we show their respective eigenvalues
(1te
are bounded below by, min s22
1 [|wllla|
.22
|
1
} and min S2
[19]
1
[1 }
Thus, it holds that
1
V A(0) = min
Ω
.2
Ω
Izd-
1
|
Restricting Q be all 0 with |0| S N for some N E 00), it holds that A(0) is strongly convex with respect to 0 on Ω.
With part d) proven, this concludes the proof of Theorem 5.1 for SymMoLLinR. We now prove the same for SygMol.ogE, referring to this section where necessary.
B.4. Convergence Guarantees of EM for SymMoLogE and SymMoLinE
In this section, we provide the proof's of Corollary B.I. Building on prior work (Lu et al., 2018), we contextualize convergence properties of MD for SymMoLinE and SymMoLogE. Before presenting the result, we briefly review key concepts. We say 01 is initialized in a locally convex region of (()) if there exists a convex set @ Ω containing 0* 0* such that for all Φ, 0 E O
((Ø) 2 L()) + (V -)
(46)
Furthermore, θ is called a-strongly convex relative to h if
[(Φ) [()) (VL(0), Φ 0) + (d.(0
(47)
The corollary that follows provides conditions for convergence EM to (1) stationary point in the KL divergence, (2) the true parameters at sub-linear rate, and (3) the true parameters at linear rate. We further note that the proof is adapted from (Kunstner et al., 2021, Proposition 2, Corollary and Corollary 3) and (Lu al., 2018, Theorem 3.1), provide it here for completeness.
Corollary B.1 (Convergence of EM). For SymMoLinE, SymMoLogE with mirror map A()) given as (31) (32) respectively, and denoting D (0', - x. y, z; +1++)))))0))))))))))) (lp(,, y, 2; 0¹), the EM iterates (0 Jeer satisfy:
1) Stationarity. For no additional conditions,
L(01) L((").
min (1 04)
(48)
EE[T]
T
2) Sub-linear Rate to 0- 0² is initialized in θ, a locally conver region of L(0) containing 0", then
(( L())) < DA(0", 0²)
(49)
T
3) Linear Rate to 0-, If 0 is initialized in e Ω2 locally strongly convex region of ((0) relative A(0) that contains 0", then
[( -(0-) < <1= C(0)))
(50)
Proof. The proof is divided into three parts that correspond to each of the three sub-results of the corollary.
Part 1): Stationarity.
Given Theorrmm 5.1, this proof follows from identical arguments to that of (Kunstner et al., 2021, Proposition 2). We write it below for completeness.
Recall from Theorem 5.1 that +1+1 is obtained as the minimizer of the convex objective, (15):
(VL(0'), 0 0²) (0,00
23

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
As such, differentiating and setting equal to holds that 0*+1 satisfies
VC(0') = VA(0") VA(0++1)
(51)
Further, by the above together with relative smoothness, it holds that
C(0'+)) v ((0") (VL(0"), 0* 0³)
,0)
L(0") + (V AV(")) VA(0*+1),0** -0') + (x(0+t), C(0') - (V 0¹) +((0+1))- A(0')
C( (0'
Thus we have shown that
DA 0 S L((') c(
(52)
The claim then follows from taking the mean over T iterations:
T
T
((0') )(0++))- C(0)) C(OT)
[(0') - L(
mi 1 5 Σ DA +0))) KT
T
Part 2): Sub-linear Rate to 0"
Given Theorem 5.1, , this proof follows from identical arguments to that of Kunstner et al. (2021, Corollary 1) and Lu et al.
(2018, Theorem 3.1). We write it below for completeness.
Here, we assume that ((0) is convex on the set θ In part 1), we used (51) to show,
(VL(0")) 0²+1 -0)) + (0 0') 00'0000*')
where the right hand side is non-positive since the Bregman divergence is non-negative if the inducing function A is convex which it is. Now, starting from relative smoothness, we see that
L( L((*) (VC(0'), 0 0³) (A ,0))
0*)
L(O') + (VC(0'), g++ 0- + 0- 0¹) (x* .00 L()') + (VL(('), g++1 0 0*) + (VL(A"), 0" - 00)) + (Dx
2 [(07) (VC())), 00+1 8t)
0t)
61) [(0") - (VA(0') VA(0*+1) 0t 0") + (0 0*)
where i) follows from convexity of (()) on the set O. Subsequently, we apply the 3-point lemma, Da(0",0") ( 0t+1) + (0' gc+l, VA(0++1) A(O')) + (6 0¹), and obtain,
((( s L()) (0", 0¹) 2,00000*+))
(53)
Finally, the result follows from summing the left and right hand side over T iterations:
T <
T(C(0T) 2(0')) S
Σ
Da(0", 0¹) DD'' 0++1) 5 DD(00,
E=1
Part 3): Linear Rate to 0"
Given Theorem 5.1, this proof follows from identical arguments to that of Kunstner et al. (2021, Corollary 3) and Lu et al.
(2018, Theorem 3.1). We write it below for completeness.
In addition to convexity, we now assume that L(0) is a-strongly convex relative to A(0) on the set O. Specifically, we have that for any o, E θ
(()) 2 L(Φ) + (VC()), Φ) + aD Φ).
(54)
24

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Using the three point lemma again, have
Da(O", , = -a(0", + (0" 0', VA(0)) VA(O++1)) + ,
(5)) DA(0", ,0) + (VL()"), 0- - 0¹) + DA(0', 1
(44
S DD + - aD'a (1 (0',)) (0Dx000 00) + (D'00' ,
-L(0') DD(00' ,
(52
S (11=))
0'+ + -(00)((-CCCC + -(0') [((
< (1 - (0*,)
I (1 D' (0",0''))
Finally, from (53), we see that
((0+11) -(0") s Da(0", (0)(000')= 1000000+
$ (1 a)" (0" -Da(0", 1
(1 a) (0"00,)))
B.5. Satisfiability of Conditions from Corollary B.1
The above result raises an important question: when does a locally convex or relatively strongly convex region of [()) containing 0* exist? Interestingly, this is closely tied to the Signal-to-Noise Ratio (SNR). Before exploring this connection, we first introduce the concept of the Missing Information Matrix (MIM) introduced in (Orchard & Woodbury, 1972).
The MIM relates the level of information the pair (xx y) holds about the latent expert label given parameters 0, and it is formally defined as
M 1 o. l
(55)
Here, Ie is the Fisher information matrix of the complete data distribution and I y. the Fisher information matrix of the conditional distribution of the latent unobserved variable given the observed ones, denoted by
82
I I : = -x,
log log 2, y: 0)
002
-V2Q(010) V²A(0)
82 log P(z (z, v. 0)
I l. -Ex,r
00
-VH(Ø0)) 6=0.
Thus, also holds that the MIM is a function of A()) and H(pl0).
M(0) = 2 {
(56)
Due to the pairwise independence of X and its rotational invariance, there exists an orthonormal matrix R such that = RI-1 R. is positive semi-definite and diagonal and J= R RT is symmetric positive semi-definite (see æ.2.yle
Lemma B.3). As such, the MIM in (55) is also symmetric and positive semi-definite: M() = R ARR JR = R AJR. Note that the MIM quantifies the difficulty of estimating parameters when only x, are observed. To understand its significance, consider the following: large eigenvalues of M indicate that æ, y contain little information about the true value of the latent variable z, making estimation more difficult. Conversely, small eigenvalues suggest that æ, provide enough information to effectively constrain the possible values of z. Thus, the MIM can be seen as analogous to the Signal-to-Noise Ratio (SNR).
In the theorem below, we show how the eigenvalues of M (0) are related to the satisfiability of the conditions for Corollary B.1 regarding the relative strong convexity of L(O) with respect to the mirror map A()).
25

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Theorem B.2. For SynMoLinE and SymMoLogE and their respective mirror mappings (31) and (32), the objective [()) a-strongly convex relative to the mirror map A(O) the convex set @ if and only if
Amaxx (M) (1 a)) for all 0 E G.
(57)
Proof. Recall that [(0) is strongly convex relative to A(0) θ if for all 0, E θ it holds that
[()) (()) (VL(0), 0) a 0).
For (() and A()) twice continuously differentiable, it was shown by (Lu et al., 2018) that this is equivalent to the following bound on the Hessian:
=
Now, using (()) = Q(010) - (o]0), we see that
V V²C(0) v2(Q(010) (0000--H(OP)))) H(
V2Q(010) H
A²A(0) {²H()0))
Therefore, since v²) is symmetric positive definite (proven in Appendix B.1), our condition simplifies to
(1 ( (1-a)v-A(0) = v2 (=)()))
- (1-a) ž V"AA0))"""1((00))000 M(0).
Finally, for from unit spherical Gaussian distribution, we know that M M(0) is symmetric positive-definte (see Lemma B.3). As a result, the above inequality is equivalent to the following bound on the eigenvalues of the MIM:
- 2 Amax (M(())
We now provide the simple Lemma that the MIM is symmetric matrix for SymMoLinE and SymMoLogE.
Lemma B.3 (M) is symmetric). For SymMoLinE and SymMologEE, the MIM is a symmetric matrix, i.e. M()) T
Proof. Recall the assumption that æ is sampled from a unit spherical Gaussian distribution: x ~ No Ia). As such, for any orthonormal x d matrix R, we know that RX ~ (0, Ia); this is called rotational invariance of the Gaussian distribution. Thus, consider orthonormal Ra E Radxd, such that R u e1 ull2 where e, E Rd the 0 vector with a at index j. Now, we can observe that for any w Rdd. Ex
xx
(LteoTuy)
diagonalizable by an orthonormal matrix R:
2 u
@ R. Rus
Ru
Ex
F.
R Ex
Ryex R
(1 +
ex u)²
(1 (1+ c Ry Rawe
Teilllll
En,,
T.l
(1 + e @TTeallalla)
#1 1ul2
E E
0
(1+e=1lu12)*
0
F1 Iulla
EE R.
(14e9)
Where in the above, 1) R = since R = R for orthonormal matrices and 2) non diagonal elements evaluate to because for all * j. the 0 mean random variables X, independent from X;-
26

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Finally, we put the above together show (0) is symmetric. We define the block diagonal orthonormal matrix RM
RM == (Po
Rww 0
R3
From the above, it follows that both for SymMoLinE and SymMoLogE, RM V²)(0) Ryy is diagonal matrix. We can now use this change of basis matrix show the MIM a symmetric matrix:
( (2) "V' H(0))))))
RMM [RMVAA()) RM) (RM ((100))) RM
RRM (RM (O((0))) [RMVAA(0) "R3d RM
-M
B.6. Correspondence Between the MIM and SNR for SymMoLinE and SymMoLogE
The above result states if (0) I2 for all 0 Z and 0* E Ω then the EM updates will converge linearly to 0 This offers unified framework for analyzing EM for MoE, linking the rate of convergence to classical statistical metric. To determine whether EM achieves linear or sub-linear convergence, we need to understand the behavior of ( and ( which indicates the existence and size of the local region where EM enjoys such convergence.
Theorem B.4. For SymMoLinE, the eigenvalues of I=) 2.a.10 belong to the set
( 2..loo)) (0([| @(1
(58)
and 1.o is given as the expectation over Y)) of function that decreasing as a function of 0|]|2
exp
0
2yæ
2
2ya
Ia Ex.y
(59)
2
+ exp
22yrr
Similarly, For SymMoLogE, the eigenvalues of I-1 0,=.110 belong the set
(I 201)) = (0( ((11)))))))))))))
(11131)), 0(13312))
(60)
and I 1, o is given as the expectation (X, Y) of a function that is decreasing a function of [|]
x
x
exp
y/23
y/2
1/ Ex.y
(61)
1 exp
y/
Proof. We divide the proof into two parts. In the first part, we consider the SymMoLinE setting and, in the second part, we consider the SymMoLogE setting.
Part a): SynMoLinE.
For I 1.00 recall that is has the following form:
Ex
2.
0
I v
[1t==
0
Ia
27

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
By Lemma B.5, that I y/ can be diagonalized into the following form:
A1
1.y. 0 RM
RM
0
0
A2
0
0
0
Ia
where RM is an orthonormal rotation matrix and A
and A2
θ (TTo1E)
Therefore, I L. has three
|
eigenvalues given as {o [
1 |
1 } follows that [
has the form
m.
1/\1
0
0
I-1 0.y.2
= RM
0
0
0
RM-
0
1/A2
0
0
Ia/
Therefore, [_] =. has three eigenvalues given as {O (1|201
(1/
(11 12)
Now, for g. first derive compact form for the conditional distribution of the latent variable, p(2 (2, y: 0). From simple Bayes rule and algebraic manipulation, we see that
(2 y: 0) plylæ, 2; 0 1)pppp() 2| ()((2)
p(x, y: 0)
(y-za) 1)2 2
expt ** Tw}
exp
22
1te
(v- g)
exp wx)
(y+0 T g)²
22
exp
2
w
+
V2x
exp
1
exp
(v
(V-24)
18)2
+ =+1
2 x w
exp
(v
[8)a
+ x: T
+ exp
(u+x
2yy T B + =+1
xr T w)
exp
2
exp
[yxx
x w) exp
(-yx (3}
exp
1+1 (2yx' B + x w)
exp
{2yæ
w
exp
2ya
w
exp
+1
2yyx
Now that we have this simplified form, we are able to derive (59) for L.j..y,0:
I Is = -Ex,y EE m
log p(2 (c, y: 0)
002
0²
( 22yz
.0)))
=-Ex.y Ex. F æ.u. e
- log
exp
.
802
2yyr
82
Ex y Ze.u0
log
+ exp
0
002
2yaz
0
Ex.y
log
+ exp
002
2yy
x
exp
0
2yfæ
2ya
2yaa
Ex,y
1 + exp
0
2yyæ
28

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
This expression depends only the random variable (X,Y) Y) and the parameter iterate 0.
Part b): SymMoLogE. Recall that for SogM Laayya has the following form:
Ex
a. T
ca
(1+
1,
Ex
(1+c=
T0)²
By Lemma B.5, that La yl0 can be diagonalized into the following form:
A1
0
0
0
0
0
A2
0
0
0
1, RM
0
0
0
RM
0
0
0
0
A4
where RM is orthonormal rotation matrix and \ Therefore, Im has four eigenvalues given
A2
A3
and h1 0
|
|wi||
|33
[3
{0) |
follows that I-1
2,z,y has the form
1/A1
0
0
0
0
0
1/22
1
0
0
0
[-1
2.2,0/ RM
0
0
0
0
RM-
0
0
0
1/X3
0
0
0
1/Aa
Therefore, 1- 0.2.y/0 has four eigenvalues given as (e (I/w [ θ
(
(1(w]12) e (1B11) θ (113112))
Now, for L= we first derive more compact form for the conditional distribution of the latent variable, P(s 0). From simple Bayes rule and algebraic manipulation, we see that
P(y () () 0 p)p(2:
P(2|x,y0 y:
p(x, y:
exp( y==1 99 exp (1 T
1tee o
1tee
exp 1e (3)
exp in T m0)
exp -+t TTT))
1t
1
1tea
1t
exp
1=+1
33 +
w
2
exp
v+1
w
exp
2
exp
2++1 [yx) B + x
()))
exp
(yar
/3 + a
exp
2+1
y2
exp
1/220
[3
29

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Now that we have this simplified form, are able derive (61) for 2.
0
12/0.y00 Exx y Ez y.
log
P(z y/
002
0²
+
1
-EEx, Y E y.
log
+ exp
.
002
2
ya
8
ac
log
+
exp
002
y/
8²
= Ex
002
log
1
+
exp
1/c
ar
@
<:
ae
exp
y/
y/
= Ex,y
+
exp
[/2]
This expression depends only the random variable (X, Y) and the parameter iterate 0.
Lemma B.S. For (O, Ia) and u E Rd and |ual2 22, the symmetric positive definite matrix
e
Ex
wx
(62)
(1 + c u
diagonalizable by an orthonormal matrix R2 E Raxd and has two eigenvalues, À1, X2 2 that satisfy
\1 θ
(63)
[
A2 θ
(64)
|2//2
Proof. Recall that Gaussian random variables are rotationally invariant. Specifically, for orthonormal matrix R E Rddd and X ~ (0, Ia), it follows that RX (0, Ia). Moreover, R R = RR = Ig- Using this notion, we will 1) diagonalize (62), then 2) evaluate the eigenvalues of (62) as the diagonal elements.
Consider the orthonormal rotation matrix R2, E Rddd that is such that R., u = ei ||u/|2 where ej the jth canonical vector of Rdd Using this change of basis matrix, we can now obtain the diagonal matrix,
)) R, = Ex Ryxx R
ex n. Ru
Ru
Ex
T.
(1 ex u))
(1 ++e* e R Rure
ER
(1 (1+ e @oTelllull))
E m
0
(1t22
Fz 1ill (1+e+11ullla)
0
ER,,
It has only two cigenvalues given in closed form as
3 21 ||ul
3 211ulz ||u/|
A1 Ex,
p(on) )dār
+ t. [|alll))
(1 + th [|ulla]
5 ||//|
5 |u]la
A Ex,
(1 eăt ||ulla p(0) )dā
(1 + قل [||ella)
30

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
The of the proof is spent evaluating tight lower and upper bounds A1 A2 in terms of ||
Part a): Bounds for Å1-
For ž (O, 1) the probability density function is bounded above by 1: p(ā,) S 1. Then we can upper bound Å as follows:
2ffe:llullz p((i))d
A1
11 + e [ulla)
00
00
2fee:/lull=
d
(1 + e mmlla)
2
27.0 - 2, dā
0
4
|
O
For the lower bounds, we will use the fact that - -1Iull2@+a33 can lower bound follows:
2 -11u1//22/2
for x
| oc
and ||2u|22 2 v2. Then,
A1
p( (i) )d
(1 + edillulla)
iff
22 2
dē
(20di||ull2) )²
30
=2 f/c -2/llu//2-
4 dar
2 2m
2
1|u 12a² 12 dā
IiC
2
1ula
Vieef(@i1l ( 1|20|12) - 221/e|1220 -
1/2T|
1u12
1
Ω
Therefore, holds that \ =
|
Part b): Bounds for d2-
For ž (0, 1) the probability density function is bounded above by 1: p(ā.) S 1. Then we can upper bound A2 as follows:
F1llul/a
A2
pl ( )dā
11 + e [ulla))
dā
(1 + e [ulla
2x 21l|ul dāc
2
1
||00|] kuol|2
O
[z/l22
31

---

will constrain Learning Mixtures of Experts with EM: A Mirror Descent Perspective
For the lower bounds, we will use the fact that e (1lu]]2a +0f/2) we can lower bound 22 follows:
/ /2 for z E [u
0 and ||2||2 2 22. Then,
00
e @51//ulla
A2
(1 + e (
@Blllullz
2
dā
J
(26211lulia)
2m
1
Bi B:|1uu12--
-2 dā
2v 2/2T
2m
J
1
00
1| || d
2 2/2x
2T
(Vrerf(ē11 || (2|2)) কमं.
2T 4
eel2
2 Ω
|2/12
Therefore, it holds that A2 = O
Tuula
B.7. Existence of Locally Convex Region
In this section, we further discuss the consequences of Corollary B.1, Theorem B.2, and Theorem B.4.
To begin, we will discuss the sufficient condition for Gradient Descent to converge linearly to the true parameters O* and compare it to that of EM. For Gradient Descent, is well understood that if 0² is initialized convex θ that contains 0" and where (()) is strongly convex, i.e.
LiL( ž aIza for all O E O,
(65)
then the parameter iterates converge linearly to 0", However, as we have shown for SymMoLinE and SygMoLogE, the sufficient condition for EM converge linearly to 0r is slightly different. Instead, we require θ to satisfy that C(0) is strongly convex relative A()) i.e.,
V2L(0) 2 A(0)
for all 0 E e.
(66)
Interestingly, if it holds that A()) is 1-smooth, we see that (66) is weaker than (65), i.c.
2 A(0) < Iea and (65) holds (66).
(67)
But more interestingly, as long as A(0) is y-smooth for some μ > O, it will hold that any set θ satisfying (65) for some > 0 will also satisfy (66) with ā = > O. We note that the converse holds for A(0) strongly convex. In summary, it then holds that, for SymMoLinE and SogMoLogE, EM's sufficient conditions for a linear rate is strictly weaker than that of Gradient Descent when the mirror map A(A) is convex, but not strongly convex.
Next, we will further diseuss the implications of Theorem B.4. In the theorem, we obtain clear lower and upper bounds for the eigenvalues of I.y.= However, it is not clear how to do the same for Is|a,y.0 the trick to rotate the axis with an orthonormal matrix R to simplify the expression will not work here because the distribution of the vector (æ, yæ) is invariant to rotation. Still, there some things that can say for special cases. For this discussion, ourselves to SygMoLoggE. However, the same lines of reasoning also apply to SymMoL First we recall from Theorem B.4 that
a
exp
0
y
1/c3
1/c
I 1sw = EX,Y
1 1
exp
0
yya
32

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
From here, one easy way to approach bounding the above is 1) recall that any outer product of the form uu has single eigenvalue given || and 2) the inner product between two vectors is equal the product of their norms and the cosine of the angle between them, i.e. s u || cos(ds,u). Thus, denoting to be the angle between 0 and the vector (z, yz)" we obtain
e (1+y") a cosl) (|2112211
I ll..v.
3 Ex
Ind
[+ 1 1110112 c00 cos(o)
= Ex
e
I
Denoting :: (11 + y | write the above expectation
1/2
[0 s²e- -00=0 co)(()))))))))))))))))))))))))) p(s, ))dade.
8
Subsequently, the idea is bound p(s, σ) = pls (o))s) in a way that makes integration easy.
The case where w. B E IR is fairly easy. Under this scenario, æ ~ (O,,1) 1) and I l.y.= can be upper-bounded as
x² yx²
ca(w+a0)
I 1/0,y.00 = Ex,
y2a²
(1 + ex(w+rl1)
$ Ex,y
[x2(1 y²)e [=(++(8)] Iz.
Now, recall that E {-1,1}, P(ul=) < 1, = 2012/2
2n 1el lel and see that
V2
22
1s Ex
2x ( -la [1([w-99] P(y 1|xz) +
- (ar P(y = -1xx))
< Ex
2x
e
[a(w-891
+ ---1z(++-))))
2²
A11 +e - ( p(z)dæ |2(
e
x²
00 (lw++++) do (-=(((=
88
11 + | - -3)
(1 + | + 81)3
0
(1 + wl - -1)3
+1+ w + 31)3
Subsequently, together with the fact that 1= 0. I max {0 (1/w1) (11B13) [ }. it holds that the eigenvalues of the MIM are upper-bounded by
fo { 3
( )
3. )²+
8 3 +
3
W
max
+
I+ w +
I 3
1+ w + 3
This special case is closely related the case where parallel to w; similar approach will work.
33

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
C. Additional Experiments
C.1. Experiment on Grayscale CIFAR-10
In this section, we consider an additional mini-batch experiment on grayscale CIFAR-IO with 5-component MoE. The MoE be trained consists of individual experts that each consist of single hidden layer MLP with hidden dimension 100 and ReLU activation. The gating function also consists of single hidden layer MLP with hidden dimension 100 and ReLU activation. We randomly initialize each linear layer have that unit-norm and execute the algorithms the same datasets and with the same initializations. For Gradient EM, the only additional code needed over GD is to define the EM Loss function appropriately, and then perform Gradient Step the Gating parameters and the Expert parameters separately as describe in A Algorithm 2. For EM, for each iteration, we perform several gradient steps an inner loop to approximately recover the solutions to the sub-problems described (9). We report our findings for the mini-batch iteration of the respective algorithms in Figure
We report the respective final test accuracy and cross-entropy loss values after 50 epochs of EM, Gradient EM and GD for fitting 5-component MoE the CIFAR-10. We that EM boasts much improved final test accuracy of 41.6%. Meanwhile, Gradient EM also registers improvement over GD at 37. 37.4%. Finally, GD obtained final accuracy of 34.5%. While these accuracies are themselves not good, the challenge was to find an optimal partitioning of the data and utilize very weak experts. Figure 4, we report the progress made on the accuracy for the test set over the 50 epochs, averaged over 25 instances. As was observed in our synthetic experiment and Fashion MNIST experiments, EM takes considerably less iterations to fit the mixture than both Gradient EM and GD, where the former also takes considerably less time to fit the mixture than GD. To validate our results fürther, we perform paired I-test. For EM and Gradient EM compared GD, obtain a T-statistic 2 22 indicating that the difference in final accuracy is statistically significant (p-value 0.000).
Gradient Descent
Accuracy Performance
Gradient EM
EM
10
3 x 10-1
Accuracy
x 10
10-1
10
20
30
40
50
Iterations
Figure 4, Convergence in predicted label Accuracy for different optimization methods.
34

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
D. Algorithms
In this section, we provide explicit formulations of EM, Gradient EM and Gradient Descent for the context MoE optimization and EM for deep and sparse MoE.
D.I. EM for MoE
(EM): EM takes a structured approach to minimizing the objective (()) in (7). Each iteration of EM decomposed into two steps follows. The first step is called "expectation": For current parameter estimate 0', we compute the expectation of the complete-data likelihooo with respect the latent variables, using the current parameter estimates O and denote it by Q O*), i.e.,
Q(010') = -Ex,y Ez pog 2; 0]]]
(68)
Then, in the second step called "maximization", we simply minimize the objective Q(0 (') (or maximize -Q(0)0") with respect and obtain our new parameter as
0" = argmin Q(0
(69)
BEN
For MoE described in Section 2, log p(y, 2|=; log p(ul2, x: 3) log p(z w). It follows that the EM objective (8) is linearly separable in the parameters B and w. Thus, we can rewrite Q(0)<) (b) the sum of two functions that depend only on B and w. respectively. Subsequently, the EM update (9) is obtained as the concatenation 0*+1 = (w*+1 1
where
w e+1 argmin Ex,y BZj [ogg(= w]]]]
wER'
ge argmin Ex.y Ez|a [logplyl> x; B)]
BERd
Algorithm 1 EM for MoE
Input: Initial 0² E Ωt data: (X,Y) ~ p(a,3; y: 0-)
for t T do
0-Update: Obtain 0=+1 as @++1 arg mingen Q(0 0t)
end for Output: HTH = (wT, ,
D.2. Gradient EM for MoE
Gradient EMM Whereas EM performs the global minimization of the EM objective given in (68), Gradient EM obtains its next parameter iterate as the concatenation of two gradient updates on the sub objectives,
Ex,y
[log P(z [[]]w:))] w)] E o (1ogp(ylz,æ;;3) [)]
(70)
Ex.y
(71)
where the EM objective is given as the summation of (70) and (71).
35

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
Algorithm 2 Gradient EM for MOE
Input: Initial 0¹ E Ωt data: (X,Y) ~ p(a,y, yi 0)), step-size: 71, 72 E (0, ox).
for i 1,...,,7 T: do
3-Update: Obtain g
8
B 3 MBBxyyy w
log p(112,
dg
w-Update: Obtain w
w'+1 w + +221Ex E2200,000
log P(z)a;
Ow
end for Output: 0 = (w=,B²)
D.3. Gradient Descent for MoE
Gradient Descent. Gradient descent is given the global minimizer of the first order approximation of L(0) at A plus quadratic regularizer, i.e.,
(()') + (VC((')) 0 - + 27
0 1/
Differentiating, and solving for equality at yields the well known gradient update.
Algorithm 3 Gradient Descent for MoE
Input: Initial 0² E Ωt data: (X,Y) ~ p(a, yi 0*), step-size: 7 E IR+
for to T
0-Update:
0++1 = 0 7 (('))
end for Output: HTH = (wT,8")
D.4. EM for Deep and Sparse MoE
We begin by formalizing the concepts of Deep and Sparse Mixtures of Experts (MoE), extending the classical MoE framework.
Deep MoE: A deep MoE is composition of 2 2 MoE blocks, denoted as MоE, MoE2, MoE, stacked sequentially. Each block MoE; consists of gating function gi (a;; ws;) and set of experts (Jaj(w;; = The input æ is processed through each MoE block in sequence, producing intermediate representations h1, h2,..., h4, where:
h an (x; wi (1)
h; Σ gi (hi-
for = 2,...,
The final output is hr-
Sparse MoE: The Sparse MoE is a variant of the MoE that is popular in deep MoE applications where, at each MoE block, only a small subset of experts (typically one or few) are activated per input both during training and inference. This is achieved via deterministic or stochastic selection mechanism (c.g., top-k gating). resulting in a sparse latent variable (21, .2, 21) E [ that encodes the sequence of selected experts across the blocks.
36

---

Learning Mixtures of Experts with EM: A Mirror Descent Perspective
EM for Deep and Sparse MoE: We describe EM-like algorithm for training deep and sparse MoE models. The key idea is to treat the expert selection sequence -21 21) latent variable and optimize the expected complete-data log-likelihood. We let 0 (....... wi 11,1. BI,x) denote all model parameters.
In the classical non-sparse MoE setting where we do not choose subset of experts to go through at each layer, the latent variables can only be resolved at the last layer. The EM surrogate then given as follows:
Q(010") EExy [e [log pyke, z,; P(i 0)l
(72)
where the given probability functions is given at the last layer as
plula, -1,..., =1; 0) p (y h1 P(zi,..., 21 20; 0) p(=(((1( -1:10;
In the sparse MoE setting where p experts are chosen at each layer, we re-define the latent variables Z; to be defined over the set of all possible ordered combinations of experts that could have been chosen out of the l available experts. The latent space embeddings is now given to h4 == Eje: 9: (: w (h), The EM surrogate is given by
Q(010") Exy EZ vn pogplulæ -1s 21; 2); x|=;0)]]
(73)
where the given probability functions are decomposed per layer
ply 22, 21,... ... 24; 0) p (y) P(2
21 x 0) P(z 0) ... 211
In practice, constructing the above EM surrogate this would require to make combinatorial number of forward passes through the model evaluate the surrogate loss function. To remediate this, will use the strategy only evaluate the loss on the greedily chosen sequence through the model is the case for GD-type solutions for Sparse MoE.
37


Reviews:
[{"review_id": "rTaYmYgDGX", "reviewer": "Reviewer_KP5h", "summary": "This paper discusses the relationship between the EM algorithm and MoE models. In particular, a relationship is shown between the EM update of the experts and router parameters and Mirror Descent (MD) with a specific expression for the Bregman divergence regulariser.", "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": "We begin by extending our thanks to the reviewer for the very constructive feedback in highlighting important strengths and limitations of our work. We address the main points raised below:\n- **The limitations of the work are not mentioned**\n\n Thanks for your feedback. We would like to mention that the main limitations of our theoretical results are that they hold for the realizable setting and extension to agnostic setting require further investigation. Also, currently we only investigated the full-batch EM, while we believe extension to mini-batch case is possible, as illustrated by our new experiments. We will include these points in the revised paper, and in particular, will mention them in the revised conclusion. \n\n- **More details about the experimental setup (needed to reproduce the results: e.g., batch size, learning rate used) are required in Appendix C.**\n\nThat is a valid point. In the revised paper, we will add all details necessary to reproduce our empirical results to Appendix C. Thanks for your comment.\n\n- **Question 1:**  \n  **Are there still settings where standard GD (and momentum-based variants, such as Adam) has advantages over EM-based optimization in the context of MoE?**  \n\nThis is a very good point. We have theoretically shown that EM requires relative convexity to hold in expectation and within a neighborhood of the true parameters to ensure linear convergence. These sufficient conditions differ from those of GD, which requires strong convexity within a neighborhood around the true parameters for linear convergence. The key insight is that we characterize the conditions under which EM enters and remains in the region where linear convergence holds, whereas it remains unclear under what conditions GD iterates satisfy the local neighborhood assumptions required for linear convergence. That said, this does not imply that EM is necessarily superior to GD or vice versa. We believe there are settings where EM may converge linearly with higher probability than GD and its variants, and vice versa.\n\n- **Question 2:**  \n  **To what extent are the assumptions A₁ - A₃ verified in the experiments?**  \n\n  For the synthetic experiments, assumptions A₁-A₃ are verifiably true; the conditions of Theorem 5.1 hold exactly. As for the experiment on FMNIST, we are not in the realizable setting, so it is safe to assume the assumptions do not hold. Still, the empirical results suggest that EM is robust to non-realizability. We will add this point to the revised paper. Thanks for raising this point. \n\n- **Question 3:**  \n  **What are the practical implications for the future of MoE of this connection with MD?**  \n\n  An example of a practical implication we rigorously showed in Appendix B.6 and B.7 is that it is preferable to initialize the gating parameters of the MoE with a much smaller norm relative to that of the expert parameters. This finding corroborates similar **empirical** findings reported on page 10 of *Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, 2022*. These corroborating results suggest that initializing the gating parameters with a small norm is a universally good heuristic when training MoE (We will add this to the discussion).  \n\n- **Question 4:**  \n  **Did the asymptotic convergence error of GD in Figures 1 and 2 reach that of EM and Gradient EM?**  \n\n  This is a very good question. Interestingly, we did not observe that this was the case, but we cannot deterministically conclude that it wouldn't have if we kept training for many more iterations. \n\n- **Question 5:**  \n  **How do you explain that the Single Expert in Table 1 trained without Random Inversion reaches a better accuracy (83.2%) compared to the three 2-Component MoLogE methods in Table 2?**  \n\n  Firstly, we clarify that results in Table 2 are obtained on the mixed dataset with random inversions whereas the single expert with accuracy 83.2% was trained on the homogeneous dataset. There are three reasons for this:  \n\n  1. **The mixed dataset is more difficult to learn than the homogeneous one.** This is observed as we report the single expert having a test accuracy of 10.2% after training on the mixed dataset, whereas training on the homogeneous dataset yielded a test accuracy of 83.2%.  \n  2. **Information sharing:** Each expert now trains on half as many images as in the base case.  \n  3. **Task difficulty:** Learning both expert and gating parameters to correctly separate the inverted and non-inverted images—and then correctly label each—is a more difficult task than labeling all inverted or non-inverted images. In fact, we observed in our experiments that the best-performing MoE model was learning to route samples based on whether they were inverted or non-inverted.\n\n**Lastly, we thank the reviewer for any additional comments and encourage the reviewer to read \"Summary of Planned Revisions Based on All Reviewers' Feedback\" in rebuttal to reviewer z9yn. Have a nice day!**", "claims_and_evidence": "The paper claims to extend the result of Kunstner et al. (2021) beyond exponential family mixtures, which is that an EM update of MoE can be interpreted as a MD with KL divergence regularizer. This includes Mixture of Linear and Logistic Experts.", "methods_and_evaluation_criteria": "The paper demonstrates the relevance of the result by conducting experiments on both synthetic data and Fashion MNIST, where the goal is to compare the convergence of vanilla GD, gradient-based EM and EM. It is shown that (i) EM converges faster than gradient-based EM, and (ii) EM outperforms the vanilla GD algorithm in terms of final performance (e.g. accuracy or cross-entropy), which makes sense.", "theoretical_claims": "The main theoretical claim is that EM is a projected Mirror Descent algorithm for general MoE, extending the result of Kunstner et al. (2021), which focuses on exponential family mixtures. This conducted in Convergence Guarantees is adapted for General MoE in Section 4.1. The result is then applied to SymMoLinE and SymMoLogE, providing a closed-form expression of the Mirror Map in this special case.", "experimental_designs_or_analyses": "The experiments are conducted on both synthetic data with SymMoLinE and SymMoLogE, and Fashion MNIST with random color image inversion to model the need for 2 experts. It shows the effectiveness of the EM algorithm in convergence speed and performance in those settings.", "relation_to_broader_scientific_literature": "The paper is mainly built upon Kunstner et al. (2021), which derived a connection between EM and MD, and Makkuva et al. (2019), Becker et al. (2020) which relates EM and MoE. It extends Kunstner et al. (2021) to general MoE that includes non-exponential family mixtures.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "Strengths: The paper is well-written and is easy to read. The scope of the theoretical contribution compared to Kunstner et al., (2021) is clear, and the results are relevant in the context of MoE, as it helps to improve their interpretation. The experimental results also highlight the practical advantages of EM-based optimization.\n\nWeaknesses: Limitations of the work are not mentioned. Question 1: Are there still settings where standard GD (and momentum-based variants, such as Adam) has advantages over EM-based optimization in the context of MoE? Furthermore, more details about the experimental setup (needed to reproduce the results: e.g., batch size, and learning rate used) are required in Appendix C.", "other_comments_or_suggestions": "Typo: It should be $p(y | x, z ; \\theta)$ instead of $p(y | x ; \\theta)$ in the second equality when deriving the decomposition (10) in section 3 (L 203). \n\nIt would be great to improve the connection between the experimental and theoretical results. Question 2: To what extent are the assumptions of $A_1-A_3$ verified in the experiments? Question 3: What are the practical implications for the future of MoE of this connection with MD?", "questions_for_authors": "Question 4: Did the asymptotic convergence error of the GD in Figures 1 and 2 reach the one of the EM and Gradient EM?\nQuestion 5: How do you explain that the Single Expert in Table 1 trained without Random Inversion reaches a better accuracy ($83.2  \\\\%$) compared to the three 2-Component MoLogE methods in Table 2?", "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": "EDDWzMnr6G", "reviewer": "Reviewer_yx5r", "summary": "This paper studies the relationship between EM for general MoE with projected mirror descent algorithms. Based on the equivalence between EM and mirror descent, this work provides non-asymptotic convergence rates for training MoE with EM.", "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": "We begin by extending our thanks to the reviewer for the very constructive feedback in highlighting important strengths and limitations of our work. We address the main points raised below:\n- **Concerns:**  \n  **1) The assumptions on the distribution are restricted to the exponential family. 2) The EM method might potentially introduce additional computational costs and is more difficult to implement compared to GD.**  \n  1) Whilst we agree with the reviewer that our results only hold for realizable settings where the conditional distribution belongs to an exponential family, we maintain that the theoretical results obtained in this paper constitute a significant improvement over the previous literature, extending the connection between EM and MD from simpler model classes like Mixture of Gaussians to a more practically relevant class today like that of General MoE described in Section 2.  \n  2) To address this concern, the empirical results in Figures 1, 2, and 3 also include the Gradient EM algorithm (\"GradEM\") which is a variant of the EM algorithm with per-iteration cost identical to that of GD. Even then, we still observe an improvement over GD, though less significant.  \n  \n  **Planned Revision:**  \n  *We will include a thorough discussion of the limitations, discussing the scalability of the EM algorithm to large-scale and sparse MoE.*  \n  \n- **Suggestion:**  \n  **Author provide a more technical highlight of the challenges or analytical innovations involved in establishing equivalence in MoE, compared to [Kunstner et al., 2021].**  \n This is a very good point. We clarify that in Kunstner et al. Theorem 1, the assumption that $p(y,z,x|\\theta)$ is an exponential family of distributions allows to decompose $Q(\\theta|\\theta^t)$ to show direct equivalence between EM and MD. In our results, we relax this assumption, only requiring the conditional distribution $p(y,z|x, \\theta)$ to belong to an exponential family of distributions (this includes more complicated mixtures like MoE). With this relaxation, we can only decompose $Q(\\theta|\\theta^t)$ inside the expectation with respect to x. Then, we utilize a  decomposition on $\\nabla L$ to relate the expression inside the expectation to the point-wise (in x) iterations of MD. Finally, this allows we show direct equivalence between the iterations of EM and a projection over the iterations of MD.  \n  \n  **Planned Revision:**  \n  *We will make space in the main text to provide a proof sketch of Theorem 4.1 that will highlight the proof techniques that allow us to bridge this key gap in the literature. See our proposed proof sketch in the screenshot: https://prnt.sc/wMDRIY7neary.*  \n  \n- **Questions:**  \n  **1) In line 147, could an additional parameter be introduced to consider the distribution $p(z|x;w) \\propto \\exp\\{\\gamma x^\\top w\\}$?**  \n  We clarify that our results hold for any probability functions $p(z|x;w)$ and $p(y|z,x;\\beta)$ so long as $p(y,z|x,\\theta)$ belongs to an exponential family. Since the norm of w is unbounded in our setting, introducing the new parameter $\\gamma$ does not provide a generalization of the problem.  \n  \n  **2) The analysis in the paper relies on the convexity of $L$. How can the obtained results be further interpreted in the context of non-convex optimization?**  \n\n  We clarify that none of our results explicitly require that the objective function $L$ is convex, as such an assumption does not generally hold for MoE. Instead, we require weaker conditions; Theorem 4.2 requires that the algorithm is initialized inside some neighborhood around the true parameter $\\theta^*$ that satisfies inequality (20) (or (21)), wherein a similar notion to convexity is maintained in expectation over feature variable x (similarly for Corollary B.1). We emphasize that the form of convergence results we obtain is consistent with previous works on EM such as [Kunstner et al. 2021, Balakrishnan et al. 2017].  \n  \n  If it is assumed that $p(y|x,z;\\beta)$ is Gaussian, our result in Theorem 4.1 can be readily applied to the setting where the experts are neural networks. In appendix B, we uncovered an interpretable result for SymMoLogE and SymMoLinE that, when training MoE with EM, one should initialize the gating parameter with a norm much smaller than that of the expert parameters, a finding that was observed empirically on page 10 of [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, 2022].\n  \n  **Planned Revision:**  \n  *We will enhance the discussion of the class of models for which our results hold and how our conditions for convergence differ from standard assumptions of convexity. We will also add a discussion on how to scale our theory to more complex experts such as neural networks.*\n\n**Lastly, we thank the reviewer for any additional comments and encourage the reviewer to read \"Summary of Planned Revisions Based on All Reviewers' Feedback\" in rebuttal to reviewer z9yn. Have a nice day!**", "claims_and_evidence": "Yes, this work provides solid proofs to the theorems and claims.", "methods_and_evaluation_criteria": "This paper mainly focus on theoretically understanding the relationship between the EM method and the mirror descent method. Their theory demonstrates that the EM method is theoretically effective for training MoE.", "theoretical_claims": "I briefly checked the proof of Theorem 4.2, the main theorem.", "experimental_designs_or_analyses": "My expertise in MoE experiments is limited, making it difficult to assess their validity. However, I believe the experimental results can support the theoretical claims to some extents as a theoretical paper.", "relation_to_broader_scientific_literature": "The key contributions in this paper primarily follow [Kunstner et al., 2021]. This work extends the results for MoE model and considers the mixtures of regression setting.", "essential_references_not_discussed": "I am not aware of any missing references.", "other_strengths_and_weaknesses": "Strengths:\n\n- This paper is well-written and easy-to-follow. It clearly presents its contributions, and the relevant theoretical assumptions are well-stated.\n- This work provides a new approach for analyzing the convergence rate of MoE training when using EM method.\n- When the $\\beta$-update in the EM method has a closed-form solution (as is the case for linear experts, mentioned in the paper), the EM approach does not require tuning learning rates, potentially offering a more robust learning algorithm.\n\nWeakness:\n\n- The assumptions on the distribution are restricted to the exponential family, which may be overly idealized compared to realistic applications.\n- The EM method might potentially introduce additional computational costs and is more difficult to implement compared to the gradient descent method, limiting the applicability of the proposed method in the paper.", "other_comments_or_suggestions": "I suggest that the author provide a more technical highlight of the challenges or analytical innovations involved in establishing equivalence in MoE, compared to [Kunstner et al., 2021]. This could enhance the paper’s impact and make it more accessible for future research to build upon.", "questions_for_authors": "- In line 147, could an additional parameter be introduced to consider the distribution $p(z \\mid x; w) \\propto \\exp(\\gamma x^\\top w)$?\n- The analysis in the paper relies on the convexity of the function L. How can the obtained results be further interpreted in the context of non-convex optimization in neural networks?", "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": "zR0W0dPfdL", "reviewer": "Reviewer_z9yn", "summary": "This paper focuses on integrating MoE optimization and EM algorithm. The authors first proved the theoretical guarantees of EM algorithm for training MoE models. Then, the authors focus on the special case of mixture of 2 linear or logistic experts and analyze the guarantees for the linear convergence. Next, the authors conduct experiments on small-scale data to demonstrate that EM algorithm can outperforms the GD algorithm.", "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": "We sincerely thank the reviewer for their constructive feedback in highlighting key strengths and limitations of our work. Below, we address the main points raised:\n\n- **The provided results are synthetic and small-scale datasets. It is hard to distinguish whether the proposed method is still effective on large-scale datasets. Moreover, this paper uses two special cases to conduct the experiments, which is not so convincing. Moreover, the authors just compare EM algorithms with GD algorithm, which also limits the contribution of this paper.**  \n\n  We agree with the reviewer that the paper does not feature large-scale experiments which could help motivate practical use of EM for large-scale MoE training. As this paper is a purely theoretical work within the broad EM and MoE literature that provides correspondence between EM and MD and subsequent SNR conditions for linear convergence, we chose to restrict the empirical experiments to supporting our theoretical claims on the two sub-cases considered in Appendix B.  \n\n  Nonetheless, based on your review, we scaled our experiments from training a mixture of 2 logistic experts model on FMNIST to training a mixture of 5 lightweight CNN models on CIFAR-10 (620k parameters, batch_size=2500, 50 epochs, and 25 instances). We report that compared to GD, EM had an average final test accuracy of **41.6%**, which is **7.1%** higher than GD. Similarly, we report that compared to GD, Gradient EM had an average final test accuracy of **37.4%**, which is **2.9%** higher than GD, while GD had a final test accuracy of **34.5%**. A screenshot of the resulting plot can be found at: [https://prnt.sc/ZXBLKD2e8bN2](https://prnt.sc/ZXBLKD2e8bN2).  \n\n  Moreover, as per reviewer kaEY, we are also providing a paired t-test to check for statistical significance.  \n  - For **EM vs GD**, we report a **T-statistic of 25.15** (p-value of **0.000**).  \n  - For **Gradient EM vs GD**, we report a **T-statistic of 22.88** (p-value of **0.000**).  \n\n  We conclude that our results are **statistically significant**.  \n\n  **Planned Revision:**  \n  *We will include our CIFAR experiments in Appendix C to highlight preliminary results on scaling EM to 'large'-scale MoE.*  \n\n- **An additional notation table is helpful to follow the proof of this paper.**  \n  **Planned Revision:**  \n  *We will include a notation section and table of contents at the beginning of the Appendix.*\n\n# Summary of Planned Revisions Based on All Reviewers' Feedback\n\nAs part of our commitment to addressing the reviewers' valuable suggestions and concerns, we have outlined a set of planned revisions. We sincerely appreciate the constructive feedback and believe that our responses, along with these well-defined and feasible revisions, will demonstrate the paper’s readiness for acceptance into the conference.  \n\n1) **Readability:**  \n   To enhance clarity and accessibility, we will:  \n   A) streamline the main text by moving portions of the background on EM theory and MoE parameter estimation to the supplementary material,  \n   B) introduce a dedicated **notation section and table of contents** in the Appendix to improve readability, and  \n   C) use the additional space in the main text to provide a **proof sketch of Theorem 4.1** (see our reply to Reviewer yx5r to see our suggested proof sketch), highlighting key proof techniques that address previous gaps in the literature.  \n\n2) **Discussions:**  \n   We will:  \n   A) provide a thorough discussion regarding the **scalability of EM to large-scale and sparse MoE** as well as challenges in extending our **theoretical framework to mini-batch EM**, and  \n   B) expand our discussion on the class of models for which our results hold (general MoE), clarifying how our convergence conditions in Theorem 4.2 and Corollary B.1 compare to standard convexity assumptions for GD. Specifically, we will elaborate on their relationship to $\\lambda$-(relative) convexity in the discussion following Theorem 4.2.  \n\n3) **Experiments:**  \n   To strengthen our empirical analysis, we will include **error bars in the plots and perform a statistical significance test** (e.g., a paired t-test) to support the claim that EM provides a meaningful improvement over GD (we have obtained these results, please see our reply to reviewer kaEY and z9yn).  \n\n   In Appendix C, we will:  \n   A) add a discussion connecting theoretical results to the experiments such as verifiability of the linear rate and satisfiability of assumptions $A_1$-$A_3$,  \n   B) include **initialization experiments** (see our reply to reviewer kaEY),  \n   C) provide our **experiments on mini-batch training of a mixture of 5 lightweight CNNs on CIFAR-10** (see our reply to reviewer z9yn), and  \n   D) provide all necessary implementation details in Appendix C (such as learning rate, precise model architecture, number of instances, etc.).  \n\nWe are looking forward to hearing your valuable feedback on our responses and proposed revisions.", "claims_and_evidence": "This paper provides detailed theoretical proof and experimental evaluation on special cases.", "methods_and_evaluation_criteria": "The provided results are synthetic and small-scale datasets. It is hard to distinguish whether the proposed method is still effective on large-scale datasets. Moreover, this paper use two special cases to conduct the experiments, which is not so convincing. Moreover, the authors just compare EM algorithms with GD algorithm, which also limits the contribution of this paper.", "theoretical_claims": "The theoretical claims in this paper are convincing. I have checked all the claims in the main pages.", "experimental_designs_or_analyses": "As mentioned in “Methods And Evaluation Criteria”, the authors just select two special case to conduct experiments. Moreover, the used datasets are synthetic and small scale. Furthermore, the compared method is only GD algorithm. All of these settings limit the contribution and convincing of this paper.", "relation_to_broader_scientific_literature": "The key contribution can provide some insights for training better MoE methods or designing MoE-based methods.", "essential_references_not_discussed": "The reference is sufficient", "other_strengths_and_weaknesses": "Please refer to the “Methods And Evaluation Criteria” and “Experimental Designs Or Analyses” parts", "other_comments_or_suggestions": "1.\tAn additional notation table is helpful to follow the proof of this paper. \n\n2.\tExperiments on large-scale datasets or diverse datasets can enhance the convincing of this paper.", "questions_for_authors": "Please refer to the “Methods And Evaluation Criteria” and “Experimental Designs Or Analyses” parts", "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": "X1Vxg3RiXs", "reviewer": "Reviewer_kaEY", "summary": "In this paper, the authors discussed the relationship between the EM algorithm and mirror descent in the context of the mixture of experts (MOE) learning. In the beginning, the authors proposed an overview of the EM-based parameter learning procedure in the MOE. On this basis, the authors then introduced the projected mirror descent (PMD), bridged the relationship between EM-based MOE learning and PMD-based MOE learning, and derived the convergence property of the EM-based MOE with the help of PMD. After that, they extend their proposed theorem to the Symmetric Mixture of Linear Experts (SymMoLinE) as a special case. Finally, various experiments are conducted to demonstrate the efficacy of the proposed approach.", "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": "We sincerely thank the reviewer for their constructive feedback in highlighting key strengths and limitations of our work. Below, we address the main points raised:\n\n- **Theory:**  \n  **The reviewer raised concerns about the nomenclature, use of KL vs L2 regularizer, the relationship between the MIM and the generalized result, and whether $\\lambda$-convexity can be used to simplify and generalize the convergence condition.**  \n  We agree that adding a dedicated notation section and a table of contents in the supplementary materials would improve clarity. Regarding the convergence conditions in Theorem 4.2, we acknowledge that assuming $L$ is $\\lambda$-convex relative to the mirror map $A$ would suffice to establish the claim. Since $L$ is $\\lambda$-convex almost surely for all $x$, the same holds in expectation, satisfying our convergence conditions. However, this imposes a stronger condition on $L$, which may not hold in general for MoE. Instead, we provide weaker conditions that relax this requirement to hold (1) in expectation with respect to $x$ and (2) only within a neighborhood around the true parameters. We verified these assumptions in Appendix B for special cases of MoE (SymMoLogE and SymMoLinE). A similar argument applies to the MIM conditions. Lastly, KL divergence regularizer makes sense in the context of MoE since it is invariant to the choice of parametrization while the L2 norm is not (see Kunstner et al., 2021 for an in depth discussion).\n\n  **Planned Revision:**  \n  We will include a notation section and table of contents in the Appendix and allocate space after Theorem 4.2 to discuss how our conditions relate to standard $\\lambda$-relative convexity.\n\n- **Mini-Batch:**  \n  **In deep learning, training is typically done via mini-batches, which introduce noise that may impact model performance. This issue has not been considered.**  \n  We acknowledge the limitation of focusing on full-batch EM and GD. However, our study maintains consistency by considering full-batch versions in both theoretical and empirical analyses.  \n  Theoretically, since EM in our settings is equivalent to projected Mirror Descent, any theoretical result for mini-batch projected Mirror Descent would apply to EM.  \n  Empirically, following your recommendation and that of other reviewers, we conducted an additional experiment on the mini-batch training of a mixture of 5 lightweight CNN experts on CIFAR-10. The results align with our FMNIST experiments (see our reply to reviewer z9yn for details).\n\n  **Planned Revision:**  \n  We will discuss challenges in extending our theoretical results to mini-batch EM and include our mini-batch empirical results on CIFAR-10 in Appendix C.\n\n- **Initialization:**  \n  **The reviewer suggests that initialization may significantly impact model performance and asks whether asymptotic convergence analysis could substitute for these experiments.**  \n  Theorem 5.2 and Corollary B.1 guarantee that EM will converge sub-linearly to a stationary point of $L$ regardless of initialization.  \n  Empirically, we tested the effect of initialization in the synthetic setting. We observed that for \"close-enough\" initialization, EM converged super-linearly to the true parameters, while for poor initialization, EM and GD converged sub-linearly, sometimes to suboptimal solutions. These findings align with previously known empirical results on EM (Kunstner et al., 2021; Xu and Jordan, 1996).\n\n  **Planned Revision:**  \n  We will provide our empirical results on initialization for the synthetic setting in Appendix C.\n\n- **Three key areas for improvement:**  \n  **1) Reorganizing the main text for clarity while moving relevant parts to the supplementary material for completeness.**  \n  **2) Reordering key results for better comprehension.**  \n  **3) Strengthening empirical results with additional statistical metrics.**\n\n  **Planned Revision:**  \n  1) We will move lines 185-220 (left column) to the appendix and condense Section 5.1. For completeness, we will add a dedicated Appendix section covering all relevant introductory details.  \n  2) We will provide a proof sketch for Theorem 4.2 (see our reply to reviewer yx5r).  \n  3) Following your recommendation, we re-ran the FMNIST experiment (50 iterations, 25 instances, full batch) to include error bars and a paired t-test (implemented using scipy's ``ttest_rel\"). We are also conducting this analysis for synthetic and CIFAR-10 experiments. Our statistical results are:  \n     - **EM vs GD:** T-statistic = 24.32, p-value = 0.000.  \n     - **Gradient EM vs GD:** T-statistic = 17.81, p-value = 0.000.  \n     We conclude that our results are statistically significant. See the updated Figure 3:  \n     [https://prnt.sc/p9rfbnIfycX2](https://prnt.sc/p9rfbnIfycX2).  \n\n**Lastly, we thank the reviewer for any additional comments and encourage the reviewer to read \"Summary of Planned Revisions Based on All Reviewers' Feedback\" in rebuttal to reviewer z9yn. Have a nice day!**", "claims_and_evidence": "Nearly all the claims made in the submission are supported by clear and convincing evidence. Nevertheless, the symbols seem quite strange in this manuscript. For example, in Eq. (1), the authors introduced $\\phi$ but did not state what $\\phi$ means. The parameter of some MOE models? In addition, on page 6, right column, `Additionally, while GD regularizes progress based on the Euclidean distance between iterates, EM adjusts progress based on the divergence between probability distributions across iterations. This is often more suitable for latent variable models, where small Euclidean changes may cause large shifts in the mixture distribution, and vice versa.`, although the reviewer buys in this statement, the authors should add extra citations or add related figures to back up this phenomenon.", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria make sense for demonstrating the derivation of convergence analysis, the authors also provide experiments on the Fashion MINST dataset to demonstrate the efficacy of the proposed approach.", "theoretical_claims": "The reviewer attempts to check the derivation of the theorem. But there remain some issues to be addressed.\n1. At the beginning, the authors are suggested to add a nomenclature to better facilitate the derivation.\n2. On this basis, some symbols should be rectified. For example, $\\mathbb{E}_{Z|\\boldsymbol{x},y;\\theta^t}$ should it be $q(Z|\\boldsymbol{x},y;\\theta^t)$? where the expectation operator is applied for some probability density function with the input $\\boldsymbol{x}$ and $y$?\n3. In appendix B.6, to the reviewer's understanding, the introduction of MIM is for proving the locally convex or relatively strongly convex region of $\\mathcal{L}(\\theta)$. Can we extend this condition to the generalized MOE model?\n4. Can we directly use the concept of $\\lambda$-convex during the proof to simplify and generalize the convergence condition?\n\n---\nReferences:  \n[1]. { Euclidean, Metric, and Wasserstein } Gradient Flows: an overview", "experimental_designs_or_analyses": "1. In the context of deep learning, our training is always conducted by a mini-batch paradigm. Thus, during the evolution of mirror descent, the parameter may be perturbed by some noise. Specifically, according to the reviewer's understanding, the results should be:\n$\\theta_{t+1} = \\theta_t - lr\\nabla_\\theta\\mathcal{L}(\\theta) + \\mathcal{N}(0,\\sigma^2)$. The noise may improve the model performance to some extent, and the authors have not considered this issue.\n2. To the reviewer's understanding, the change of initial value may affect the model performance greatly during model training. However, this part of the experiments has not been investigated. It would be better to add some asymptotic convergence analysis to sidestep conducting this experiment?", "relation_to_broader_scientific_literature": "The MOE is of great importance in the era of LLM, and the authors proposed a novel understanding of why the GD-based MOE learning procedure does not take effect and how to improve it from the perspective of mirror descent.", "essential_references_not_discussed": "The reviewer thought that the GD-based approach is mainly focused on the first-order parameter learning procedure, and the GD utilized in the manuscript is a non-mini-batch. Thus two related works may be considered:\n1. Optimizer with Second Order Information like reference [1].\n2. Minibatch sampling effect for parameter learning [2]. \n---\nReferences:  \n[1]. Optimizing Neural Networks with Kronecker-factored Approximate Curvature (ICML 2015)\n[2]. Stochastic Gradient Hamiltonian Monte Carlo (ICML 2014)", "other_strengths_and_weaknesses": "### Strengths\n1. The topic is related to the ICML conferences.\n2. The proposed approach is interesting.\n3. The derivation is rigorous.\n\n### Weaknesses\n1. The organization of the main content should be revised. For example, the parameter learning of MOE should be more compact, this part can be moved to supplementary material. \n2. The illustration of the convergence analysis can be reformulated. For example, an introduction of SymMoLinE can be given then the generalization to other MOE structures can be given in the latter content.\n3. Figures 2 and 3 have not proposed the error bars.\n4. The results in Tables 1 and 2 have not included the statistically significant test.", "other_comments_or_suggestions": "See the abovementioned chat window.", "questions_for_authors": "See the abovementioned chat window.", "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: This paper revisits the idea of learning mixtures of experts (MoEs) using the classical expectation-maximization (EM) algorithm. The main contribution of the paper is making the connection between the EM for MoEs and projected mirror descent with unit step size and a KL regularizer; this connection enables the development of convergence rates leveraging theory from mirror decent. The paper concludes with empirical studies comparing to MoEs trained using gradient descent, with the goal of highlighting that EM is more suitable than gradient descent for fitting specific models. 

Overall reviewers were positive about the paper and unanimously have recommended acceptance. The reviewers agreed taht the paper is well-written, the scope of the theoretical contribution is clear and its relationship to related work is well-discussed. The reviewers also appreciated the experimental results as a means to highlight practical advantages of EM-based optimization. Thus, I recommend accepting this paper.

Some points that the reviewers would like to see discussed in the revision:
* Limitations of the approach
* More discussion comparing to the setting of gradient descent
* More discussion regarding the empirical setup for reproducibility purposes
* Some revisions to streamline the theoretical results discussed

Overall, most of these points were addressed in the author rebuttal, and the authors provided a very thorough plan for revising the paper according to the reviewers' comments.
Decision: Accept (poster)