专家混合模型MoE
专家混合模型MoE——Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer 大得离谱的神经网络:稀疏门控混合专家层
研究背景与动机-Why
- Dense-LLM 传统密集模型:传统的LLM主要通过扩大训练数据量和增加模型大小规模来提高预测准确率
- 无论遇到什么问题(比如修水管还是写法律文件),整个团队的每个专家都必须到场,发表一点意见,这非常低效
- 同时,随着模型规模和训练数据量的增加,会导致训练成本爆炸式增长
研究方法与创新-How
条件计算
- 条件计算:一种计算范式,核心思想是:神经网络并非对每个输入都激活和使用全部的参数进行运算(即调用整个模型),而是根据输入的具体内容(即条件),动态地、稀疏地选择网络中的一部分路径或参数进行计算
- 有一个聪明的“项目经理”(门控网络)。他先看一下问题是什么(输入),然后只呼叫最相关的几个专家(比如水管工和电工)来解决问题。其他专家(比如律师和医生)就可以休息,不参与当前的计算。
- 条件:当前输入的数据,对于神经网络的某一层,条件就是传递到该层的输入值,所以,条件是输入数据,而不是预设的、固定的指令,而是动态变化的,对于每个输入都不同
- 优势:在显著增加模型总参数,扩大模型规模的同时,几乎不增加计算成本和推理时间,因为计算成本只和激活的专家有关,而不是所有专家
混合专家模型-MoE
本文提出了一种混合专家模型,模型规模是提升模型性能的关键因素之一。在有限的计算资源预算下,用更少的训练步数训练一个更大的模型,往往比用更多的步数训练一个较小的模型效果更佳。
MoE 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下,可以显著扩大模型或数据集的规模。特别是在预训练阶段,与稠密模型相比,混合专家模型通常能够更快地达到相同的质量水平。
MoE模型的理念:如果有一个包括了多个领域知识的复杂问题:该使用什么样的方法来解决呢?最简单的办法就是把各个领域的专家集合到一起来攻克这个任务,比如事先要把不同的任务先分离出来,这样才便于分发给不同领域的专家,让他们来帮忙处理,最后再汇总结论。它由多个专业化的子模型(即“专家”)组合而成,每一个“专家”都在其擅长的领域内做出贡献。而决定哪个“专家”参与解答特定问题的,是一个称为“门控网络”的机制。
混合专家模型-MoE的构成:
稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个),每个专家本身是一个独立的神经网络。在实际应用中,这些专家通常是前馈网络 (FFN),但它们也可以是更复杂的网络结构,甚至可以是 MoE 层本身,从而形成层级式的 MoE 结构。 门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如,在下图中,“More”这个令牌可能被发送到第二个专家,而“Parameters”这个令牌被发送到第一个专家。有时,一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点,因为路由器由学习的参数组成,并且与网络的其他部分一同进行预训练。
门控:门控不是输入值,本身不是条件,而是“根据条件做决策的机制”
Thinking:我在最开始看到这个模型的第一时间就想到了Hadoop中的MapReduce,和分布式系统非常类似,先由Map进行任务shuffle切片成不同的Slice,通过处理后再由Reduce汇总,但是我仔细对比了一下,发现了有几点不同:
MapReduce的任务通常是按照数据分片进行静态分配,而MoE模型输入的token是动态分配的,不同的样本可以映射到不同的专家集合,且是可以随着训练不断优化的,同时,这实现了一种负载均衡,避免了有的专家过于专业被样本压爆,有的专家不太专业没有token输入导致出现饥饿现象的路由不均
研究结果与未解决问题-What’s Next
- 训练挑战: 虽然 MoE 能够实现更高效的计算预训练,但它们在微调阶段往往面临泛化能力不足的问题,长期以来易于引发过拟合现象。
- 推理挑战: MoE 模型虽然可能拥有大量参数,但在推理过程中只使用其中的一部分,这使得它们的推理速度快于具有相同数量参数的稠密模型。然而,这种模型需要将所有参数加载到内存中,因此对内存的需求非常高。以 Mixtral 8x7B MoE 为例,需要足够的 VRAM 来容纳一个 47B 参数的稠密模型。之所以是 47B 而不是 8 x 7B = 56B,是因为在 MoE 模型中,只有 FFN 层被视为独立的专家,而模型的其他参数是共享的。此外,假设每个token只使用两个专家,那么推理速度 (以 FLOPs 计算) 类似于使用 12B 模型 (而不是 14B 模型),因为虽然它进行了 2x7B 的矩阵乘法计算,但某些层是共享的。
延申-Extension
堆叠的LSTM,也称为深度LSTM,是指将多个LSTM层一个接一个地堆叠起来,形成一个更深度的循环神经网络架构。
堆叠LSTM:一栋高楼。每一层楼(LSTM层)都接收来自下一层的信息,并进行更高级、更抽象的处理。低楼层处理细节(如单个单词的特征),高楼层处理宏观信息(如句子整体的语义或情感)。
堆叠分层的核心思想:通过增加深度来学习不同层级的抽象特征
底层LSTM:接近输入,捕获序列中底层级、局部的模式和依赖关系
高层LSTM:接近输出,接收底层LSTM处理后的,已经部分抽象化的序列信息,并在此基础上捕获更高层级、全局的模式例如:在一个NLP任务中,底层的LSTM可能负责学习单词的形态、词性以及相邻单词之间短期依赖关系,而高层LSTM负责理解整个句子的主旨、感情倾向或段落之间的逻辑关系
softmax函数:将一个未归一化的数值(通常是一个神经网络的输出)转化为一个概率分布(所有的输出值都在0-1之间,并且概率的输出值和为1)
通常位于神经网络的最后一层,以预测每个类别的概率,然后选择概率最大的类别作为最终的输出。W_g的本质是一个可学习(通过模型的训练,W_g的参数可以不断调整,使得各个专家能够逐渐专业化)的“路由字典”
形状为[input_dim,n_experts]的可训练矩阵,即[输入向量x的维度,模型中专家的总数量]MoE模型在通信开销上的优化:
为了控制通信开销,将MoE中每个专家分配的批量控制在一个较高的水准,确保所有专家都有足够大的批量来处理,避免了某些专家因为批量太小而“饿死”(类似于操作系统中的饥饿算法),从而提高计算设备的利用率,同时使得每次通信传输的数据包更大,次数更少,分摊了固定通信开销,提高了通信效率
同时,通过负载均衡技术,避免例如专家1收到10个token,而专家2收到1014个token,导致系统总耗时受限于最慢的环节,专家1的大部分时间都被浪费了,因为它大部分时间都在等待通信
通过负载均衡,使得门控网络将数据均匀路由,每个专家都收到512个token,方案B的总耗时(512C + L)远小于方案A的总耗时(1014C + L),因为计算负载被均衡了,通信次数也做到了最小化
关于MoE模型的一篇综述-A Survey on Mixture of Experts in Large Language Models
发展历程
MoE
参数介绍
- fi:各个专家神经网络,通常是一个线性-ReLU 线性网络,最终生成fi(x; Wi)
- Wi:专家神经网络中传入的参数
- G:门控网络,由Θ作为输入参数,输出为G(x; Θ),并由门控功能的设计,MoE层分成了Dense MoE和Sparse MoE
Dense MoE
- Dense MoE密集专家混合层在每次迭代中激活所有的专家网络,{f1, . . . , fN }
-
- g(x; Θ):softmax函数操作前的门控值
Sparse MoE
虽然密集MoE通常可以获得更高的预测精度,但是也会显著增加计算开销。
为了解决这个问题,引入Sparse MoE,稀疏的门控MoE层,目的是在每次向前传递时只激活选定的专家子集,而不是汇总所有专家的输出,从而实现一种稀疏性。
与Dense MoE的不同点在于,引入了TopK函数,其中传入要选择的top的k的数量与topk的门控值,并添加了一个噪声Rnoise。
TopK(·, k)函数只保留向量前 k 个条目原始值,而将所有其他条目设置为 -∞。经过Softmax函数操作后,被赋值为-∞的条目概率近似为零。
超参数 k 根据具体应用进行选择,常见的选择是 k = 1 或 k = 2
虽然Sparse Gate G(x; Θ)能够在不增加相应计算成本的条件下极大地扩展模型的参数空间,但是它可能导致负载均衡问题。这种问题是指专家之间的token分配不均,有的专家过于专业被样本压爆,有的专家不太专业没有token输入导致出现饥饿现象的路由不均。为了解决这个问题,每个MoE层都包含了一个辅助损失函数,来促进每批专家之间token的even distribution平均分配。
其中,Di代表了专家i分到token的比例,Pi代表通过门控机制,专家i分配概率的比例。为了确保可以实现负载均衡,even distribution,Lload-balancing函数值应该是最小的。
例如,最理想条件,每个专家分配的token数量与概率是相同的,此时,Di和Pi的取值都是1/N,那么函数值则为:
为了在不相应增加计算需求的情况下有效扩展模型参数,专家混合(MoE)架构成为一种可行的解决方案。混合专家架构利用一系列专业模型和门控机制,动态选择适当的 “专家网络 “来处理给定的输入。这使得模型能够根据需要分配计算资源,这一概念被称为条件计算。
Algorithm Design MoE的算法设计
门控设计
Sparse MoE Token Choice Gating 门控
N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean, “Outrageously large neural networks: The sparselygated mixture-of-experts layer,” arXiv preprint arXiv:1701.06538, 2017.中,有必要对前k位专家的输入进行门控,来提高MoE的计算效率。基本原理:同时咨询多位专家对给定输入的意见,网络可以有效地权衡和整合他们各自的贡献,来提高效率。为了增加专家的数量,该论文中采用了两级分层的MoE,来减少在专家人数众多的情况下的分支因子。
W. Fedus, B. Zoph, and N. Shazeer, “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,” Journal of Machine Learning Research, vol. 23, no. 120, pp. 1–39, 2022.中,top-1门控策略也可以产生有竞争力的结果。
A. Yang, J. Lin, R. Men, C. Zhou, L. Jiang, X. Jia, A. Wang, J. Zhang, J. Wang, Y. Li et al., “M6-t: Exploring sparse expert models and beyond,” arXiv preprint arXiv:2105.15082, 2021.中,提出了一种名为“expert prototyping”的新变种,即把专家分成k组,然后在每组中应用top-1选取,实验结果的好坏依次为:分成4组的top-1的expert prototyping、1个top-4门控、1个top-16门控、1个top-1门控。
Token-Choice Auxiliary Loss 门控的辅助损失函数
N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton,and J. Dean, “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer,” arXiv preprint arXiv:1701.06538,2017.
为了促进专家之间token的平均分配,引入Auxiliary Loss损失函数
首先定义单个专家的重要性Importance,其中B是一个训练批次(batch),g(x;Θ)是对单个固定专家给出的分配权重/概率(每个token都会有一个值)。把同一个专家在这个训练批次batch上所有的token的权重求和,就得到该专家在这个batch上的Importance。
然后定义辅助损失函数Loss,将其定义为重要性集合的变异系数(Coefficient of Variation, CV)的平方
直观含义:把每个专家的Importance看成一个向量,计算这个向量的变异系数CV,CV越大说明专家之间的负载差异越大。然后把CV的平方作为损失,加到总损失里去最小化,就会促使gating网络把token分配更均匀。
Expert Capacity for Token-Choice Gating 令牌选择门控的专家容量
- [25]
- 结合专家容量的限制,定义了专家可以处理的令牌数量的阈值。
- 但是这可能会因为token数据量过大,导致token溢出,即多余的token没有指定的专家处理。