作者：Hanbo Cheng等

解讀：AI生成未來

亮點直擊

系統(tǒng)性分析與統(tǒng)一視角：對軌跡蒸餾（TD）進行了系統(tǒng)性分析，揭示了其本質(zhì)是一種有損壓縮過程。這一視角解釋了為何TD方法雖然能有效保留全局結(jié)構(gòu)，卻不可避免地會犧牲精細細節(jié)。

創(chuàng)新的分層蒸餾框架：重新審視軌跡蒸餾和分布蒸餾的角色，提出一個新穎的分層蒸餾（Hierarchical Distillation, HD）框架。該框架協(xié)同利用兩種方法的優(yōu)勢，先構(gòu)建結(jié)構(gòu)，再優(yōu)化細節(jié)。

專為細節(jié)優(yōu)化的判別器：為配合HD框架，設(shè)計了自適應(yīng)加權(quán)判別器（Adaptive Weighted Discriminator, AWD），一種新穎的對抗機制，專為優(yōu)化高質(zhì)量生成器的局部瑕疵而設(shè)計，從而顯著提升最終生成質(zhì)量。

圖 1.50 步教師 SANA和我們的 1 步高清方法的生成質(zhì)量比較。本文方法達到了與多步驟教師相當(dāng)?shù)馁|(zhì)量。

解決的問題

本文旨在解決擴散模型推理延遲過高的問題，特別是在單步或少步生成場景下。現(xiàn)有的加速方法主要分為兩類：

軌跡蒸餾（Trajectory-based Distillation, TD）：此類方法能很好地保留生成內(nèi)容的全局結(jié)構(gòu)，但由于其“有損壓縮”的特性，會犧牲高頻細節(jié)，導(dǎo)致保真度下降。

分布蒸餾（Distribution-based Distillation）：此類方法理論上可以達到更高的保真度，但常常受困于模式崩潰（mode collapse）和訓(xùn)練不穩(wěn)定的問題，尤其是在初始分布不佳的情況下。 本文的工作旨在克服這兩類方法的固有缺陷，將它們的優(yōu)勢結(jié)合起來，實現(xiàn)高保真、高效率的單步生成。

提出的方案

本文提出了一個名為分層蒸餾（Hierarchical Distillation, HD）的兩階段框架，其核心思想是“先搭骨架，再填血肉”。

第一階段：結(jié)構(gòu)化初始化

利用基于軌跡蒸餾的方法（具體采用MeanFlow），將一個多步教師模型的結(jié)構(gòu)先驗知識“注入”到學(xué)生模型中。此階段的目標不是生成最終結(jié)果，而是為學(xué)生模型提供一個結(jié)構(gòu)合理、接近真實數(shù)據(jù)流形的“草圖”或高質(zhì)量的初始分布。這有效穩(wěn)定了后續(xù)的訓(xùn)練過程。

第二階段：分布優(yōu)化

將第一階段預(yù)訓(xùn)練好的模型作為生成器，進行分布匹配（Distribution Matching, DM）的微調(diào)。

此階段引入對抗性訓(xùn)練來恢復(fù)在第一階段丟失的高頻細節(jié)并避免模式崩潰。

為解決傳統(tǒng)判別器在面對高質(zhì)量生成器時難以提供有效監(jiān)督信號的問題，本文設(shè)計了自適應(yīng)加權(quán)判別器（AWD）。AWD通過注意力機制動態(tài)地為特征圖上的不同空間位置（token）分配權(quán)重，使判別器能更專注于局部瑕疵，從而為生成器的細節(jié)優(yōu)化提供更精準的指導(dǎo)。

應(yīng)用的技術(shù)點

軌跡蒸餾（Trajectory Distillation, TD）：以MeanFlow作為實現(xiàn)方式，用于第一階段的結(jié)構(gòu)化初始化。

分布匹配蒸餾（Distribution Matching Distillation, DMD）：用于第二階段的細節(jié)優(yōu)化，旨在將生成分布與真實數(shù)據(jù)分布對齊。

對抗性訓(xùn)練（Adversarial Training）：在第二階段引入，以穩(wěn)定訓(xùn)練并減輕模式崩潰問題。

自適應(yīng)加權(quán)判別器（Adaptive Weighted Discriminator, AWD）：本文提出的核心技術(shù)之一，通過可學(xué)習(xí)的查詢嵌入（query embedding）和注意力機制來動態(tài)加權(quán)特征，從而聚焦于局部偽影的判別。

達到的效果

本文的方法在多個任務(wù)上均取得了當(dāng)前最優(yōu)（SOTA）的性能。

在ImageNet 256×256的類條件生成任務(wù)上，本文的單步模型達到了2.26的FID分數(shù)，這一成績不僅在單步模型中領(lǐng)先，甚至可以媲美其250步的教師模型（FID為2.27）。

在MJHQ-30K高分辨率文生圖基準測試上，單步和兩步模型的FID和CLIP分數(shù)均優(yōu)于現(xiàn)有的其他蒸餾方法，證明了其強大的泛化能力。

方法

本節(jié)介紹分層蒸餾（HD）框架的技術(shù)細節(jié)首先進行理論分析，統(tǒng)一主流的軌跡蒸餾（TD）方法，揭示它們共同的局限性，以此作為我們方法的動機。隨后，詳細介紹我們流水線的第一階段，其中基于MeanFlow的TD階段為學(xué)生模型注入了強大的結(jié)構(gòu)先驗。最后，描述了第二階段，在這一階段，對這個良好初始化的模型應(yīng)用分布匹配，對其進行優(yōu)化以實現(xiàn)高保真度的結(jié)果。

軌跡蒸餾的統(tǒng)一視角

本節(jié)進行理論分析以闡明軌跡蒸餾（TD）的建模目標。通過數(shù)學(xué)推導(dǎo)，證明了幾種主流TD方法的目標，包括一致性模型（CM/sCM）和漸進式蒸餾（PGD），可以統(tǒng)一在平均速度估計的共同框架下�；�于這一觀察，識別出大多數(shù)TD方法固有的一個共同局限。

命題 1.連續(xù)一致性模型隱式地建模了區(qū)間【0，t】上的平均速度。

證明. 一致性模型的核心原則是強制網(wǎng)絡(luò)輸出  在任何給定的PF-ODE軌跡上保持一致性。這種一致性約束的微分形式可以表示為（詳細推導(dǎo)見附錄6）：

其中，在極限dt→0時，這個離散關(guān)系產(chǎn)生了微分形式：

回想一下方程(5)中瞬時速度和平均速度之間的關(guān)系。具體來說，對于從  開始的區(qū)間，它變?yōu)椋?/p>

這揭示了當(dāng) dt→0 時，CM網(wǎng)絡(luò)的輸出  被隱式地訓(xùn)練來建模區(qū)間  上的平均速度，即 。

命題 2.隨著蒸餾步驟數(shù)趨于無窮，漸進式蒸餾（PGD）收斂于對整個區(qū)間 上的平均速度進行建模。

證明. 漸進式蒸餾（PGD）是一個迭代過程，它在N輪中將一個 步的教師模型蒸餾成一個步的學(xué)生模型。在每一輪k中，學(xué)生模型被訓(xùn)練來預(yù)測其教師模型（）在兩個連續(xù)時間步上的輸出的平均值。

經(jīng)過N輪蒸餾后，最終的單步學(xué)生模型的輸出可以表示為原始多步教師模型在個離散時間步上的輸出的算術(shù)平均值（詳細推導(dǎo)見附錄6）：

其中 是離散的時間步。原始的教師模型被訓(xùn)練來近似瞬時速度，即 當(dāng)蒸餾輪數(shù)  時，該和收斂為一個積分：

根據(jù)定義，方程(13)的右側(cè)是整個區(qū)間  上的平均速度 。這表明PGD也隱式地試圖建模平均速度。

統(tǒng)一視角的影響. 統(tǒng)一分析產(chǎn)生了兩個關(guān)鍵的見解，它們構(gòu)成了我們所提出方法的理論基礎(chǔ)。（1）我們的證明揭示了主流TD方法，無論其具體公式如何，都共享一個共同的本質(zhì)：它們訓(xùn)練一個單一的神經(jīng)網(wǎng)絡(luò)  來近似一個動態(tài)的平均速度函數(shù) 。作為一個軌跡段上的積分，這個平均速度函數(shù)封裝了來自多步教師模型PF-ODE路徑的豐富高頻動態(tài)信息。強迫一個有限容量的學(xué)生模型在單步中完美復(fù)制一個復(fù)雜的函數(shù)，從信息論的角度引入了一個根本性的瓶頸。這為所有單步TD方法都不可避免地遭受細粒度細節(jié)損失提供了理論解釋。我們在5.1節(jié)中為這一主張?zhí)峁┝诉M一步的經(jīng)驗驗證。（2）因為所有TD方法都共享建模平均速度這一基本目標，我們選擇MeanFlow作為它們的代表性實現(xiàn)。這個決定不僅是基于其最先進的性能。更重要的是，MeanFlow明確且直接地對平均速度進行建模，提供了一個數(shù)學(xué)上優(yōu)雅且穩(wěn)健的實現(xiàn)。

第一階段：通過TD進行結(jié)構(gòu)化初始化

如前所述，從零開始應(yīng)用分布匹配蒸餾（DMD）進行單步生成面臨訓(xùn)練不穩(wěn)定和模式崩潰的問題。一個主要原因是生成分布和真實數(shù)據(jù)分布之間缺乏重疊。為了解決這個問題，我們引入了一個結(jié)構(gòu)化初始化階段。利用軌跡蒸餾（TD）來有效地將多步教師模型積累的豐富結(jié)構(gòu)先驗注入到學(xué)生模型中。這確保了在分布匹配階段開始之前，學(xué)生模型就已經(jīng)具備了捕捉目標分布宏觀結(jié)構(gòu)和布局的強大能力�；�于前面的分析，采用MeanFlow作為我們TD階段的蒸餾目標。盡管MeanFlow最初是為從零開始訓(xùn)練模型而提出的，但我們認為將其重新用作蒸餾框架可以提供一個方差更低的學(xué)習(xí)信號。從零開始訓(xùn)練時，模型從數(shù)據(jù)和噪聲的隨機配對中學(xué)習(xí)，其中每個樣本都呈現(xiàn)一個獨特的、高方差的目標。相比之下，蒸餾利用了一個已經(jīng)收斂到從噪聲到數(shù)據(jù)的固定、確定性映射的預(yù)訓(xùn)練教師模型。來自教師的這種指導(dǎo)確保了學(xué)習(xí)目標在訓(xùn)練期間是一致的，從而降低了梯度信號的方-差，并導(dǎo)致一個更穩(wěn)定和高效的初始化階段。

圖 2.分層蒸餾 (HD) 管道。我們的方法包括兩個主要階段：(1) 結(jié)構(gòu)化初始化：基于 MeanFlow 的方法為學(xué)生灌輸基礎(chǔ)結(jié)構(gòu)信息。(2) 分布細化：第二階段采用專為 HD 框架設(shè)計的自適應(yīng)加權(quán)判別器 (AWD)，恢復(fù)高頻細節(jié)。SN "和 "LN "分別指譜規(guī)范[22]和層規(guī)范

這個階段的流程圖如圖2的上半部分所示。在我們的蒸餾框架中，將（最初在標準MeanFlow中由線性插值導(dǎo)出的）真實瞬時速度場  替換為預(yù)訓(xùn)練教師模型  的輸出。這直接引導(dǎo)學(xué)生學(xué)習(xí)教師的軌跡動態(tài)。具體來說，我們使用教師模型上的無分類器指導(dǎo)（CFG）來定義瞬時速度場：

其中是指導(dǎo)尺度， 分別代表條件和無條件輸入。通過將這個由教師定義的速度場代入MeanFlow的訓(xùn)練目標（方程(6)），我們構(gòu)建了我們的蒸餾損失。這個階段的產(chǎn)物是一個被賦予了教師結(jié)構(gòu)先驗的學(xué)生生成器。盡管其在最少步數(shù)下的保真度不完美，但它為后續(xù)的分布匹配和優(yōu)化提供了一個適定（well-posed）的初始化。

第二階段：分布優(yōu)化

在通過第一階段模型  初始化后，生成器  接著進入第二階段的分布優(yōu)化，恢復(fù)那些在僅從教師軌跡學(xué)習(xí)時固有丟失的高頻細節(jié)。我們采用基于DMD的策略來將  的單步輸出分布與真實數(shù)據(jù)分布對齊。由于  的初始分布已經(jīng)占據(jù)了數(shù)據(jù)流形上的一個有利區(qū)域，與真實分布有顯著重疊，DMD的訓(xùn)練過程更加穩(wěn)定和高效。其主要任務(wù)從“盲目探索”轉(zhuǎn)變?yōu)?ldquo;有針對性地優(yōu)化細節(jié)”。對于DMD中的分數(shù)網(wǎng)絡(luò)，盡管MeanFlow學(xué)生模型本身可以預(yù)測瞬時速度，我們?nèi)匀皇褂妙A(yù)訓(xùn)練的教師模型  來初始化真實和偽造分數(shù)分支。這可以防止?jié)撛诘恼`差累積，并提供更準確的速度場估計。損失函數(shù)如方程(8)所示。

為了進一步穩(wěn)定訓(xùn)練并減輕模式崩潰的風(fēng)險，引入了對抗性訓(xùn)練策略。我們引入一個判別器D，它不在高維像素空間中進行判別，而是在教師模型  的特征空間中操作，這遵循了[29]的做法�？偟膶�抗性損失由生成器和判別器損失 ,  組成：其中  表示特征提取函數(shù)，它接收圖像、時間和條件作為輸入，并返回教師模型  的中間特征。這里， 是一個帶噪聲的真實圖像，而  是一個生成的樣本。

最終，總損失公式為：

自適應(yīng)加權(quán)判別器. 經(jīng)過TD初始化后，學(xué)生模型已經(jīng)捕捉到了目標分布的整體結(jié)構(gòu)。不完美之處不再是全局性的，而是表現(xiàn)為微妙的、局部化的偽影。這使得依賴全局平均池化（GAP）的傳統(tǒng)判別器在很大程度上失效。為了應(yīng)對這一挑戰(zhàn)，我們設(shè)計了自適應(yīng)加權(quán)判別器（AWD），如圖2底部所示。我們的判別器不是為所有令牌（token）分配統(tǒng)一的權(quán)重，而是采用一個可學(xué)習(xí)的查詢嵌入和一個注意力機制來動態(tài)地加權(quán)特征圖上的不同令牌。因此，判別器可以專注于最可能包含偽影的局部區(qū)域，為生成器提供更精確和有效的梯度。

通過這個分層框架訓(xùn)練出的最終學(xué)生模型，能夠在最少步數(shù)下生成與多步教師模型質(zhì)量相媲美的圖像，同時保持多樣性。

實驗

實驗部分首先通過一個二維玩具實驗驗證了其核心理論假設(shè)：軌跡蒸餾（TD）存在信息瓶頸，其性能上限受限于學(xué)生模型的能力。實驗表明，增加模型容量能顯著提升單步學(xué)生模型的性能，但即使容量增加50倍以上，也無法完美復(fù)制多步教師模型的軌跡，證明了僅靠TD不足以實現(xiàn)最優(yōu)的單步生成質(zhì)量，必須有后續(xù)的優(yōu)化階段。

圖 3.軌跡蒸餾 (TD) 性能與模型大小的關(guān)系。TD 性能的上限隨著模型參數(shù)數(shù)量的增加而增加。

在主要對比實驗中，本文在ImageNet 256×256和文生圖（MJHQ-30K）兩個基準上驗證了HD框架的有效性。

在ImageNet上，HD單步模型的FID達到了2.26，優(yōu)于所有其他單步方法（如MeanFlow的3.43，DMD的6.63），并且?guī)缀跖c250步的教師模型（FID 2.27）持平，實現(xiàn)了巨大的推理加速（約70倍）而幾乎沒有性能損失。

在文生圖任務(wù)上，HD在單步和兩步設(shè)置下的FID和CLIP分數(shù)均優(yōu)于包括SDXL-LCM、DMD2和MeanFlow在內(nèi)的現(xiàn)有方法，證明了其方法的普適性和卓越性能。

消融實驗進一步剖析了HD框架各個組件的貢獻。結(jié)果表明：

TD初始化至關(guān)重要：與沒有TD初始化的模型相比，經(jīng)過TD初始化的模型性能有顯著提升（FID從5.49提升到3.09），證實了提供一個高質(zhì)量的結(jié)構(gòu)先驗是成功的關(guān)鍵。

AWD的有效性：與使用標準全局平均池化（GAP）的判別器相比，本文提出的自適應(yīng)加權(quán)判別器（AWD）將FID從3.09進一步降低到2.26，證明了AWD在聚焦局部瑕疵、指導(dǎo)模型優(yōu)化細節(jié)方面的優(yōu)越性。

結(jié)論

本工作為軌跡蒸餾（TD）提出了一個統(tǒng)一的理論公式，識別出一個共同的“平均速度”建模目標，該目標導(dǎo)致了一個信息論瓶頸。這一分析揭示了為何TD擅長于全局結(jié)構(gòu)，卻在根本上難以處理精細細節(jié)。受此啟發(fā)，提出了一個新穎的分層蒸餾（HD）框架，該框架協(xié)同地結(jié)合了TD和分布匹配。本文方法首先利用TD作為一個強大的初始化器，從教師模型中注入豐富的結(jié)構(gòu)先驗，為學(xué)生模型建立一個適定（well-posed）的起點。隨后，通過分布匹配來優(yōu)化這個強大的初始模型。為了增強這一階段，引入了一個量身定制的對抗性訓(xùn)練過程，并配備了本文提出的自適應(yīng)加權(quán)判別器（AWD）。通過動態(tài)地關(guān)注良好初始化模型的局部偽影，它為細節(jié)優(yōu)化提供了更精確的指導(dǎo)。大量的實驗表明，本文的單步學(xué)生模型顯著優(yōu)于現(xiàn)有的蒸餾方法，并實現(xiàn)了與其多步教師模型相當(dāng)?shù)谋Ｕ娑�。通過診斷并克服TD的瓶頸，本工作為少步乃至單步高保真度生成提供了一個有效的新范式。

參考文獻

[1] From Structure to Detail: Hierarchical Distillation for Efficient Diffusion Model

       原文標題 : 一步直接封神！單步擴散媲美250步教師模型！中科大&字節(jié)發(fā)布圖像生成“分層蒸餾術(shù)”