【深度学习笔记】MoE的介绍以及调研

【深度学习笔记】MoE的介绍以及调研

MoE介绍

MoE（Mixture of Experts）模型与标准Dense模型的主要区别在于：

• 在相同参数规模下，MoE模型的预训练速度和推理速度更快。
• MoE模型需要更高的内存来装载。
• MoE模型的训练更具挑战性。

与标准Transformer模型相比，MoE模型的主要结构变化是将FFN层替换为多个专家（expert）层，并添加了一个门控（gate）层。门控层负责将每个token路由到不同的专家，最终将专家的输出合并为模型的输出。

直观上，使用多个专家可以使得每个专家专注于特定的任务或领域。然而，这种结构的训练相对复杂，需要确保所有专家都能得到充分学习，避免模型退化为一个Dense模型。

MoE训练方法

目前主要的MoE相关研究包括Mixtral、DeepSeekMoE和上海人工智能研究院的Llama-MoE。这些研究采用了不同的训练和初始化策略。

从头训练

DeepSeekMoE采用了从头训练的方法，并进行了多尺寸实验。研究在2B（等效0.3B）、16B（等效7B）和145B（等效22B）等不同规模上进行了验证。DeepSeekMoE采用了共享专家和多个领域专家的策略，以提供综合性和多样性的知识获取方式。

对于2B的MoE模型，具体参数设置如下：

• 9层Transformer
• 1280的hidden size
• 10个注意力头，维度为128
• 64个expert（1个共享专家+63个领域专家）
• 每个expert的大小为1/4个FFN大小
• 整个MoE层相当于16倍的FFN大小

训练细节包括：

• 参数初始化标准差为0.006
• AdamW优化器，beta1=0.9，beta2=0.95
• weight decay 0.1
• batch size token 4M
• 学习率调度：2000个step warmup到1.08e-3，80%step时缩小到0.316，90%step时继续缩小到0.316的平方

实验结果显示，在2B规模时，MoE模型达到了与Dense模型相当的性能。在16B规模下，使用40%的计算量达到了Llama2-7B的性能。在145B规模下，使用28.5%的计算量达到了DeepSeek-67B的性能。

从现有模型扩展

Mixtral和Llama-MoE采用了从现有模型初始化并微调的策略，但选择了不同的FFN初始化方向。

向上扩展

Mixtral使用了向上扩展的策略，将Mistral-7B扩展为一个47B等效于12B模型的MoE模型。具体做法是将FFN复制8份，其他参数完全复用。Mixtral在32k的上下文长度上进行了微调，整体性能接近Llama2-70B和GPT-3.5。由于使用滑动窗口机制，Mixtral在长文本处理上具有优势，例如在密钥检索任务中达到了100%的准确率。

向下扩展

Llama-MoE采用了向下扩展的策略，将Llama2-7B缩小为一个7B等效于3~3.5B的MoE模型。实验结果显示，Llama-MoE-v1的性能低于Llama2-7B，但优于同激活量的Dense模型。

训练细节包括：

• 在Llama2-7B的基础上训练
• 使用112块A100 80G GPU
• global batch token 15M
• 学习率使用cosine调度，从2e-4到2e-5
• 总计训练200B token（13.6k step）

技术认知

expert认知

对于expert的初始化，不同研究采用了不同的策略。Llama-MoE等向下拆分的结构需要解决如何拆分expert的问题。研究验证了以下几种方法：

• Independent Random：将原始FFNs的参数随机分割成多个专家网络，每个专家网络独立且不重叠。
• Independent Clustering：基于聚类来构建专家网络，试图找到权重之间的自然分组。
• Sharing Inter：在专家网络之间共享部分中间神经元。
• Sharing Inner：在专家网络之间共享更多的中间神经元。

实验结果显示，独立随机拆分的效果最好，但需要在15B token之后才能显现出来。聚类方法初期效果较好，但长期效果不佳，可能是因为降低了expert的表达能力。

expert区分

研究还验证了不同expert是否有效以及它们捕捉的信息类型。Mixtral分析了不同层的token分配情况，发现浅层网络接近随机分配，而深层网络则更倾向于选择特定的expert。Llama-MoE也发现不同的任务或数据类型会选择不同的expert。

实现细节

Mixtral使用megablocks实现MoE，而Llama-MoE基于transformers开发，并在GitHub上开源。实现中需要注意Deepspeed版本（0.13.1以上）以及一些潜在的bug，特别是在x!=y的MoE模型训练中。关于gate的初始化，实验表明不需要冻结gate，因为在大约20B参数时，冻结与不冻结的loss差异已经很小。

参考文献

1. https://huggingface.co/blog/moe
2. Llama-MoE: https://github.com/pjlab-sys4nlp/llama-moe/blob/main/docs/LLaMA_MoE.pdf
3. DeepSeekMoE: http://arxiv.org/abs/2401.06066
4. Mixtral: http://arxiv.org/abs/2401.04088
5. Switch Transformers: https://arxiv.org/abs/2101.03961