擴(kuò)散模型已成為現(xiàn)代文本到圖像 (T2I) 生成技術(shù)的核心，能夠生成高質(zhì)量圖像，但其迭代式推理過程導(dǎo)致生成速度緩慢。多數(shù)模型通常需要20–50 個(gè)去噪步驟，這嚴(yán)重制約了其在實(shí)時(shí)應(yīng)用中的部署。
現(xiàn)有的蒸餾技術(shù)旨在加速擴(kuò)散模型的采樣過程，然而，這些方法往往會(huì)引入穩(wěn)定性問題，在極低步數(shù)下出現(xiàn)質(zhì)量下降，并可能導(dǎo)致顯著的內(nèi)存需求。
Nvidia 提出的 SANA-Sprint 是一種混合蒸餾框架，它整合了連續(xù)時(shí)間一致性模型 (sCM)和潛在對(duì)抗擴(kuò)散蒸餾 (LADD)，旨在實(shí)現(xiàn)以下目標(biāo)：
- 無步訓(xùn)練，并支持靈活的1–4 步推理
- 卓越的速度與質(zhì)量平衡，單步推理即可達(dá)到FID 7.59GenEval 0.74的指標(biāo)。
- 在 H100 GPU 上實(shí)現(xiàn) 0.1 秒生成 1024×1024 圖像，速度比FLUX-Schnell 快 10 倍，同時(shí)保持更高圖像質(zhì)量。
本文將深入探討 SANA-Sprint 實(shí)現(xiàn)上述性能的技術(shù)原理。
傳統(tǒng)蒸餾方法在超低步數(shù)推理中的局限性
擴(kuò)散模型依賴于隨機(jī)微分方程 (SDE) 或常微分方程 (ODE)進(jìn)行圖像生成，該過程通常需要多個(gè)步驟。盡管存在多種步數(shù)縮減技術(shù)，但每種方法都存在其固有的局限性：
- 基于 GAN 的蒸餾方法（例如，LADD)可以加速推理過程，但容易遭受模式崩潰泛化能力不足的問題。
- 一致性模型 (CM)能夠?qū)崿F(xiàn)快速采樣，但在超低步數(shù) (少于 4 步)的情況下，由于軌跡截?cái)嗾`差，語義對(duì)齊性能會(huì)顯著下降
- 變分分?jǐn)?shù)蒸餾 (VSD)需要額外訓(xùn)練輔助擴(kuò)散模型，這會(huì)顯著增加GPU 內(nèi)存占用和計(jì)算開銷
SANA-Sprint 通過整合 sCM 和 LADD 到統(tǒng)一框架中，克服了上述挑戰(zhàn)，從而在確保快速推理的同時(shí)，實(shí)現(xiàn)了高圖像質(zhì)量。
基于無訓(xùn)練一致性變換的預(yù)訓(xùn)練模型重用
擴(kuò)散模型通常采用流匹配或基于分?jǐn)?shù)的學(xué)習(xí)方法進(jìn)行訓(xùn)練，而一致性模型 (CM)則基于TrigFlow 參數(shù)化。為了實(shí)現(xiàn)無需重新訓(xùn)練的快速蒸餾，SANA-Sprint 引入了一種數(shù)學(xué)變換，可以將預(yù)訓(xùn)練的流匹配模型轉(zhuǎn)化為 TrigFlow 模型。
該變換確保了以下關(guān)鍵特性：
- 時(shí)域映射的無縫銜接：實(shí)現(xiàn)了從流匹配模型的 [0,1] 區(qū)間TrigFlow 模型的 [0, π/2] 區(qū)間的平滑轉(zhuǎn)換。
- 信噪比 (SNR) 的一致性：在模型適配過程中，保持了信噪比的穩(wěn)定，確保圖像保真度。
- 模型輸出的正確參數(shù)化：保證了轉(zhuǎn)換后模型輸出的速度場與 TrigFlow 框架的公式保持一致。
通過上述變換，預(yù)訓(xùn)練模型可以直接應(yīng)用于 SANA-Sprint 框架，無需額外的重新訓(xùn)練，從而顯著提升了效率。
解決大規(guī)模一致性模型訓(xùn)練不穩(wěn)定性問題
將一致性模型擴(kuò)展到更高分辨率和更大模型規(guī)模時(shí)，常常會(huì)面臨訓(xùn)練不穩(wěn)定性的挑戰(zhàn)，這主要是由于梯度爆炸現(xiàn)象引起的。SANA-Sprint 通過以下兩項(xiàng)關(guān)鍵技術(shù)來穩(wěn)定訓(xùn)練過程：
密集時(shí)間嵌入以抑制梯度爆炸
- 傳統(tǒng)擴(kuò)散模型通常使用乘法因子（例如，1000 * t）來縮放時(shí)間嵌入，這種方法會(huì)放大時(shí)間導(dǎo)數(shù)梯度，容易導(dǎo)致訓(xùn)練崩潰。
- SANA-Sprint 采用歸一化時(shí)間嵌入方法，確保時(shí)間步長表示的均勻分布，從而有效提升訓(xùn)練穩(wěn)定性和樣本質(zhì)量
- 這種方法使得模型能夠更快收斂，并生成更清晰銳利的圖像
QK 歸一化實(shí)現(xiàn)穩(wěn)定的自注意力和交叉注意力機(jī)制
- 隨著模型規(guī)模的擴(kuò)大 (參數(shù)量從 0.6B 增至 1.6B)，梯度范數(shù)變得不穩(wěn)定 (>1?3)，導(dǎo)致訓(xùn)練失敗。
- SANA-Sprint 在注意力層的 Query 和 Key (QK) 組件中引入 RMS 歸一化，在不改變模型架構(gòu)的前提下，有效穩(wěn)定了梯度。
- 僅需5,000 次微調(diào)迭代，即可顯著降低訓(xùn)練不穩(wěn)定性，從而為大規(guī)模擴(kuò)散模型的穩(wěn)定蒸餾奠定基礎(chǔ)。
結(jié)合一致性模型與對(duì)抗監(jiān)督
傳統(tǒng)一致性模型主要依賴局部軌跡學(xué)習(xí)，這導(dǎo)致其收斂速度較慢，并且在單步生成中容易丟失細(xì)節(jié)信息。SANA-Sprint 通過引入基于 GAN 的對(duì)抗監(jiān)督機(jī)制 (LADD)(Latent Adversarial Diffusion Distillation)，對(duì)一致性模型進(jìn)行了增強(qiáng)：
- 使用凍結(jié)的教師模型提取高層潛在空間表征，以強(qiáng)制模型學(xué)習(xí)數(shù)據(jù)分布的一致性。
- 引入多頭判別器學(xué)習(xí)特征層面的差異，避免了像素空間直接比對(duì)可能導(dǎo)致的問題。
- 采用鉸鏈損失函數(shù)，提升了訓(xùn)練穩(wěn)定性和生成樣本的真實(shí)感
該技術(shù)顯著提升了單步圖像生成質(zhì)量，有效保留了傳統(tǒng)一致性模型難以捕捉的高頻細(xì)節(jié)。
評(píng)估與結(jié)果
SANA-Sprint 在速度和質(zhì)量方面均達(dá)到了新的技術(shù)水平。相較于 FLUX-Schnell，SANA-Sprint 的推理速度提升了 10 倍，同時(shí)能夠生成更高質(zhì)量的圖像。在單步推理下，SANA-Sprint 取得了7.59 的 FID 值和 0.74 的 GenEval 值，性能超越了需要多步推理的模型。即使在RTX 4090 等消費(fèi)級(jí) GPU上，SANA-Sprint 也能在0.31 秒內(nèi)生成 1024×1024 像素的圖像，使得高質(zhì)量 AI 圖像生成技術(shù)更加普及。在H100 GPU上，文本到圖像生成僅需0.1 秒，ControlNet 任務(wù)耗時(shí)0.25 秒，實(shí)現(xiàn)了近乎實(shí)時(shí)的視覺反饋。
總結(jié)
與需要20 步以上的傳統(tǒng)擴(kuò)散模型不同，SANA-Sprint 僅需1-4 步即可生成高質(zhì)量圖像，且無需額外的訓(xùn)練過程。單步推理速度極快，非常適合實(shí)時(shí)應(yīng)用場景。兩步生成能夠在保證速度 (低于 0.25 秒)的前提下，有效提升圖像細(xì)節(jié)。四步生成則在質(zhì)量和效率之間實(shí)現(xiàn)了最佳平衡。
該論文在數(shù)學(xué)原理上具有一定的復(fù)雜性，但其技術(shù)方案堪稱杰出，非常值得深入閱讀和研究。SANA-Sprint 的工作有望推動(dòng)Flow Matching DiT 模型的下游優(yōu)化，進(jìn)而實(shí)現(xiàn)更快、更低成本的圖像生成。
蒸餾推理技術(shù)的進(jìn)步，使得高質(zhì)量圖像生成技術(shù)更加普惠化。
https://avoid.overfit.cn/post/c9690cdfa56046e7833462825ef93352
作者：Pietro Bolcato
熱門跟貼