最近,扩散模型(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 社区的发展。
以上是加速擴散模型,最快1步驟產生SOTA級圖片,位元組Hyper-SD開源了的詳細內容。更多資訊請關注PHP中文網其他相關文章!