加速擴散模型，最快1步生成SOTA級圖片，字節Hyper-SD開源了
最近，擴散模型（Diffusion Model）在圖像生成領域取得了顯著的進展，為圖像生成和視頻生成任務帶來了前所未有的發展機遇。盡管取得了令人印象深刻的結果，擴散模型在推理過程中天然存在的多步數迭代去噪特性導致了較高的計算成本。近期出現了一系列擴散模型蒸餾算法來加速擴散模型的推理過程。這些方法大致可以分為兩類：i) 軌跡保持蒸餾；ii) 軌跡重構蒸餾。然而，這兩類方法會分別受到效果天花板有限或者或輸出域變化這兩個問題的限制。
為了解決這些問題，字節跳動技術團隊提出了一種名為 Hyper-SD 的軌跡分段一致性模型。Hyper-SD 的開源也得到了Huggingface首席執行官 Clem Delangue的肯定。
該模型是一種新穎的擴散模型蒸餾框架，結合了軌跡保持蒸餾和軌跡重構蒸餾兩種策略的優點，在壓縮去噪步數的同時保持接近無損的性能。與現有的擴散模型加速算法相比，該方法取得了卓越的加速效果。經過大量實驗和用戶評測的驗證，Hyper-SD 在 SDXL 和 SD1.5 兩種架構上都能在 1 到 8 步生成中實現 SOTA 級別的圖像生成性能。
- 項目主頁：https://hyper-sd.github.io/
- 論文鏈接：https://arxiv.org/abs/2404.13686
- Huggingface 鏈接：https://huggingface.co/ByteDance/Hyper-SD
- 單步生成 Demo 鏈接：https://huggingface.co/spaces/ByteDance/Hyper-SDXL-1Step-T2I
- 實時畫板 Demo 鏈接：https://huggingface.co/spaces/ByteDance/Hyper-SD15-Scribble
引言
現有用于擴散模型加速的蒸餾方法大致可以分為兩大類：軌跡保持蒸餾和軌跡重構蒸餾。軌跡保持蒸餾技術旨在維持擴散對應的常微分方程（ODE）的原始軌跡。其原理是通過迫使蒸餾模型和原始模型產生相似的輸出來減少推理步驟。然而需要注意的是，盡管能夠實現加速，由于模型容量有限以及訓練擬合過程中不可避免的誤差，這類方法可能導致生成質量下降。相比之下，軌跡重構方法則直接利用軌跡上的端點或真實圖像作為監督的主要來源，忽略了軌跡的中間步驟，能夠通過重建更有效的軌跡來減少推理步驟的數量，并在有限的步驟內探索模型的潛力，將其從原始軌跡的約束中解放出來。然而，這通常會導致加速模型與原始模型的輸出域不一致，從而得到不理想的結果。
本論文提出了一種結合軌跡保持和重構策略優點的軌跡分段一致性模型（簡稱 Hyper-SD）。具體而言，該算法首先引入軌跡分段一致性蒸餾，在每個段內強制保持一致性，并逐漸減少段的數量以實現全時一致性。這一策略解決了由于模型擬合能力不足和推理誤差累積導致的一致性模型性能次優的問題。隨后，該算法利用人類反饋學習（RLHF）來提升模型的生成效果，以彌補加速過程中模型生成效果的損失，使其更好地適應低步數推理。最后，該算法使用分數蒸餾來增強一步生成性能，并通過統一的 LORA 實現理想化的全時間步數一致擴散模型，在生成效果上取得了卓越的成果。
方法
1. 軌跡分段一致性蒸餾
一致性蒸餾（CD）[24] 和一致性軌跡模型（CTM）[4] 都旨在通過一次性蒸餾將擴散模型轉換為整個時間步范圍 [0，T] 的一致性模型。然而，由于模型擬合能力的限制，這些蒸餾模型往往達不到最優性。受到 CTM 中引入的軟一致性目標的啟發，我們通過將整個時間步范圍 [0, T] 劃分為 k 段并逐步執行分段一致模型蒸餾來細化訓練過程。
在第一階段，我們設置 k=8 并使用原始擴散模型來初始化
和
。起始時間步
是從
中均勻隨機采樣的。然后，我們對結束時間步
進行采樣，其中
計算如下：
訓練損失計算如下：
其中
通過公式 3 進行計算，
表示學生模型的指數滑動平均(EMA)。
隨后，我們恢復上一階段的模型權重并繼續訓練
，逐漸將 k 減少到 [4,2,1]。值得注意的是，k=1 對應于標準 CTM 訓練方案。對于距離度量 d，我們采用了對抗性損失和均方誤差 (MSE) 損失的混合。在實驗中，我們觀察到，當預測值和目標值接近時（例如，對于 k=8, 4），MSE 損失更為有效，而對抗性損失則隨著預測和目標值之間的差異增加而變得更加精確（例如，對于 k=2, 1)。因此，我們在整個訓練階段動態增加對抗性損失的權重并減少 MSE 損失的權重。此外，我們還集成了噪聲擾動機制來增強訓練穩定性。以兩階段軌跡分段一致性蒸餾（TSCD）過程為例。如下圖所示，我們第一階段在
和
時間段內執行獨立一致性蒸餾 ，然后基于之前的兩段一致性蒸餾結果，進行全局一致性軌跡蒸餾。
完整的算法流程如下：
2. 人類反饋學習
除了蒸餾之外，我們進一步結合反饋學習以提高加速擴散模型的性能。具體來說我們通過利用人類審美偏好和現有視覺感知模型的反饋來提高加速模型的生成質量。對于審美反饋，我們利用 LAION 審美預測器和 ImageReward 中提供的審美偏好獎勵模型來引導模型生成更具美感的圖像，如下所示：
其中
是審美獎勵模型，包括 LAION 數據集和 ImageReward 模型的審美預測器，c 是文本提示，
與ReLU函數一起作為鉸鏈損失 。除了來自審美偏好的反饋之外，我們注意到嵌入有關圖像的豐富先驗知識的現有視覺感知模型也可以作為良好的反饋提供者。根據經驗，我們發現實例分割模型可以指導模型生成結構合理的物體。具體來說，我們首先將潛在空間中圖像
上的噪聲擴散到
，之后，類似于 ImageReward，我們執行迭代去噪，直到 特定時間步
并直接預測
。隨后，我們利用感知實例分割模型通過檢查真實圖像實例分割標注與去噪圖像的實例分割預測結果之間的差異來評估結構生成的性能，如下所示：
其中
是實例分割模型（例如 SOLO）。實例分割模型可以更準確地捕獲生成圖像的結構缺陷并提供更有針對性的反饋信號。值得注意的是，除了實例分割模型之外，其他感知模型也適用。這些感知模型可以作為主觀審美的補充反饋，更多地關注客觀生成質量。因此，我們用反饋信號優化擴散模型可以定義為：
3. 一步生成強化
由于一致性損失的固有限制，一致性模型框架內的一步生成并不理想。正如 CM 中分析的那樣，一致性蒸餾模型在引導位置
處的軌跡端點
方面表現出卓越的準確性。因此，分數蒸餾是一種合適且有效的方法來進一步提升我們的 TSCD 模型的一步生成效果。具體來說，我們通過優化的分布匹配蒸餾（DMD）技術來推進一步生成。DMD 通過利用兩個不同的評分函數來增強模型的輸出：來自教師模型分布
和來自假模型的
。我們將均方誤差 (MSE) 損失與基于分數的蒸餾結合起來，以提高訓練穩定性。在這個過程中，前面提到的人類反饋學習技術也被集成進來，用來微調我們的模型以有效地生成具有保真度的圖像。
通過集成這些策略，我們的方法不僅能夠實現在 SD1.5 和 SDXL 上都實現卓越的低步數推理效果（并且無需 Classifier-Guidance），同時能夠實現理想的全局一致性模型，無需針對每個特定的步數訓練 UNet 或者 LoRA 實現統一的低步數推理模型。
實驗
在 SD1.5 和 SDXL 上和目前現有的各種加速算法的定量比較，可以看到 Hyper-SD 顯著優于當前最先進的方法
此外，Hyper-SD 能夠用一個模型來實現各種不同低步數的推理，上面的定量指標也顯示了我們方法在使用統一模型推理時的效果。
在 SD1.5 和 SDXL 上的加速效果可視化直觀地展示了 Hyper-SD 在擴散模型推理加速上的優越性。
大量的 User-Study 也表明 Hyper-SD 相較于現有的各種加速算法的優越性。
Hyper-SD 訓練得到的加速 LoRA 能夠很好地兼容不同的風格的文生圖底模。
同時，Hyper-SD 的 LoRA 也能適配現有的 ControlNet，實現低步數下高質量的可控圖像生成。
總結
論文提出了 Hyper-SD，一個統一的擴散模型加速框架，可以顯著提升擴散模型的在低步數情況下的生成能力，實現基于 SDXL 和 SD15 的新 SOTA 性能。該方法通過采用軌跡分段一致性蒸餾，增強了蒸餾過程中的軌跡保存能力，實現接近原始模型的生成效果。然后，通過進一步利用人類反饋學習和變分分數蒸餾提升模型在極端低步數下的潛力，從而產生了更優化、更高效的模型生成效果。論文還開源了用于 SDXL 和 SD15 從 1 到 8 步推理的 Lora 插件，以及專用的一步 SDXL 模型，旨在進一步推動生成式 AI 社區的發展。
本文轉自機器之心 ，作者：機器之心
