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:

• Theory:
  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.
  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).

  Planned Revision:
  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.

• Mini-Batch:
  In deep learning, training is typically done via mini-batches, which introduce noise that may impact model performance. This issue has not been considered.
  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.
  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.
  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).

  Planned Revision:
  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.

• Initialization:
  The reviewer suggests that initialization may significantly impact model performance and asks whether asymptotic convergence analysis could substitute for these experiments.
  Theorem 5.2 and Corollary B.1 guarantee that EM will converge sub-linearly to a stationary point of $L$ regardless of initialization.
  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).

  Planned Revision:
  We will provide our empirical results on initialization for the synthetic setting in Appendix C.

• Three key areas for improvement:
  1) Reorganizing the main text for clarity while moving relevant parts to the supplementary material for completeness.
  2) Reordering key results for better comprehension.
  3) Strengthening empirical results with additional statistical metrics.

  Planned Revision:

  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.
  2. We will provide a proof sketch for Theorem 4.2 (see our reply to reviewer yx5r).
  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:
     • EM vs GD: T-statistic = 24.32, p-value = 0.000.
     • Gradient EM vs GD: T-statistic = 17.81, p-value = 0.000.
       We conclude that our results are statistically significant. See the updated Figure 3:
       https://prnt.sc/p9rfbnIfycX2.

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!

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:

• 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.

  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.

  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.

  Moreover, as per reviewer kaEY, we are also providing a paired t-test to check for statistical significance.

  • For EM vs GD, we report a T-statistic of 25.15 (p-value of 0.000).
  • For Gradient EM vs GD, we report a T-statistic of 22.88 (p-value of 0.000).

  We conclude that our results are statistically significant.

  Planned Revision:
  We will include our CIFAR experiments in Appendix C to highlight preliminary results on scaling EM to 'large'-scale MoE.

• An additional notation table is helpful to follow the proof of this paper.
  Planned Revision:
  We will include a notation section and table of contents at the beginning of the Appendix.

Summary of Planned Revisions Based on All Reviewers' Feedback

As 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.

1. Readability:
   To enhance clarity and accessibility, we will:
   A) streamline the main text by moving portions of the background on EM theory and MoE parameter estimation to the supplementary material,
   B) introduce a dedicated notation section and table of contents in the Appendix to improve readability, and
   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.

2. Discussions:
   We will:
   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
   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.

3. Experiments:
   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).

   In Appendix C, we will:
   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$,
   B) include initialization experiments (see our reply to reviewer kaEY),
   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
   D) provide all necessary implementation details in Appendix C (such as learning rate, precise model architecture, number of instances, etc.).

We are looking forward to hearing your valuable feedback on our responses and proposed revisions.

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:

• The limitations of the work are not mentioned

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.

• More details about the experimental setup (needed to reproduce the results: e.g., batch size, learning rate used) are required in Appendix C.

That 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.

• 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?

This 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.

• Question 2:
  To what extent are the assumptions A₁ - A₃ verified in the experiments?

  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.

• Question 3:
  What are the practical implications for the future of MoE of this connection with MD?

  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).

• Question 4:
  Did the asymptotic convergence error of GD in Figures 1 and 2 reach that of EM and Gradient EM?

  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.

• Question 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?

  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:

  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%.
  2. Information sharing: Each expert now trains on half as many images as in the base case.
  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.

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!

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:

• Concerns:
  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.

  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.
  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.

  Planned Revision:
  We will include a thorough discussion of the limitations, discussing the scalability of the EM algorithm to large-scale and sparse MoE.

• Suggestion:
  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 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.

  Planned Revision:
  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.

• Questions:
  1) In line 147, could an additional parameter be introduced to consider the distribution $p(z|x;w) \propto \exp\{\gamma x^\top w\}$?
  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.

  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?

  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].

  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].

  Planned Revision:
  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.

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!