Branch-Train-MiX: 可以大幅度降低训练成本的新型混合专家模型架构

在人工智能的广阔范畴中，大型言语模型（LLMs）的研讨始终是技能前沿的热门话题。随着技能的不断前进，咱们对LLMs的期望也在不断进步，不只期望它们能够处理日常的言语使命，还期待它们在专业范畴如编程、数学推理和世界常识等方面展示出更高的才能。在这样的背景下，Meta的研讨团队提出了一种名为Branch-Train-MiX（BTX）的立异办法，旨在高效地进步LLMs在多个专业范畴的才能。

论文标题：Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

组织：FAIR at Meta

论文链接：arxiv.org/pdf/2403.07…

1. Motivation

传统的分布式LLMs练习办法本钱昂扬、通讯本钱大、同步练习软弱等问题。这些问题限制了模型的练习功率，阻止了模型在特定范畴的运用。本文所介绍的Branch-Train-MiX便是一个针对于垂直范畴的LLMs的低代价练习办法。

2. 立异的办法

BTX办法的中心在于它的两阶段练习战略：

• 首要是分支练习（Branch-Train），从一个预练习的基础模型出发，练习多个专家模型，这些专家模型是在不同的数据集上相互隔离的进行并行练习的；
• 其次是混合专家（MiX），将这些专家模型的前馈参数混合到Mixture-of-Experts（MoE）层中，并经过MoE微调阶段学习token级别的路由。这种办法不只进步了练习功率，还坚持了模型的一致性，使得终究的模型能够像任何其他规范LLM相同进行微调或运用。

下面我将具体解说这些阶段：

2.1 Branch & Train: 并行异步专家练习

Branch（分支）:

BTX办法从一个预练习的种子模型（seed model） 开端，更简略的说，咱们首要要有一个预练习完的LLM。然后咱们将这个种子模型复制多份，每个复制的模型都会练习为一个独立的专家模型（expert LLM）。这个进程称为“分支”，因为它就像代码的分支相同，从一个相同的基础代码中，改动不同的部分，然后形成不同的版别。

Train（练习）:

接下来，每个专家模型在各自的特定范畴数据集前进行练习。这些数据集分别对应于不同的常识范畴，比如咱们将复制三个种子模型，然后分别在数学数据集，代码数据集，维基百科数据集前进行练习，每个数据集对应不同的范畴，每个专家模型经过这种方式学习并专精于其对应范畴的特定常识。

关键点是，这些专家模型是并行且异步地练习的，即它们独登时在不同的核算资源上运转，相互之间无需任何同步。这种练习方式大幅度压低了通讯本钱，并进步了练习的吞吐量。

2.2 MiX: 结合独立专家成为混合专家模型

在一切专家模型异步练习完成后，BTX办法进入混合阶段。在这个阶段，专家模型的前馈子层（feedforward sublayers）被兼并到一个一致的Mixture-of-Experts（MoE）模块中。这个MoE模块在模型的每一层都包含一切专家的前馈子层，并经过一个路由器网络（router network）动态挑选在给定输入令牌时应运用哪个专家的前馈子层。

首要阐明一下概念：

前馈子层（Feedforward Sublayers）

在深度学习的Transformer架构中，前馈子层（一般表明为FF层）是网络中的一个根本组件，它对输入数据进行非线性变换。具体来说，每个Transformer层包含一个自注意力子层（Attention Sublayer）和一个前馈子层。前馈子层一般由两个线性变换组成，中心夹着一个非线性激活函数（如ReLU或GELU）。在BTX办法中，每个专家模型的前馈子层被规划为处理特定范畴的数据，这意味着它们学习到了与各自范畴相关的特征表明。

Mixture-of-Experts（MoE）模块

MoE模块是一种特别的网络结构，它答应模型在处理不同输入时动态地挑选和激活最合适的专家子网络。MoE的中心思维是将大型网络分解为多个较小的专家网络，每个专家网络担任处理特定类型的输入数据。这种办法能够明显进步模型的核算功率，因为在任何给定时刻，只要一小部分专家网络被激活，然后削减了无效核算。

路由器网络（Router Network）

路由器网络是MoE结构中的关键组件，担任依据输入数据的特征动态挑选要激活的专家网络。在BTX办法中，路由器网络接纳输入令牌（token），并输出一个路由概率分布，指示每个专家网络处理该令牌的相对重要性。路由器的输出一般经过softmax函数转换为概率分布，以保证挑选的专家网络总数与MoE模块中的专家数量相匹配。

MiX（混合）进程:

在BTX办法的MiX阶段，研讨者首要将一切专家模型的前馈子层逐层的兼并到一个MoE模块中。这个MoE模块在每个Transformer层中包含一切专家的前馈子层，并经过路由器网络来决定对于每个输入令牌应该激活哪些专家的前馈子层。这种路由机制答应模型在处理特定使命时，运用在练习阶段学习到的范畴特定常识。

具体来说，假定咱们输入了一个token $x$，每个专家的前馈子层都会核算FFl i(x)，可是只要那些被路由器网络选中的专家子层才会被实际用于核算终究输出。这经过引进一个路由函数$g$和一个线性变换WlW_lWl​来完成，公式如下：

FFMoEl(x)=∑i=1Ngi(Wlx)FFil(x)mathrm{FF}_{mathrm{MoE}}^l(x)=sum_{i=1}^N g_ileft(W_l xright) mathrm{FF}_i^l(x)FFMoEl​(x)=∑i=1N​gi​(Wl​x)FFil​(x)

• FFil(x)FF_i^{l}(x)FFil​(x)表明第$i$个范畴专家模型在第$l$层的前馈子层输出。
• WlW_lWl​是第$l$层的线性变换矩阵。
• gig_igi​是第$i$个专家的路由函数。
• $x$是输入向量。
• $N$是专家的总数。

这里的gi(Wlx)g_ileft(W_l xright)gi​(Wl​x)是一个稀疏的路由函数，它决定了每个token应该由哪些专家模型处理。一般运用Top-k路由，其间k=2，意味着对于每个token，只会挑选两个最相关的专家子层进行核算。

除了前馈子层外，模型的其他部分（如自注意力层和嵌入层）经过简略地均匀各个专家模型的权重来兼并。这个进程不引进新的参数，保证了模型的一致性和进一步微调的才能。

简而言之，所谓的混合进程便是：将自注意力层和嵌入层经过均匀值进行兼并，而前馈层则在每一层都增加一个路由函数，模型在中心核算时在每一层都可能会挑选不同的前馈子层进行核算。

MoE-finetuning（MoE微调）:

在混合阶段之后，模型会进行MoE微调，这是一个继续练习的进程，旨在优化路由器的挑选战略，并调整模型权重以进步全体功能。这一阶段，模型在一切用于练习专家模型的数据前进行微调，使得路由器能够学习如安在不同范畴间有效地混合专家的前馈子层。

4. 试验剖析

4.1 试验设置

试验根据Llama-2 7B模型。研讨者们创建了三个专家模型，分别在数学、编程和维基百科数据集前进行练习。此外，还包含了原始的Llama-2 7B模型作为一个“通才”专家。这些专家模型被兼并到一个MoE模型中，并在一切数据源前进行了进一步的微调。

4.2 首要结果

4.2.1 全体功能

• 专家模型的专业化：试验结果显现，每个专家模型在其特长的范畴内体现出色。例如，数学专家在GSM8K和MATH等数学使命上取得了明显的前进，而编程专家在HumanEval和MBPP等编程使命上体现优异。
• BTX模型的全面进步：BTX模型在一切范畴使命上都有所进步，尤其是在数学和编程使命上，接近或超过了专门化模型的功能。这表明BTX办法能够有效地结合专家模型的专业化常识，一起坚持或进步在其他使命上的功能。

• 与基线的比较：BTX模型与其他基线模型比较，包含Llama-2 7B、Dense模型、Sparse upcycling和Branch-Train-Merge (BTM) 办法，展示了更好的功能。特别是在数学和编程范畴，BTX模型的功能明显优于BTM模型，这表明MoE微调阶段学习令牌级路由的优点。

4.2.2 核算功率

• 与Dense和BTM的比较：尽管BTX在MoE练习阶段运用了较少的练习预算，但其在一般才能上的进步与Dense和BTM比较更为明显。
• 与Sparse upcycling的比较：作为一种特别的BTX（没有专家练习阶段），Sparse upcycling在相同的或更大的核算预算下体现不如BTX，这表明专家并行练习的核算功率优势。

4.3 融化研讨和剖析

4.3.1 BTX练习的融化研讨

• 路由办法：研讨了不同的路由办法，包含Switch、Sample Top-1和Top-2，以及是否运用负载平衡（load balancing）。结果表明，Top-2路由与负载平衡相结合在大多数使命上体现最佳。
• 练习战略：测试了不同的BTX练习变体，如不运用负载平衡、冻住专家的前馈模块等。这些融化试验有助于了解不同规划挑选对模型功能的影响。

4.3.2 路由剖析

• 路由决议计划：剖析了不同范畴使命的令牌路由决议计划，发现Top-2路由与负载平衡保证了专家间的负载分布更均匀。
• 专家参与度：在数学和编程使命中，GSM8K使命倾向于挑选编程和通才专家，而MATH使命更多依赖于数学专家。这表明不同使命依据其内容和难度挑选了最合适的专家。

5 结论

本文提出的Branch-Train-MiX (BTX) 办法的首要意图是为了解决大型言语模型（LLMs）在多个专业范畴内进行高效练习的问题。具体来说，BTX旨在完成以下几点：

1. 进步多范畴才能：经过练习专家模型来增强LLMs在特定范畴（如编程、数学推理、世界常识等）的功能。
2. 下降通讯本钱：经过并行异步练习专家模型，削减在分布式练习中坚持多个模型副本同步所需的通讯开支。
3. 进步练习功率：并行练习答应模型运用更多的核算资源，然后加快练习速度。
4. 增强模型通用性：经过MoE结构整合专家模型，坚持一个一致的模型，以便进行进一步的微调和运用。
5. 坚持核算功率：尽管BTX模型具有更多的参数，但它经过稀疏激活的方式坚持了较低的推理核算本钱。

相对于其他办法，BTX的优势包含：

1. 专家练习的并行性：与同步练习比较，BTX的分支练习阶段能够异步进行，这意味着每个专家模型能够独立于其他模型练习，然后进步了练习进程的鲁棒性和功率。
2. 细粒度的MoE混合：BTX不是简略地混合专家模型的终究输出，而是在Transformer的每个层内进行更细粒度的混合，这答应模型在更深层次上运用专家常识。
3. 路由机制的优化：BTX运用路由器网络来动态挑选最合适的专家前馈子层，这种根据令牌的路由办法比传统的根据使命的路由办法愈加灵活和高效。
4. 一致的微调才能：与Branch-Train-Merge (BTM) 办法比较，BTX经过MoE微调阶段学习令牌级路由，终究得到一个能够像任何规范LLM相同进行微调的一致模型。
5. 功能与功率的平衡：BTX在坚持较低推理核算本钱的一起，完成了在多个专业范畴内的功能进步，这使得它在核算功率和使命功能之间取得了良好的平衡。
6. 削减过拟合危险：因为BTX办法在多个范畴上练习专家模型，这有助于模型学习到更泛化的特征表明，然后削减在特定范畴数据上过拟合的危险。