文章鏈接：https://arxiv.org/pdf/2511.08930

亮點(diǎn)直擊

• 系統(tǒng)性分析與統(tǒng)一視角：對(duì)軌跡蒸餾（TD）進(jìn)行了系統(tǒng)性分析，揭示了其本質(zhì)是一種有損壓縮過(guò)程。這一視角解釋了為何TD方法雖然能有效保留全局結(jié)構(gòu)，卻不可避免地會(huì)犧牲精細(xì)細(xì)節(jié)。
• 創(chuàng)新的分層蒸餾框架：重新審視軌跡蒸餾和分布蒸餾的角色，提出一個(gè)新穎的分層蒸餾（Hierarchical Distillation, HD）框架。該框架協(xié)同利用兩種方法的優(yōu)勢(shì)，先構(gòu)建結(jié)構(gòu)，再優(yōu)化細(xì)節(jié)。
• 專為細(xì)節(jié)優(yōu)化的判別器：為配合HD框架，設(shè)計(jì)了自適應(yīng)加權(quán)判別器（Adaptive Weighted Discriminator, AWD），一種新穎的對(duì)抗機(jī)制，專為優(yōu)化高質(zhì)量生成器的局部瑕疵而設(shè)計(jì)，從而顯著提升最終生成質(zhì)量。

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

解決的問(wèn)題

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

1. 軌跡蒸餾（Trajectory-based Distillation, TD）：此類方法能很好地保留生成內(nèi)容的全局結(jié)構(gòu)，但由于其“有損壓縮”的特性，會(huì)犧牲高頻細(xì)節(jié)，導(dǎo)致保真度下降。
2. 分布蒸餾（Distribution-based Distillation）：此類方法理論上可以達(dá)到更高的保真度，但常常受困于模式崩潰（mode collapse）和訓(xùn)練不穩(wěn)定的問(wèn)題，尤其是在初始分布不佳的情況下。 本文的工作旨在克服這兩類方法的固有缺陷，將它們的優(yōu)勢(shì)結(jié)合起來(lái)，實(shí)現(xiàn)高保真、高效率的單步生成。

提出的方案

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

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

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

• 第二階段：分布優(yōu)化
• 將第一階段預(yù)訓(xùn)練好的模型作為生成器，進(jìn)行分布匹配（Distribution Matching, DM）的微調(diào)。
• 此階段引入對(duì)抗性訓(xùn)練來(lái)恢復(fù)在第一階段丟失的高頻細(xì)節(jié)并避免模式崩潰。
• 為解決傳統(tǒng)判別器在面對(duì)高質(zhì)量生成器時(shí)難以提供有效監(jiān)督信號(hào)的問(wèn)題，本文設(shè)計(jì)了自適應(yīng)加權(quán)判別器（AWD）。AWD通過(guò)注意力機(jī)制動(dòng)態(tài)地為特征圖上的不同空間位置（token）分配權(quán)重，使判別器能更專注于局部瑕疵，從而為生成器的細(xì)節(jié)優(yōu)化提供更精準(zhǔn)的指導(dǎo)。

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

• 軌跡蒸餾（Trajectory Distillation, TD）：以MeanFlow作為實(shí)現(xiàn)方式，用于第一階段的結(jié)構(gòu)化初始化。
• 分布匹配蒸餾（Distribution Matching Distillation, DMD）：用于第二階段的細(xì)節(jié)優(yōu)化，旨在將生成分布與真實(shí)數(shù)據(jù)分布對(duì)齊。
• 對(duì)抗性訓(xùn)練（Adversarial Training）：在第二階段引入，以穩(wěn)定訓(xùn)練并減輕模式崩潰問(wèn)題。
• 自適應(yīng)加權(quán)判別器（Adaptive Weighted Discriminator, AWD）：本文提出的核心技術(shù)之一，通過(guò)可學(xué)習(xí)的查詢嵌入（query embedding）和注意力機(jī)制來(lái)動(dòng)態(tài)加權(quán)特征，從而聚焦于局部偽影的判別。

達(dá)到的效果

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

• 在ImageNet 256×256的類條件生成任務(wù)上，本文的單步模型達(dá)到了2.26的FID分?jǐn)?shù)，這一成績(jī)不僅在單步模型中領(lǐng)先，甚至可以媲美其250步的教師模型（FID為2.27）。
• 在MJHQ-30K高分辨率文生圖基準(zhǔn)測(cè)試上，單步和兩步模型的FID和CLIP分?jǐn)?shù)均優(yōu)于現(xiàn)有的其他蒸餾方法，證明了其強(qiáng)大的泛化能力。

方法

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)

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

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

在主要對(duì)比實(shí)驗(yàn)中，本文在ImageNet 256×256和文生圖（MJHQ-30K）兩個(gè)基準(zhǔn)上驗(yàn)證了HD框架的有效性。

• 在ImageNet上，HD單步模型的FID達(dá)到了2.26，優(yōu)于所有其他單步方法（如MeanFlow的3.43，DMD的6.63），并且?guī)缀跖c250步的教師模型（FID 2.27）持平，實(shí)現(xiàn)了巨大的推理加速（約70倍）而幾乎沒(méi)有性能損失。
• 在文生圖任務(wù)上，HD在單步和兩步設(shè)置下的FID和CLIP分?jǐn)?shù)均優(yōu)于包括SDXL-LCM、DMD2和MeanFlow在內(nèi)的現(xiàn)有方法，證明了其方法的普適性和卓越性能。

消融實(shí)驗(yàn)進(jìn)一步剖析了HD框架各個(gè)組件的貢獻(xiàn)。結(jié)果表明：

1. TD初始化至關(guān)重要：與沒(méi)有TD初始化的模型相比，經(jīng)過(guò)TD初始化的模型性能有顯著提升（FID從5.49提升到3.09），證實(shí)了提供一個(gè)高質(zhì)量的結(jié)構(gòu)先驗(yàn)是成功的關(guān)鍵。
2. AWD的有效性：與使用標(biāo)準(zhǔn)全局平均池化（GAP）的判別器相比，本文提出的自適應(yīng)加權(quán)判別器（AWD）將FID從3.09進(jìn)一步降低到2.26，證明了AWD在聚焦局部瑕疵、指導(dǎo)模型優(yōu)化細(xì)節(jié)方面的優(yōu)越性。

結(jié)論

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

本文轉(zhuǎn)自AI生成未來(lái) ，作者：AI生成未來(lái)

原文鏈接:??https://mp.weixin.qq.com/s/Jc6GLiDEt2oAoA0fs1I-Yg??