[阅读笔记20][BTX]Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

这篇论文是meta在24年3月发表的，它提出的BTX结构融合了BTM和MoE的优点，既能保证各专家模型训练时的高度并行，又是一个统一的单个模型，可以进一步微调。

这篇论文研究了以高效方法训练LLM使其获得各领域专家的能力，例如写代码、数学推理以及自然知识。现有的融合多个专家模型的方法有Branch-Train-Merge和Mixture-of-Experts，前者BTM各专家模型在不进行任何同步的情况下并行训练，大大提升了训练时的吞吐量，但是缺乏一个统一的模型，导致没法进行后续的SFT和RLHF，这两步是对齐LLM的重要步骤。后者MoE虽然是一个统一的模型，可以进行微调了，但是训练时是各专家模型是完全同步的，并且由于all-to-all通信，随着专家数量增加通信成本也在增加。

这篇论文提出的BTX就是融合了BTM和MoE的优点，弥补了二者的缺点，具体来说，BTX的各个专家模型可以异步的独立训练，大大提高了模型训练时的数据吞吐量，另外BTX是一个统一的模型，所以之后可以对其进行微调。
实现分为三步，首先是Branch，这里取了四个一样的种子模型LLaMA-2 7B，然后其中三个分别在数学数据集、代码数据集、维基百科上预训练，最后剩下的保留LLaMA原始权重。前三个专家模型分别具有数学推理能力、代码能力、世界知识，最后一个专家模型作为通才专家，将通用知识迁移到模型中。第二步是Train，这三个领域专家分别在各自领域数据集（Llemma、CodeLlama、Wikipedia）上预训练，这个过程是并行且互不干扰的。第三步是Mix，也就是将这四个专家模型进行混合，这一步在下一张ppt会详细说明。

具体融合四个专家的过程其实就是把这四个专家的前馈层进行合并，也就是将同层次的四个前馈合并为一个MoE层，下图公式展示了如何合并，整个MoE层输出是各前馈层的加权和。这里使用了Top2路由，对于输入x使用投影矩阵Wl进行投影，然后取值最大的两个专家模型进行混合，混合比例由softmax计算得到。
对于模型的其他部分，例如注意力层、embedding层，BTX混合各专家的方法是直接把对应的模型参数取平均，这个方法比较粗暴，作者给出的解释是这个做法基于一个假设：自注意力层比前馈层更通用化。另外后续微调阶段还会对这些参数进行调整，所以问题不大。

预训练说完了，接下来就是微调了。微调使用的数据仍然来自训练时数据，作者对用于训练四个专家的数据集进行采样得到微调用的数据，采样概率数学是30%，代码是40%，维基百科是10%，LLaMA-2是19%。
下图是训练三个专家模型时使用的三个专业领域数据集及其采样比例。

接下来是结果展示。左图是各模型的训练代价和平均性能对比，x轴是训练所需要的GPU天数，圆圈大小是推理时激活参数数量。右图是不同领域的性能差异，可以看到BTX在各方面都得到了很大的改善，尤其是代码领域，已经接近专用模型CodeLlama了。

这张图反映了每个token来自于哪两个专家，共有四个专家，所以有六种组合，分别用不同的颜色来代表。如果是由领域内专家生成的token，则标上下划线。对于数学、代码、知识问答这三个领域的三个输出，大部分token都是由各自领域专家生成。