Mixture-of-Experts (MoE)

The following explanation of MoE models comes from the DeepSeekMoE paper.

A standard Transformer language model is constructed by stacking $L$ layers of standard Transformer blocks, where each block can be represented as follows:

$$\begin{array}{r}{\text{u}}_{1:T}^l=\text{Self-Attn}\left({\text{h}}_{1:T}^{l-1}\right)+{\text{h}}_{1:T}^{l-1}\\ {\text{h}}_t^l=\text{FFN}\left({\text{u}}_t^l\right)+{\text{u}}_t^l\end{array}$$

where $T$ denotes the sequence length, $\text{Self-Attn}\left(\cdot \right)$ denotes the self-attention module, $\text{FFN}\left(\cdot \right)$ denotes the Feed-Forward Network (FFN), ${\text{u}}_{1:T}^l\in {\mathbb{R}}^{T\times d}$ are the hidden states of all tokens after the $l$-th attention module, and ${\text{h}}_t^l\in {\mathbb{R}}^d$ is the output hidden state of the $t$-th token after the $l$-th Transformer block. For brevity, we omit the layer normalisation in the above formulations.

A typical practice to construct an MoE language model usually substitutes FFNs in a Transformer with MoE layers at specified intervals. An MoE layer is composed of multiple experts, where eah expert is structurally identical to a standard FFN. Then, each token will be assigned to one or two experts. If the $l$-th FFN is substituted with an MoE layer, the computation for its output hidden state ${\text{h}}_t^l$ is expressed as:

$$\begin{array}{rl}{\text{h}}_t^l& =\sum _{i=1}^N\left(g_{i,t}{\text{FFN}}_i\left({\text{u}}_t^l\right)\right)+{\text{u}}_t^l\\ g_{i,t}& =\{\begin{array}{ll}{s}_{i,t}& s_{i,t}\in \text{Topk}(\{s_{j,t}|1\le j\le N\},K)\\ 0& \text{otherwise}\end{array}\\ s_{i,t}& ={\text{Softmax}}_i\left({\text{u}}_t^{l^T}{\text{e}}_i^l\right)\end{array}$$

where $N$ denotes the total number of experts, ${\text{FFN}}_i(\cdot )$ is the $i$-th expert $\text{FFN}$, $g_{i,t}$ denotes the gate value for the $i$-th expert, $s_{i,t}$ denotes the token-to-expert affinity, $\text{Topk}(\cdot ,K)$ denotes the set comprising $K$ highest affinity scores among those calculated for the $t$-th token and all $N$ experts, and ${\text{e}}_i^l$ is the centroid of the $i$-th expert in the $l$-th layer. Note that $g_{i,t}$ is sparse, indicating that only $K$ out of $N$ gate values are non-zero. This sparsity property ensures computational efficiency within an MoE layer, i.e., each token will only be assigned to and computed in only $K$ experts. Also, in the above formulations, we omit the layer normalisation operation for brevity.

References

• DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models