分开训练,合并使用:基于混合专家模型的分模块后训练
BAR(分支-适配-路由)是一种逐步训练语言模型能力的方案:独立训练领域专家,将它们合并为一个混合专家模型,并能在不影响其他模块的情况下升级任一专家。
在预训练之后,语言模型需要经过一系列中期和后训练阶段才能变得实用——学习遵循指令、推理问题、可靠地调用工具等。然而,在这些阶段之后更新或扩展模型往往具有挑战性。最可靠的选择是从头开始重新训练,但成本高昂且需要完全访问原始训练设置。进一步训练新数据更便宜,但可能导致模型丢失已有能力。而且,由于后训练通常涉及多个阶段(每个阶段有自己的数据和目标),添加新技能意味着需要重新运行或调整每个阶段以避免破坏已有成果。
为了解决这些问题,艾伦人工智能研究所(Ai2)的研究团队提出了BAR(Branch-Adapt-Route,分支-适配-路由)方法,一种模块化后训练的方案。与一次训练一个处理所有数据的单一模型不同,BAR独立训练领域专家——每个专家通过自己的完整训练流程——并通过混合专家(MoE)架构将它们组合成一个统一模型。每个专家都可以独立开发、升级或替换,而无需影响其他专家。该团队发布了配方、技术报告以及用于验证该方法的模型检查点。
BAR的动机源于之前的FlexOlmo工作。FlexOlmo表明,基于MoE的模块化训练在预训练阶段效果良好:可以从共享基础分支,训练特定领域的FFN专家同时冻结所有共享层,然后合并回。然而,研究人员发现这种方案不适用于后训练。原因在于:预训练主要更新知识表示,这些知识主要存在于FFN层;而后训练则引入行为变化(如新的输出格式、推理模式和安全约束),这些变化需要改变共享参数(如注意力层、嵌入层和语言模型头)。例如,当研究人员在强化学习与验证奖励(RLVR)中直接使用FlexOlmo方法时,奖励曲线完全平坦;模型在所有共享参数冻结的情况下根本学不会。这促使他们专门为后训练开发了新的配方。
BAR包含三个阶段。阶段一:独立专家训练。每个领域专家实例化为包含两个专家的MoE:一个冻结的“锚定”专家(保留基础模型的FFN权重)和一个可训练的专家。专家根据领域需求经历相应的训练阶段。在实验中,数学和代码专家经历中期训练、有监督微调(SFT)和RLVR;工具使用和安全专家仅使用SFT。关键的技术贡献是跨阶段的渐进式解冻计划:中期训练时所有共享层冻结;SFT时解冻嵌入层和语言模型头,这对于引入新特殊令牌(如工具使用的函数调用格式)是必要的;RLVR时解冻所有共享参数,包括注意力层。此外,每个专家还训练混合了领域特定和通用SFT数据的数据集,因为仅领域SFT会严重损害通用能力。
阶段二:专家合并。训练后,所有专家合并为一个单一的MoE模型。跨专家运行中分化的共享参数(由于在SFT或RLVR期间解冻)被简单地平均。研究发现,与任何单个专家相比,这种平均化在领域特定评估中几乎没有引入可测量的性能损失。阶段三:路由器训练。最后,在MoE内部训练路由器,同时冻结所有其他专家和共享权重。研究人员发现,仅用SFT数据的分层5%样本就足以实现有效路由,使得这一阶段快速且廉价。
实验表明,BAR模型在19个基准测试中超越了所有不需要从零重新训练中期训练的基线方法。平均而言,BAR优于仅进行后训练重训练的模型(49.1 vs 47.8),尤其在数学(+7.8)和代码(+4.7)方面提升显著。这归因于模块化训练的结构性优势:在单一流水线中,后期对数学和代码的RL可能会削弱早期SFT阶段学到的安全能力,而模块化训练完全避免了这一点。值得注意的是,中期训练后的稠密模型合并效果极差(整体6.5),而BAR在无需中期训练的情况下仍然大幅领先。
BAR最实用的特性之一是专家可以独立升级。研究展示了两种升级方式:用更高质量数据和RL训练的代码专家替换原有专家,使代码性能提升16.5点,而其他领域几乎不变;在已有数学专家基础上添加RL训练,使数学性能提升13点。在这两种情况下,只有受影响的专家和轻量级路由器需要重新训练,而传统单一流水线则需要重新训练整个模型。这使得BAR的领域更新成本呈线性增长,而单一模型更新成本则几乎是二次方增长。
研究还得出一些实用经验:后训练需要比预训练更大的灵活性,渐进式解冻至关重要;仅使用领域SFT数据是不够的,必须混合通用SFT数据;解冻后的权重平均效果出奇地好;并非所有专家都需要在推理时激活,激活4/5专家即可达到几乎相同的性能。未来工作包括从原生稀疏架构开始,而非将稠密模型转换为稀疏模型,这可能提高模块化方法的效率和可扩展性。