• 論文標題：On the Design of One-step Diffusion via Shortcutting Flow Paths
• 論文地址: https://openreview.net/forum?id=k6q8rRYVQR
• 代碼開源: https://github.com/EDAPINENUT/ExplicitShortCut/
• 項目主頁: https://edapinenut.github.io/explicitshortcut-project-page/

前言

近期，基于捷徑化概率流路徑（shortcut probability flow trajectory）并從頭訓練的一步擴散生成模型，展現(xiàn)出強大的實證有效性。然而，這類方法的提出通常建立在較為復雜的理論推導之上，并且往往與具體實現(xiàn)細節(jié)高度耦合。這帶來一個直接的問題：究竟哪些設(shè)計是方法成立的本質(zhì)要素，哪些又只是可以靈活替換的實現(xiàn)組件。

對于研究人員，這種復雜性進一步體現(xiàn)在多個層面：一方面是深奧的理論背景，繁復的方法形式推導、對應(yīng)的學習目標設(shè)計，另一方面，是應(yīng)用實現(xiàn)上的技術(shù)細節(jié)，包括時間采樣器、課程學習方法等。面對這些內(nèi)容，研究人員往往很難迅速抓住不同方法之間真正共享的核心思想，也因此缺少一種足夠直觀的設(shè)計范式，來理解這類一步生成模型到底 “為什么這樣設(shè)計”。

為了解決這個問題，西湖大學研究人員首先從一個更實用的角度出發(fā)，嘗試為這類 shortcut 模型提出一個統(tǒng)一的設(shè)計框架。其指出，無論是離散時間還是連續(xù)時間的一步生成方法，其背后都遵循著同一個基本原則：用單步參數(shù)化預(yù)測去逼近一個兩步擴散構(gòu)建的目標。在此基礎(chǔ)上，該工作進一步給出了這一設(shè)計范式成立的一般性理論說明。這樣的統(tǒng)一視角，使得原本耦合在一起的生成方法模塊能夠被拆解開來，以供研究人員更清晰地理解各個組件之間是如何協(xié)同作用的，以及在整體方法設(shè)計中究竟還保留了多少可調(diào)整的自由度。

得益于這一框架，本文對當前具有代表性的離散時間與連續(xù)時間一步擴散生成方法進行更細粒度的組件級分析，包括路徑選擇、時間采樣器等關(guān)鍵模塊。通過這一分析，本文進一步定位了影響訓練 — 推理穩(wěn)定性以及逼近誤差的一個關(guān)鍵來源：損失函數(shù)設(shè)計中，學習目標里的速度場監(jiān)督帶來的方差。基于這一發(fā)現(xiàn)，本文對現(xiàn)有的一步擴散生成方法進行了改進，并提出了 ESC（ExplicitShortCut） 變體。該方法在訓練穩(wěn)定性和生成質(zhì)量上都取得了進一步提升。在 ImageNet 256×256 上，基于 SiT-XL/2 架構(gòu)，ESC 在 1-NFE 設(shè)置下達到 FID50k 2.53，超越了此前所有從頭訓練的捷徑擴散模型。

如何從頭訓練一個一步擴散模型？

圖 1 離散時間與連續(xù)時間的一步擴散模型在理想學習與實際學習中的物理圖景 (a) 展示了從 N (0,1) 到高斯混合分布的邊緣速度場。(b)(c) 對應(yīng)理想學習情形，此時 xr 采樣自 PF-ODE 的同一條軌跡，因此構(gòu)造的學習目標可作為正確監(jiān)督信號；(d)(e) 則展示實際學習情形，此時目標偏離原始軌跡，從而導致模型預(yù)測也相應(yīng)發(fā)生漂移。

表 1 代表性的一步生成模型中，具體的組件選擇

如何選擇設(shè)計組件？

圖 2 不同一步生成模型在訓練過程中 FID50k 的比較：(a) 表示在 CIFAR-10 上的無條件生成（Uncond.）；(b) 表示在 CIFAR-10 上的類別條件生成（Cond.）；(c) 表示在 ImageNet-256×256 上的 classifier-free guidance（CFG）訓練。

Q1: 路徑選擇：線性 vs 余弦？

一般而言，線性路徑擴散如 Rectified Flow、DiT 等工作在理論上更加簡潔，實踐上更方便以使用各種訓練和采樣技巧，而余弦路徑擴散如 EDM 等工作往往被認為像素空間生成中通常因固定方差的隨機過程而被視為更穩(wěn)定。

本文在此探尋路徑選擇對模型的影響。根據(jù)圖 2 (a),(b) 中 sCT (-Cosine) 以及 sCT-Linear 的實驗結(jié)果表明，在線性路徑與余弦路徑的對比中，線性路徑在一步生成的設(shè)定下通常更有競爭力。一種合理的解釋是：線性路徑作為條件路徑時，對應(yīng)的邊緣速度場具有更低的凸運輸代價，因此軌跡彎曲程度更小。這意味著，由兩步采樣映射構(gòu)造出來的監(jiān)督目標更接近理想情形，從而讓 “一步去逼近兩步” 這件事本身更容易成立。換句話說，在一步生成的問題設(shè)定里，線性路徑并不只是實現(xiàn)更方便，而是更符合這一類模型的近似需求。

Q2: 損失類型：離散 vs 連續(xù)？

前面討論過，當 s→t 的時候，損失函數(shù)就會將一步擴散推至連續(xù)時間模型。本文在此探尋連續(xù)與離散模型的優(yōu)劣勢。在相同訓練設(shè)置和統(tǒng)一代碼實現(xiàn)下，根據(jù)圖 2 (a), (b), (c) 中的 MeanFlow、sCT 這樣的連續(xù)時間上的一步擴散模型與 CT、SCD、IMM 等離散時間方法的對比，可以得到連續(xù)時間一步擴散模型整體上明顯優(yōu)于離散時間變體。

進一步，本文給出了推理誤差分析，其核心結(jié)論是：對于線性路徑，連續(xù)時間方法的誤差上界相比離散時間方法少了一項與局部 Lipschitz 常數(shù)、時間步長以及監(jiān)督方差耦合的額外誤差項；而這項額外項很可能會直接放大訓練不穩(wěn)定性和一步推理誤差。也就是說，連續(xù)時間方法更強，并不只是因為 “參數(shù)化更平滑”，而是因為它在目標逼近誤差的結(jié)構(gòu)上更占優(yōu)。這也解釋了為什么同樣是一步生成擴散，基于連續(xù)時間的模型往往表現(xiàn)得更穩(wěn)、更準。

Q3: 時間采樣：固定終點時間 vs 隨機化

換句話說，固定終止時間更像是在優(yōu)化 “早期易學性”，而隨機化終止時間則是在優(yōu)化 “最終生成與表達能力”。

如何改進現(xiàn)有難題？

基于上述理論和實證分析，我們采用連續(xù)時間線性路徑基線（MeanFlow with SiT-B/2），并通過三個方向進行改進：

1. 即插即用的理想速度（Plug-in velocity）

第一個改進方向是用 plug-in velocity 替代直接使用單樣本對應(yīng)的條件速度作為監(jiān)督信號。其動機來自前面的誤差分析：在訓練時，真正難以處理的不是目標形式本身，而是引入條件速度而帶來的高方差。本文指出，如果能夠使用基于整個經(jīng)驗分布構(gòu)造的 “理想速度”，那么速度項的方差可以顯著下降；但這需要對整個訓練集求和，在大規(guī)模數(shù)據(jù)集上幾乎不可行。為此，作者采用在一個 mini-batch 上的 plug-in velocity，將一個批次內(nèi)多個樣本的條件速度以不同加權(quán)的形式組合起來，從而在僅引入較小偏差的代價下，將監(jiān)督方差從依賴單樣本的水平降到與批次大小相關(guān)的更低水平。這種替換直接帶來更穩(wěn)定的訓練監(jiān)督和更低的一步推理誤差，是對核心瓶頸的直接修正。

2. 漸進時間采樣器（Gradual time sampler）

第二個改進方向是設(shè)計一個漸進式時間采樣器，讓訓練從更容易的目標開始，再平滑過渡到完整的 MeanFlow 式時間采樣。這個設(shè)計直接回應(yīng)了前面對 “固定終止時間還是隨機化” 的分析：固定 r=0 的監(jiān)督更像去噪任務(wù)，早期更容易優(yōu)化，而覆蓋完整區(qū)間的隨機時間采樣雖然更難，卻更有助于模型學習全局 “捷徑模式”�；谶@一點，本文采用了一個 漸進式課程策略：在訓練初期，以較高概率采樣 r=0，讓模型先獲得更穩(wěn)定、更直接的監(jiān)督；隨著訓練推進，再逐步衰減這一概率，最終完全過渡到 MeanFlow 的采樣方式。這樣做的好處是同時兼顧了前期收斂速度和后期全局建模能力，本質(zhì)上是在易學性與最終性能之間建立了一條更平滑的訓練路徑。

3. 優(yōu)化技巧 (Other techniques)

第三個方向則是吸收已有連續(xù)時間一步擴散方法中已經(jīng)驗證有效的訓練技巧，并將其系統(tǒng)整合到新的框架中。包括一些已經(jīng)在原始工作中驗證過的優(yōu)化手段，例如 variational adaptive loss weighting，進一步提升優(yōu)化過程的穩(wěn)定性、緩解訓練早期的脆弱性，并改善不同監(jiān)督項之間的平衡。此外，在 classifier-free guidance 的訓練場景下，本文額外引入類別一致的批次處理技術(shù)（class-consistent mini-batching）：通過同類采樣避免 mini-batch 平均過程過度稀釋類別信號，以達到大規(guī)模訓練中真正 “跑得穩(wěn)、訓得動、效果好” 的工程化方案。

綜合上述提出的技術(shù)，本文在SiT-B/2的網(wǎng)絡(luò)架構(gòu)上，進行了實驗消融，說明幾種改進的結(jié)合可以穩(wěn)定地提升模型的生成能力。

大規(guī)模實驗評估結(jié)果

我們在 ImageNet-256x256 的潛空間中使用 SiT-XL/2（約 676M 參數(shù)）評估 ESC。遵循 MeanFlow 在分類器自由引導下的訓練協(xié)議，ESC 從頭訓練 240 個 epoch（約 120 萬次迭代），ESC+ 延長至 480 個 epoch（約 240 萬次迭代）。

在 ImageNet-256x256 的 1-NFE 設(shè)置下，ESC 將 FID 從 3.43 提升至 2.85，ESC+ 進一步提升至 2.53。這一結(jié)果不僅大幅超越了此前所有從頭訓練的捷徑模型，甚至超過了 MeanFlow 的兩步結(jié)果（2-NFE，F(xiàn)ID50k 2.93）。

此外，我們還可以從中觀察到：

1)類別一致的批次處理可以提高收斂速度。

2)即插即用的理想速度幾乎不增加計算開銷 (千分之九)，同時提高穩(wěn)定性。

3) 性能提升在更大的骨干網(wǎng)絡(luò)上更明顯，表明降低方差對于模型穩(wěn)定性的提升，隨著模型容量增長，效果越來越明顯

ESC 在 ImageNet-256x256 上展現(xiàn)出卓越的一步生成能力（SiT-XL/2 架構(gòu)），且能夠生成高質(zhì)量、多樣化的圖像。（FID:2.85）

ExplicitShortCut (ESC)通過統(tǒng)一的設(shè)計空間分析，系統(tǒng)性地理解了一步擴散模型的關(guān)鍵設(shè)計選擇。通過改進目標構(gòu)建和降低監(jiān)督方差，ESC 在 ImageNet-256x256 上實現(xiàn)了 FID 2.53 的 SOTA 一步生成性能，也為之后設(shè)計高效的一步擴散生成提供了關(guān)鍵見解。

然而，目前的從頭訓練的一步生成模型仍然面臨著極大的挑戰(zhàn)，比如在連續(xù)情形下 jvp 函數(shù)的使用以及其導致的顯存開銷、對 CFG 的依賴導致設(shè)計不夠靈活、超參數(shù)空間龐大等，即便如此，一步生成的擴散模型仍然是未來的研究熱門，正如何凱明在 CVPR 上的演講所說：

“生成模型正處在一個類似「AlexNet 之前」的階段：從多步推理，走向端到端的一步生成。” 一步擴散不是終點，它是下一個 AIGC 時代的開始。

作者及團隊簡介

本文第一作者為林海濤，來自西湖大學，共同一作為胡佩炎，來自中科院數(shù)學所。其研究方向關(guān)注生成模型，一步擴散與蒸餾，離散擴散模型等，以第一作者在 ICML、ICLR、NeurIPS 等人工智能頂級會議或期刊上發(fā)表論文超過十篇；本文的通訊作者為西湖大學李子青、吳泰霖，以及深勢科技柯國霖。李子青實驗室關(guān)注人工智能及生命科學交叉領(lǐng)域。吳泰霖實驗室關(guān)注生成式 AI（包括擴散模型、多智能體等）及其在聚變、水下具身智能、虛擬細胞領(lǐng)域的應(yīng)用，長期招募希望在以上領(lǐng)域做出有影響力工作的同學（博士生、博后、助理研究員、實習生）。