DeepSeekMoE

The DeepSeekMoE is a mixture-of-experts architecture that involves two principal strategies:

• finely segmenting the experts into $mN$ ones and activating $mK$ from them, allowing for a more flexible combination of activated experts
• isolating $K_s$ experts as shared ones, aiming at capturing common knowledge and mitigating redundancy in routed experts

Fine-Grained Expert Segmentation

The idea here is that many smaller experts is better than few large experts. The rationale is that large experts are:

• be more likely to encounter a wide variety of knowledge
• therefore more likely to learn to generalise to this wide knowledge

With many small experts this wide variety of knowledge has more potential to be decomposed into the experts, allowing each expert to retain a high level of expert specialisation.

DeepSeekMoE:

1. splits each expert $\text{FFN}$ into $m$ smaller experts by reducing the $\text{FFN}$ intermediate hidden dimension to $\frac{1}{m}$ times its original size
2. increases the total experts $m$ times to keep the same memory cost, as illustrated Figure 2(b) where $N$ is doubled
3. number of activated experts $m$ times to keep the same computational cost, as illustrated in Figure 2(b) where $K$ is doubled

The output of an MoE layer with fine-grained expert segmentation can be expressed as:

$$\begin{array}{rl}{\text{h}}_t^l& =\sum _{i=1}^{mN}\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 mN\},mK)\\ 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 the total number of expert parameters is equal to $N$ times the number of parameters in a standard $\text{FFN}$, and $mN$ denotes the total number of fine-grained experts. With the fine-grained expert segmentation strategy, the number of non-zero gates will also increase to $mK$.

Shared Expert Isolation

Certain tokens may require each expert to have some common knowledge. In conventional routing, multiple experts might acquire this common knowledge, resulting in redundancy in expert parameters. This redundancy can be mitigated by using dedicated shared experts to capture and consolidate this common knowledge.

To achieve this, $K_s$ experts are isolated to serve as shared experts. Each token will deterministically be assigned to these shared experts. The number of experts activated by the router is decreased by $K_s$ in order to maintain a constant computational cost, as illustrated in Figure 2(c) where $K$ is decreased by the number of shared experts.

Therefore, an MoE layer in the complete DeepSeekMoE architecture is formulated as follows:

$$\begin{array}{rl}{\text{h}}_t^l& =\sum _{i=1}^{K_s}{\text{FFN}}_i\left({\text{u}}_t^l\right)+\sum _{i=K_s+1}^{mN}\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_{l,t}|K_s\le j\le mN\},mK-K_s)\\ 0& \text{otherwise}\end{array}\\ s_{i,t}& ={\text{Softmax}}_i\left({\text{u}}_t^{l^T}{\text{e}}_i^l\right)\end{array}$$

Finally, in DeepSeekMoE the:

• number of shared experts is $K_s$
• total number of routed experts is $mN-K_s$
• number of non-zero gates is $mK-K_s$

Load Balance Consideration

There are two main routing considerations when training an MoE model:

• If the same experts are always chosen, other experts are prevented from learning. This is called routing collapse
• If experts are distributed across multiple devices, load imbalance can exacerbate computation bottlenecks

In practice, a:

• small expert-level balance factor is used to mitigate the risk of routing collapse
• large device-level balance factor is used to promote balanced computation across devices

Expert-Level Balance Loss

$$\begin{array}{rl}{\mathcal{L}}_{ExpBal}& ={\alpha }_1\sum _{i=1}^{N^{\prime }}{f}_i{P}_i\\ f_i& =\frac{N^{\prime }}{K^{\prime }T}\sum _{t=1}^{T}1(\text{Token }t\text{ selects Expert }i)\\ P_i& =\frac{1}{T}\sum _{t=1}^T{s}_{i,t}\end{array}$$

where ${\alpha }_1$ is a hyperparameter called expert-level balance factor, $N^{\prime }$ is equal to $mN-K_s$ and $K^{\prime }$ is equal to $mK-K_s$ for brevity. $1(\cdot )$ denotes the indicator function.

Device-Level Balance Loss

If we partition all routed experts into $D$ groups $\{{\mathcal{E}}_1,,{\mathcal{E}}_2,\dots ,{\mathcal{E}}_D\}$, and deploy each group on a single device, then device-level balance loss is computed as follows:

$$\begin{array}{rl}{\mathcal{L}}_{DevBal}& ={\alpha }_2\sum _{i=1}^D{f}_i^{\prime }{P}_i^{\prime }\\ f_i^{\prime }& =\frac{1}{|{\mathcal{E}}_i|}\sum _{j\in {\mathcal{E}}_i}{f}_j\\ P_i^{\prime }& =\sum _{j\in {\mathcal{E}}_i}{P}_j\end{array}$$

where $\alpha 2$ is a hyperparameter called device-level balance factor.

References

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