Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Mixtures of Experts with EM: A Mirror Descent Perspective

Quentin Fruytier, Aryan Mokhtari, Sujay Sanghavi

TL;DR

The paper establishes a theoretical connection between EM for Mixtures of Experts (MoE) and projected Mirror Descent (MD) on the conditional log-likelihood, proving that EM updates are equivalent to a one-step MD step with a KL regularizer when the complete-data distribution lies in an exponential family. This MD view yields convergence guarantees to stationary points and, under favorable initialization, linear convergence to the true parameters, with rates described via the Missing Information Matrix (MIM) and the signal-to-noise ratio (SNR). The authors further show that for symmetric two-expert MoEs (SymMoLinE and SymMoLogE), EM reduces to MD without projection and provide explicit conditions for linear convergence. They extend the framework to deep and sparse MoE via a principled EM-like approach and validate the theory with synthetic and real-data experiments where EM outperforms gradient-based methods in convergence speed and accuracy. Overall, the work provides a unifying, theory-grounded perspective on training MoE with EM and suggests practical benefits and future directions for scalable EM-based MoE optimization.

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

TL;DR

The paper establishes a theoretical connection between EM for Mixtures of Experts (MoE) and projected Mirror Descent (MD) on the conditional log-likelihood, proving that EM updates are equivalent to a one-step MD step with a KL regularizer when the complete-data distribution lies in an exponential family. This MD view yields convergence guarantees to stationary points and, under favorable initialization, linear convergence to the true parameters, with rates described via the Missing Information Matrix (MIM) and the signal-to-noise ratio (SNR). The authors further show that for symmetric two-expert MoEs (SymMoLinE and SymMoLogE), EM reduces to MD without projection and provide explicit conditions for linear convergence. They extend the framework to deep and sparse MoE via a principled EM-like approach and validate the theory with synthetic and real-data experiments where EM outperforms gradient-based methods in convergence speed and accuracy. Overall, the work provides a unifying, theory-grounded perspective on training MoE with EM and suggests practical benefits and future directions for scalable EM-based MoE optimization.

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

Paper Structure

This paper contains 30 sections, 9 theorems, 170 equations, 4 figures, 2 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Mixture of Experts
  4. EM for Mixtures of Experts
  5. EM for Symmetric Mixture of 2-Experts
  6. EM for Deep and Sparse MoE
  7. Main Result
  8. Convergence Guarantees for General MoE
  9. Discussion of Main Result
  10. Special case of SymMoLinE and SymMoLogE
  11. Experiments
  12. Conclusion
  13. EM, Projected Mirror Descent, and General MoE.
  14. EM is Projected Mirror Descent for General MoE.
  15. Convergence Results for EM applied to General MoE.
  16. ...and 15 more sections

Key Result

Theorem 4.1

For General MoE, there exists a natural reparameterization $\bm \theta_{\bm x} \in \{\eta(\cdot, \bm \theta): \theta \in \Omega\}$ with and a mirror map $A(\bm \theta_{\bm x})$ such that the EM update in eqn:EM Iteration simplifies and is equivalent to the expectation moment projection, where for each $\bm x$, $\bm \tilde{\bm \theta}^{t+1}_{\bm x}$ is obtained from the following MD step, with $

Figures (4)

  • Figure 1: Convergence of objective errors $\mathcal{L}(\bm \theta^t) - \mathcal{L}(\bm \theta^*)$ and $\mathcal{L}(\bm \theta^t) - \mathcal{L}(\bm \theta^T)$ in Fig \ref{['fig:sub1_stat_err_lik']} and Fig \ref{['fig:sub2_opt_err_lik']}, respectively, averaged over $50$ instances when fitting a SymMoLinE.
  • Figure 2: This figure shows the progress made towards the true parameters, $\frac{\Vert \bm \beta^t - \bm \beta^*\Vert_2}{\Vert\bm \beta^*\Vert_2}$ and $\frac{\Vert \bm w^t - \bm w^*\Vert_2}{\Vert \bm w^*\Vert_2}$ in figures \ref{['fig:sub1_beta_progress']} and \ref{['fig:sub2_w_progress']} respectively, averaged over $50$ instances when fitting a SymMoLinE
  • Figure 3: Test accuracy and objective function, $\frac{1}{n}\sum_{i=1}^n \mathbbm{1}_{\hat{y_i} = y_i}$ and $\mathcal{L}(\bm \theta^t)$ in \ref{['fig:sub1_accuracy_progress']} and \ref{['fig:sub2_cross_entropy_progress']}, respectively, averaged over $25$ instances for a $2$-component MoLogE train on Random Invert FMNIST.
  • Figure 4: Convergence in predicted label Accuracy for different optimization methods.

Theorems & Definitions (19)

  • Theorem 4.1
  • proof : Proof sketch
  • Theorem 4.2: Convergence of EM
  • Theorem 5.1
  • proof
  • proof
  • Lemma A.1: cosine law for Bregman divergence
  • proof
  • proof
  • Corollary B.1: Convergence of EM
  • ...and 9 more