Learning Mixtures of Experts with EM

Mixtures of Experts (MoE)[JJ94] are a crucial class of parametric latent variable models that have gained significant popularity in deep learning for their ability to reduce both training and inference costs [CDWGL22]. MoE are particularly effective when the feature space can be divided and processed by specialized models, known as experts. Instead of relying on a 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 expert to focus on mapping its designated inputs to outputs using a separate, optimized model. By leveraging multiple specialized experts rather than a monolithic model, MoE can achieve greater scalability and flexibility. This modular 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 Mixture of Experts span across various domains, including language translation, speech recognition, recommendation systems, and more [FZS22, MZYCHC18, Hin+12].

Training an MoE model involves training both (a) the parameters in the individual experts, and (b) the gating function that routes inputs to the appropriate expert. Typically, these are both learnt jointly by first formulating the final loss function as applied to the ensembled 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 MoE using a 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. It then solves each problem in isolation and in parallel, and then collates the outputs to arrive at the updated set of gate and expert parameters.

For our theoretical results, we consider two simple instances of MoE models: Mixture of Linear Experts (where each expert is a linear regressor) and Mixture of Logistic Experts (where each expert is a logisitic regressor). The router in each case is a linear softmax.

Main Contributions:

• 1)

  A recent but cornerstone result in EM [KKS22] showed that for exponential families, EM corresponds to Mirror Descent with unit step size and KL divergence regularizer. MoE models (inlcuding the ones we focus on) are generally not exponential families; however, we show that nonetheless this correspondence still holds.

• 2)

  Armed with this perspective, we study our linear and logistic MoE when there are two experts. We provide sufficient conditions for EM to converge to the true model parameters at sub-linear or linear rates; we also characterize under what parameter regimes these conditions will hold. To our knowledge, no similar such characterizations are known for SGD-based methods.

• 3)

  Finally we show, both on synthetic and small scale proof of concept real world datasets, that EM outperforms gradient descent both in terms of convergence rate and the achieved performance.

Next, we formally describe the setting under consideration for the Mixture of $k$ Experts (MoE).

Data Generation Model: First, the input or feature vector variable $\bm{x}\in\mathbb{R}^{d}$ is sampled based on a probability density function $p(\bm{x})$. Second, given the feature vector $\bm{x}$, a latent variable $z\in[k]$, responsible for routing $\bm{x}$ to the appropriate expert is sampled with probability mass function $P(z|\bm{x})$. Finally, given the pair $(\bm{x},z)$, the output or label $y$ is generated according to the probability distribution $p(y|\bm{x},z)$. Hence, the complete data distribution is $p(\bm{x},y,z)=p(y|\bm{x},z)p(z|\bm{x})p(\bm{x})$ and the joint input-output probability distribution can be written as

This is indeed the most general form of data generation under mixture of $k$-experts, but here we focus on the case where $\bm{x}$ is sampled from a unit spherical Gaussian distribution, i.e. $\bm{x}\sim\mathcal{N}(\bm{0},\bm{I}_{d})$, and $P(z|\bm{x})=P(z|\bm{x};\bm{w}^{*})$ is parameterized by a set of vectors $(\bm{w}_{1}^{*},\dots,\bm{w}_{k}^{*})$ and can be written as

In other words, the probability mass function of the latent variable is the softmax function between the inner product of $\bm{x}$ with the parameter vectors concatenated as $\bm{w}^{*}=(\bm{w}^{*}_{1},...,\bm{w}^{*}_{k})\in\mathbb{R}^{d\times k}$.

Regarding $p(y|\bm{x},z)$, the probability distribution of the output variable $y$ conditioned on the input variable $x$ and latent routing variable $z\!=\!i$ (i.e., expert $i$), we also assume that it is parameterized by a vector $\bm{\beta}_{i}^{*}$ and we have $p(y|\bm{x},z\!=\!i)=p(y|\bm{x},z\!=\!i;\bm{\beta}_{i}^{*})$. Thus, for the concatenated vector $\bm{\beta}^{*}=(\bm{\beta}^{*}_{1},...,\bm{\beta}^{*}_{k})\in\mathbb{R}^{s\times k}$, then $p(y|\bm{x},z)=p(y|\bm{x},z;\bm{\beta}^{*})$.

To benefit the understanding of the reader in later sections, we focus on two special cases of experts. In the first case, we have

as the density function of the normal distribution. This is equivalent to assume that the output $y=\bm{x}^{\top}\beta_{i}^{*}+\epsilon$ where $\epsilon$ is additive zero mean Gaussian noise with unit variance. For future reference, we will refer to this special case as a Mixture of Linear Experts.

In the second case, we assume a logistic model for each expert. Thus, the density function $p(y|\bm{x},z=i;\bm{\beta}_{i}^{*})$ can be written as

and indeed $P(y=-1|\bm{x},z\!=\!i;\bm{\beta}_{i}^{*})=1-P(y=1|\bm{x},z\!=\!i;\bm{\beta}_{i}^{*})$. For future reference, we will refer to this special case as a Mixture of Logistic Experts.

Maximum Likelihood Loss: Given the assumed data distribution of $(\bm{x},y)$, our goal is to find the set of feasible parameters $\bm{\beta}\in\mathbb{R}^{d\times k}$ and $\bm{w}\in\mathbb{R}^{d\times k}$ that maximize the log likelihood function. For ease of notation we define $\bm{\theta}:=(\bm{\beta},\bm{w})\in\mathbb{R}^{2d\times k}$ as the concatenation of all the parameters. Specifically, from (2), the expected log likelihood is given by

Given the fact that only the last term of the sum depends on parameters $\bm{\theta}=(\bm{w},\bm{\beta})^{\top}$ the negative log-likelihood objective function, ${\cal{L}}(\bm{\theta})$, that we aim to minimize can be written as

Note that as we will discuss in detail, for both mixtures of linear or logistic experts, the above objective function is known to be non convex with respect to $\bm{\theta}$. 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 Expectation-Maximization (EM) algorithm for solving this optimization problem.

Next, we present the Expectation-Maximization (EM) algorithm for MoE. EM takes a structured approach to minimizing the objective $\mathcal{L}(\bm{\theta})$ given in (6). Each iteration of the EM algorithm is decomposed into two steps as follows. The first step is called “expectation”: For current parameter estimate $\bm{\theta}^{t}$, we compute the expectation of the complete-data log-likelihood with respect to the hidden variables, using the current parameter estimates $\bm{\theta}^{t}$ and denote it by $Q(\bm{\theta}|\bm{\theta}^{t})$, i.e.,

Then, in the second step called “maximization”, we simply minimize the objective $Q(\bm{\theta}|\bm{\theta}^{t})$ (or maximize $-Q(\bm{\theta}|\bm{\theta}^{t})$) with respect to $\bm{\theta}\in\Omega$ and obtain our new parameter iterate as

Since $\log p(y,z|\bm{x};\bm{\theta})=\log p(y|z,\bm{x};\beta)+\log p(z|\bm{x};\bm{w})$, it follows that the EM objective (7) is linearly separable in the parameters $\bm{\beta}$ and $\bm{w}$. Thus, we can re-write $Q(\bm{\theta}|\bm{\phi})$ as the sum of $2$ functions that depend only on $\bm{\beta}$ and $\bm{w}$ respectively. Subsequently, the EM update (8) is obtained as the concatenation $\bm{\theta}^{t+1}=(\bm{w}^{t+1},\bm{\beta}^{t+1})^{\top}$ where

are both strongly convex minimization problems for the special cases of mixture of linear experts and mixture of logistic experts.

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 function, $\mathcal{L}(\bm{\theta})$, can be decomposed into the difference between the EM objective and another function that is bounded below by $0$. Specifically, for any $\bm{\theta},\bm{\phi}\in\Omega$,

Denoting $H(\bm{\theta}|\bm{\phi}):=-\mathbb{E}_{\bm{X},Y}\mathbb{E}_{Z|\bm{x},y,\bm{\phi}}\left[\log p(z|\bm{x},y;\bm{\theta})\right]$, it follows that

where $H(\bm{\theta}|\bm{\phi})$ is bounded below by $0$. Thus, $Q(\bm{\theta}|\bm{\phi})$ acts as an upper bound on the negative log-likelihood.

Next, applying Jensen’s inequality shows that $H(\bm{\theta}|\bm{\phi})$ is minimized at $\bm{\theta}=\bm{\phi}$ where $H(\bm{\theta}|\bm{\theta})=0$. Consequently, it follows that the gradient of the negative log-likelihood objective matches that of the surrogate EM objective at $\bm{\phi}=\bm{\theta}$, i.e., $\nabla\mathcal{L}(\bm{\theta})=\nabla Q(\bm{\theta}|\bm{\theta})$. This suggests that any stationary point of $\mathcal{L}(\bm{\theta})$ is also a stationary point of the EM algorithm and vice versa.

In the previous section, we derived the EM update for both MoE and SMoE as the solution to minimizing the EM objective in (7). We further demonstrated that this solution can be decomposed into the concatenation of the respective solutions to two minimization sub-problems. In this section, we will show that this update is exactly equivalent to performing a single step of the Mirror Descent (MD) update on $\mathcal{L}(\bm{\theta})$ with a unit step size and the KL divergence as a regularizer.

To illustrate more clearly that minimizing $Q(\bm{\theta}|\bm{\theta}^{t})$ in (7) corresponds to a MD step on the loss $\mathcal{L}(\bm{\theta})$ at the point $\bm{\theta}^{t}$, we first provide a brief overview of the core concept behind the MD update.

In most gradient-based methods, the next iteration is obtained by minimizing an upper bound of the objective function. For example, in Gradient Descent (GD), the next iterate is found by minimizing the first-order Taylor expansion of the objective at $\bm{\theta}^{t}$, with a squared norm regularizer. Specifically, for minimizing the loss $\mathcal{L}(.)$, the GD update with step size $\eta$ at $\bm{\theta}^{t}$ is equivalent to minimizing the following quadratic function:

This function indeed serves as an upper bound for $\mathcal{L}(.)$ if $\eta\leq 1/L$, where $L$ is the Lipschitz constant of $\nabla\mathcal{L}(\bm{\theta})$.

Mirror Descent solves a similar sub-problem where instead of a squared norm regularizer, we employ the Bregman Divergence regularizer: The Bregman divergence induced by a differentiable, convex function $h:\mathbb{R}^{d}\to\mathbb{R}$, measures the difference between $h(\bm{\theta})$ and its first-order approximation at $\bm{\theta}^{t}$,i.e.,

Thus, the iterations of MD are derived by minimizing the following expression:

As mentioned, this scheme is reasonable if the function approximation using the Bregman divergence serves as an upper bound for the function $\mathcal{L}(\bm{\theta})$. This can be ensured when the step size $\eta$ is sufficiently small, and the condition of relative smoothness is satisfied.

In the upcoming theorem, we formally establish that for the Symmetric Mixture of Linear Experts (SMoLinE) and the Symmetric Mixture of Logistic Experts (SMoLogE) models, minimizing the EM objective function defined in equation (8) is exactly equivalent to minimizing the subproblem associated with a single step of MD, as defined in equation (14), with a step size of $\eta=1$. This result demonstrates that the EM update in these specific models is essentially performing a MD step on $\cal L$. The proof of this result is quite technical and is deferred to Appendix A.

As the loss is is 1-relative smooth with respect to $A$, this validates the choice of $\eta=1$ for the Mirror Descent Update, and subsequent convergence results from the MD literature.

This section presents robust convergence guarantees for EM on SMoLinE and SMoLogE, starting with a review of Mirror Descent guarantees for context. We then examine how these conditions relate to concrete statistical metrics and the true model parameters.

In this section, we empirically validate our theoretical results by comparing the performance of EM with Gradient EM (Algorithm 1), and Gradient Descent (Algorithm 2). Recall that EM for MoE obtains its next parameter iterate 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 sub-minimization problems of EM. This differs from Gradient Descent (GD) that obtains its next parameter iterate as the gradient update on the negative log-likelihood objective $\mathcal{\bm{\theta}}$. We evaluate these methods on both a synthetic dataset and the real-world Fashion MNIST dataset, consistently showing significant improvements for EM and Gradient EM over GD. Note that our aim is not to achieve state-of-the-art accuracy, but to demonstrate that EM is more suitable than GD for fitting specific models.

Synthetic Dataset. We created the synthetic dataset so as to simulate a population setting of SMoLinE. We sampled $10^{3}$ data points from an SMoLinE with known additive unit Gaussian noise (i.e. $\mathcal{N}(0,1)$) and true parameters $\bm{\beta}^{*},\bm{w}^{*}\in\mathbb{R}^{10}$ that sastisfy $\|\bm{\beta}^{*}\|_{2}=\|\bm{w}^{*}\|_{2}=4$. Subsequently, we run 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 neighborhood of the true parameters, and are consistent across all benchmarks.

Figure 1 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 SMoLinE parameters. Once again, EM requires significantly fewer iterations to fit the mixture compared to 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 theoretical results for MoE and is consistent with findings for Mixtures of Gaussians and Mixtures of Linear Regression in high SNR scenarios.

Validation Experiment on Fashion MNIST. For the small scale proof of concept experiment on Fashion MNIST [XRV17], we alter the dataset so as to simulate a mixture of $2$ Logistic Experts. To do this, we simply perform an inversion transformation on the images at random with probability $\frac{1}{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 can be seen 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 an accuracy of $10.2\%$ on the altered dataset. This suggests a $2$-component MoLogE is appropriate for fitting the altered dataset, so long as the ground truth partitioning is linear in image space.

The $2$-component MoLogE to be trained consists of one Linear gating layer of dimension $2\times 28\times 28$, and $2$ logistic experts of dimension $10\times 28\times 28$ each. We randomly initialize each linear layer to be unit norm and execute the algorithms on 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 a Gradient Step on the Gating parameters and the Expert parameters separately as describe in Algorithm 1. For EM, for each iteration, we perform several gradient steps in an inner loop to approximately recover the solutions to the sub-problems described in (8). We report our findings in Table 2 and Figure 3.

In Table 2, we report the respective final test accuracy and cross-entropy loss values after $100$ iterations of EM, Gradient 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 original 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, averaged 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.

In this paper, we addressed the problem of Mixtures of Experts (MoE), focusing on the use of the EM method to solve these problems. We first showed that the EM update for MoE could be interpreted as a Mirror Descent step on the log-likelihood with a unit step size and a KL divergence regularizer, extending the result of [KKS22] beyond the exponential family. Building on this, we characterized different convergence rates for EM in this setting under various assumptions about the log-likelihood function and specified when these assumptions held. We also compared EM with gradient descent empirically and showed that EM outperformed gradient descent in both convergence rate and final performance.