最近，擴散模型（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 社群的發展。