最近，擴(kuò)散模型（Diffusion Model）在圖像生成領(lǐng)域取得了顯著的進(jìn)展，為圖像生成和視頻生成任務(wù)帶來了前所未有的發(fā)展機(jī)遇。盡管取得了令人印象深刻的結(jié)果，擴(kuò)散模型在推理過程中天然存在的多步數(shù)迭代去噪特性導(dǎo)致了較高的計(jì)算成本。近期出現(xiàn)了一系列擴(kuò)散模型蒸餾算法來加速擴(kuò)散模型的推理過程。這些方法大致可以分為兩類：i) 軌跡保持蒸餾；ii) 軌跡重構(gòu)蒸餾。然而，這兩類方法會(huì)分別受到效果天花板有限或者或輸出域變化這兩個(gè)問題的限制。

為了解決這些問題，字節(jié)跳動(dòng)技術(shù)團(tuán)隊(duì)提出了一種名為 Hyper-SD 的軌跡分段一致性模型。Hyper-SD 的開源也得到了Huggingface首席執(zhí)行官 Clem Delangue的肯定。

該模型是一種新穎的擴(kuò)散模型蒸餾框架，結(jié)合了軌跡保持蒸餾和軌跡重構(gòu)蒸餾兩種策略的優(yōu)點(diǎn)，在壓縮去噪步數(shù)的同時(shí)保持接近無損的性能。與現(xiàn)有的擴(kuò)散模型加速算法相比，該方法取得了卓越的加速效果。經(jīng)過大量實(shí)驗(yàn)和用戶評(píng)測的驗(yàn)證，Hyper-SD 在 SDXL 和 SD1.5 兩種架構(gòu)上都能在 1 到 8 步生成中實(shí)現(xiàn) SOTA 級(jí)別的圖像生成性能。

• 項(xiàng)目主頁：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
• 實(shí)時(shí)畫板 Demo 鏈接：https://huggingface.co/spaces/ByteDance/Hyper-SD15-Scribble

引言

現(xiàn)有用于擴(kuò)散模型加速的蒸餾方法大致可以分為兩大類：軌跡保持蒸餾和軌跡重構(gòu)蒸餾。軌跡保持蒸餾技術(shù)旨在維持?jǐn)U散對(duì)應(yīng)的常微分方程（ODE）的原始軌跡。其原理是通過迫使蒸餾模型和原始模型產(chǎn)生相似的輸出來減少推理步驟。然而需要注意的是，盡管能夠?qū)崿F(xiàn)加速，由于模型容量有限以及訓(xùn)練擬合過程中不可避免的誤差，這類方法可能導(dǎo)致生成質(zhì)量下降。相比之下，軌跡重構(gòu)方法則直接利用軌跡上的端點(diǎn)或真實(shí)圖像作為監(jiān)督的主要來源，忽略了軌跡的中間步驟，能夠通過重建更有效的軌跡來減少推理步驟的數(shù)量，并在有限的步驟內(nèi)探索模型的潛力，將其從原始軌跡的約束中解放出來。然而，這通常會(huì)導(dǎo)致加速模型與原始模型的輸出域不一致，從而得到不理想的結(jié)果。

本論文提出了一種結(jié)合軌跡保持和重構(gòu)策略優(yōu)點(diǎn)的軌跡分段一致性模型（簡稱 Hyper-SD）。具體而言，該算法首先引入軌跡分段一致性蒸餾，在每個(gè)段內(nèi)強(qiáng)制保持一致性，并逐漸減少段的數(shù)量以實(shí)現(xiàn)全時(shí)一致性。這一策略解決了由于模型擬合能力不足和推理誤差累積導(dǎo)致的一致性模型性能次優(yōu)的問題。隨后，該算法利用人類反饋學(xué)習(xí)（RLHF）來提升模型的生成效果，以彌補(bǔ)加速過程中模型生成效果的損失，使其更好地適應(yīng)低步數(shù)推理。最后，該算法使用分?jǐn)?shù)蒸餾來增強(qiáng)一步生成性能，并通過統(tǒng)一的 LORA 實(shí)現(xiàn)理想化的全時(shí)間步數(shù)一致擴(kuò)散模型，在生成效果上取得了卓越的成果。

方法

1. 軌跡分段一致性蒸餾

一致性蒸餾（CD）[24] 和一致性軌跡模型（CTM）[4] 都旨在通過一次性蒸餾將擴(kuò)散模型轉(zhuǎn)換為整個(gè)時(shí)間步范圍 [0，T] 的一致性模型。然而，由于模型擬合能力的限制，這些蒸餾模型往往達(dá)不到最優(yōu)性。受到 CTM 中引入的軟一致性目標(biāo)的啟發(fā)，我們通過將整個(gè)時(shí)間步范圍 [0, T] 劃分為 k 段并逐步執(zhí)行分段一致模型蒸餾來細(xì)化訓(xùn)練過程。

在第一階段，我們?cè)O(shè)置 k=8 并使用原始擴(kuò)散模型來初始化 

  和

。起始時(shí)間步

是從

中均勻隨機(jī)采樣的。然后，我們對(duì)結(jié)束時(shí)間步

進(jìn)行采樣，其中

計(jì)算如下：

訓(xùn)練損失計(jì)算如下：

其中

通過公式 3 進(jìn)行計(jì)算，

表示學(xué)生模型的指數(shù)滑動(dòng)平均(EMA)。

隨后，我們恢復(fù)上一階段的模型權(quán)重并繼續(xù)訓(xùn)練

，逐漸將 k 減少到 [4,2,1]。值得注意的是，k=1 對(duì)應(yīng)于標(biāo)準(zhǔn) CTM 訓(xùn)練方案。對(duì)于距離度量 d，我們采用了對(duì)抗性損失和均方誤差 (MSE) 損失的混合。在實(shí)驗(yàn)中，我們觀察到，當(dāng)預(yù)測值和目標(biāo)值接近時(shí)（例如，對(duì)于 k=8, 4），MSE 損失更為有效，而對(duì)抗性損失則隨著預(yù)測和目標(biāo)值之間的差異增加而變得更加精確（例如，對(duì)于 k=2, 1)。因此，我們?cè)谡麄€(gè)訓(xùn)練階段動(dòng)態(tài)增加對(duì)抗性損失的權(quán)重并減少 MSE 損失的權(quán)重。此外，我們還集成了噪聲擾動(dòng)機(jī)制來增強(qiáng)訓(xùn)練穩(wěn)定性。以兩階段軌跡分段一致性蒸餾（TSCD）過程為例。如下圖所示，我們第一階段在 

 和 

 時(shí)間段內(nèi)執(zhí)行獨(dú)立一致性蒸餾 ，然后基于之前的兩段一致性蒸餾結(jié)果，進(jìn)行全局一致性軌跡蒸餾。

完整的算法流程如下：

2. 人類反饋學(xué)習(xí)

除了蒸餾之外，我們進(jìn)一步結(jié)合反饋學(xué)習(xí)以提高加速擴(kuò)散模型的性能。具體來說我們通過利用人類審美偏好和現(xiàn)有視覺感知模型的反饋來提高加速模型的生成質(zhì)量。對(duì)于審美反饋，我們利用 LAION 審美預(yù)測器和 ImageReward 中提供的審美偏好獎(jiǎng)勵(lì)模型來引導(dǎo)模型生成更具美感的圖像，如下所示：

其中

是審美獎(jiǎng)勵(lì)模型，包括 LAION 數(shù)據(jù)集和 ImageReward 模型的審美預(yù)測器，c 是文本提示，

與ReLU函數(shù)一起作為鉸鏈損失 。除了來自審美偏好的反饋之外，我們注意到嵌入有關(guān)圖像的豐富先驗(yàn)知識(shí)的現(xiàn)有視覺感知模型也可以作為良好的反饋提供者。根據(jù)經(jīng)驗(yàn)，我們發(fā)現(xiàn)實(shí)例分割模型可以指導(dǎo)模型生成結(jié)構(gòu)合理的物體。具體來說，我們首先將潛在空間中圖像

上的噪聲擴(kuò)散到

，之后，類似于 ImageReward，我們執(zhí)行迭代去噪，直到 特定時(shí)間步

并直接預(yù)測

。隨后，我們利用感知實(shí)例分割模型通過檢查真實(shí)圖像實(shí)例分割標(biāo)注與去噪圖像的實(shí)例分割預(yù)測結(jié)果之間的差異來評(píng)估結(jié)構(gòu)生成的性能，如下所示：

其中

是實(shí)例分割模型（例如 SOLO）。實(shí)例分割模型可以更準(zhǔn)確地捕獲生成圖像的結(jié)構(gòu)缺陷并提供更有針對(duì)性的反饋信號(hào)。值得注意的是，除了實(shí)例分割模型之外，其他感知模型也適用。這些感知模型可以作為主觀審美的補(bǔ)充反饋，更多地關(guān)注客觀生成質(zhì)量。因此，我們用反饋信號(hào)優(yōu)化擴(kuò)散模型可以定義為：

3. 一步生成強(qiáng)化

由于一致性損失的固有限制，一致性模型框架內(nèi)的一步生成并不理想。正如 CM 中分析的那樣，一致性蒸餾模型在引導(dǎo)位置

處的軌跡端點(diǎn)

方面表現(xiàn)出卓越的準(zhǔn)確性。因此，分?jǐn)?shù)蒸餾是一種合適且有效的方法來進(jìn)一步提升我們的 TSCD 模型的一步生成效果。具體來說，我們通過優(yōu)化的分布匹配蒸餾（DMD）技術(shù)來推進(jìn)一步生成。DMD 通過利用兩個(gè)不同的評(píng)分函數(shù)來增強(qiáng)模型的輸出：來自教師模型分布

和來自假模型的

。我們將均方誤差 (MSE) 損失與基于分?jǐn)?shù)的蒸餾結(jié)合起來，以提高訓(xùn)練穩(wěn)定性。在這個(gè)過程中，前面提到的人類反饋學(xué)習(xí)技術(shù)也被集成進(jìn)來，用來微調(diào)我們的模型以有效地生成具有保真度的圖像。

通過集成這些策略，我們的方法不僅能夠?qū)崿F(xiàn)在 SD1.5 和 SDXL 上都實(shí)現(xiàn)卓越的低步數(shù)推理效果（并且無需 Classifier-Guidance），同時(shí)能夠?qū)崿F(xiàn)理想的全局一致性模型，無需針對(duì)每個(gè)特定的步數(shù)訓(xùn)練 UNet 或者 LoRA 實(shí)現(xiàn)統(tǒng)一的低步數(shù)推理模型。

實(shí)驗(yàn)

在 SD1.5 和 SDXL 上和目前現(xiàn)有的各種加速算法的定量比較，可以看到 Hyper-SD 顯著優(yōu)于當(dāng)前最先進(jìn)的方法

此外，Hyper-SD 能夠用一個(gè)模型來實(shí)現(xiàn)各種不同低步數(shù)的推理，上面的定量指標(biāo)也顯示了我們方法在使用統(tǒng)一模型推理時(shí)的效果。

在 SD1.5 和 SDXL 上的加速效果可視化直觀地展示了 Hyper-SD 在擴(kuò)散模型推理加速上的優(yōu)越性。

大量的 User-Study 也表明 Hyper-SD 相較于現(xiàn)有的各種加速算法的優(yōu)越性。

Hyper-SD 訓(xùn)練得到的加速 LoRA 能夠很好地兼容不同的風(fēng)格的文生圖底模。

同時(shí)，Hyper-SD 的 LoRA 也能適配現(xiàn)有的 ControlNet，實(shí)現(xiàn)低步數(shù)下高質(zhì)量的可控圖像生成。

總結(jié)

論文提出了 Hyper-SD，一個(gè)統(tǒng)一的擴(kuò)散模型加速框架，可以顯著提升擴(kuò)散模型的在低步數(shù)情況下的生成能力，實(shí)現(xiàn)基于 SDXL 和 SD15 的新 SOTA 性能。該方法通過采用軌跡分段一致性蒸餾，增強(qiáng)了蒸餾過程中的軌跡保存能力，實(shí)現(xiàn)接近原始模型的生成效果。然后，通過進(jìn)一步利用人類反饋學(xué)習(xí)和變分分?jǐn)?shù)蒸餾提升模型在極端低步數(shù)下的潛力，從而產(chǎn)生了更優(yōu)化、更高效的模型生成效果。論文還開源了用于 SDXL 和 SD15 從 1 到 8 步推理的 Lora 插件，以及專用的一步 SDXL 模型，旨在進(jìn)一步推動(dòng)生成式 AI 社區(qū)的發(fā)展。

本文轉(zhuǎn)自機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/dqDqlWv1xe-8zayeJCGq8A??