Mixture of experts

The meta-pi network, reported by Hampshire and Waibel,[2] uses ${\displaystyle f(x)=\sum _{i}w(x)_{i}f_{i}(x)}$ as the output. The model is trained by performing gradient descent on the mean-squared error loss ${\displaystyle L:={\frac {1}{N}}\sum _{k}\|y_{k}-f(x_{k})\|^{2}}$. The experts may be arbitrary functions.

In their original publication, they were solving the problem of classifying phonemes in speech signal from 6 different Japanese speakers, 2 females and 4 males. They trained 6 experts, each being a "time-delayed neural network"[3] (essentially a multilayered convolution network over the mel spectrogram). They found that the resulting mixture of experts dedicated 5 experts for 5 of the speakers, but the 6th (male) speaker does not have a dedicated expert, instead his voice was classified by a linear combination of the experts for the other 3 male speakers.

The adaptive mixtures of local experts [4][5] uses a gaussian mixture model. Each expert simply predicts a gaussian distribution, and totally ignores the input. Specifically, the ${\displaystyle i}$-th expert predicts that the output is ${\displaystyle y\sim N(\mu _{i},I)}$, where ${\displaystyle \mu _{i}}$ is a learnable parameter. The weighting function is a linear-softmax function:

$${\displaystyle w(x)_{i}={\frac {e^{k_{i}^{T}x+b_{i}}}{\sum _{j}e^{k_{j}^{T}x+b_{j}}}}}$$

The mixture of experts predict that the output is distributed according to the probability density function:

$${\displaystyle f_{\theta }(y|x)=\ln \left[\sum _{i}{\frac {e^{k_{i}^{T}x+b_{i}}}{\sum _{j}e^{k_{j}^{T}x+b_{j}}}}N(y|\mu _{i},I)\right]=\ln \left[(2\pi )^{-d/2}\sum _{i}{\frac {e^{k_{i}^{T}x+b_{i}}}{\sum _{j}e^{k_{j}^{T}x+b_{j}}}}e^{-{\frac {1}{2}}\|y-\mu _{i}\|^{2}}\right]}$$

It is trained by maximal likelihood estimation, that is, gradient ascent on ${\displaystyle f(y|x)}$. The gradient for the ${\displaystyle i}$-th expert is

$${\displaystyle \nabla _{\mu _{i}}f_{\theta }(y|x)={\frac {w(x)_{i}N(y|\mu _{i},I)}{\sum _{j}w(x)_{j}N(y|\mu _{j},I)}}\;(y-\mu _{i})}$$

and the gradient for the weighting function is

$${\displaystyle \nabla _{[k_{i},b_{i}]}f_{\theta }(y|x)={\begin{bmatrix}x\\1\end{bmatrix}}{\frac {w(x)_{i}}{\sum _{j}w(x)_{j}N(y|\mu _{j},I)}}(f_{i}(x)-f_{\theta }(y|x))}$$

For each input-output pair ${\displaystyle (x,y)}$, the weighting function is changed to increase the weight on all experts that performed above average, and decrease the weight on all experts that performed below average. This encourages the weighting function to learn to select only the expects that make the right predictions for each input.

The ${\displaystyle i}$-th expert is changed to make its prediction closer to ${\displaystyle y}$, but the amount of change is proportional to ${\displaystyle w(x)_{i}N(y|\mu _{i},I)}$. This has a Bayesian interpretation. Given input ${\displaystyle x}$, the prior probability that expert ${\displaystyle i}$ is the right one is ${\displaystyle w(x)_{i}}$, and ${\displaystyle N(y|\mu _{i},I)}$ is the likelihood of evidence ${\displaystyle y}$. So, ${\displaystyle {\frac {w(x)_{i}N(y|\mu _{i},I)}{\sum _{j}w(x)_{j}N(y|\mu _{j},I)}}}$ is the posterior probability for expert ${\displaystyle i}$, and so the rate of change for the ${\displaystyle i}$-th expert is proportional to its posterior probability.

In words, the experts that, in hindsight, seemed like the good experts to consult, are asked to learn on the example. The experts that, in hindsight, were not, are left alone.

The combined effect is that the experts become specialized: Suppose two experts are both good at predicting a certain kind of input, but one is slightly better, then the weighting function would eventually learn to favor the better one. After that happens, the lesser expert is unable to obtain a high gradient signal, and becomes even worse at predicting such kind of input. Conversely, the lesser expert can become better at predicting other kinds of input, and increasingly pulled away into another region. This has a positive feedback effect, causing each expert to move apart from the rest and take care of a local region alone (thus the name "local experts").

Hierarchical mixtures of experts[6][7] uses multiple levels of gating in a tree. Each gating is a probability distribution over the next level of gatings, and the experts are on the leaf nodes of the tree. They are similar to decision trees.

For example, a 2-level hierarchical MoE would have a first order gating function ${\displaystyle w_{i}}$, and second order gating functions ${\displaystyle w_{j|i}}$ and experts ${\displaystyle f_{j|i}}$. The total prediction is then ${\displaystyle \sum _{i}w_{i}(x)\sum _{j}w_{j|i}(x)f_{j|i}(x)}$.

The mixture of experts, being similar to the gaussian mixture model, can also be trained by the expectation-maximization algorithm, just like gaussian mixture models. Specifically, during the expectation step, the "burden" for explaining each data point is assigned over the experts, and during the maximization step, the experts are trained to improve the explanations they got a high burden for, while the gate is trained to improve its burden assignment. This can converge faster than gradient ascent on the log-likelihood.[7][8]

The choice of gating function is often a softmax gating. Other than that, [9] proposed using gaussian distributions, and [8] proposed using exponential families.

Instead of performing a weighted sum of all the experts, in hard MoE [10] only the highest ranked expert is chosen. That is, ${\displaystyle f(x)=f_{\arg \max _{i}w_{i}(x)}(x)}$. This can accelerate training and inference time.[11]

The experts can use more general forms of multivariant gaussian distributions. For example, [6] proposed ${\displaystyle f_{i}(y|x)=N(y|A_{i}x+b_{i},\Sigma _{i})}$, where ${\displaystyle A_{i},b_{i},\Sigma _{i}}$ are learnable parameters. In words, each expert learns to do linear regression, with a learnable uncertainty estimate.

One can use different experts than gaussian distributions. For example, one can use Laplace distribution,[12] or Student's t-distribution.[13] For binary classification, it also proposed logistic regression experts, with

$${\displaystyle f_{i}(y|x)={\begin{cases}{\frac {1}{1+e^{\beta _{i}^{T}x+\beta _{i,0}}}},&y=0\\1-{\frac {1}{1+e^{\beta _{i}^{T}x+\beta _{i,0}}}},&y=1\end{cases}}}$$

where ${\displaystyle \beta _{i},\beta _{i,0}}$ are learnable parameters. This is later generalized for multi-class classification, with multinomial logistic regression experts.[14]

The sparsely-gated MoE layer,[17] published by researchers from Google Brain, uses feedforward networks as experts, and linear-softmax gating. Similar to the previously proposed hard MoE, they achieve sparsity by a weighted sum of only the top-k experts, instead of the weighted sum of all of them. Specifically, in a MoE layer, there are feedforward networks ${\displaystyle f_{1},...,f_{n}}$, and a gating network ${\displaystyle w}$. The gating network is defined by ${\displaystyle w(x)=\mathrm {softmax} (\mathrm {top} _{k}(Wx+{\text{noise}}))}$, where ${\displaystyle \mathrm {top} _{k}}$ is a function that keeps the top-k entries of a vector the same, but sets all other entries to ${\displaystyle -\infty }$. The addition of noise helps with load balancing.

The choice of ${\displaystyle k}$ is a hyperparameter that is chosen according to application. Typical values are ${\displaystyle k=1,2}$. The ${\displaystyle k=1}$ version is also called the Switch Transformer.[18]

As demonstration, they trained a series of models for machine translation with alternating layers of MoE and LSTM, and compared with deep LSTM models.[19] Table 3 shows that the MoE models used less inference time compute, despite having 30x more parameters.

Vanilla MoE tend to have issues of load balancing: some experts are consulted often, while other experts rarely or not at all. To encourage the gate to select each expert with equal frequency (proper load balancing) within each batch, each MoE layer has two auxiliary loss functions. This is improved by [18] into a single auxiliary loss function. Specifically, let ${\displaystyle n}$ be the number of experts, then for a given batch of queries ${\displaystyle \{x_{1},x_{2},...,x_{T}\}}$, the auxiliary loss for the batch is

$${\displaystyle n\sum _{i=1}^{n}f_{i}P_{i}}$$

Here, ${\displaystyle f_{i}={\frac {1}{T}}\#({\text{queries sent to expert }}i)}$ is the fraction of time where expert ${\displaystyle i}$ is ranked highest, and ${\displaystyle P_{i}={\frac {1}{T}}\sum _{j=1}^{T}w_{i}(x_{j})}$ is the fraction of weight on expert ${\displaystyle i}$. This loss is minimized at ${\displaystyle 1}$, precisely when every expert has equal weight ${\displaystyle 1/n}$ in all situations.

In sparsely-gated MoE, only the top-k experts are queried, and their outputs are weighted-summed. There are other methods.[20]

In Hash MoE,[21] routing is performed deterministically by a hash function, fixed before learning begins. For example, if the model is a 4-layered Transformer, and input is a token for word "eat", and the hash of "eat" is ${\displaystyle (1,4,2,3)}$, then the token would be routed to the 1st expert in layer 1, 4th expert in layer 2, etc. Despite its simplicity, it achieves competitive performance as sparsely gated MoE with ${\displaystyle k=1}$.

In soft MoE, suppose in each batch, each expert can process ${\displaystyle p}$ queries, then there are ${\displaystyle n\times p}$ queries that can be assigned per batch. Now for each batch of queries ${\displaystyle \{x_{1},x_{2},...,x_{T}\}}$, the soft MoE layer computes an array ${\displaystyle w_{i,j,k}}$, such that ${\displaystyle (w_{i,j,1},...,w_{i,j,T})}$ is a probability distribution over queries, and the ${\displaystyle i}$-th expert's ${\displaystyle j}$-th query is ${\displaystyle \sum _{k}w_{i,j,k}x_{k}}$.[22] However, this does not work with autoregressive modelling, since the weights ${\displaystyle w_{i,j,k}}$ over one token depends on all other tokens'.[23]

Other approaches include solving it as a constrained linear programming problem,[24] making each expert choose the top-k queries it wants (instead of each query choosing the top-k experts for it),[25] using reinforcement learning to train the routing algorithm (since picking an expert is a discrete action, like in RL).[26]

Suppose there are ${\displaystyle n}$ experts in a layer. For a given batch of queries ${\displaystyle \{x_{1},x_{2},...,x_{T}\}}$, each query is routed to one ore more experts. For example, if each query is routed to one expert as in Switch Transformers, and if the experts are load-balanced, then each expert should expect on average ${\displaystyle T/n}$ queries in a batch. In practice, the experts cannot expect perfect load balancing: in some batches, one expert might be underworked, while in other batches, it would be overworked.

Since the inputs cannot move through the layer until every expert in the layer has finished the queries it is assigned, load balancing is important. As a hard constraint on load balancing, there is the capacity factor: each expert is only allowed to process up to ${\displaystyle c\cdot T/n}$ queries in a batch. [20] found ${\displaystyle c\in [1.25,2]}$ to work in practice.

MoE layers are used in very large Transformer models, for which learning and inferring over the full model is too costly. In Transformer models, the MoE layers are often used to select the feedforward layers (typically a linear-ReLU-linear network), appearing in each Transformer block after the multiheaded attention. This is because the feedforward layers take up an increasing portion of the computing cost as models grow larger. For example, 90% of parameters in PALM-540B are in feedforward layers.[27]

A series of large language models from Google used MoE. GShard[28] uses MoE with up to top-2 experts per layer. Specifically, the top-1 expert is always selected, and the top-2th expert is selected with probability proportional to that experts' weight according to the gating function. Later, GLaM[29] demonstrated a language model with 1.2 trillion parameters, each MoE layer using top-2 out of 64 experts. Switch Transformers[18] use top-1 in all MoE layers.

The NLLB-200 by Meta AI is a machine translation model for 200 languages.[30] Each MoE layer uses a hierarchical MoE with two levels. On the first level, the gating function chooses to use either a "shared" feedforward layer, or to use the experts. If using the experts, then another gating function computes the weights and chooses the top-2 experts (see Figure 19).[31]

MoE large language models can be adapted for downstream tasks by instruction tuning.[32]

Generally, MoE are used when dense models have become too costly. As of 2023, the largest models tend to be large language models. Outside of those, Vision MoE[33] is a Transformer model with MoE layers. They demonstrated it by training a model with 15 billion parameters.