CVPR | LoRA遇上RoPE！WaDi：面向單步圖像生成權(quán)重方向感知蒸餾

本文作者王雷，南開大學(xué)，PCA Lab成員

盡管擴(kuò)散模型（如 Stable Diffusion，SD）在圖像生成領(lǐng)域表現(xiàn)出色，但其緩慢的推理速度限制了實(shí)際部署。近期工作通過將多步擴(kuò)散蒸餾為單步生成器來加速推理。

為了更好地理解蒸餾機(jī)制，南開大學(xué) PCA Lab 團(tuán)隊(duì)分析了單步學(xué)生模型與多步教師模型之間 U-Net/DiT 權(quán)重的變化規(guī)律。分析表明，權(quán)重方向上的變化顯著超過權(quán)重范數(shù)上的變化，這揭示了方向是蒸餾過程中的關(guān)鍵因素。受此啟發(fā)，團(tuán)隊(duì)提出了權(quán)重方向低秩旋轉(zhuǎn)（Low-rank Rotation of weight Direction，LoRaD）—— 一種專為單步擴(kuò)散蒸餾設(shè)計(jì)的參數(shù)高效適配器。

LoRaD 通過可學(xué)習(xí)的低秩旋轉(zhuǎn)矩陣對(duì)預(yù)訓(xùn)練權(quán)重的方向進(jìn)行建模。團(tuán)隊(duì)進(jìn)一步將 LoRaD 集成到變分得分蒸餾（Variational Score Distillation，VSD）中，提出了權(quán)重方向感知蒸餾（Weight Direction-aware Distillation，WaDi）—— 一種新穎的單步蒸餾框架。

WaDi 在 COCO 2014 和 COCO 2017 上取得了最先進(jìn)的 FID 分?jǐn)?shù)，而可訓(xùn)練參數(shù)僅占 U-Net/DiT 全量參數(shù)的約 10%。此外，蒸餾后的單步模型展現(xiàn)出強(qiáng)大的通用性和可擴(kuò)展性，能夠良好地泛化到可控生成、關(guān)系反演、高分辨率合成等多種下游任務(wù)。

• 論文鏈接：https://arxiv.org/abs/2603.08258
• 代碼鏈接：https://github.com/gudaochangsheng/WaDi
• Project: https://gudaochangsheng.github.io/WaDi-Page/
• Demo: https://huggingface.co/spaces/gudaochangsheng/WaDi-1.5
• 講解視頻: https://www.youtube.com/watch?v=j6CuQxynJcA

圖 1. 使用我們提出的方法 WaDi（即 SD 2.1）一步生成的圖像。

引言

擴(kuò)散模型（DMs）在圖像生成領(lǐng)域受到了廣泛關(guān)注，在文本到圖像生成、文本到視頻生成以及圖像到視頻生成等任務(wù)中均有廣泛應(yīng)用。然而，擴(kuò)散模型依賴多步采樣，導(dǎo)致計(jì)算成本高、推理速度慢。

為此，近期蒸餾方法將采樣步數(shù)壓縮至數(shù)步甚至一步。有趣的是，在蒸餾過程中，團(tuán)隊(duì)發(fā)現(xiàn)權(quán)重范數(shù)在各層間保持相對(duì)穩(wěn)定，而在將權(quán)重重參數(shù)化為范數(shù)與方向時(shí)，方向則呈現(xiàn)出更大的變化幅度。

受權(quán)重重參數(shù)化的啟發(fā)，團(tuán)隊(duì)采用類似的分解方式來分析擴(kuò)散蒸餾中的權(quán)重變化。為此，團(tuán)隊(duì)研究了最先進(jìn)（SOTA）單步模型（如 DMD2 和 Pixart-α DMD）與其對(duì)應(yīng)多步模型（如 SD 1.5 和 Pixart-α）之間的權(quán)重更新。

如圖 2 (a) 所示，在基于 U-Net 的架構(gòu)中，各層權(quán)重范數(shù)幾乎保持穩(wěn)定，均值和標(biāo)準(zhǔn)差（STD）分別約為 0.1% 和 0.2%。相比之下，權(quán)重方向的變化則明顯更大，均值為 2.2%，標(biāo)準(zhǔn)差為 2.1%，對(duì)應(yīng)為范數(shù)變化的 22 倍和 10 倍。在基于 DiT 的架構(gòu)中也觀察到類似規(guī)律（見圖 2 (a) 右）。

這些觀察表明，權(quán)重方向可能攜帶了蒸餾中更豐富、更敏感的信息。

此外，方向上的變化是否具有結(jié)構(gòu)化規(guī)律？為此，團(tuán)隊(duì)對(duì)殘差矩陣（單步與多步方向矩陣之差）進(jìn)行奇異值分解（SVD），發(fā)現(xiàn)僅保留 30% 的秩即可恢復(fù) 93% 的信息，突顯了其低秩本質(zhì)（見圖 2 (b)）。

圖 2. 我們方法的動(dòng)機(jī)分析。(a) 一步學(xué)生模型與教師模型之間的權(quán)重范數(shù)和方向差異。更多細(xì)節(jié)和補(bǔ)充示例見補(bǔ)充材料 E。(b) DMD2 殘差矩陣的 SVD 分析。(c) 將一步模型的范數(shù)替換為多步模型的范數(shù)影響很?、佗埽惶鎿Q方向會(huì)嚴(yán)重降低生成質(zhì)量②⑤）。(d) 與 (c) 對(duì)應(yīng)的定性示例。(e) LoRaD 示意圖。

為了量化這兩個(gè)分量的影響，團(tuán)隊(duì)通過有選擇地將單步模型的范數(shù)或方向替換為多步模型的對(duì)應(yīng)值，進(jìn)行受控消融實(shí)驗(yàn)（見圖 2 (d)）。如圖 2 (c) 所示，替換范數(shù)對(duì)性能影響微乎其微（如 DMD2：FID +0.7，CLIP 不變），而替換方向則導(dǎo)致嚴(yán)重退化（如 DMD2：FID +241.3，CLIP -0.18）。

這些發(fā)現(xiàn)表明，方向重建是蒸餾中性能提升的核心因素，而范數(shù)變化的影響相對(duì)次要。一種可能的解釋是：用教師權(quán)重初始化學(xué)生模型對(duì)齊了初始范數(shù)，訓(xùn)練過程中的權(quán)重衰減進(jìn)一步約束了范數(shù)漂移；而蒸餾信號(hào)則主要通過調(diào)整權(quán)重方向來減少表征差異。

綜合來看，這些結(jié)果表明方向重建是蒸餾性能提升的核心驅(qū)動(dòng)因素。

上述蒸餾方法大致可分為兩類：全量微調(diào)（FT）基于低秩適配（LoRA）的微調(diào)。然而，二者在優(yōu)化范數(shù)和方向時(shí)均直接更新模型參數(shù)，導(dǎo)致范數(shù)和方向的變化相互耦合，增加了優(yōu)化難度。此外，F(xiàn)T 和 LoRA 均面臨收斂慢、不穩(wěn)定和過擬合等問題，進(jìn)一步增加了優(yōu)化的復(fù)雜性。

為此，團(tuán)隊(duì)提出了權(quán)重方向低秩旋轉(zhuǎn)（LoRaD）（見圖 2 (e)），通過可學(xué)習(xí)的旋轉(zhuǎn)矩陣調(diào)整預(yù)訓(xùn)練權(quán)重的方向。鑒于方向變化的結(jié)構(gòu)化特性（即低秩特性），旋轉(zhuǎn)角度被參數(shù)化為兩個(gè)低秩矩陣的乘積，以進(jìn)一步減少可學(xué)習(xí)參數(shù)數(shù)量。團(tuán)隊(duì)將 LoRaD 集成到變分得分蒸餾（VSD）中，提出了權(quán)重方向感知蒸餾（WaDi），一種新穎的單步文本到圖像蒸餾框架。

在 COCO 2014 和 COCO 2017 數(shù)據(jù)集上的實(shí)驗(yàn)表明，WaDi 取得了最先進(jìn)的 FID 分?jǐn)?shù)，超越所有現(xiàn)有單步生成方法。這一成果僅通過優(yōu)化方向?qū)崿F(xiàn)，將蒸餾難度降低，同時(shí) U-Net 可訓(xùn)練參數(shù)僅約占 10%，極大提升了參數(shù)效率。

此外，團(tuán)隊(duì)將 WaDi 應(yīng)用于可控生成、關(guān)系反演、高分辨率合成和圖像定制化等下游任務(wù)，展示了其加速能力和廣泛適用性。本文貢獻(xiàn)總結(jié)如下：

• 團(tuán)隊(duì)對(duì)多步與單步生成模型之間 U-Net 權(quán)重變化進(jìn)行了深入分析，將權(quán)重方向調(diào)整確定為單步蒸餾的關(guān)鍵驅(qū)動(dòng)因素，為高效蒸餾提供了新的理論視角；
• 團(tuán)隊(duì)提出了一種新穎的單步文本到圖像蒸餾框架 WaDi，采用 LoRaD 通過低秩旋轉(zhuǎn)建模權(quán)重方向，有效引導(dǎo)學(xué)生模型對(duì)齊教師分布；
• WaDi 在 COCO 數(shù)據(jù)集和多個(gè)下游任務(wù)上進(jìn)行了評(píng)估，定性和定量結(jié)果均表明 WaDi 在顯著提升推理效率的同時(shí)取得了實(shí)質(zhì)性的圖像質(zhì)量提升。

相關(guān)工作

擴(kuò)散模型

擴(kuò)散模型在圖像生成領(lǐng)域表現(xiàn)卓越，但像素空間計(jì)算開銷大。為提升效率，Rombach 等人提出了潛在擴(kuò)散模型（LDM），將去噪過程遷移至潛在空間。然而，現(xiàn)有基于文本引導(dǎo)的方法由于多步生成仍然較慢。盡管大多數(shù)方法采用 U-Net 骨干，擴(kuò)散 Transformer（DiT）以 Transformer 替代 U-Net 以獲得更好的可擴(kuò)展性，推動(dòng)了文本到圖像生成的進(jìn)步。盡管有所改進(jìn)，迭代去噪仍是一個(gè)緩慢的過程。近期，許多加速方法相繼出現(xiàn)。

擴(kuò)散模型加速

現(xiàn)有加速方法可分為無訓(xùn)練和基于訓(xùn)練兩類。無訓(xùn)練加速方法主要通過緩存減少冗余計(jì)算，或采用高階求解器減少采樣步數(shù)。然而這兩類方法的加速效果有限，因此基于訓(xùn)練的方法受到了更多關(guān)注。

基于訓(xùn)練的加速方法大致可分為四類：一致性蒸餾（CD）、漸進(jìn)蒸餾（PD）、擴(kuò)散 - GAN 蒸餾和變分得分蒸餾（VSD）。CD 在軌跡層面學(xué)習(xí)一致性以加速采樣，但圖像保真度往往較低。PD 分階段降低步數(shù)，引入顯著訓(xùn)練開銷。擴(kuò)散 - GAN 蒸餾（如 Diffusion2GAN）通過將多步擴(kuò)散蒸餾為 GAN 來提升保真度。VSD 采用雙教師策略實(shí)現(xiàn)分布對(duì)齊，SwiftBrush 實(shí)現(xiàn)了單步無圖像生成，SwiftBrushv2 利用模型集成進(jìn)一步改善，DMD 引入回歸損失進(jìn)一步提升性能，DMD2 將 VSD 擴(kuò)展至少步生成，并支撐了近期文本到視頻加速框架。

然而，現(xiàn)有基于訓(xùn)練的方法通常使用 FT 或 LoRA，這可能增加優(yōu)化難度。團(tuán)隊(duì)發(fā)現(xiàn)方向變化在蒸餾中通常更具影響力，因此提出 WaDi，利用 LoRaD 專注于建模方向旋轉(zhuǎn)。

方法

團(tuán)隊(duì)首先簡要回顧變分得分蒸餾（VSD），它是本工作的基礎(chǔ)。受權(quán)重方向變化在蒸餾中發(fā)揮關(guān)鍵作用這一觀察的啟發(fā)，緊接著引入權(quán)重方向低秩旋轉(zhuǎn)（LoRaD）模塊（更多理論說明見補(bǔ)充材料 D）。最后，團(tuán)隊(duì)將 LoRaD 集成到 VSD 中，形成我們提出的蒸餾框架 —— 權(quán)重方向感知蒸餾（WaDi）。

2.權(quán)重方向低秩旋轉(zhuǎn)

圖 3. （左）權(quán)重方向低秩旋轉(zhuǎn)（LoRaD）模塊的詳細(xì)結(jié)構(gòu)。LoRaD 使用可學(xué)習(xí)的低秩旋轉(zhuǎn)角來旋轉(zhuǎn)預(yù)訓(xùn)練權(quán)重方向。（右）權(quán)重方向感知蒸餾（WaDi）框架概覽。

注：旋轉(zhuǎn)不影響范數(shù)，因此無需顯式分離范數(shù)矩陣。

實(shí)驗(yàn)

1.實(shí)驗(yàn)設(shè)置

評(píng)估數(shù)據(jù)集與指標(biāo)。團(tuán)隊(duì)在 COCO 2014 和 COCO 2017 數(shù)據(jù)集上系統(tǒng)評(píng)估 WaDi 的零樣本文本到圖像生成能力，分別隨機(jī)采樣 30k 和 5k 張圖像。為全面評(píng)估生成質(zhì)量，團(tuán)隊(duì)使用 Fréchet Inception Distance（FID）衡量圖像保真度，CLIP 分?jǐn)?shù)評(píng)估文本 - 圖像語義對(duì)齊。FID 使用 Inception V3 作為特征提取器，CLIP 分?jǐn)?shù)基于 ViT-G/14 模型。團(tuán)隊(duì)還采用精確率和召回率評(píng)估保真度與多樣性，并在 Human Preference Score v2（HPSv2）基準(zhǔn)上評(píng)估文本 - 圖像對(duì)齊質(zhì)量。詳見補(bǔ)充材料 G.1。

實(shí)現(xiàn)細(xì)節(jié)。遵循先前方法，WaDi 中的學(xué)生模型采用與教師相同的架構(gòu)，并用教師權(quán)重初始化。WaDi 在從 JourneyDB 數(shù)據(jù)集中采樣的 140 萬條提示詞上進(jìn)行訓(xùn)練。訓(xùn)練過程中，學(xué)生模型的學(xué)習(xí)率（LR）設(shè)為 1e-4，虛假模型的學(xué)習(xí)率設(shè)為 1e-2。團(tuán)隊(duì)使用 AdamW 作為優(yōu)化器，批大小為 128（每 GPU 16 張）。無分類器引導(dǎo)（CFG）系數(shù)設(shè)為 1.5，訓(xùn)練進(jìn)行 2 個(gè) epoch。團(tuán)隊(duì)在三種不同骨干上蒸餾學(xué)生模型，分別為 SD 1.5、SD 2.1 和 PixArt-α（256×256）。對(duì)于 SD 1.5 和 SD 2.1，學(xué)生模型的 LoRaD 秩設(shè)為 256，而對(duì)于 PixArt-α 設(shè)為 128。所有虛假模型的 LoRaD 秩統(tǒng)一設(shè)為 32。詳見補(bǔ)充材料 F.1。

2.與最先進(jìn)方法的比較

定量結(jié)果。團(tuán)隊(duì)在 COCO 2014 數(shù)據(jù)集上與三種骨干（SD 1.5、SD 2.1、PixArt-α）的 SOTA 零樣本單步生成方法全面評(píng)估 WaDi。為保證公平比較并考慮計(jì)算約束，團(tuán)隊(duì)遵循 TiUE 的設(shè)置，統(tǒng)一使用 140 萬條提示詞復(fù)現(xiàn) WaDi、DMD2、SiD-LSG 和 SwiftBrushv2。如表 1 所示，WaDi 在所有骨干上均取得最佳 FID 和召回率分?jǐn)?shù)，表現(xiàn)出優(yōu)越的保真度和多樣性。在 CLIP 和精確率上也排名第一或第二，顯示出強(qiáng)大的文本 - 圖像對(duì)齊和感知質(zhì)量。

值得注意的是，SD 1.5、SD 2.1 和 PixArt-α 可訓(xùn)練參數(shù)分別僅占 9.74%、10.92% 和 13.30%，突顯了 WaDi 的參數(shù)高效性。這些改進(jìn)源于團(tuán)隊(duì)提出的 LoRaD，通過低秩旋轉(zhuǎn)重參數(shù)化權(quán)重更新，實(shí)現(xiàn)穩(wěn)定高效的蒸餾。

表 1. WaDi 與其他方法在零樣本 COCO 2014 結(jié)果上的定量比較。? 表示我們復(fù)現(xiàn)的結(jié)果，? 表示使用官方預(yù)訓(xùn)練模型得到的結(jié)果?！?」表示未知。最佳和次佳分?jǐn)?shù)分別用粗體和下劃線標(biāo)出。「Image-free」指在沒有真實(shí)圖像監(jiān)督的情況下進(jìn)行訓(xùn)練。

定性結(jié)果。圖 4 展示了 WaDi 與 SOTA 單步生成方法在 SD 1.5 和 SD 2.1 骨干上的定性比較。在不同提示詞下，WaDi 一致地生成視覺連貫、語義對(duì)齊的結(jié)果。例如，在第一、二行中，WaDi 更好地保留了結(jié)構(gòu)和風(fēng)格保真度，捕捉到清晰的細(xì)節(jié)和鮮艷的色彩，無偽影或失真。在第三、四行中，它能準(zhǔn)確跟隨涉及特定主體（如 sphynx cat、corgi、shiba inu）和場景（如劇院、服裝）的提示詞，而其他方法往往漏掉關(guān)鍵屬性或生成不真實(shí)的形狀。值得注意的是，在最后一行，WaDi 生成了空間構(gòu)圖和背景細(xì)節(jié)一致的復(fù)雜場景（如狗看電視），展示了相對(duì)于其他基線更優(yōu)越的整體理解能力。詳見補(bǔ)充材料 G.5。

圖 4. 與其他方法的定性比較，其中 ? 表示團(tuán)隊(duì)復(fù)現(xiàn)的結(jié)果。

3.下游任務(wù)

可控生成。ControlNet 是一種廣泛使用的可控生成模型，通過將空間條件整合到 SD 中實(shí)現(xiàn)精細(xì)控制。如圖 5 所示，將 WaDi 應(yīng)用于 ControlNet 可顯著提升推理效率，將推理時(shí)間縮短 86.26%，同時(shí)保持圖像質(zhì)量，忠實(shí)遵循空間條件，提示詞遵從度與 ControlNet 相當(dāng)。

圖 5. 使用或不使用 WaDi 的 ControlNet [74] 質(zhì)量結(jié)果。

關(guān)系反演。Reversion 是首個(gè)在 SD 中通過關(guān)系提示詞引導(dǎo)特定對(duì)象關(guān)系合成的方法。將 WaDi 集成到 Reversion 中可顯著加速推理。如圖 6 所示，WaDi 將推理時(shí)間縮短 88.89%，生成與關(guān)系提示詞對(duì)齊的高保真圖像，質(zhì)量接近原始多步 Reversion。

圖 6. 使用或不使用 WaDi 的 Reversion [22] 質(zhì)量結(jié)果。

圖像定制化。Dreambooth 是一種開創(chuàng)性的個(gè)性化文本到圖像框架，通過對(duì) U-Net 進(jìn)行微調(diào)將目標(biāo)主體綁定到稀有標(biāo)記。為增強(qiáng)參數(shù)效率，團(tuán)隊(duì)將 LoRaD 集成到 Dreambooth 中，并與 Dreambooth（FT）和 LoRA 進(jìn)行比較。如圖 7 所示，原始 DreamBooth 會(huì)捕捉主體但記憶訓(xùn)練圖像，降低提示詞敏感性。LoRA 緩解了過擬合，但降低了主體保真度和圖像質(zhì)量。相比之下，LoRaD 在保持提示詞遵從度的同時(shí)維持了主體保真度，實(shí)現(xiàn)了更好的平衡。團(tuán)隊(duì)將此 DreamBooth 實(shí)驗(yàn)僅作為說明性示例，而非擴(kuò)散微調(diào)的全面研究。

圖 7. 使用或不使用 LoRaD 的 DreamBooth 質(zhì)量結(jié)果。

4.用戶研究

為評(píng)估圖像質(zhì)量和文本 - 圖像對(duì)齊，團(tuán)隊(duì)邀請(qǐng) 57 名參與者進(jìn)行了用戶研究，涵蓋零樣本生成和下游任務(wù)。如圖 9 所示，結(jié)果清楚地表明團(tuán)隊(duì)的方法優(yōu)于現(xiàn)有基線。詳見補(bǔ)充材料 F.5。

圖 9. 與其他方法相比的用戶研究結(jié)果。

5.消融研究

表 2 在 COCO 2017 的 VSD 損失下比較了五種適配器類型。LoRaD 以最少 83.8M 可訓(xùn)練參數(shù)（比 LoRA/DoRA 少約 31%，比 FT 少約 90%）取得最低 FID（20.86）和競爭性 CLIP 分?jǐn)?shù)（0.31）。它還取得最高的方向均值（2.89，而 FT 為 2.21%，LoRA/DoRA 變體為 ≤0.92%），表明在緊湊參數(shù)化下具有更廣泛、更有效的更新方向空間。與 DoRA 和 DoRA（凍結(jié)范數(shù)）不同，后兩者通過 LoRA 式加性更新對(duì)歸一化權(quán)重后跟動(dòng)態(tài)重歸一化進(jìn)行方向優(yōu)化，LoRaD 直接將預(yù)訓(xùn)練權(quán)重參數(shù)化為低秩正交旋轉(zhuǎn)，保留范數(shù)并純?cè)诜较蚩臻g中操作。總體而言，LoRaD 展示了良好的質(zhì)量 - 效率權(quán)衡。

表 2. 在 COCO 2017 數(shù)據(jù)集上，WaDi（SD 1.5，VSD 損失）中適配器類型影響的消融實(shí)驗(yàn)。「NM」和「DM」分別表示所有層的范數(shù)均值和方向均值。

團(tuán)隊(duì)在 COCO 2014 上進(jìn)行了秩配置的消融研究。如表 3 所示，團(tuán)隊(duì)得出三個(gè)主要觀察：1）增大學(xué)生秩可持續(xù)改善性能。將秩從設(shè)置 A 提升到 C，F(xiàn)ID 從 13.64 降至 10.79，表明更高的秩使學(xué)生能夠更好地捕捉教師分布，提升生成質(zhì)量。2）超過閾值后繼續(xù)增大秩會(huì)導(dǎo)致收益遞減。對(duì)比設(shè)置 C 和 D，進(jìn)一步增大秩導(dǎo)致 FID 退化（12.75 vs. 10.79），CLIP 也從 0.31 降至 0.30，表明過大的秩可能導(dǎo)致過擬合。3）虛假模型的秩對(duì)保真度的影響大于對(duì)齊度。調(diào)整虛假模型秩（設(shè)置 C、E、F）會(huì)改變 FID 但 CLIP 基本穩(wěn)定，表明保真度對(duì)容量更敏感而對(duì)齊度較為穩(wěn)健。

綜上，設(shè)置 C 在模型容量和性能之間取得了良好的權(quán)衡，與圖 8 的定性結(jié)果一致。詳見補(bǔ)充材料 G.2、G.4。

表 3. 在 COCO 2014 數(shù)據(jù)集上，WaDi（SD 1.5，VSD 損失）中秩影響的消融實(shí)驗(yàn)。

圖 8. 不同設(shè)置下的一步圖像生成結(jié)果。

結(jié)論

本文提出了權(quán)重方向感知蒸餾（WaDi），一種高效的單步文本到圖像蒸餾框架。通過對(duì)多步與單步模型權(quán)重變化的深入分析，團(tuán)隊(duì)發(fā)現(xiàn)權(quán)重方向的變化是蒸餾中的關(guān)鍵機(jī)制，而范數(shù)變化的作用相對(duì)次要。

基于這一洞見，團(tuán)隊(duì)引入權(quán)重方向低秩旋轉(zhuǎn)（LoRaD）模塊，以參數(shù)高效的方式建模方向調(diào)整。

大量實(shí)驗(yàn)表明，WaDi 在圖像質(zhì)量和推理速度上顯著優(yōu)于現(xiàn)有單步方法—— 包括 DMD、SiD-LSG 和 SwiftBrush。此外，蒸餾后的模型可無縫適配多種下游任務(wù)，展示了強(qiáng)大的泛化能力和實(shí)際適用性。本工作為高效擴(kuò)散模型蒸餾提供了新穎的理論視角和實(shí)踐方案。