中文

BAM!就是如此:专家混合模型的简单高效参数升级复用

机器学习 2024-08-20 v2

摘要

专家混合(MoE)框架因其相较于稠密模型的优越性能,已成为大型语言模型的热门架构。然而,从头大规模训练MoE的成本极其高昂。现有方法通过独立预训练多个稠密专家模型并用其初始化MoE来缓解这一问题。这一方法通过使用专家的前馈网络(FFN)来初始化MoE的专家层,同时合并其他参数。然而,该方法将稠密模型参数的重用限制在FFN层,从而限制了将这些模型"升级复用"为MoE时的优势。我们提出BAM(Branch-Attend-Mix),一种简单而有效的方法来解决这一缺陷。BAM通过不仅使用稠密模型的FFN来初始化MoE层,还充分利用专家的注意力参数,将它们初始化为注意力混合(MoA)层的软变体,从而充分利用专用稠密模型。我们探索了两种升级复用注意力参数的方法:1)从稠密模型中初始化包含所有注意力参数的独立注意力专家,以获得最佳模型性能;2)在所有专家间共享键和值参数,以提高推理效率。为进一步提高效率,我们为MoE采用了并行注意力Transformer架构,使得注意力专家和FFN专家可以并行计算。我们在590百万到20亿参数范围的种子模型上的实验表明,在相同的计算和数据约束下,BAM在困惑度和下游任务性能上均超越了基线。

关键词

引用

@article{arxiv.2408.08274,
  title  = {BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts},
  author = {Qizhen Zhang and Nikolas Gritsch and Dwaraknath Gnaneshwar and Simon Guo and David Cairuz and Bharat Venkitesh and Jakob Foerster and Phil Blunsom and Sebastian Ruder and Ahmet Ustun and Acyr Locatelli},
  journal= {arXiv preprint arXiv:2408.08274},
  year   = {2024}
}