清华,arxiv,发表于2025年1月
https://github.com/THUDM/MoELoRA_Riemannian
用于微调基础模型的更强大低秩专家混合方法
摘要
为简化基础模型的微调过程,低秩适配器(LoRAs)已在包括指令调整和领域适应等多个领域广泛应用。LoRA的基本概念是将满秩矩阵分解为两个低秩矩阵的乘积,这减少了存储消耗并加速了训练过程。此外,为解决LoRA表达能力有限的问题,专家混合(MoE)被引入以整合多个LoRA适配器。LoRA专家的集成在多个下游场景中带来了明显的性能提升。然而,LoRAs混合(MoE-LoRA)在调优和推理过程中仍表现出较低的稳健性。受将LoRA训练为子空间投影仪的黎曼预条件器的启发,我们提出了一种新的MoE-LoRA训练策略,通过多空间投影来稳定和促进其特征学习过程。对随机梯度下降(SGD)和自适应矩估计权重衰减(AdamW)优化器的实验验证了我们方法的有效性。源代码可在https://github.com/THUDM/MoELoRA_Riemannian获取。
1 引言
参数高效微调(PEFT)技术为基础模型(FMs)的微调提供了一种经济高效的解决方案(Zhang等人,2025a)。其中,低秩适应(LoRA)因其通用性和简单性而成为一种流行的技术。具体而言,LoRA引入了可训练的低秩矩阵A和B来更新基础模型的内部模块,公式为\(X = W + BA\)。从某种意义上说,它们的乘积近似于对预训练权重的满秩更新。虽然LoRA显著减少了可训练参数的数量,但它也存在两个局限性:表示能力有限和梯度次优。
- 局限性1:表示能力有限。低秩矩阵的一个自然问题是表示能力较弱,尤其是在复杂任务中。为解决这一问题,一种直接的解决方案是将多个LoRA模块集成到专家混合框架中,即MoE-LoRA。图1(左)展示了一个简单的MoE-LoRA框架。这些努力在许多场景中显著提高了LoRA的性能,如视觉语言任务、多任务学习、持续学习等。简而言之,MoE-LoRA的发展路径大致可分为两类:(i)为特定领域设计专用的MoE-LoRA框架,如MOELoRA(Liu等人,2023)和MoCLE(Gou等人,2023)。(ii)通过架构、更新和损失约束在技术上改进MoE-LoRA,如MoLA(Gao等人,2024)和HydraLoRA(Tian等人,2024)。然而,这些努力大多没有考虑到训练MoE-LoRA的不稳定性和低效性。
- 局限性2:梯度次优。困扰LoRA的另一个问题是梯度次优。这是因为低秩矩阵A和B共同构成了一个具有一定曲率的商流形空间,导致流形内最优梯度与满秩最优梯度不一致。这进一步导致了LoRA的训练过程次优。为缓解这一问题,Zhang等人(Zhang & Pilanci,2024)通过黎曼梯度预条件器增强了LoRA的梯度,公式为\(\nabla_{A} L = (B^{T}B)^{-1}\nabla_{A} L\)和\(\nabla_{B} L = \nabla_{B} L(AA^{T})^{-1}\)。经过数学推导,这些预条件器有助于构建两个梯度投影仪,确保更新是根据满秩梯度在A的行空间和B的列空间上的投影进行的,即\(X_{new} = X - \eta[Proj_{col(B)}(\nabla_{X} L)^{T} + Proj_{row(A)}(\nabla_{X} L)]\)。
通过对局限性1和局限性2的综合分析,一个自然的问题出现了:基于LoRA的结构如何在保证解决这两个局限性的同时,进一步逼近完全微调?受MoE-LoRA和梯度预条件方法的启发,这个问题的一个直接答案是将两种方法结合起来,同时克服表示能力和次优性的局限。
具体来说,每个LoRA专家的梯度可以通过各自的黎曼预条件器进行优化。然而,我们认为在MoE-LoRA中对专家进行加权求和的过程为每个LoRA专家的流形引入了基于门的缩放,从而根据各自的门值\(g_{i}\)改变了它们的曲率。
我们在图1的右侧说明了这一现象,该图绘制了一个2 - 专家MoE-LoRA在\(g_{1}<g_{2}\)条件下的示例。具体来说,在专家1和专家2各自的空间中,由\(B_{1}A_{1}\)和\(B_{2}A_{2}\)构建的流形最初具有相同的曲率,因为它们的低秩矩阵具有相同的秩。
然而,在乘以门值后,流形\(g_{1}B_{1}A_{1}\)被更显著地缩放,因此在MoE全空间中比\(g_{2}B_{2}A_{2}\)具有更大的曲率。结果,专家1在全局最优下降和流形内最优下降之间表现出更高的差异。这一现象表明,需要进一步优化每个专家的预条件器,以考虑门值的影响。
在本文中,我们提出了一种简单而有效的解决方案,通过各自的门值\(g_{i}\)以轻量级的方式进一步缩放每个专家的梯度。我们改进的MoE-LoRA梯度更新过程如下:
\[\begin{array}{r}
X_{new }=X-\eta \sum_{i=1}^{N_{Expert }} g_{i} Proj_{col\left(B_{i}\right)}\left(\nabla_{X} \mathcal{L}\right)^{T} \\
-\eta \sum_{i=1}^{N_{Expert }} g_{i} Proj_{row\left(A_{i}\right)}\left(\nabla_{X} \mathcal{L}\right)
\end{array}\]
图1 MoE-LoRA整体架构及其梯度更新过程的深入分析。该图左侧展示了LoRAs混合的流程,它固定前馈神经网络(FFN)的预训练权重,并与路由门一起训练一系列LoRA适配器。右侧展示了MoE-LoRA的更新方式。
具体来说,我们绘制了一个2 - 专家MoE-LoRA在\(g_{1}<g_{2}\)条件下的示例,这导致流形\(g_{1}B_{1}A_{1}\)进一步扭曲。这里我们简单省略了固定的预训练权重,并假设\(X = g_{1}E_{1} + g_{2}E_{2}\)以便于展示。因此,对于随机步骤t,我们绘制一个状态点\(\frac{1}{2} X^{(t)}\),它等于\(\frac{g_{1}^{(t)}B_{1}^{(t)}A_{1}^{(t)} + g_{2}^{(t)}B_{2}^{(t)}A_{2}^{(t)}}{2}\),因此作为t时刻两个流形状态的中心点。该图说明\(g_{1}B_{1}A_{1}\)具有更高的曲率,因此其局部最优下降和全局最优下降投影更加明显。这表明需要与门相关的预条件器。
我们总结我们的贡献如下:
- 我们将LoRAs混合结构与黎曼预条件器相结合,以缓解LoRA的表示能力有限和次优性问题。
- 我们分别为MoE-LoRA的门值重新缩放梯度预条件提出了理论和工程解决方案。
- 我们在一系列基础模型下实现并检验了我们的重新缩放方法,展示了其在各种任务中的有效性。
2 相关工作
2.1 LoRA及其变体
LoRA(Hu等人,2021)将满秩矩阵分解为两个低秩矩阵的乘积,这被广泛认为是参数高效微调的有效解决方案。研究提出了几种变体来改进LoRA:在初始化方面,PISSA(Meng等人,2024)利用奇异值分解(SVD)来获得W的主奇异分量,而MiLoRA(Wang等人,2024a)利用次要奇异值和向量。
LoRA-Pro(Wang & Liang,2024)和LoRA-GA(Wang等人,2024c)近似初始梯度的方向,使其与完全微调的方向一致。LoRA +(Hayou等人,2024)引入了学习率分离策略,\(\eta_{B}>\eta_{A}\)。ResLoRA(Shi等人,2024)和SIBO(Wen等人,2024)通过引入残差路径加速收敛并减轻过平滑问题。
DoRA(Liu等人,2024b)将权重向量分解为方向和幅度,仅使用其方向分量。rsLoRA(Kalajdzievski,2023)提出了秩稳定缩放因子\(\lambda_{t}=r_{t}^{1 / 2}\)以确保稳定的梯度更新。为防止过拟合,BiLoRA(Qiang等人,2024)采用双层优化策略,而其他方法则实现了随机失活(dropout)机制(Wang等人,2024b;Lin等人,2024)。
2.2 LoRAs混合
MoE已成为解决复杂任务的关键框架。通过整合多个专家模块,它可以根据特定输入动态选择合适的专家(Jacobs等人,1991)。早期研究,如LoRAMoE(Dou等人,2024)和MixLoRA(Li等人,2024),率先通过为全局和下游任务集成LoRA专家引入了MoE-LoRA架构。此后,MoE-LoRA在一系列领域展示了其有效性,如持续学习(Dou等人,2024;Yang等人,2024)、视觉语言多模态任务(Gou等人,2023;Chen等人,2024)和多任务应用(Liu等人,2023)。
最近的研究集中在通过架构改进和改进训练策略来增强MoE-LoRA。例如,MoLA(Gao等人,2024)在不同层分配不同数量的专家,而MixDA(Diao等人,2023)引入多个领域自适应模块以支持多领域知识。还提出了其他方法(Wu等人,2024a;Liu等人,2023;Wu等人,2024b;Gou等人,2023;Wang等人,2022)来加强MoE-LoRA。为了促进MoE-LoRA的训练,Luo等人(Luo等人,2024)通过引入对比损失解决了随机路由问题。
同时,MoV(Zadouri等人,2023)选择将轻量级向量与稀疏选择机制相结合,以实现高效的专家分配。其他方法,包括(Dou等人,2024;Li等人,2024;Zhu等人,2023),专注于专家之间的负载平衡。然而,据我们所知,仍然缺乏专门针对MoE-LoRA模型的梯度优化工作。
2.3 梯度预条件器
在大多数深度学习场景中,梯度下降算法通过计算基于梯度的更新来更新模型参数。为了加速优化过程,引入了梯度预条件的概念。先进技术如Adagrad(Duchi等人,2011)通过累积平方梯度\(G_{t}=\sum_{i=1}^{t} g_{i}^{2}\)动态调整学习率,并通过\(\Delta \theta_{t}=-\eta G_{t}^{-1 / 2} \cdot g_{t}\)更新模型。Adam(Kingma,2014)通过结合动量和偏差校正扩展了这种方法,通过对角预条件器缩放梯度,并以\(\Delta \theta_{t}=-\eta \frac{m_{t}}{\sqrt{v_{t}}+\epsilon}\)的形式进行更新,其中\(v_{t}=\beta_{2} v_{t-1}+(1-\beta_{2}) g_{t}^{2}\)。
AdamW(Loshchilov,2017)进一步在Adam中引入了权重衰减。最近的研究为不同预条件策略下的缩放梯度下降方法提供了理论支持。其核心思想是通过向梯度引入缩放矩阵来调整更新的方向和幅度。Tong等人(Tong等人,2021)证明了缩放梯度下降方法的局部收敛性。Jia等人(Jia等人,2024)扩展了这项工作,证明了缩放梯度下降对于最小二乘矩阵分解问题\(\left\|A B^{T}-Y\right\|_{F}^{2} / 2\)的全局收敛性,表明该方法在不同条件数下都能实现全局收敛。还出现了其他缩放梯度下降的变体,如Zhang等人提出的两种正则化策略(Zhang等人,2023;2024)。
在高维设置中,缩放梯度下降已进一步扩展到张量优化(Tong等人,2022;Ma等人,2023)。Mishra等人(Mishra等人,2013;Mishra & Sepulchre,2016)也将黎曼原理应用于涉及低秩矩阵的优化。考虑到数据的流形几何结构,引入了黎曼度量\(g_{p}(v, w)\)来指导沿流形的梯度更新。最近,Zhang等人(Zhang & Pilanci,2024)将黎曼预条件器的思想引入到LoRA中,通过在低秩矩阵的梯度上附加一个\(r × r\)的预条件器。结果,与传统的梯度优化器如SGD和AdamW相比,他们提高了LoRA的微调性能。
3 方法
我们详细阐述我们的动机,并详细说明我们对黎曼预条件方法针对MoE-LoRA所做的修改。我们还介绍了理论基础和工程解决方案。
3.1 LoRA专家中的黎曼预条件器
作为预备知识,我们首先简要介绍黎曼预条件器(Zhang & Pilanci,2024)。假设预训练模型的权重为W,其加性低秩分量为B和A,令\(X = W + BA\)表示整个权重矩阵,令c和n分别表示损失函数和学习率。对于普通梯度下降方法,梯度更新过程通过公式(1)至(4)描述,其中从(2)到(3)的推导依赖于忽略学习率的二阶项。显然,(4)中的\(B \nabla_{A} L+\nabla_{B} L A\)是X的理想快速傅里叶变换(FFT)梯度的近似。
\[\begin{array}{rlrl}
X_{new } & =W+B_{new } A_{new } & (1) \\
& =W+\left(B-\eta \nabla_{B} \mathcal{L}\right)\left(A-\eta \nabla_{A} \mathcal{L}\right) & & (2) \\
& \approx W+B A-\eta B \nabla_{A} \mathcal{L}-\eta \nabla_{B} \mathcal{L} A & (3) \\
& =X-\eta\left(B \nabla_{A} \mathcal{L}+\nabla_{B} \mathcal{L} A\right) &
\end{array}\]
随后,根据推导链式法则以及\(X = W + BA\)这一简单事实,我们直接得到\(\nabla_{A} L=(\nabla_{A} X)(\nabla_{X} L)=B^{T}(\nabla_{X} L)\),同样\(\nabla_{B} L=(\nabla_{X} L) A^{T}\)。因此,(4)可以转换为:
\[X_{new }=X-\eta\left[B B^{T}\left(\nabla_{X} \mathcal{L}\right)+\left(\nabla_{X} \mathcal{L}\right) A^{T} A\right], (5)\]
这实际上与FFT更新公式\(X_{new }=X-\eta \nabla_{X} L\)在不同方向上更新模型。这种现象的发生是因为由\(BA\)构建的X的扭曲子空间导致其流形内的最优梯度下降与全矩阵X的最优梯度下降不一致。为了解决这种不一致性,Zhang等人(Zhang & Pilanci,2024)通过以下方式缩放A和B的梯度:
\[\begin{aligned}
\nabla_{A} \mathcal{L} & =\left(B^{T} B\right)^{-1} \nabla_{A} \mathcal{L} \\
\nabla_{B} \mathcal{L} & =\nabla_{B} \mathcal{L}\left(A A^{T}\right)^{-1}
\end{aligned}\]
使得(5)可以表示为:
\[\begin{gathered}
X_{new }=X-\eta\left[B\left(B^{T} B\right)^{-1} B^{T}\left(\nabla_{X} \mathcal{L}\right)\right. \\
\left.+\left(\nabla_{X} \mathcal{L}\right) A^{T}\left(A A^{T}\right)^{-1} A\right] \\
=X-\eta\left[Proj_{col(B)}\left(\nabla_{X} \mathcal{L}\right)^{T}\right. \\
\left.+Proj_{row (A)}\left(\nabla_{X} \mathcal{L}\right)\right]
\end{gathered}\]
其中流形内的更新是根据全矩阵梯度在A的行空间和B的列空间上的投影进行的。因此,与未缩放的下降步骤相比,它更好地逼近了完全微调。受这项工作的启发,将其解决方案扩展到MoE-LoRA的一种直接方法是通过(6)分别缩放每个LoRA专家的梯度。然而,在MoE-LoRA中,公式\(X = W + BA呈现为不同的形式:
\[X = W + \sum_{i = 1}^{N_{Expert}}g_{i}B_{i}A_{i}\]
其中,\(N_{Expert}\)表示激活专家的数量,\(g_{i}\)表示特定专家\(i\)的门值。结果,它不仅在公式(1) - (4)中为每个专家\(i\)引入了门值\(g_{i}\),而且在公式(5)中也为每个专家\(i\)引入了额外的门值\(g_{i}\),因为推导链式法则\(\nabla_{B_{i}}L = g_{i}(\nabla_{X}L)A_{i}^{T}\)和\(\nabla_{A_{i}}L = g_{i}B_{i}^{T}(\nabla_{X}L)\)。为了进一步说明,我们正式推导整个结果。注意,门值是通过具有复杂非线性操作的softmax计算得到的,因此为了便于推导近似,我们将它们视为常数。按照(6)中的传统黎曼预条件器,我们有:
\[\begin{aligned}
X_{new} &= W + \sum_{i = 1}^{N_{Expert}}g_{i}(B_{i} - \eta\nabla_{B_{i}}\mathcal{L})(A_{i} - \eta\nabla_{A_{i}}\mathcal{L})\\
&\approx X - \eta\sum_{i = 1}^{N_{Expert}}g_{i}(B_{i}\nabla_{A_{i}}\mathcal{L} + \nabla_{B_{i}}\mathcal{L}A_{i})\\
&= X - \eta\sum_{i = 1}^{N_{Expert}}g_{i}[B_{i}(B_{i}^{T}B_{i})^{-1}\nabla_{A_{i}}\mathcal{L}\\
&\quad + \nabla_{B_{i}}\mathcal{L}(A_{i}A_{i}^{T})^{-1}A_{i}]
\end{aligned}\]
\[\begin{aligned}
&= X - \eta\sum_{i = 1}^{N_{Expert}}g_{i}[g_{i}B_{i}B_{i}(B_{i}^{T}B_{i})^{-1}B_{i}^{T}(\nabla_{X}\mathcal{L})\\
&\quad + g_{i}(\nabla_{X}\mathcal{L})A_{i}^{T}(A_{i}A_{i}^{T})^{-1}A_{i}]\\
&= X - \eta\sum_{i = 1}^{N_{Expert}}g_{i}^{2}Proj_{col(B_{i})}(\nabla_{X}\mathcal{L})^{T}\\
&\quad - \eta\sum_{i = 1}^{N_{Expert}}g_{i}^{2}Proj_{row(A_{i})}(\nabla_{X}\mathcal{L})
\end{aligned}\]
其中推导步骤(9)表示传统的黎曼预条件器缩放。应该解释的是,(10)由全矩阵梯度在\(A\)专家的行空间和\(B\)专家的列空间上的投影集合组成。
3.2 重新缩放预条件器
公式(10)呈现了梯度投影集合的平方值加权和。一般来说,激活的专家越多,每个专家的门值越小,从而导致组合梯度更小;另一方面,专家之间越平衡,组合梯度也越小,因为基本不等式定理\(\sum_{i}x_{i}^{2} \geq \frac{(\sum_{i}x_{i})^{2}}{n} = \frac{1}{n}\)在\(x_{i}\)相等时取等号。结果,由于这些平方门值,全矩阵\(X\)的梯度会被低估。从流形和曲率的角度来看,我们解释为将(8)中的\(g_{i}\)视为流形缩放器,它减小了\(B_{i}A_{i}\)的大小,从而可能增加其曲率。然而,传统的黎曼预条件器没有考虑到流形缩放器\(g_{i}\),因为它是为单个LoRA适配器设计的。为了缓解这个平方问题,我们假设对黎曼预条件器进行进一步的重新缩放步骤:
\[\begin{aligned}
\nabla_{A_{i}}\mathcal{L} &= \frac{(B_{i}^{T}B_{i})^{-1}\nabla_{A_{i}}\mathcal{L}}{g_{i}}\\
\nabla_{B_{i}}\mathcal{L} &= \frac{\nabla_{B_{i}}\mathcal{L}(A_{i}A_{i}^{T})^{-1}}{g_{i}}
\end{aligned}\]
它被引入以在公式(9)的推导中替换(6),从而消除变量\(g_{i}\),并在最终公式(10)中仅保留\(g_{i}\)的一次幂。通过这种变换,最终的多专家投影集合与单个LoRA适配器的投影具有相同的尺度,如公式(12)所示。因此,MoE - LoRA的训练将避免被低估。
\[\begin{aligned}
X_{new} = X & - \eta\sum_{i = 1}^{N_{Expert}}g_{i}Proj_{col(B_{i})}(\nabla_{X}\mathcal{L})^{T}\\
& - \eta\sum_{i = 1}^{N_{Expert}}g_{i}Proj_{row(A_{i})}(\nabla_{X}\mathcal{L})
\end{aligned}\]
3.3 工程近似
虽然公式(11)为MoE - LoRA提供了一种消除低估的方法,但它是不可实现的,因为每个LoRA模块对于每个批次样本的每个单个令牌都存在各自的\(g_{i}\)。实际上,在训练过程中,反向传播总是在对批次中每个样本的每个单个令牌的所有损失进行平均后运行。因此,当我们优化一个LoRA模块时,不可能重构和重新缩放每个单个令牌贡献的各自梯度。作为替代,我们设计了一种对(11)和(12)的工程近似方法,在模型前向传播期间用其平方根\(\sqrt{g_{i}}\)替换每个门值\(g_{i}\)。因此,只有在(6)的预条件器下才能实现公式(12),因为公式(10)中门值\(g_{i}^{2}\)的二次项现在自然地变成了线性项\(g_{i}\)。
用\(\sqrt{g_{i}}\)替换\(g_{i}\)同时会对前向传播造成破坏,因为平方根的和不等于1。一种可能的解决方案是重新归一化这些平方根,使其和为1。然而,这会在正向和反向传播期间为专家分配的权重之间引入不一致性。因此,我们提出另一种策略来兼顾两个方面,即手动分配公式(8)中的可优化和不可优化组件,以满足公式(8)中的前向传播和公式(12)中的反向传播的要求。在前向传播过程中,所提出的策略简单表示为:
\[X = \hat{W} + \sum_{i = 1}^{N_{Expert}}\hat{\sqrt{g_{i}}}B_{i}A_{i} + (g_{i} - \sqrt{g_{i}})\hat{B}_{i}\hat{A}_{i}, \quad(13)\]
其中\(\hat{p}\)表示\(p\)不需要梯度,这也意味着\(p\)应该在整个神经网络中与梯度跟踪分离。通过这样分解可优化和不可优化组件,低秩矩阵\(A\)和\(B\)能够按照(12)进行优化。此外,通过在前向传播中保留可优化的\(g_{i}\)项,并将所有\(\sqrt{g_{i}}\)视为不受优化影响的常数,保留了门(\(g = G(x)\))的传统训练行为。另外,这种修改对原始前向计算过程仅引入了最小的开销。
4 实验
我们进行了一系列对比实验,以评估MoE - LoRA在各种下游任务中的性能,包括问答任务、通用语言理解评估基准(GLUE Benchmark)和视觉 - 语言任务。具体而言,我们的实验主要涉及两种类型的实验候选对象:(1)使用黎曼缩放优化器独立更新专家的MoE - LoRA;(2)使用黎曼缩放优化器更新,并结合我们提出的重新缩放技术(工程近似)的MoE - LoRA。我们分别在SGD和AdamW优化器上实现这两种方法。作为进一步的参考,我们还展示了与先前MoE - LoRA基线(如MoLA(Gao等人,2024))的比较以及集成的可能性。最后,为了支持我们的理论基础,我们在没有黎曼预条件器支持的经典优化器下,通过评估我们的前向传播修改进行了消融研究。
4.1 实验设置
在大多数实验中,除非另有说明,我们构建一个由总共20个专家组成的LoRAs模块混合体,每个专家的秩为4,每次激活前10个专家。此外,在消融部分还讨论了一系列其他MoE架构设置。我们以Llama - 3.2 - 3B(Touvron等人,2023)、GLM - 4 - 9B(GLM等人,2024)和LLaVAv1.5 - 7B(Liu等人,2024a)为基础模型进行实验。
在训练过程中,我们遵循线性衰减学习率调度器。与其他可训练组件相比,我们为门模块分配相对较小的学习率,以实现稳定的训练行为。门的较小学习率有助于防止模型经历突然和不稳定的路由变化。为了进一步稳定,我们还将其最大梯度范数限制在1.0。我们仔细为各种任务分配不同的初始学习率,试图确保所有模型在合理的运行时间内达到最佳性能。
我们将专家数量、前\(k\)个和每个专家的秩分别表示为\(n\)、\(k\);对于使用传统黎曼预条件优化器的实验候选对象,我们将它们表示为\(RSGD_{n, k, r}\)和\(RAdamW_{n, k, r}\),其中前面的\(R\)代表黎曼(Riemannian);而那些结合了我们基于门的重新缩放方法的候选对象分别表示为\(gRSGD_{n, k, r}\)和\(gRAdamW_{n, k, r}\),其中前面的\(g\)表示我们通过门值重新缩放梯度。
4.2 问答评估
我们在几个问答基准上评估我们提出的方法,包括科学问答(ScienceQA)(Lu等人,2022)、常识问答(CommonsenseQA)(Talmor等人,2018)、开卷问答(OpenBookQA)(Mihaylov等人,2018)和社会交互问答(SIQA)(Sap等人,2019)。这些问答数据集涵盖了不同的领域和类型,如科学、社会交互、常识和开卷考试等。
我们以Llama3.2 - 3B为基础模型实现所有实验候选对象。对于SGD优化器,我们为每个LoRA专家设置初始学习率为\(3×10^{-5}\);对于AdamW优化器,我们使用\(1×10^{-5}\)的初始学习率。我们运行所有实验,直到它们达到稳定的性能,特别确保每对比较的候选对象(即独立使用黎曼预条件的MoE - LoRA和使用我们提出的重新缩放方法的MoE - LoRA)经过相同的训练步骤,以确保它们具有公平的可比性。
具体而言,根据数据集的复杂程度,我们在所有问答评估中选择800步或1400步这两种设置,除了在常识问答上的\(RAdamW\)和\(gRAdamW\),我们将其训练到2000步,以更清楚地区分两个可比的候选对象。我们在表1中展示评估的性能。
可以观察到:(1)结合我们方法的黎曼预条件优化器在每个问答基准上都取得了更好的性能,尽管改进程度有所不同;(2)总体而言,与AdamW相比,我们的方法对黎曼预条件SGD的贡献更大:我们将\(RSGD\)的性能提高了约8.5%,而将\(RAdamW\)的性能提高了约1.5%。
表1 以Llama - 3.2 - 3B为基础模型,在四个问答数据集上的问答评估。我们基于门的重新缩放方法在SGD和AdamW方面均优于传统的黎曼预条件优化器。每对比较的候选对象经过相同的训练步骤,直到它们都达到良好的稳定性能。
除了在最终性能上的改进,我们还见证了在我们的优化下收敛速度的提升。为了清楚地展示这一点,我们在图2中绘制了在SGD优化器下四个问答数据集的损失下降曲线和指标变化。可以清楚地看到,\(gRSGD\)在训练和评估损失以及准确率指标方面都比\(RSGD\)收敛得更快。
图2 以Llama - 3.2 - 3B为基础模型的\(RSGD_{20,10,4}\)和\(gRSGD_{20,10,4}\) MoE - LoRA的收敛性能。我们绘制了前500步的训练和评估损失以及准确率指标。
4.3 GLUE基准上的性能
为了全面检验我们方法的有效性,我们在GLUE基准(Wang,2018)上进行了一系列下游评估,该基准是用于评估模型在自然语言理解方面性能的资源集合。我们首先以Llama - 3.2 - 3B为基础模型,在GLUE的所有评估中运行,并在表2中展示基准结果。
对于大多数SGD实验,我们为LoRA专家设置初始学习率为\(3×10^{-5}\),除了在WNLI任务中,我们将其初始学习率设置为\(3×10^{-6}\);对于AdamW实验,我们从\(\{3×10^{-5}, 1×10^{-5}\}\)中选择初始学习率。对于大多数数据集,我们训练2000步,除了一些AdamW实验,由于它们似乎已经收敛甚至过拟合,我们在大约1000步时提前停止。
表2展示了在Llama - 3.2 - 3B下我们的方法在各种下游应用中的有效性以及整体评估结果。在整体性能方面,我们的方法分别将\(RSGD\)和\(RAdamW\)的性能提高了11.9%和3.0%。
随后,我们将实验扩展到更大的基础模型GLM - 4 - 9B。由于9B模型在少样本学习方面更强大,对于一些数据集,如SST - 2等,我们分别为SGD和AdamW设置较低的学习率,如\(3×10^{-6}\)和\(1×10^{-6}\),以确保可以看到明显的损失下降阶段。
我们为每对竞争候选对象训练相同数量的步骤。表2也展示了使用GLM4 - 9B训练MoELoRA的不同优化策略的性能。结果仍然证明了我们的总体优势。特别是,我们将\(RSGD\)的平均性能提高了约4.3%,将\(RAdamW\)的平均性能提高了约0.7%。
表2 以Llama - 3.2 - 3B和GLM - 4 - 9B为基础模型,在九个任务上的GLUE基准评估。我们基于门的重新缩放方法在黎曼预条件SGD和AdamW方面对GLUE基准做出了总体改进。
4.4 LLaVA上的性能
除了文本基准,我们进一步在计算机视觉领域评估我们基于门的重新缩放方法。具体而言,我们为著名的视觉 - 语言基础模型LLaVAv1.5 - 7B(Chen等人,2024)实现了一种MoE - LoRA架构。我们在LLaVA - v1.5 - 7B的视觉和文本模块中都引入了可训练的MoELoRA适配器。
为了进行评估,我们使用了Visual7W(Zhu等人,2016)和VMCBench(Zhang等人,2025b)数据集,这两个数据集都由多模态样本组成,每个样本包含一个选择题和一张相关图像。问题可以通过理解提供的图像来回答。Visual7W是Visual Genome(Krishna等人,2017)数据集的一个子集,而VMCBench是由20个现有的视觉问答(VQA)数据集创建的基准。对于VMCBench,由于其测试集未标记,我们仅使用其开发集。我们从1000个标记样本中选取900个作为训练样本,其余100个用于评估。表3展示了所有实验候选对象的结果。我们的方法始终展示出明显的改进,特别是对于SGD。
表3 以LLaVA - v1.5 - 7B为基础模型,训练1000步后在Visual7W和VMCBench上的性能。(对于VMCBench,我们使用100个样本进行评估,因此准确率最多为两位小数。这就是为什么我们在这里以百分比列出所有数字,以便更直观地展示。)
4.5 与MoE - LoRA基线的比较和集成
然后,我们将我们的方法与现有的MoE - LoRA基线进行比较和集成。我们与两个基线进行比较:(1)纯LoRAs混合(Liu等人,2023),我们将其表示为MoELoRA,并使用令牌级路由;(2)MoLA(Gao等人,2024),这是一种MoE - LoRA变体,专门聚焦于为不同层分配不同数量的专家,并证明更高层需要更多的LoRA专家。需要注意的是,我们提出的基于门的重新缩放方法可以与大多数MoE - LoRA变体集成,因为它们并不冲突。以MoLA为例,我们可以通过在更高层实现更多专家,并使用黎曼预条件器和基于门的重新缩放方法进行训练,将我们的方法与MoLA集成。我们重现了MoELoRA和MoLA,实现了集成,并在表4中展示了它们的性能。我们以Llama - 3.2 - 3B为基础模型,并遵循MoLA的配置,这意味着我们将每个专家的秩设置为4,top - k设置为2,所有层的专家总数设置为140。通过这种方式,MoELoRA和我们的方法为每层分配5个专家,而MoLA分别为底层、中下层、中上层和顶层分配2、4、6和8个专家。在表4中,我们将这种特殊的分配策略表示为(2, 4, 6, 8),而平均分配为(5, 5, 5, 5),其中每个数字涵盖Llama - 3.2 - 3B下的七层。在MoLA架构的背景下,我们的方法仍然提供了性能提升。
表4 基线比较和集成。前三行提供了纯MoE - LoRA、MoLA和我们基于门重新缩放的黎曼预条件方法之间的比较。最后两行分别提供了集成了传统和基于门重新缩放预条件方法的MoLA。所有候选对象均使用SGD优化器训练多达2000步。
4.6 消融研究
理论依赖性:尽管我们提出的方法基于黎曼预条件器,但重要的是要注意,我们的工程实现并不固有地要求与黎曼预条件器共存。原因是我们的修改仅专注于改变MoE - LoRA的前向传播惯例。这就引出了一个关键问题,即我们的修改在不依赖黎曼预条件上下文的情况下,对提升MoE - LoRA性能的独立有效性。理想情况下,由于传统的未预条件优化器不能保证在低秩空间中投影全矩阵梯度,对于它们来说,通过用\(\sqrt{g_{i}}\)替换\(g_{i}\)来归一化专家梯度之和应该是微不足道的。为了证实这一点,我们通过将基于门的修改与传统的未预条件SGD优化器集成进行消融研究。图3所示的损失下降曲线表明,直接在纯SGD优化器上应用我们的方法没有帮助,这反而证明了我们的改进与黎曼预条件算法高度耦合。
各种MoE架构:为了证明我们提出的方法可以推广到各种LoRA混合设置,我们构建了不同的MoE - LoRA架构进行进一步探索,包括专家数量、每个专家的秩和top - k数量的变化。具体来说,我们在ScienceQA数据集上测试了七种结构条件,所有候选对象都使用相同的初始学习率训练800步。表5展示了结果,表明在大多数MoE结构情况下,我们的方法都能取得更好的性能。此外,从SGD性能中还可以观察到,专家数量或每个专家的秩的变化对我们方法的有效性影响有限,而较大的top - k大致表现出更高的提升。这一观察结果与我们的理论分析一致,即更多的激活专家会导致每个专家的门值更小,从而为我们的修改留出更大的发挥空间。
图3 以Llama3.2 - 3B为基础模型,在传统和黎曼预条件SGD优化以及两者都集成基于门的重新缩放方法下,ScienceQA训练损失的曲线。
表5 以Llama - 3.2 - 3B为基础模型,在各种MoE架构下,传统和基于门重新缩放的黎曼优化器在ScienceQA上的准确率和提升。
5 结论
我们引入黎曼梯度预条件器来训练低秩专家混合(MoE - LoRA)。与直接将黎曼预条件器应用于每个专家的梯度以追求局部最优不同,我们认为在前向传播中,将专家\(B_{i}A_{i}\)乘以其各自的门值\(g_{i}\)会导致专家\(i\)构建的流形进一步缩放。为了缓解这一问题,为MoE - LoRA设计的黎曼预条件器应进行修改以纳入门值。为了近似这一概念,我们提出了一种工程解决方案,将前向传播变量分解为可优化和不可优化的组件。在各种下游任务上的实验证明了我们的方法相对于传统黎曼预条件器的性能提升。消融研究进一步证明了我们的理论基础和通用性。