基于LoRA构建稀疏混合专家模型(MoE)的方法(MixLoRA)-探析(微调)

发表于 分类于 阅读次数： 本文字数： 4k 阅读时长 ≈ 13 分钟

介绍

MixLoRA是一种用于优化大规模语言模型(LLMs)微调的新方法，结合了LoRA(Low-Rank Adaptation)和专家混合(Mixture of Experts, MoE)技术。大规模语言模型的微调通常需要大量的计算资源和显存。LoRA通过引入低秩适配器，显著减少了微调时的参数数量和显存需求。然而，LoRA在多任务学习场景中的性能仍有提升空间。专家混合模型(MoE)在多任务学习中表现出色，但其资源需求对普通消费者级GPU来说是一个挑战。

主要特点：

  • 架构设计：1.MixLoRA在冻结的预训练密集模型的前馈网络块中插入多个LoRA专家模块；2.使用常见的top-k路由器（如Mixtral）或top-1开关路由器（如Switch Transformers）来动态选择合适的专家。
  • 性能提升：1.MixLoRA在多任务学习场景下的准确率比现有的参数高效微调(PEFT)方法提高了约9%；2.通过独立的注意力层LoRA适配器增强模型性能。
  • 资源效率：1.引入辅助负载平衡损失来解决路由器的不平衡问题；2.新的高吞吐框架在训练和推理过程中减少了40%的GPU显存消耗和30%的计算延迟。

大型语言模型(LLM)的指令微调已在自然语言处理(NLP)领域中取得了很好的成绩。随着参数规模的增加，LLM已被证明能够识别复杂的语言模式，从而实现强大的跨任务泛化能力。然而指令微调将会使其在计算资源和下游任务的性能之间做出权衡。为了减少全参数微调过程所需的计算和内存资源，引入了参数高效微调(PEFT)方法，其中，低秩自适应是流行的PEFT方法中的一种。MixLoRA通过结合LoRA和专家混合技术，提供了一种高效的微调大规模语言模型的方法，显著提升了多任务学习的性能，同时大幅减少了计算资源和显存的需求。

PEFT，大型语言模型(LLM)在NLP任务中表现出了卓越的能力。在此之后，指令微调进一步使LLM能够理解人类意图并遵循指令，成为聊天系统的基础。但是，随着LLM规模的扩大，对其进行微调成为一个耗时且占用大量内存的过程。为了缓解这个问题，研究探索了不同的方法：参数高效微调(PEFT)、蒸馏、量化、修剪等。LoRA利用低秩矩阵分解线性层权重，是最流行的PEFT方法之一，不仅提高了模型性能，而且不会在推理过程中引入任何额外的计算开销。例如，VeRA结合可学习的缩放向量来调整跨层的共享冻结随机矩阵对。此外，FedPara专注于联邦学习场景的低秩Hadamard积。Tied-Lora实现了权重绑定以进一步减少可训练参数的数量。AdaLoRA使用奇异值分解(SVD)来分解矩阵并且修剪不太重要的奇异值以简化更新。DoRA将预训练权重分解为两个分量，即幅度和方向，并在微调期间利用LoRA进行方向更新，从而有效减少可训练参数的数量。

专家混合：专家混合(MoE)的概念可以追溯到1991年，它引入了一种新颖的监督学习方法，涉及多个网络（专家），每个网络专门处理一组训练示例。MoE的现代版本通过合并稀疏激活的专家来修改Transformer块内的前馈子层，从而能够在不增加计算量的情况下大幅增加模型宽度。LLaVA-MoLE有效地将token路由到Transformer层内的特定领域专家，从而缓解数据冲突并实现与普通LoRA基线一致的性能提升。对于其他基于MoE的架构，MoRAL解决了将LLM适应新领域/任务的挑战。LoRAMoE使用路由器网络集成LoRA，以缓解知识遗忘。PESC使用MoE架构将密集模型转换为稀疏模型，从而降低了计算成本和GPU内存要求。MoE-LoRA提出了一种新颖的参数高效的MoE方法，该方法具有分层专家分配(MoLA)，适用于Transformer的模型。MoCLE可根据指令集群激活定制任务的模型参数。MixLoRA将LoRA作为随机专家进行集成，从而降低了计算成本，同时扩展了模型容量并增强了LLM的泛化能力。

MixLoRA架构

LoRA仅调整额外的自适应参数并替换原始权重更新。LoRA块由两个矩阵组成，和，其中和表示LLM预训练权重的维度()，参数表示LoRA的隐藏维度，其中。然后，更新后的权重，通过以下公式计算：

混合专家(MoE)。MoE架构最初通过GShard引入到语言模型中。MoE层由n个专家组成，表示为，以及路由器。给定隐藏状态的MoE层输出由以下公式确定：

其中，表示路由器输出到第位专家，是第位专家的结果。MixLoRA由两个主要部分构成。第一部分添加了LoRA的vanilla Transformer块构建稀疏MoE块。第二部分利用top-k路由器将来自各种任务（例如ARC、OBQA、PIQA等）的每个token分配给不同的专家模块。假设输入文本为，标签为。让表示第个大型语言模型(LLM)层的第个token的输出隐藏状态，其中是LLM层的总数，是隐藏维度。大型语言模型由堆叠的多头自注意力(MSA)和前馈神经网络(FFN)组成。每个块内都应用了层归一化(LN)和残差连接。普通transformers块中第个LLM层的输出ℓ 通过以下方式计算：

MixLoRA Forward。MixLoRA基于LoRA专家构建的。MixLoRA利用LoRA在微调期间存储每个专家的更新参数，而不是仅使用LoRA来构建每个专家。这种方法使MixLoRA与现有的预训练MoE模型更加一致。在MixLoRA的MoE块中，这些专家的基本权重由密集模型的单个前馈网络(FFN)共享，以提高训练和推理效率：

其中是FFN层的预训练权重，由共享，表示我们用不同token和任务选择特定LoRA专家的Top-K路由器，表示MixLoRA模块中的第k个LoRA专家。MixLoRA的作用是取代公式中密集模型的FFN层，其关键概念是通过路由器为每个token选择不同的专家，其中每个专家由不同的LoRA和原始FFN层组成。

Top-K路由器。MoE层中的Top-K路由器将每个token分配给最合适的专家。路由器是一个线性层，它计算输入token被路由到每个专家的概率：。在稀疏转换器块中，此路由器根据输入token激活最合适的LoRA专家。它利用softmax激活函数来模拟专家的概率分布。路由器的权重是路由网络的可训练参数。我们的设计采用了top-2门控路由器，它为每个输入token 从个可用的中选择最佳的两个专家：

在推理过程中，top-k门控路由器会动态地为每个token选择最佳的k个专家。通过这种机制，混合专家和路由器协同工作，使专家能够开发不同的能力并有效地处理不同类型的任务。专家负载平衡。专家负载不平衡是MoE面临的重大挑战。这是因为某些专家往往会更频繁地被top-k路由器选择。为了解决这种负载不平衡问题，我们在训练时应用负载平衡损失来减轻专家的负载不平衡。受Switch Transformers的启发，计算辅助损失并将其添加到总损失中。给定个专家（索引为i = 1到N）和具有T个token的批次B，辅助损失计算如下：

其中是前个路由器，是分配给专家的token分数，是分配给专家的路由器概率的分数。最终损失乘以专家数量，以保持损失在专家数量变化时保持不变。此外，我们使用作为辅助损失的乘积系数，该系数足够大以确保负载平衡，同时保持足够小而不会影响主要交叉熵目标。添加带有注意力层的LoRA。MixLoRA进一步扩展了其微调功能以涵盖注意力层。先前的研究，例如ST-MoE，表明微调注意力层可以显着提高性能。为了增强MixLoRA的微调过程，将LoRA适配器集成到密集模型的注意力层中。实验结果表明，经过q、k、v和o投影微调的MixLoRA模型与仅使用稀疏的LoRA专家混合层（即MixLoRA MoE）训练的配置相比，始终能够获得更高的平均分数。

性能优化

  • 降低计算复杂度。MixLoRA中的每个专家网络包括一个共享的冻结FFN和多个LoRA，用于在微调期间存储专家的每个线性投影层更新参数。此设置导致计算复杂度根据路由器的K设置而变化。这是因为MixLoRA中的token级Top-K路由器将每个token发送给K个专家进行计算，然后聚合他们的残差并产生输出。以LLaMA为例，LLaMA的FFN块由三个线性投影权重组成：和，前向传播过程可以表示为。如图 3 (a) 所示，MixLoRA中的每个专家都有相同的前向传播过程，只是每个线性投影层都有一个单独的 LoRA，并且每个专家的输入 x 由 MixLoRA路由器预先分配。这在处理长序列输入时会带来显著的开销，对MixLoRA的性能优化提出了重大挑战：如何在保持模型精度的同时降低 MixLoRA的计算复杂度？考虑到我们已经共享了 FFN 块的权重，我们进一步共享计算结果以降低计算复杂度。如上图所示，方法不是预先分配MixLoRA块的输入序列，而是首先将输入直接并行发送到FFN块的和，然后根据MixLoRA路由器输出的路由权重对这些线性投影的输出进行切片。投影的计算复杂度无法降低，因为其计算过程取决于和投影的输出。这种方法显著降低了MixLoRA计算复杂度的三分之一。MoE块。在实验中，这种方法保持相同的模型性能的情况下token计算延迟比vanilla MixLoRA大约低30%。
  • 优化多MixLoRA训练和推理。受m-LoRA提出的多LoRA优化的启发，我们还优化了MixLoRA以实现多模型高吞吐量训练和推理。之前，我们通过消除重复计算来降低MixLoRA的计算复杂度。当使用两个或更多MixLoRA模型进行训练和推理时，这些模型的多任务输入被打包成一个批次，以提高训练吞吐量。具体来说，我们首先将批处理输入并行发送到和，然后使用MixLoRA模型的单独路由权重对这些线性投影的输出进行切片。这种方法通过共享相同的预训练模型权重，显着减少了多个MixLoRA模型的内存使用量。这种方法保持了相同的性能的情况下对每个模型的峰值GPU内存使用量减少了约45%。

结论

这里介绍了MixLoRA，这是一种使用多个LoRA的专家和一个冻结的共享FFN块的MoE方法。与传统的LoRA-MoE方法不同，MixLoRA将多个LoRA与共享FFN层融合，并在微调期间使用它们为每个专家存储更新的参数，使其与预先训练的MoE模型更加一致。它还采用了自注意力LoRA适配器和辅助负载平衡损失来提高性能并解决路由器不平衡问题。此外，还设计了一个高性能框架来优化MixLoRA中多个LoRA的专家的计算过程，以进行训练和推理。因此，该框架将MixLoRA的计算复杂度降低了30%，并在训练或推理多个MixLoRA模型时节省了大约40%的GPU内存使用量。评估表明，MixLoRA在单任务和多任务场景中的表现优于基线。对于单任务学习，MixLoRA在LLaMA-2 7B上的平均准确率比LoRA提高了6.2%，比DoRA提高了2.6%。在多任务学习中，MixLoRA的准确率明显优于LoRA 9.8%，优于DoRA 9%。