推动专家混合模型的极限:用于指令调优的极端参数高效MoE

 人工智能咨询培训老师叶梓 转载标明出处

MoE是一种在特定子任务上具有专门化子模型(专家)的神经网络架构,这些子模型可以独立优化,以提高整体性能。然而,传统的MoE在大规模应用时面临挑战,因为需要在内存中存储所有专家。这不仅增加了内存的需求,而且在完全微调(full fine-tuning)时计算成本极高。为了克服这些限制,Cohere for AI的研究团队提出了一种极端参数高效的MoE方法。该方法通过将MoE架构与轻量级专家相结合,实现了在只更新不到1%的参数的情况下,与全参数微调相媲美的性能。这种方法不仅在未见任务上具有良好的泛化能力,而且不依赖于任何先验任务知识。

论文链接:https://arxiv.org/pdf/2309.05444

项目链接:GitHub - for-ai/parameter-efficient-moe

方法

研究者们将任务集划分为训练任务(Ttrain)和评估任务(Teval),即。基础预训练模型首先在训练任务上进行微调,然后在评估任务上的每个未见任务进行零样本评估。标准方法是微调所有模型参数,这导致计算和内存成本高。

研究者们将专家混合(MoE)架构推向参数效率的极端,使用参数高效微调(PEFT)方法。PEFT 方法通过限制权重更新至有限数量的参数,解决了在大规模完全微调大型语言模型(LLM)时出现的与更新大量参数相关的挑战。

(IA)3 方法:引入三个新向量 ​, ,,这些向量重新调整自注意力中的关键和价值激活,以及逐位置前馈层中的中间激活:其中Q,K,V 是自注意力的查询、键和值投影矩阵, 是预训练模型中前馈层的冻结权重。由于 (IA)3 只更新,它在参数效率上极为高效。对于3亿参数的 T5 模型,它仅更新了总参数的0.018%。

LORA 方法:优化了大型语言模型中密集层的低秩分解。对于预训练权重矩阵 和输入激活 ,LORA 将 W_0​ 分解为两个低秩矩阵: 其中 ​,秩。在微调期间,所有预训练权重都冻结,只更新A 和 B 矩阵。LORA 适配可以用于每个 Transformer 块中的所有线性层,包括自注意力的查询 Q、键 K、值 V 以及前馈层 W_1​ 和 W_2​ 的输出。与 (IA)3 不同,LORA 适配在参数使用上提供了更多的灵活性。

研究者们提出了一种极端参数高效的 MoE 框架,该框架利用轻量级“适配器”作为专家,建立在预训练密集模型之上。MoE 层由路由器网络R 和一组n 专家​ 组成,每个专家 E_i​ 是一个参数化函数。路由器网络通常由一个密集层和随后的 softmax 函数组成,这一层接收中间令牌表示 x 作为输入,并根据门控分数​ 组合每个专家的输出:

MoV 架构如图 2 左侧所示,MoV 架构的概览突出了软合并,其中只有向量和路由器针对每个多头注意力块进行更新。右侧提供了实现 MoV 层核心的 JAX 风格伪代码。

通过 MoV 和 MoLORA 方法,只更新一小部分参数,在训练和推理时都有多个实际好处,后者是 MoE 架构独有的。这些增益如下:

  • 训练效率:极端参数高效的 MoE 公式导致内存显著减少。在训练期间冻结大多数参数减少了计算模型参数梯度的计算开销,也减少了存储模型优化器状态的内存需求。
  • 推理效率:MoV 和 MoLORA 方法的固有结构模块化允许在推理时显著减少内存需求。对于传统的 MoE 模型,需要在内存中存储许多完整的前馈块(甚至是基于特定架构的模型的完整副本),这是一项昂贵的工作。而使用研究者们的方法,无论确切类型如何,除了轻量级参数高效专家外,只需要在内存中存储模型主干的单一副本。这导致推理时内存需求的显著减少。

想要掌握如何将大模型的力量发挥到极致吗?2024年10月26日叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具。

留言“参加”即可来叶老师的直播间互动,1小时讲解让您轻松上手,学习如何使用 Llama Factory 微调模型。

实验 

实验设置:

  • 基础预训练模型:使用了不同大小的T5模型,参数量从770M到11B不等,包括T5 v1.1+LM适配。
  • 优化器:所有实验都使用Adafactor优化器进行微调,学习率设置为3×10−43×10−4。
  • 序列长度:输入序列长度设置为1024,目标序列长度设置为256。
  • 批量大小和训练步数:对于所有参数高效的MoE变体,使用32的批量大小在500K步内进行微调。

基线比较:

  • 与完全微调的模型(T0基线)和标准的参数高效微调方法((IA)3和LORA)进行比较。
  • T0基线:通过使用较大的批量大小和学习率进行微调,获得了更好的性能。具体来说,使用256的批量大小和1×10−31×10−3的学习率进行10K步的微调。
  • (IA)3和LORA:对于秩为4的(IA)3和LORA,使用相同的训练超参数,即学习率为3×10−43×10−4,批量大小为32,训练500K步。

评估指标:

  • 在8个未见数据集上测试了他们的方法和基线方法,这些数据集包括ANLI、HellaSwag、WinoGrande和5个Super Glue数据集,涵盖了指代消解、自然语言推理、多项选择问答、故事完成和词义消歧等不同任务。
  • 计算每个评估数据集在不同提示模板下的中位数准确率,并报告每个数据集的结果以及所有数据集的平均结果。

基础设施:

  • 所有实验都在TPU v4机器上进行,最多使用256个pod切片。
  • 使用SeqIO和T5X框架进行训练、评估和推理,这些框架支持跨TPU核心的数据和模型并行处理。

消融研究:

  • 路由输入: 令牌与句子嵌入的比较,探讨了任务表示形式的指令嵌入如何影响路由和下游泛化。
  • 路由策略: 软合并与离散top-k路由策略的比较,研究了在参数高效的MoE中哪种路由策略最有效。

研究者们首先将他们提出的MoE模型(包括MoV和MoLORA)与标准的PEFT方法((IA)3和LORA)进行了比较。这些比较是基于使用T5-3B作为基础模型的零样本(zero-shot)性能测试。表1展示了这些方法在不同任务上的平均中位数结果。

表1中数据显示,MoV和MoLORA在多个任务上相较于(IA)3和LORA显示出显著的性能提升。特别是,使用30个专家的MoV实现了相比其密集对应模型14.57%的性能提升。这一改进在所有未见任务上都是一致的,并且仅通过增加每个专家0.018%的参数更新实现。对于LORA,配备15个专家的MoLORA实现了平均中位数得分5.70%的提升,虽然这一提升相比MoV来说不那么显著,但仍然证明了混合学习的优势。

研究者们进一步比较了MoE方法与全参数微调(Full-FT)的性能。表1显示,即使是只更新0.32%的参数的MoV-10,也能达到与全参数微调相当的性能。当增加专家数量时,MoV和MoLORA甚至略微超过了全参数微调的性能。

研究者们探讨了MoE方法在不同规模的基础模型(从770M到11B参数)上的表现。图1(右图)展示了MoV在不同模型规模上与(IA)3和全参数微调的性能比较。结果显示,MoV在所有评估的模型规模上均保持比标准PEFT更高的性能,并与全参数微调取得了相当的结果。

研究者们还研究了专家数量对未见任务性能的影响。图4(中间图)展示了在所有模型规模上,增加专家数量通常会提高未见任务的性能。例如,在770M和11B参数的模型上,使用60个专家的MoV方法实现了最佳性能。

研究者们比较了MoV使用不同路由策略的性能。图4(右图)显示,软合并(soft merging)策略显著优于离散的top-k路由策略。这表明在极端参数高效的MoE设置中,软合并能够实现更优的性能。

为了理解专家如何在不同任务中被激活,研究者们分析了MoV中5个专家在最后解码器块的前馈层的平均路由概率。图5展示了不同任务(包括训练任务和未见评估任务)中专家的平均路由概率。结果表明,专家的专业化是内在的,并且可以从训练任务迁移到未见任务。

最后,研究者们对MoE架构对超参数的敏感性进行了广泛的消融研究。他们发现,较大的批量大小更可能导致MoE模型坍塌为单一专家,而较小的学习率有助于稳定训练,防止快速、不平衡的更新抑制多样性,从而导致次优解决方案。

通过这些详细的分析,研究者们展示了他们提出的MoE模型在参数效率和性能上的显著优势,以及在不同设置下的有效性和鲁棒性。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值