仅需一行代码！上交大联合腾讯优图提出 SaRA: 让‘懒惰‘参数重获新生的微调魔法

论文地址: https://arxiv.org/pdf/2409.06633
项目主页: https://sjtuplayer.github.io/projects/SaRA/
开源代码: https://github.com/sjtuplayer/SaRA

引言

想象一下，如果毕加索突然决定改行做一名厨师。他可能会把调味料撒得满桌都是。然后宣称，这是一道"立体主义风味的番茄酱意面"。这就是我们今天要讨论的主角——扩散模型的工作方式。它们先把图像搞得一团糟,然后再慢慢把它们变成艺术品。是的~这就是数字世界版的"破坏再重建"疗法!

所以，哪怕是诸如毕加索这样的艺术大师，也是需要学习如何正确使用厨具一样，扩散模型其实也需要经过精心调教才能在特定任务上大显身手。这个过程就叫做"微调"。初次听起来是不是就像是给模型做 SPA 一样。而传统的微调方法就像是给模型全身按摩，虽然效果不错，但往往会让你的钱包大喊"OMG！住手"!

说到这，让我们来聊聊这些微调方法的主要缺点：

• 首先，它们消耗的计算资源比一个热爱网购的青少年还要多。
• 其次，实现这些方法的复杂度堪比组装宜家家具——没有说明书的那种。
• 最后,这些方法往往需要很长时间,就像等待披萨外卖,你可能会怀疑送餐员是不是在送餐的途中睡着了。

不过，幸运的是，就在我们即将放弃希望,准备接受"要么富有,要么等待"的残酷现实时,一位名叫 SaRA 的英雄悄然登场。它不仅解决了上述所有问题,还做到了"一行代码改变世界"。这简直就像发现了一种只需要动一动手指就能完成全身运动的健身方法!

那么，这个神奇的 SaRA 到底是什么？它又是如何做到既省钱又高效的呢？👇，让我们带着这些疑问，跟随上海交通大学DMCV团队及腾讯优图的研究员们一同深入探索这块奇妙世界吧!

SaRA

SaRA 本质上是一种针对预训练扩散模型的高效微调方法。通过微调预训练扩散模型中的无效参数，赋予模型对下游任务的处理能力。SaRA 能够显著节省计算显存开销与代码复杂度，仅修改一行训练代码即可实现微调过程。该方法的核心创新在于：

• 参数重要性分析：SaRA首先对预训练模型中的参数重要性进行分析，发现预训练扩散模型中绝对值最小的10%至20%的参数在生成过程中的作用微乎其微。并且这些参数的无效性并非模型固有属性，而是由于训练过程中的不稳定性导致。
• 稀疏低秩训练：基于上述发现，SaRA 提出利用这些暂时无效的参数，通过优化稀疏权重矩阵来学习特定任务的知识。为了避免过拟合，SaRA 进一步采用了基于核范数的低秩稀疏训练方案，有效约束了学习过程中的参数秩。
• 渐进式参数调整策略：SaRA 设计了一种参数重调整策略，通过在微调过程中重定向可训练参数，确保几乎所有参数都能有效地贡献于新任务的学习。
• 非结构化反向传播策略：SaRA 提出了一种新颖的反向传播策略，显著降低了微调过程中的内存成本，相较于LoRA可节省约45%的训练显存占用。

最后，通过在多个下游任务上进行的广泛实验验证，SaRA 实现了包括基模型能力提升、下游数据微调、图像定制化、可控视频生成等。实验结果表明 SaRA 不仅能够提升基础模型在原始任务的生成能力，在下游任务中，能兼顾下游任务的学习以及预训练先验的维护，实现优越的模型微调效果！

参数重要性分析

预训练模型中的无效参数

对于"参数"这个名词，想必各位 AIer 都耳熟能详。在后深度学习时代，模型的参数量是越做越大。

但是，根据模型剪枝理论，其实并非所有参数都对模型的输出有积极的影响。

因此，作者首先研究了多个版本的预训练 Stable Diffusion（包括1.4,1.5,2.0,与3.0）中，绝对值权重较小的参数对生成结果的影响。通过将绝对值权重小于的参数置为0,后，让模型根据GPT-4o生成的1000个文本，生成对应的1000张图像，计算生成图像的FID指标以及CLIP Score：

如图1所示，我们可以发现，其实将模型10%~20%的参数置为0时（无效参数），模型的生成结果并没有受到影响，甚至有些情况下还能略微提升，证明了这些小参数在预训练Stable Diffusion模型中的无效性。

无效参数的潜在有效性

导致无效参数的原因可能有两个：

1. 由于模型结构设计的原因，这些参数天生就是冗余、无效的参数，因此无法在训练过程中起到作用；
2. 由于模型训练过程中的随机性，导致这些参数恰好在训练结束的时候停留在0附近。

因此，作者进一步对参数无效的原因展开研究。选取了Stable Diffusion在FFHQ的预训练模型，标记了初始权重最小的1%参数，将该模型继续在FFHQ上训练，并在训练过程中实时跟踪这1%参数的变化：

如图 2所示，可见，随着训练的进行，大部分初始的小参数（蓝色线条）逐渐跳出了1%的阈值，而初始大于1%阈值的参数，大部分跌入了1%以内，并且小于该阈值的参数总量始终维持在1%左右，证明了在训练过程中，所有参数都以一定的概率跳出或者跌入1%阈值中，说明初始的小参数是由训练过程的随机性导致的。

因此，这些暂时无效的参数具有潜在的学习能力，可以在微调过程中利用这些暂时无效的参数，赋予模型针对下游任务的生成能力。

方法

为了充分利用这些"暂时无效"的参数，本文提出了一种渐进式稀疏低秩训练方法。

该方法的核心思想是：通过对这些无效参数进行微调，使其在下游任务中发挥作用。

具体来说，SaRA 首先确定一个阈值，将小于该阈值的参数视为暂时无效的参数。然后，通过优化一个稀疏权重矩阵，使得这些参数能够学习到新任务的知识。

为了避免在训练过程中出现过拟合，SaRA引入了基于核范数的低秩约束。此外，SaRA还提出渐进式的训练策略，最大化地利用所有当前无效的参数。最后，研究人员还设计了一种非结构化梯度回传策略，有效降低了模型微调过程中的显存占用。

稀疏矩阵训练

SaRA致力于微调预训练模型中暂时无效的参数(即权重绝对值小于一定阈值${\theta }_t$的参数)，使预训练的扩散模型适应下游任务，同时尽可能地保留原始的模型先验。具体而言，首先为初始参数P计算一个稀疏掩码M，满足：

$$|p|<{\theta }_t,\forall p\in P_0\odot M,$$

SaRA基于该稀疏掩码来更新初始无效的参数$P\odot M$, 同时保持初始有效参数$P\odot (1-M)$不变。在训练期间，对于所有参数的梯度$\nabla P$, 利用该稀疏掩码M来保留所需要更新参数的梯度，并更新相应的可训练参数$P_M=P\odot M$:

$$\begin{array}{rl}& \nabla P_M=\nabla P\odot M+0\odot (1-M),\\ & P_{new}=P-\lambda \cdot \nabla P_M.\end{array}$$

基于核范数的低秩约束

稀疏参数矩阵$P_M$可能会具有较高的秩，从而导致过于强大的表征能力，使得模型在下游任务训练过程中出现过拟合的问题。为了缓解这个问题，SaRA在稀疏矩阵上引入了基于核范数的低秩约束，以约束稀疏参数矩阵的秩。低秩约束的一种直接方法是最小化稀疏参数矩阵的秩$Rank(P_M)$。然而，由于矩阵秩求解的非凸性，$Rank(P_M)$难以实现可微的计算过程。因此，SaRA使用矩阵的核范数来估计稀疏参数矩阵的秩：
$${\Vert P_M\Vert }_{\ast }=\sum _i{\sigma }_i\left(P_M\right),$$
其中${\sigma }_i(P_M)$是参数矩阵$P_M$的第$i$个奇异值，通过最小化该式子，可以实现对参数矩阵的低秩约束。

为了计算核范数$\Vert P_M{\Vert }_{\ast }$,对参数矩阵进行奇异值分解$P_M=U\Sigma V^T$,其中$U$和$V$是正交矩阵,$\Sigma$是包含奇异值${\sigma }_i(P_M)$的对角矩阵.$\Vert P_M{\Vert }_{\ast }$关于$P_M$的次梯度,可以由下式求出:
∂ ∥ P M ∥ ∗ = { U V T + W ∣ W ∈ R m × n , U T W = 0 , W V = 0 , ∥ W ∥ 2 ≤ 1 } . \partial \| P_M\|_*=\{UV^T+W|W \in R^{m\times n}, U^TW=0,WV=0,\| W\|_2\leq 1\}. ∂∥PM​∥∗​={UVT+W∣W∈Rm×n,UTW=0,WV=0,∥W∥2​≤1}.
基于核范数梯度的推导,可以保证低秩约束的可微性,从而实现基于核范数的低秩约束损失:
$$L_{rank}=\Vert P_M{\Vert }_{\ast }=\sum _i{\sigma }_i(P_M).$$

渐进式参数调整

在模型的微调过程中，由于训练的随机性，仍然会存在部分参数停留在阈值以下，尤其是微调过程的总轮次往往较少，导致最终存在一部分的参数仍然无效。

如上图 2 所示，初始的小参数在训练初期会快速跳出阈值，而后期的趋势逐渐放缓，当微调轮次较少时，可训练参数中可能存在15%的参数仍然无效。

因此，SaRA提出渐进式的参数调整策略，在微调的前半阶段，首先对初始的无效参数进行训练，使得大部分的无效参数跳出阈值，而在后半阶段，再次对剩余的无效参数进行训练，使其快速跳出阈值。

通过这种分阶段的渐进式训练策略，SaRA可以更有效地利用这些无效参数，提高模型在新任务上的性能，最终在1%的threshold下保证99.98%的参数得到了合理的利用。

非结构化反向传播策略

目前, 基于LoRA的方法和参数选择的高效微调方法都对计算资源造成了沉重的负担:

1. 对于基于LoRA的方法, 由于LoRA模块的引入, 它不需要存储模型参数的梯度, 但却需要额外的内存成本来存储LoRA模块中的中间变量, 如图3(a)所示。
2. 对于参数选择的方法, 一个一直困扰的问题是: 它们需要与全参数微调相同甚至更多的计算资源（尤其是GPU显存）。虽然它们只对模型参数的子集进行微调, 但它们保留了整个参数矩阵P的梯度, 因为主流的深度学习库(如PyTorch和TensorFlow)只支持对整个参数矩阵的梯度反向传播。

因此, 以往的基于参数选择的方法必须在整个参数矩阵P上执行梯度反向传播, 然后使用预先计算的掩码矩阵M, 通过∇PM=∇P⊙M屏蔽不需要的参数梯度, 并通过Pnew=P−λ·∇PM实现整体参数的更新（如图3(b)所示）。

这种方法需要存储所有模型参数的梯度和额外的掩码矩阵, 导致比全参数微调更大的计算资源需求。因此, 为了解决这些问题, SaRA提出了非结构化梯度回传策略, 通过将可训练参数从参数矩阵中剥离, 结合非结构化前向和反向传播函数, 实现训练过程中显存的大幅降低。

具体地，SaRA 首先将训练模型的所有参数变为非叶节点，并通过系数矩阵 M，获取可学习参数 $P_{learn}=P[M]\in R^{|\mathbf{M}|}$ ，将可学习参数作为真实的叶节点。定义非结构化映射函数 UM(·)，在前向过程中将可学习参数映射到模型参数中：

$$P=UM(P,P_{learn},M),where\{\begin{array}{ll}P[\mathbf{M}]& =P_{learn},\\ P[1-\mathbf{M}]& =P[1-\mathbf{M}].\end{array}$$

在反向过程中，定义非结构化反向传播函数 UB(·)，将模型参数的梯度自动回传至可训练参数：
$\nabla P_{learn}=UB(\nabla P,M)=\nabla P[\mathbf{M}].$
由于模型参数成为了非叶子节点，$P_{learn}$ 为训练过程中的实际叶节点，因此，梯度流经模型参数到达 $P_{learn}$ 后，模型参数的梯度会自动消除，最终存储的梯度仅有 $P_{learn}$ 的梯度，从而节省大量的计算显存。

实验效果

为了验证方法的有效性，SaRA在多个主流与下游的任务上进行了验证，包含基模型提升、下游数据集微调、图像定制化与可控视频生成。

基模提升

SaRA主要致力于将预训练模型中的无效参数利用起来，赋予模型更强大的生成能力，这与一般微调方法仅针对下游任务设计的理念不尽相同。

因此，SaRA可以用来提升预训练模型在原本任务上的生成能力。实验选取了在ImageNet、FFHQ、CelebA-HQ上预训练的Stable Diffusion，利用SaRA在相应数据集对模型进行进一步的微调，以完全利用模型中的无效参数，结果如图4所示，可以看出，SaRA能够稳定地提升基模型的生成能力（降低约5%的FID）。

下游数据集微调

在下游数据集微调实验中，将SaRA应用于多个不同的数据集，并在不同StableDiffusion版本（1.5,2.0,3.0）与参数规模（50M，20M，5M）下进行了训练。

数据集包括BarbieCore, Cyberpunk, Elementfire, Expedition, Hornify五个风格。结果如图5所示，可见SaRA学习到了最丰富的目标数据特征，同时能够保持语义与文本的一致性。

此外，实验还计算了生成数据的FID与文本-图像的CLIP Score，以及一个归一化指标VLHI同时衡量FID与CLIP Score，定量结果如表1所示，可见，SaRA在不同版本的Stable Diffusion以及不同的参数量下，均取得了最好的表现。

图像定制任务

图像定制化通过从少量几张图像中学习到共有的对象特征，然后将该对象生成到新的图片中。DreamBooth作为一种主流的图像定制化方法，需要微调扩散模型实现对目标对象特征的捕捉，因此，SaRA可以直接用于DreamBooth的微调过程。

实验比较了不同微调方法在DreamBooth上的表现，定性结果如图6所示，可见，SaRA在成功捕捉目标对象特征的同时，还较好地维护了生成图像语义与文本的一致性。

表2计算了不同方法在三种定制化数据上的定量表现，可以看出，SaRA同时兼顾了特征捕捉与图文一致性，展现了在定制化任务中的优秀表现。

视频生成任务

SaRA不仅在图像生成任务中大展身手，在视频生成任务中也同样能取得较好的结果。实验将不同微调方法应用在视频生成模型AnimateDiff上，在不同运镜数据集下进行微调（镜头放大、缩小、右移）。

结果如图7所示，其他的微调方法在视频生成任务中展现出一定的过拟合与内容崩溃的问题，相较之下，SaRA在微调过程中展现出丰富的运镜理解能力，同时较好地维护了模型的先验，保证了生成视频的质量以及与文本的一致性。

计算资源比较

SaRA引入了非结构化梯度回传策略，有效解决了基于参数选择的微调方法中遇到的显存开销大的问题。

图8的实验比较了LT-SFT（一种基于参数选择的方法）、LoRA以及SaRA在Stable Diffusion 2.0用不同Batch size训练过程中的显存开销与训练时间。

可以看出，SaRA比LT-SFT减少了固定的9.2GB显存占用（对应所有参数的梯度占用空间），在Batch Size较小时(<=4)节省了45%的显存。

而LoRA随着Batchsize的增大，显存占用急速上升。SaRA在Batch Size=16时比LoRA节省了52%的显存占用，并且节省了49%的训练时间。

训练参数分析

一个好的微调方法,在微调过程中,应该能够学习到更多的任务相关的信息,同时最大化保留预训练权重的先验知识。因此，作实验分析了 SaRA 与 LoRA 在 Expedition 数据集上微调后的学习到的参数 ΔP 与预训练权重 P 之间的关系。

表 3 通过 F 范数量化了 ΔP 前 r 维子空间与预训练权重 P 子空间的相关性，可见 SaRA 学习到的参数与 P 相关性更小，说明相较于 LoRA 学习到了更多的下游任务的知识。

此外，还计算了放大因子 Amplification Factor，量化了 ΔP 对 P 中未强调的特征方向的放大倍数，同样证明了 SaRA 对新知识更强的学习能力。

图9计算了SaRA与LoRA训练后的模型参数与预训练参数的前r维子空间的相似性，可以看出，SaRA的相似性在95%以上，而LoRA维持在80%附近，证明了融合SaRA训练参数的模型，能够更好地维护预训练权重的先验知识。

总结

本文提出了 SaRA，一种新颖的参数高效微调方法，该方法充分利用了预训练模型中绝对值最小的无效参数。作者提出了一种基于核范数的低秩损失，以约束学习到的稀疏矩阵的秩，从而避免模型过拟合。此外，设计了一种渐进式参数调整策略，进一步提高了微调参数的有效性。最后，提出了一种新型的非结构化反向传播方法，大大节省了参数微调过程中的内存开销。

大量实验证明了本文方法的有效性，它在保持预训练模型的先验信息的同时，实现了最佳的拟合能力。此外，作者高效地封装了本文的方法，使其只需修改一行代码即可实现，这大大增强了代码在其他模型和任务中的易用性和适应性。

写在最后

如果你对本文有任何的观点或疑问，欢迎评论区下方留言讨论。同时也欢迎对前沿AI技术感兴趣的同学扫添加微信好友: cv_huber，备注“交流学习”即可。