Branch-Train-MiX: Meta开源一个融合多个领域专家模型成一个MOE模型的训练框架

        
          
https://arxiv.org/pdf/2403.07816.pdf  

      

Branch-Train-MiX (BTX) ，提高大型语言模型（LLMs）在多个专业领域（如编程、数学推理和世界知识）的能力。BTX方法的核心思想是结合了Branch-Train-Merge (BTM) 方法和Mixture-of-Experts (MoE) 架构的优势，同时减少了它们的不足。

BTX方法分为三个主要步骤：

1. Branch & Train ：

• 从一个预训练好的种子模型（seed model）开始，创建多个副本。
• 这些副本（称为专家模型，Expert LLMs）在不同的数据子集上独立训练，每个数据子集对应一个特定的知识领域，如数学、编程或维基百科。
• 这种训练方式是并行的，每个专家模型可以异步训练，减少了通信成本，提高了训练吞吐量。

• MiX ：

• 训练完成后，将这些专家模型的feedforward sublayers合并到一个MoE模块中，形成一个统一的MoE模型。
• 在每个Transformer层中，使用一个router network 来选择每个 Token 应该使用哪个专家的feedforward sublayers。
• self-attention sublayers 和其他模块的权重则通过简单平均的方式来合并。

• MoE Finetuning ：

• 合并后的模型在所有训练数据上进行进一步的微调，以便 router network 学习如何在不同的专家之间进行Token级路由。
• 这个过程称为MoE finetuning，它允许模型在测试时动态选择最合适的专家来处理特定的输入。

BTX方法的优势在于：

• 专家训练阶段是并行和异步的，减少了通信成本，提高了训练效率。
• 最终的BTX模型是一个统一的神经网络，可以像任何其他标准LLM一样进行微调。
• BTX模型在推理时的FLOPs（浮点运算次数）不会显著增加，因为它是稀疏激活的，尽管参数数量大大增加。

此外，文章还探讨了BTX的一些变体，如负载平衡（load balancing）、不同的路由方法（如Switch、Soft routing和Sample Top-1）以及专家的分割（splitting）和混合（blending）等策略，以进一步提高模型的性能和效率。通过实验，BTX在多个任务上展示了优于BTM和其他基线模型的性能，证明了其在继续预训练中的有效性和计算效率。