PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, Wenfeng Liang

arXiv:2401.06066v1cs.CL

TL;DR

Conventional MoE architectures face high computational costs and difficulty achieving expert specialization. DeepSeekMoE addresses these challenges through fine-grained expert segmentation and shared expert isolation, achieving comparable performance to larger or denser models with substantially fewer computations.

  • Problem

    Conventional MoE architectures face high computational costs and challenges in ensuring that experts acquire non-overlapping, focused knowledge.

  • Method

    DeepSeekMoE uses fine-grained expert segmentation and shared expert isolation to pursue greater expert specialization.

  • Results

    DeepSeekMoE achieves comparable performance with DeepSeek 7B and LLaMA2 7B using about 40% of computations, while its 145B model is comparable with DeepSeek 67B using 28.5% of computations.

  • Takeaways & Limitations

    DeepSeekMoE demonstrates higher expert specialization and approaches the upper-bound performance of MoE models while retaining computational advantages over prevailing architectures.

Abstract

from arXiv · show

In the era of large language models, Mixture-of-Experts (MoE) is a promising architecture for managing computational costs when scaling up model parameters. However, conventional MoE architectures like GShard, which activate the top-$K$ out of $N$ experts, face challenges in ensuring expert specialization, i.e. each expert acquires non-overlapping and focused knowledge. In response, we propose the DeepSeekMoE architecture towards ultimate expert specialization. It involves two principal strategies: (1) finely segmenting the experts into $mN$ ones and activating $mK$ from them, allowing for a more flexible combination of activated experts; (2) isolating $K_s$ experts as shared ones, aiming at capturing common knowledge and mitigating redundancy in routed experts. Starting from a modest scale with 2B parameters, we demonstrate that DeepSeekMoE 2B achieves comparable performance with GShard 2.9B, which has 1.5 times the expert parameters and computation. In addition, DeepSeekMoE 2B nearly approaches the performance of its dense counterpart with the same number of total parameters, which set the upper bound of MoE models. Subsequently, we scale up DeepSeekMoE to 16B parameters and show that it achieves comparable performance with LLaMA2 7B, with only about 40% of computations. Further, our preliminary efforts to scale up DeepSeekMoE to 145B parameters consistently validate its substantial advantages over the GShard architecture, and show its performance comparable with DeepSeek 67B, using only 28.5% (maybe even 18.2%) of computations.

1. Introduction

DeepSeekMoE targets limited expert specialization in conventional MoE models through fine-grained expert segmentation and shared expert isolation. Experiments from 2B to 145B parameters report strong performance and computational advantages over comparable MoE and dense baselines.

  • Motivation: MoE architectures help scale language-model parameters while keeping computational costs modest, but existing designs can suffer from knowledge hybridity and redundancy.These issues limit expert specialization, defined as each expert acquiring focused, non-overlapping knowledge.
  • Architecture: DeepSeekMoE combines fine-grained expert segmentation with shared expert isolation to improve expert specialization.The architecture is explicitly designed to address knowledge hybridity and redundancy in existing MoE models.
  • 2B Validation: DeepSeekMoE 2B matches GShard 2.9B despite using fewer expert parameters and less computation, and nearly approaches its dense counterpart’s performance with equal total parameters.GShard 2.9B has 1.5× the expert parameters and computation; the dense counterpart establishes an upper bound for MoE models.
  • Scaling: With only about 40% of computations, DeepSeekMoE 16B achieves comparable performance with DeepSeek 7B and LLaMA2 7B.LLaMA2 7B has approximately 2.5 times the activated parameters, while DeepSeekMoE 16B also supports a chat model after supervised fine-tuning.

2. Preliminaries: Mixture-of-Experts for Transformers

A Transformer language model stacks blocks containing self-attention and feed-forward networks. MoE replaces selected FFNs with multiple expert FFNs and uses sparse routing to assign each token to only a subset of experts.

  • Transformer Structure: A standard Transformer language model stacks L Transformer blocks, each containing self-attention followed by a feed-forward network.The hidden states after attention and after the Transformer block are represented separately in the formulation.
  • MoE Layers: An MoE layer substitutes selected Transformer FFNs with multiple structurally identical expert FFNs.Each token is assigned to one or two experts in the conventional architectures described.
  • Routing: The routing computation combines expert outputs using sparse gate values determined by token-to-expert affinities and top-K selection.TopK selects the K highest affinity scores among N experts for each token.
  • Routing: Because only K of N gate values are nonzero, each token is computed by only K experts, preserving MoE computational efficiency.The sparse gate values determine which expert outputs contribute to the token representation.
  • DeepSeekMoE Comparison: Figure 2 compares conventional top-2 routing, fine-grained expert segmentation, and shared expert isolation while holding expert parameters and computational costs constant.The third configuration integrates both DeepSeekMoE strategies.

3. DeepSeekMoE Architecture

DeepSeekMoE targets expert specialization through fine-grained expert segmentation and shared expert isolation, while addressing routing imbalance with expert- and device-level balance losses.

  • 3.1. Fine-Grained Expert Segmentation: Fine-grained expert segmentation splits each expert into smaller experts and activates more of them while keeping expert parameters and computation constant.Each expert is reduced to 1/m of its original intermediate hidden dimension, while the number of activated experts increases m-fold.
  • 3.1. Fine-Grained Expert Segmentation: 4,426,165,368 potential combinations replace 120 under N=16 when four-way segmentation changes top-2 routing to top-8 routing.The larger combination space is intended to support more accurate and targeted knowledge acquisition.
  • 3.2. Shared Expert Isolation: Shared expert isolation deterministically assigns every token to K_s shared experts and reduces routed activations by K_s to preserve computation.The shared experts capture common knowledge, while routed experts can avoid reproducing that knowledge in their own parameters.
  • 3.3. Load Balance Consideration: Expert-level balance loss mitigates routing collapse, where only a few experts may be selected and others receive insufficient training.The expert-level balance factor is denoted α1, with N′ = mN − K_s and K′ = mK − K_s.
  • 3.3. Load Balance Consideration: Device-level balance loss prioritizes balanced computation across devices rather than imposing strict expert-level balance constraints that could compromise performance.The authors use a small expert-level balance factor and a larger device-level balance factor to pursue these distinct objectives.

4. Validation Experiments

Validation experiments compare DeepSeekMoE with dense and MoE baselines under controlled settings, then test its architectural components and expert specialization. DeepSeekMoE matches or surpasses larger-computation alternatives while showing benefits from shared experts, finer segmentation, and less redundant routed experts.

  • Experimental Setup: DeepSeekMoE uses a multilingual corpus and HAI-LLM-based distributed training with tensor, data, pipeline, and expert parallelism.The validation setup uses 100B training tokens and NVIDIA A100 or H800 GPU clusters.
  • Baseline Comparisons: Under equal total and activated parameters, DeepSeekMoE demonstrates overwhelming advantages over GShard.The validation compares models trained on the same corpus and with the same training hyperparameters.
  • Baseline Comparisons: DeepSeekMoE achieves comparable performance with GShard×1.5 despite GShard×1.5 using 1.5× the expert parameters and computation.At larger scale, DeepSeekMoE can even distinctly outperform GShard×1.5.
  • Baseline Comparisons: DeepSeekMoE nearly approaches Dense×16, the stated upper bound for MoE models in model capacity.Dense×16 is constructed with 16 shared experts, each matching a standard FFN in parameter count.
  • Ablation Studies: Isolating one shared expert improves performance across a majority of benchmarks compared with GShard.This ablation supports the contribution of shared expert isolation to stronger model performance.
  • Ablation Studies: Progressively finer expert segmentation consistently enhances overall model performance under matched total and activated parameters.Experts are segmented into 32 or 64 total experts, including one shared expert and 31 or 63 routed experts.
  • Expert Specialization: Disabling the shared expert raises Pile loss from 1.808 to 2.414 despite maintaining the same computational cost.Activating one additional routed expert does not replace the shared expert’s function.
  • Expert Specialization: With only 4 routed experts activated, DeepSeekMoE achieves a Pile loss comparable with GShard.A separate comparison trained from scratch also finds DeepSeekMoE outperforming GShard with the same total expert parameters and half the activated expert parameters.

5. Scaling up to DeepSeekMoE 16B

DeepSeekMoE 16B is evaluated against dense and open-source models using matched 2T-token training data and broad benchmarks. It delivers comparable or better performance with roughly 40% of the computation, while enabling faster single-GPU inference.

  • Evaluation setup: The evaluation spans language modeling, reasoning, reading comprehension, code generation, closed-book question answering, math, multiple-choice, disambiguation, and Chinese benchmarks.The Open LLM Leaderboard provides an additional comparison with open-source models.
  • Internal comparison with DeepSeek 7B: With only 40.5% of computations, DeepSeekMoE 16B achieves comparable performance with DeepSeek 7B trained on the same 2T-token corpus.The matched corpus isolates the architecture comparison from training-data differences.
  • Evaluation results: DeepSeekMoE 16B shows notable strengths in language modeling and knowledge-intensive tasks, including Pile, HellaSwag, TriviaQA, and NaturalQuestions.Its math-reasoning and code-generation capabilities are also reported as stronger than LLaMA2 7B, while Chinese benchmarks favor its bilingual pretraining.
  • Deployment: With appropriate operator optimizations, DeepSeekMoE 16B can run on a single 40GB GPU and achieve nearly 2.5 times the inference speed of a 7B dense model.The deployment claim is attributed to its modest parameter count.
  • Internal comparison with LLaMA2 7B: With only 39.6% of computations, DeepSeekMoE 16B outperforms LLaMA2 7B on the majority of evaluated benchmarks.Both models are pretrained on 2T tokens; DeepSeekMoE has 245% of LLaMA2 7B’s total parameters.

6. Alignment for DeepSeekMoE 16B

The authors construct DeepSeekMoE Chat 16B through supervised fine-tuning on bilingual, multi-domain data and compare it fairly with two dense chat models. It maintains broad parity at about 40% of their computation, with strengths in code and Chinese benchmarks but weaker multiple-choice performance than DeepSeek Chat 7B.

  • Training and evaluation: DeepSeekMoE Chat 16B is fine-tuned on 1.4M bilingual examples covering math, code, writing, question answering, reasoning, and summarization.The evaluation excludes Pile and CHID while adding BBH for chat-model reasoning assessment.
  • Results: With nearly 40% of computations, DeepSeekMoE Chat 16B achieves comparable performance with 7B dense models across understanding, reasoning, reading comprehension, mathematics, and knowledge-intensive tasks.The result covers PIQA, ARC, BBH, RACE, GSM8K, MATH, TriviaQA, and NaturalQuestions.
  • Baselines: Using the same fine-tuning data, the study compares DeepSeekMoE Chat 16B with LLaMA2 SFT 7B and DeepSeek Chat 7B.The dense comparison models use about 2.5 times the FLOPs.
  • Results: DeepSeekMoE Chat 16B significantly outperforms LLaMA2 SFT 7B on HumanEval and MBPP and also surpasses DeepSeek Chat 7B on code generation.This is the strongest reported post-alignment advantage over the compared dense models.
  • Results: DeepSeekMoE Chat 16B remains behind DeepSeek Chat 7B on MMLU, CEval, and CMMLU, although supervised fine-tuning narrows the gap with DeepSeek 7B.It notably outperforms LLaMA2 SFT 7B on all Chinese benchmarks.

7. DeepSeekMoE 145B Ongoing

A preliminary 145B-parameter study compares DeepSeekMoE with similarly trained MoE and dense baselines. DeepSeekMoE 145B outperforms GShard 137B and matches DeepSeek 67B with substantially fewer computations, while half-activated routing reaches similar performance at still lower cost.

  • Setup: DeepSeekMoE 145B is trained on 245B tokens using the same corpus and tokenizer as DeepSeekMoE 16B.The study uses the same internal benchmarks as the 16B evaluation.
  • Baselines: The comparison includes DeepSeek 67B dense, GShard 137B, and DeepSeekMoE 142B with half-activated routed experts.All compared models share the same training corpus, while MoE models share training hyperparameters and are trained from scratch.
  • Results: Despite comparable total parameters and computations, DeepSeekMoE 145B significantly outperforms GShard 137B.This result is reported as evidence of the architecture’s advantage at larger scale.
  • Results: With only 28.5% of computations, DeepSeekMoE 145B achieves comparable performance with DeepSeek 67B dense.The model remains particularly strong in language modeling and knowledge-intensive tasks but has limitations on multiple-choice tasks.
  • Results: With only 18.2% of computations, DeepSeekMoE 142B with half the activated expert parameters still matches DeepSeek 67B and outperforms GShard 137B.Its performance does not lag far behind DeepSeekMoE 145B.

8. Related Work

Prior MoE language models commonly use top-1 or top-2 routing while facing challenges in training stability, fine-tuning, and expert specialization. DeepSeekMoE addresses the remaining specialization gap through finer-grained expert design.

  • MoE foundations: MoE originated with independent expert modules and was later introduced into language-model training, including large-scale LSTM-based systems.Transformer-based MoE models subsequently extended FFNs into MoE layers.
  • Routing strategies: GShard and Switch Transformer use learnable top-2 or top-1 routing to scale MoE language models to extremely large sizes.Hash Layer and StableMoE instead use fixed routing strategies.
  • Training challenges: ST-MoE targets training instability and fine-tuning difficulty in MoE models.This work complements architectural research with training-strategy improvements.
  • Expert specialization: Most previous MoE models rely on conventional top-1 or top-2 routing, leaving room to improve expert specialization.DeepSeekMoE is presented as a response to that remaining limitation.

9. Conclusion

DeepSeekMoE combines fine-grained expert segmentation with shared expert isolation to pursue stronger expert specialization. Across larger scales, it retains computational advantages while matching several larger or dense models.

  • 9. Conclusion: DeepSeekMoE targets ultimate expert specialization through fine-grained expert segmentation and shared expert isolation.The paper also reports higher specialization than GShard.
  • 9. Conclusion: 16B parameters and 2T tokens yield performance comparable with DeepSeek 7B and LLaMA2 7B using about 40% of computations.The model was also adapted into an MoE chat model through supervised fine-tuning.
  • 9. Conclusion: 145B parameters retain substantial advantages over GShard and achieve performance comparable with DeepSeek 67B using 28.5% or possibly 18.2% of computations.This larger-scale result is described as a preliminary exploration.
  • 9. Conclusion: A public DeepSeekMoE 16B checkpoint can be deployed on a single GPU with 40GB of memory.The release is intended to support research and accelerate large-scale language-model development.

Appendices

The appendices document DeepSeekMoE configurations across model sizes and define the basis for comparing relative expert size.

  • Appendices: Table 7 presents DeepSeekMoE hyper-parameters across various model sizes.The appendix uses this table to provide an overview of the configurations.
  • Appendices: Relative expert size is reported against a standard FFN.This provides the reference used for interpreting expert-size values.

B. Comparing DeepSeekMoE with Larger Models

The larger-model appendix compares DeepSeekMoE with scaled GShard and dense baselines. At 13B total parameters, DeepSeekMoE outperforms GShard×1.5 distinctly.

  • B. Comparing DeepSeekMoE with Larger Models: Comparisons include DeepSeekMoE against GShard×1.2 and GShard×1.5.These comparisons are reported in Table 8 and, at larger scale, Table 10.
  • B. Comparing DeepSeekMoE with Larger Models: Additional comparisons evaluate DeepSeekMoE against Dense×4 and Dense×16 baselines.These results are presented in Table 9.
  • B. Comparing DeepSeekMoE with Larger Models: 13B total parameters allow DeepSeekMoE to outperform GShard×1.5 distinctly.The comparison is conducted at a larger scale and reported in Table 10.

C. Training Benchmark Curves of DeepSeekMoE 16B

The appendix provides benchmark curves during training for DeepSeekMoE 16B and the dense DeepSeek 7B model.

  • C. Training Benchmark Curves of DeepSeekMoE 16B: Figure 7 shows benchmark curves during training for DeepSeekMoE 16B and DeepSeek 7B (Dense).The figure is presented for reference.
Loading 2401.06066v1…