一步直接封神!單步擴散媲美250步教師模型!中科大&字節發布圖像生成“分層蒸餾術”
作者:Hanbo Cheng等
解讀:AI生成未來
亮點直擊
系統性分析與統一視角:對軌跡蒸餾(TD)進行了系統性分析,揭示了其本質是一種有損壓縮過程。這一視角解釋了為何TD方法雖然能有效保留全局結構,卻不可避免地會犧牲精細細節。
創新的分層蒸餾框架:重新審視軌跡蒸餾和分布蒸餾的角色,提出一個新穎的分層蒸餾(Hierarchical Distillation, HD)框架。該框架協同利用兩種方法的優勢,先構建結構,再優化細節。
專為細節優化的判別器:為配合HD框架,設計了自適應加權判別器(Adaptive Weighted Discriminator, AWD),一種新穎的對抗機制,專為優化高質量生成器的局部瑕疵而設計,從而顯著提升最終生成質量。
圖 1.50 步教師 SANA和我們的 1 步高清方法的生成質量比較。本文方法達到了與多步驟教師相當的質量。
解決的問題
本文旨在解決擴散模型推理延遲過高的問題,特別是在單步或少步生成場景下。現有的加速方法主要分為兩類:
軌跡蒸餾(Trajectory-based Distillation, TD):此類方法能很好地保留生成內容的全局結構,但由于其“有損壓縮”的特性,會犧牲高頻細節,導致保真度下降。
分布蒸餾(Distribution-based Distillation):此類方法理論上可以達到更高的保真度,但常常受困于模式崩潰(mode collapse)和訓練不穩定的問題,尤其是在初始分布不佳的情況下。 本文的工作旨在克服這兩類方法的固有缺陷,將它們的優勢結合起來,實現高保真、高效率的單步生成。
提出的方案
本文提出了一個名為分層蒸餾(Hierarchical Distillation, HD)的兩階段框架,其核心思想是“先搭骨架,再填血肉”。
第一階段:結構化初始化
利用基于軌跡蒸餾的方法(具體采用MeanFlow),將一個多步教師模型的結構先驗知識“注入”到學生模型中。此階段的目標不是生成最終結果,而是為學生模型提供一個結構合理、接近真實數據流形的“草圖”或高質量的初始分布。這有效穩定了后續的訓練過程。
第二階段:分布優化
將第一階段預訓練好的模型作為生成器,進行分布匹配(Distribution Matching, DM)的微調。
此階段引入對抗性訓練來恢復在第一階段丟失的高頻細節并避免模式崩潰。
為解決傳統判別器在面對高質量生成器時難以提供有效監督信號的問題,本文設計了自適應加權判別器(AWD)。AWD通過注意力機制動態地為特征圖上的不同空間位置(token)分配權重,使判別器能更專注于局部瑕疵,從而為生成器的細節優化提供更精準的指導。
應用的技術點
軌跡蒸餾(Trajectory Distillation, TD):以MeanFlow作為實現方式,用于第一階段的結構化初始化。
分布匹配蒸餾(Distribution Matching Distillation, DMD):用于第二階段的細節優化,旨在將生成分布與真實數據分布對齊。
對抗性訓練(Adversarial Training):在第二階段引入,以穩定訓練并減輕模式崩潰問題。
自適應加權判別器(Adaptive Weighted Discriminator, AWD):本文提出的核心技術之一,通過可學習的查詢嵌入(query embedding)和注意力機制來動態加權特征,從而聚焦于局部偽影的判別。
達到的效果
本文的方法在多個任務上均取得了當前最優(SOTA)的性能。
在ImageNet 256×256的類條件生成任務上,本文的單步模型達到了2.26的FID分數,這一成績不僅在單步模型中領先,甚至可以媲美其250步的教師模型(FID為2.27)。
在MJHQ-30K高分辨率文生圖基準測試上,單步和兩步模型的FID和CLIP分數均優于現有的其他蒸餾方法,證明了其強大的泛化能力。
方法
本節介紹分層蒸餾(HD)框架的技術細節首先進行理論分析,統一主流的軌跡蒸餾(TD)方法,揭示它們共同的局限性,以此作為我們方法的動機。隨后,詳細介紹我們流水線的第一階段,其中基于MeanFlow的TD階段為學生模型注入了強大的結構先驗。最后,描述了第二階段,在這一階段,對這個良好初始化的模型應用分布匹配,對其進行優化以實現高保真度的結果。
軌跡蒸餾的統一視角
本節進行理論分析以闡明軌跡蒸餾(TD)的建模目標。通過數學推導,證明了幾種主流TD方法的目標,包括一致性模型(CM/sCM)和漸進式蒸餾(PGD),可以統一在平均速度估計的共同框架下。基于這一觀察,識別出大多數TD方法固有的一個共同局限。
命題 1.連續一致性模型隱式地建模了區間【0,t】上的平均速度。
證明. 一致性模型的核心原則是強制網絡輸出 在任何給定的PF-ODE軌跡上保持一致性。這種一致性約束的微分形式可以表示為(詳細推導見附錄6):

其中,在極限dt→0時,這個離散關系產生了微分形式:

回想一下方程(5)中瞬時速度和平均速度之間的關系。具體來說,對于從 開始的區間,它變為:

這揭示了當 dt→0 時,CM網絡的輸出 被隱式地訓練來建模區間 上的平均速度,即
。
命題 2.隨著蒸餾步驟數趨于無窮,漸進式蒸餾(PGD)收斂于對整個區間
上的平均速度進行建模。
證明. 漸進式蒸餾(PGD)是一個迭代過程,它在N輪中將一個
步的教師模型蒸餾成一個
步的學生模型。在每一輪k中,學生模型
被訓練來預測其教師模型(
)在兩個連續時間步上的輸出的平均值。
經過N輪蒸餾后,最終的單步學生模型
的輸出可以表示為原始多步教師模型在
個離散時間步上的輸出的算術平均值(詳細推導見附錄6):

其中
是離散的時間步。原始的教師模型
被訓練來近似瞬時速度,即
當蒸餾輪數
時,該和收斂為一個積分:

根據定義,方程(13)的右側是整個區間 上的平均速度 。這表明PGD也隱式地試圖建模平均速度。
統一視角的影響. 統一分析產生了兩個關鍵的見解,它們構成了我們所提出方法的理論基礎。(1)我們的證明揭示了主流TD方法,無論其具體公式如何,都共享一個共同的本質:它們訓練一個單一的神經網絡 來近似一個動態的平均速度函數 。作為一個軌跡段上的積分,這個平均速度函數封裝了來自多步教師模型PF-ODE路徑的豐富高頻動態信息。強迫一個有限容量的學生模型在單步中完美復制一個復雜的函數,從信息論的角度引入了一個根本性的瓶頸。這為所有單步TD方法都不可避免地遭受細粒度細節損失提供了理論解釋。我們在5.1節中為這一主張提供了進一步的經驗驗證。(2)因為所有TD方法都共享建模平均速度這一基本目標,我們選擇MeanFlow作為它們的代表性實現。這個決定不僅是基于其最先進的性能。更重要的是,MeanFlow明確且直接地對平均速度進行建模,提供了一個數學上優雅且穩健的實現。
第一階段:通過TD進行結構化初始化
如前所述,從零開始應用分布匹配蒸餾(DMD)進行單步生成面臨訓練不穩定和模式崩潰的問題。一個主要原因是生成分布和真實數據分布之間缺乏重疊。為了解決這個問題,我們引入了一個結構化初始化階段。利用軌跡蒸餾(TD)來有效地將多步教師模型積累的豐富結構先驗注入到學生模型中。這確保了在分布匹配階段開始之前,學生模型就已經具備了捕捉目標分布宏觀結構和布局的強大能力。基于前面的分析,采用MeanFlow作為我們TD階段的蒸餾目標。盡管MeanFlow最初是為從零開始訓練模型而提出的,但我們認為將其重新用作蒸餾框架可以提供一個方差更低的學習信號。從零開始訓練時,模型從數據和噪聲的隨機配對中學習,其中每個樣本都呈現一個獨特的、高方差的目標。相比之下,蒸餾利用了一個已經收斂到從噪聲到數據的固定、確定性映射的預訓練教師模型。來自教師的這種指導確保了學習目標在訓練期間是一致的,從而降低了梯度信號的方-差,并導致一個更穩定和高效的初始化階段。
圖 2.分層蒸餾 (HD) 管道。我們的方法包括兩個主要階段:(1) 結構化初始化:基于 MeanFlow 的方法為學生灌輸基礎結構信息。(2) 分布細化:第二階段采用專為 HD 框架設計的自適應加權判別器 (AWD),恢復高頻細節。SN "和 "LN "分別指譜規范[22]和層規范
這個階段的流程圖如圖2的上半部分所示。在我們的蒸餾框架中,將(最初在標準MeanFlow中由線性插值導出的)真實瞬時速度場 替換為預訓練教師模型 的輸出。這直接引導學生學習教師的軌跡動態。具體來說,我們使用教師模型上的無分類器指導(CFG)來定義瞬時速度場:

其中
是指導尺度,
分別代表條件和無條件輸入。通過將這個由教師定義的速度場代入MeanFlow的訓練目標(方程(6)),我們構建了我們的蒸餾損失。這個階段的產物是一個被賦予了教師結構先驗的學生生成器。盡管其在最少步數下的保真度不完美,但它為后續的分布匹配和優化提供了一個適定(well-posed)的初始化。
第二階段:分布優化
在通過第一階段模型 初始化后,生成器 接著進入第二階段的分布優化,恢復那些在僅從教師軌跡學習時固有丟失的高頻細節。我們采用基于DMD的策略來將 的單步輸出分布與真實數據分布對齊。由于 的初始分布已經占據了數據流形上的一個有利區域,與真實分布有顯著重疊,DMD的訓練過程更加穩定和高效。其主要任務從“盲目探索”轉變為“有針對性地優化細節”。對于DMD中的分數網絡,盡管MeanFlow學生模型本身可以預測瞬時速度,我們仍然使用預訓練的教師模型 來初始化真實和偽造分數分支。這可以防止潛在的誤差累積,并提供更準確的速度場估計。損失函數如方程(8)所示。
為了進一步穩定訓練并減輕模式崩潰的風險,引入了對抗性訓練策略。我們引入一個判別器D,它不在高維像素空間中進行判別,而是在教師模型 的特征空間中操作,這遵循了[29]的做法。總的對抗性損失由生成器和判別器損失 , 組成:
其中 表示特征提取函數,它接收圖像、時間和條件作為輸入,并返回教師模型 的中間特征。這里, 是一個帶噪聲的真實圖像,而 是一個生成的樣本。
最終,總損失公式為:

自適應加權判別器. 經過TD初始化后,學生模型已經捕捉到了目標分布的整體結構。不完美之處不再是全局性的,而是表現為微妙的、局部化的偽影。這使得依賴全局平均池化(GAP)的傳統判別器在很大程度上失效。為了應對這一挑戰,我們設計了自適應加權判別器(AWD),如圖2底部所示。我們的判別器不是為所有令牌(token)分配統一的權重,而是采用一個可學習的查詢嵌入和一個注意力機制來動態地加權特征圖上的不同令牌。因此,判別器可以專注于最可能包含偽影的局部區域,為生成器提供更精確和有效的梯度。
通過這個分層框架訓練出的最終學生模型,能夠在最少步數下生成與多步教師模型質量相媲美的圖像,同時保持多樣性。
實驗
實驗部分首先通過一個二維玩具實驗驗證了其核心理論假設:軌跡蒸餾(TD)存在信息瓶頸,其性能上限受限于學生模型的能力。實驗表明,增加模型容量能顯著提升單步學生模型的性能,但即使容量增加50倍以上,也無法完美復制多步教師模型的軌跡,證明了僅靠TD不足以實現最優的單步生成質量,必須有后續的優化階段。
圖 3.軌跡蒸餾 (TD) 性能與模型大小的關系。TD 性能的上限隨著模型參數數量的增加而增加。
在主要對比實驗中,本文在ImageNet 256×256和文生圖(MJHQ-30K)兩個基準上驗證了HD框架的有效性。
在ImageNet上,HD單步模型的FID達到了2.26,優于所有其他單步方法(如MeanFlow的3.43,DMD的6.63),并且幾乎與250步的教師模型(FID 2.27)持平,實現了巨大的推理加速(約70倍)而幾乎沒有性能損失。
在文生圖任務上,HD在單步和兩步設置下的FID和CLIP分數均優于包括SDXL-LCM、DMD2和MeanFlow在內的現有方法,證明了其方法的普適性和卓越性能。
消融實驗進一步剖析了HD框架各個組件的貢獻。結果表明:
TD初始化至關重要:與沒有TD初始化的模型相比,經過TD初始化的模型性能有顯著提升(FID從5.49提升到3.09),證實了提供一個高質量的結構先驗是成功的關鍵。
AWD的有效性:與使用標準全局平均池化(GAP)的判別器相比,本文提出的自適應加權判別器(AWD)將FID從3.09進一步降低到2.26,證明了AWD在聚焦局部瑕疵、指導模型優化細節方面的優越性。


結論
本工作為軌跡蒸餾(TD)提出了一個統一的理論公式,識別出一個共同的“平均速度”建模目標,該目標導致了一個信息論瓶頸。這一分析揭示了為何TD擅長于全局結構,卻在根本上難以處理精細細節。受此啟發,提出了一個新穎的分層蒸餾(HD)框架,該框架協同地結合了TD和分布匹配。本文方法首先利用TD作為一個強大的初始化器,從教師模型中注入豐富的結構先驗,為學生模型建立一個適定(well-posed)的起點。隨后,通過分布匹配來優化這個強大的初始模型。為了增強這一階段,引入了一個量身定制的對抗性訓練過程,并配備了本文提出的自適應加權判別器(AWD)。通過動態地關注良好初始化模型的局部偽影,它為細節優化提供了更精確的指導。大量的實驗表明,本文的單步學生模型顯著優于現有的蒸餾方法,并實現了與其多步教師模型相當的保真度。通過診斷并克服TD的瓶頸,本工作為少步乃至單步高保真度生成提供了一個有效的新范式。
參考文獻
[1] From Structure to Detail: Hierarchical Distillation for Efficient Diffusion Model
原文標題 : 一步直接封神!單步擴散媲美250步教師模型!中科大&字節發布圖像生成“分層蒸餾術”
請輸入評論內容...
請輸入評論/評論長度6~500個字
最新活動更多
- 1 特斯拉Optimus Gen3量產在即,哪些環節最具確定性?
- 2 OpenAI深夜王炸!ChatGPT Images 2.0實測:中文穩、細節炸,設計師慌了
- 3 AI狂歡遇上油價破百,全球股市還能漲多久? | 產聯看全球
- 4 6000億美元估值錨定:字節跳動的“去單一化”突圍與估值重構
- 5 Tesla AI5芯片最新進展總結
- 6 連夜測了一波DeepSeek-V4,我發現它可能只剩“審美”這個短板了
- 7 熱點丨AI“瑜亮之爭”:既生OpenClaw,何生Hermes?
- 8 2026,人形機器人只贏了面子
- 9 AI界的殺豬盤:9秒刪庫跑路,全員被封號,還繼續扣錢!
- 10 AI Infra產業鏈卡在哪里了?


分享













