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

本文作者王雷，南開大學，PCA Lab成員

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

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

LoRaD 通過可學習的低秩旋轉矩陣對預訓練權重的方向進行建模。團隊進一步將 LoRaD 集成到變分得分蒸餾（Variational Score Distillation，VSD）中，提出了權重方向感知蒸餾（Weight Direction-aware Distillation，WaDi）—— 一種新穎的單步蒸餾框架。

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

• 論文鏈接：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）一步生成的圖像。

引言

擴散模型（DMs）在圖像生成領域受到了廣泛關注，在文本到圖像生成、文本到視頻生成以及圖像到視頻生成等任務中均有廣泛應用。然而，擴散模型依賴多步采樣，導致計算成本高、推理速度慢。

為此，近期蒸餾方法將采樣步數壓縮至數步甚至一步。有趣的是，在蒸餾過程中，團隊發現權重范數在各層間保持相對穩定，而在將權重重參數化為范數與方向時，方向則呈現出更大的變化幅度。

受權重重參數化的啟發，團隊采用類似的分解方式來分析擴散蒸餾中的權重變化。為此，團隊研究了最先進（SOTA）單步模型（如 DMD2 和 Pixart-α DMD）與其對應多步模型（如 SD 1.5 和 Pixart-α）之間的權重更新。

如圖 2 (a) 所示，在基于 U-Net 的架構中，各層權重范數幾乎保持穩定，均值和標準差（STD）分別約為 0.1% 和 0.2%。相比之下，權重方向的變化則明顯更大，均值為 2.2%，標準差為 2.1%，對應為范數變化的 22 倍和 10 倍。在基于 DiT 的架構中也觀察到類似規律（見圖 2 (a) 右）。

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

此外，方向上的變化是否具有結構化規律？為此，團隊對殘差矩陣（單步與多步方向矩陣之差）進行奇異值分解（SVD），發現僅保留 30% 的秩即可恢復 93% 的信息，突顯了其低秩本質（見圖 2 (b)）。

圖 2. 我們方法的動機分析。(a) 一步學生模型與教師模型之間的權重范數和方向差異。更多細節和補充示例見補充材料 E。(b) DMD2 殘差矩陣的 SVD 分析。(c) 將一步模型的范數替換為多步模型的范數影響很小①④）；替換方向會嚴重降低生成質量②⑤）。(d) 與 (c) 對應的定性示例。(e) LoRaD 示意圖。

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

這些發現表明，方向重建是蒸餾中性能提升的核心因素，而范數變化的影響相對次要。一種可能的解釋是：用教師權重初始化學生模型對齊了初始范數，訓練過程中的權重衰減進一步約束了范數漂移；而蒸餾信號則主要通過調整權重方向來減少表征差異。

綜合來看，這些結果表明方向重建是蒸餾性能提升的核心驅動因素。

上述蒸餾方法大致可分為兩類：全量微調（FT）基于低秩適配（LoRA）的微調。然而，二者在優化范數和方向時均直接更新模型參數，導致范數和方向的變化相互耦合，增加了優化難度。此外，FT 和 LoRA 均面臨收斂慢、不穩定和過擬合等問題，進一步增加了優化的復雜性。

為此，團隊提出了權重方向低秩旋轉（LoRaD）（見圖 2 (e)），通過可學習的旋轉矩陣調整預訓練權重的方向。鑒于方向變化的結構化特性（即低秩特性），旋轉角度被參數化為兩個低秩矩陣的乘積，以進一步減少可學習參數數量。團隊將 LoRaD 集成到變分得分蒸餾（VSD）中，提出了權重方向感知蒸餾（WaDi），一種新穎的單步文本到圖像蒸餾框架。

在 COCO 2014 和 COCO 2017 數據集上的實驗表明，WaDi 取得了最先進的 FID 分數，超越所有現有單步生成方法。這一成果僅通過優化方向實現，將蒸餾難度降低，同時 U-Net 可訓練參數僅約占 10%，極大提升了參數效率。

此外，團隊將 WaDi 應用于可控生成、關系反演、高分辨率合成和圖像定制化等下游任務，展示了其加速能力和廣泛適用性。本文貢獻總結如下：

• 團隊對多步與單步生成模型之間 U-Net 權重變化進行了深入分析，將權重方向調整確定為單步蒸餾的關鍵驅動因素，為高效蒸餾提供了新的理論視角；
• 團隊提出了一種新穎的單步文本到圖像蒸餾框架 WaDi，采用 LoRaD 通過低秩旋轉建模權重方向，有效引導學生模型對齊教師分布；
• WaDi 在 COCO 數據集和多個下游任務上進行了評估，定性和定量結果均表明 WaDi 在顯著提升推理效率的同時取得了實質性的圖像質量提升。

相關工作

擴散模型

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

擴散模型加速

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

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

然而，現有基于訓練的方法通常使用 FT 或 LoRA，這可能增加優化難度。團隊發現方向變化在蒸餾中通常更具影響力，因此提出 WaDi，利用 LoRaD 專注于建模方向旋轉。

方法

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

2.權重方向低秩旋轉

圖 3. （左）權重方向低秩旋轉（LoRaD）模塊的詳細結構。LoRaD 使用可學習的低秩旋轉角來旋轉預訓練權重方向。（右）權重方向感知蒸餾（WaDi）框架概覽。

注：旋轉不影響范數，因此無需顯式分離范數矩陣。

實驗

1.實驗設置

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

實現細節。遵循先前方法，WaDi 中的學生模型采用與教師相同的架構，并用教師權重初始化。WaDi 在從 JourneyDB 數據集中采樣的 140 萬條提示詞上進行訓練。訓練過程中，學生模型的學習率（LR）設為 1e-4，虛假模型的學習率設為 1e-2。團隊使用 AdamW 作為優化器，批大小為 128（每 GPU 16 張）。無分類器引導（CFG）系數設為 1.5，訓練進行 2 個 epoch。團隊在三種不同骨干上蒸餾學生模型，分別為 SD 1.5、SD 2.1 和 PixArt-α（256×256）。對于 SD 1.5 和 SD 2.1，學生模型的 LoRaD 秩設為 256，而對于 PixArt-α 設為 128。所有虛假模型的 LoRaD 秩統一設為 32。詳見補充材料 F.1。

2.與最先進方法的比較

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

值得注意的是，SD 1.5、SD 2.1 和 PixArt-α 可訓練參數分別僅占 9.74%、10.92% 和 13.30%，突顯了 WaDi 的參數高效性。這些改進源于團隊提出的 LoRaD，通過低秩旋轉重參數化權重更新，實現穩定高效的蒸餾。

表 1. WaDi 與其他方法在零樣本 COCO 2014 結果上的定量比較。? 表示我們復現的結果，? 表示使用官方預訓練模型得到的結果。「-」表示未知。最佳和次佳分數分別用粗體和下劃線標出。「Image-free」指在沒有真實圖像監督的情況下進行訓練。

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

圖 4. 與其他方法的定性比較，其中 ? 表示團隊復現的結果。

3.下游任務

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

圖 5. 使用或不使用 WaDi 的 ControlNet [74] 質量結果。

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

圖 6. 使用或不使用 WaDi 的 Reversion [22] 質量結果。

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

圖 7. 使用或不使用 LoRaD 的 DreamBooth 質量結果。

4.用戶研究

為評估圖像質量和文本 - 圖像對齊，團隊邀請 57 名參與者進行了用戶研究，涵蓋零樣本生成和下游任務。如圖 9 所示，結果清楚地表明團隊的方法優于現有基線。詳見補充材料 F.5。

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

5.消融研究

表 2 在 COCO 2017 的 VSD 損失下比較了五種適配器類型。LoRaD 以最少 83.8M 可訓練參數（比 LoRA/DoRA 少約 31%，比 FT 少約 90%）取得最低 FID（20.86）和競爭性 CLIP 分數（0.31）。它還取得最高的方向均值（2.89，而 FT 為 2.21%，LoRA/DoRA 變體為 ≤0.92%），表明在緊湊參數化下具有更廣泛、更有效的更新方向空間。與 DoRA 和 DoRA（凍結范數）不同，后兩者通過 LoRA 式加性更新對歸一化權重后跟動態重歸一化進行方向優化，LoRaD 直接將預訓練權重參數化為低秩正交旋轉，保留范數并純在方向空間中操作。總體而言，LoRaD 展示了良好的質量 - 效率權衡。

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

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

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

表 3. 在 COCO 2014 數據集上，WaDi（SD 1.5，VSD 損失）中秩影響的消融實驗。

圖 8. 不同設置下的一步圖像生成結果。

結論

本文提出了權重方向感知蒸餾（WaDi），一種高效的單步文本到圖像蒸餾框架。通過對多步與單步模型權重變化的深入分析，團隊發現權重方向的變化是蒸餾中的關鍵機制，而范數變化的作用相對次要。

基于這一洞見，團隊引入權重方向低秩旋轉（LoRaD）模塊，以參數高效的方式建模方向調整。

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