修改一行代碼就能實現高效微調！上海交大&騰訊開源SaRA：兼顧原始生成和下游任務
文章鏈接：https://arxiv.org/pdf/2409.06633
項目鏈接：https://sjtuplayer.github.io/projects/SaRA/
1.引言
SaRA是一種針對預訓練擴散模型的高效微調方法。通過微調預訓練擴散模型中的無效參數，賦予模型對下游任務的處理能力。SaRA能夠顯著節省計算顯存開銷與代碼復雜度，僅修改一行訓練代碼即可實現微調過程。該方法的核心創新在于：
- 參數重要性分析：SaRA首先對預訓練模型中的參數重要性進行分析，發現預訓練擴散模型中絕對值最小的10%至20%的參數在生成過程中的作用微乎其微。并且這些參數的無效性并非模型固有屬性，而是由于訓練過程中的不穩定性導致。
- 稀疏低秩訓練：基于上述發現，SaRA提出利用這些暫時無效的參數，通過優化稀疏權重矩陣來學習特定任務的知識。為了避免過擬合，SaRA采用了基于核范數的低秩稀疏訓練方案，有效約束了學習過程中的參數秩。
- 漸進式參數調整策略：SaRA設計了一種參數重調整策略，通過在微調過程中重定向可訓練參數，確保幾乎所有參數都能有效地貢獻于新任務的學習。
- 非結構化反向傳播策略：SaRA提出了一種新穎的反向傳播策略，顯著降低了微調過程中的內存成本。
SaRA在多個下游任務上進行了廣泛的實驗驗證，包括基模型能力提升、下游數據微調、圖像定制化、可控視頻生成等。實驗結果表明SaRA不僅能夠提升基礎模型在原始任務的生成能力，在下游任務中，能兼顧下游任務的學習以及預訓練先驗的維護，實現優越的模型微調效果。
2. 參數重要性分析
2.1 預訓練模型中的無效參數
圖 1：Stable Diffusion預訓練模型參數分布與小參數對生成結果的影響
2.2 無效參數的潛在有效性
2.1中導致無效參數的原因可能有兩個：一是由于模型結構設計的原因，這些參數天生就是冗余、無效的參數，因此無法在訓練過程中起到作用，另外一個原因則可能是由于模型訓練過程中的隨機性，導致這些參數恰好在訓練結束的時候停留在0附近。因此，作者進一步對參數無效的原因展開研究。選取了Stable Diffusion在FFHQ的預訓練模型，標記了初始權重最小的1%參數，將該模型繼續在FFHQ上訓練，并在訓練過程中實時跟蹤這1%參數的變化，結果如圖 2所示，可見，隨著訓練的進行，初始的小參數（藍色線條）逐漸跳出了1%的閾值，而初始大于1%閾值的參數，大部分跌入了1%以內，并且小于該閾值的參數總量始終維持在1%左右，證明了在訓練過程中，所有參數都以一定的概率跳出或者跌入1%閾值中，說明初始的小參數是由訓練過程的隨機性導致的，因此，可以在微調過程中利用這些暫時無效的參數，賦予模型針對下游任務的生成能力。
圖 2：訓練過程中權重絕對值小于初始1%閾值θ_t的參數分布變化
3. 方法介紹
3.1 稀疏矩陣訓練
3.2 基于核范數的低秩約束
3.3 漸進式參數調整
在模型的微調過程中，由于訓練的隨機性，仍然會存在部分參數停留在閾值以下，尤其是微調過程的總輪次往往較少，導致最終存在一部分的參數仍然無效。如圖 2 所示，初始的小參數在訓練初期會快速跳出閾值，而后期的趨勢逐漸放緩，當微調輪次較少時，可訓練參數中可能存在15%的參數仍然無效。因此，SaRA提出漸進式的參數調整策略，在微調的前半階段，首先對初始的無效參數進行訓練，使得大部分的無效參數跳出閾值，而在后半階段，再次對剩余的無效參數進行訓練，使其快速跳出閾值。通過這種分階段的漸進式訓練策略，SaRA可以更有效地利用這些無效參數，提高模型在新任務上的性能。
3.4 非結構化反向傳播策略
圖3：非結構化梯度反向傳播
4.實驗效果
為了驗證方法的有效性，SaRA在多個主流與下游的任務上進行了驗證，包含基模型提升、下游數據集微調、圖像定制化與可控視頻生成。
4.1 基模型提升
SaRA主要致力于將預訓練模型中的無效參數利用起來，賦予模型更強大的生成能力，這與一般微調方法僅針對下游任務設計的理念不盡相同。因此，SaRA可以用來提升預訓練模型在原本任務上的生成能力。實驗選取了在ImageNet、FFHQ、CelebA-HQ上預訓練的Stable Diffusion，利用SaRA在相應數據集對模型進行進一步的微調，以完全利用模型中的無效參數，結果如圖4所示，可以看出，SaRA能夠穩定地提升基模型的生成能力（降低約5%的FID）。
圖4：基模型在原始任務上的微調
4.2下游數據集微調
在下游數據集微調實驗中，將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以及不同的參數量下，均取得了最好的表現。
圖5：不同微調方法在下游數據集微調的表現。
表1：不同微調方法在下游數據集微調的定量表現。
4.3 圖像定制任務
圖像定制化通過從少量幾張圖像中學習到共有的對象特征，然后將該對象生成到新的圖片中。Dreambooth作為一種主流的圖像定制化，需要微調擴散模型實現對目標特征的捕捉，因此，SaRA可以直接用于Dreambooth的微調過程。實驗比較了不同微調方法在DreamBooth上的表現，定性結果如圖6所示，可見，SaRA在成功捕捉目標對象特征的同時，還較好地維護了生成圖像語義與文本的一致性。表2計算了不同方法在三個定制化數據集上的定量表現，可以看出，SaRA同時兼顧了特征捕捉與圖文一致性，展現了在定制化任務重的優秀表現。
圖6：不同微調方法在Dreambooth上的定性表現
表2：不同微調方法在Dreambooth上的定量表現
4.4 視頻生成任務
SaRA不僅在圖像生成任務中大展身手，在視頻生成任務重也同樣能取得較好的結果。實驗將不同微調方法應用在視頻生成模型AnimateDiff上，在不同運鏡數據集下進行微調（鏡頭放大、縮小、右移）。結果如圖7所示，其他的微調方法在視頻生成任務中展現出一定的過擬合與內容崩潰的問題，相較之下，SaRA在微調過程中展現出豐富的運鏡理解能力，同時較好地維護了模型的先驗，保證了生成視頻的質量以及與文本的一致性。
圖7：不同微調方法在可控運鏡的視頻生成上的表現
4.5 計算資源比較
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%的訓練時間。
圖8：基于參數選擇的方法、LoRA、SaRA在不同batch size下的顯存開銷與訓練時間
4.6 訓練參數分析
一個好的微調方法，在微調過程中，應該能夠學習到更多的任務相關的信息，同時最大化保留預訓練權重的先驗知識。因此，作者實驗分析了SaRA與LoRA在Expedition數據集上微調后的學習到的參數ΔP與預訓練權重P之間的關系。表3通過F范數量化了ΔP前r維子空間與預訓練權重P子空間的相關性，可見SaRA學習到的參數與P相關性更小，說明相較于LoRA學習到了更多的下游任務的知識。此外，還計算了放大因子Amplification Factor，量化了ΔP對P中未強調的特征方向的放大倍數，同樣證明了SaRA對新知識更強的學習能力。
表3：SaRA與LoRA訓練參數ΔP，與預訓練參數P的關系
圖9計算了SaRA與LoRA訓練后的模型參數ΔP+P與預訓練參數P的前r維子空間的相似性，可以看出，SaRA的相似性在95%以上，而LoRA維持在80%附近，證明了融合SaRA訓練參數的模型，能夠更好地維護預訓練權重的先驗知識。
圖9：SaRA與LoRA訓練后的模型參數ΔP+P與預訓練參數P的關系
結論
本文提出了 SaRA，一種新穎的參數高效微調方法，該方法充分利用了預訓練模型中絕對值最小的無效參數。作者提出了一種基于核范數的低秩損失，以約束學習到的稀疏矩陣的秩，從而避免模型過擬合。此外，設計了一種漸進式參數調整策略，進一步提高了微調參數的有效性。最后，提出了一種新型的非結構化反向傳播方法，大大節省了參數微調過程中的內存開銷，同時也能降低其他選擇性 PEFT 方法的內存成本。大量實驗證明了本文方法的有效性，它在保持預訓練模型的先驗信息的同時，實現了最佳的擬合能力。此外，作者高效地封裝了本文的方法，使其只需修改一行代碼即可實現，這大大增強了代碼在其他模型和任務中的易用性和適應性。
本文轉自 AI生成未來 ，作者：AI生成未來
