![]()
![]()
- 論文標題: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)并從頭訓練的一步擴散生成模型,展現出強大的實證有效性。然而,這類方法的提出通常建立在較為復雜的理論推導之上,并且往往與具體實現細節高度耦合。這帶來一個直接的問題:究竟哪些設計是方法成立的本質要素,哪些又只是可以靈活替換的實現組件。
對于研究人員,這種復雜性進一步體現在多個層面:一方面是深奧的理論背景,繁復的方法形式推導、對應的學習目標設計,另一方面,是應用實現上的技術細節,包括時間采樣器、課程學習方法等。面對這些內容,研究人員往往很難迅速抓住不同方法之間真正共享的核心思想,也因此缺少一種足夠直觀的設計范式,來理解這類一步生成模型到底 “為什么這樣設計”。
為了解決這個問題,西湖大學研究人員首先從一個更實用的角度出發,嘗試為這類 shortcut 模型提出一個統一的設計框架。其指出,無論是離散時間還是連續時間的一步生成方法,其背后都遵循著同一個基本原則:用單步參數化預測去逼近一個兩步擴散構建的目標。在此基礎上,該工作進一步給出了這一設計范式成立的一般性理論說明。這樣的統一視角,使得原本耦合在一起的生成方法模塊能夠被拆解開來,以供研究人員更清晰地理解各個組件之間是如何協同作用的,以及在整體方法設計中究竟還保留了多少可調整的自由度。
得益于這一框架,本文對當前具有代表性的離散時間與連續時間一步擴散生成方法進行更細粒度的組件級分析,包括路徑選擇、時間采樣器等關鍵模塊。通過這一分析,本文進一步定位了影響訓練 — 推理穩定性以及逼近誤差的一個關鍵來源:損失函數設計中,學習目標里的速度場監督帶來的方差。基于這一發現,本文對現有的一步擴散生成方法進行了改進,并提出了 ESC(ExplicitShortCut) 變體。該方法在訓練穩定性和生成質量上都取得了進一步提升。在 ImageNet 256×256 上,基于 SiT-XL/2 架構,ESC 在 1-NFE 設置下達到 FID50k 2.53,超越了此前所有從頭訓練的捷徑擴散模型。
如何從頭訓練一個一步擴散模型?
![]()
圖 1 離散時間與連續時間的一步擴散模型在理想學習與實際學習中的物理圖景 (a) 展示了從 N (0,1) 到高斯混合分布的邊緣速度場。(b)(c) 對應理想學習情形,此時 xr 采樣自 PF-ODE 的同一條軌跡,因此構造的學習目標可作為正確監督信號;(d)(e) 則展示實際學習情形,此時目標偏離原始軌跡,從而導致模型預測也相應發生漂移。
![]()
![]()
![]()
表 1 代表性的一步生成模型中,具體的組件選擇
![]()
如何選擇設計組件?
![]()
圖 2 不同一步生成模型在訓練過程中 FID50k 的比較:(a) 表示在 CIFAR-10 上的無條件生成(Uncond.);(b) 表示在 CIFAR-10 上的類別條件生成(Cond.);(c) 表示在 ImageNet-256×256 上的 classifier-free guidance(CFG)訓練。
Q1: 路徑選擇:線性 vs 余弦?
一般而言,線性路徑擴散如 Rectified Flow、DiT 等工作在理論上更加簡潔,實踐上更方便以使用各種訓練和采樣技巧,而余弦路徑擴散如 EDM 等工作往往被認為像素空間生成中通常因固定方差的隨機過程而被視為更穩定。
本文在此探尋路徑選擇對模型的影響。根據圖 2 (a),(b) 中 sCT (-Cosine) 以及 sCT-Linear 的實驗結果表明,在線性路徑與余弦路徑的對比中,線性路徑在一步生成的設定下通常更有競爭力。一種合理的解釋是:線性路徑作為條件路徑時,對應的邊緣速度場具有更低的凸運輸代價,因此軌跡彎曲程度更小。這意味著,由兩步采樣映射構造出來的監督目標更接近理想情形,從而讓 “一步去逼近兩步” 這件事本身更容易成立。換句話說,在一步生成的問題設定里,線性路徑并不只是實現更方便,而是更符合這一類模型的近似需求。
Q2: 損失類型:離散 vs 連續?
前面討論過,當 s→t 的時候,損失函數就會將一步擴散推至連續時間模型。本文在此探尋連續與離散模型的優劣勢。在相同訓練設置和統一代碼實現下,根據圖 2 (a), (b), (c) 中的 MeanFlow、sCT 這樣的連續時間上的一步擴散模型與 CT、SCD、IMM 等離散時間方法的對比,可以得到連續時間一步擴散模型整體上明顯優于離散時間變體。
進一步,本文給出了推理誤差分析,其核心結論是:對于線性路徑,連續時間方法的誤差上界相比離散時間方法少了一項與局部 Lipschitz 常數、時間步長以及監督方差耦合的額外誤差項;而這項額外項很可能會直接放大訓練不穩定性和一步推理誤差。也就是說,連續時間方法更強,并不只是因為 “參數化更平滑”,而是因為它在目標逼近誤差的結構上更占優。這也解釋了為什么同樣是一步生成擴散,基于連續時間的模型往往表現得更穩、更準。
Q3: 時間采樣:固定終點時間 vs 隨機化
![]()
換句話說,固定終止時間更像是在優化 “早期易學性”,而隨機化終止時間則是在優化 “最終生成與表達能力”。
如何改進現有難題?
基于上述理論和實證分析,我們采用連續時間線性路徑基線(MeanFlow with SiT-B/2),并通過三個方向進行改進:
1. 即插即用的理想速度(Plug-in velocity)
第一個改進方向是用 plug-in velocity 替代直接使用單樣本對應的條件速度作為監督信號。其動機來自前面的誤差分析:在訓練時,真正難以處理的不是目標形式本身,而是引入條件速度而帶來的高方差。本文指出,如果能夠使用基于整個經驗分布構造的 “理想速度”,那么速度項的方差可以顯著下降;但這需要對整個訓練集求和,在大規模數據集上幾乎不可行。為此,作者采用在一個 mini-batch 上的 plug-in velocity,將一個批次內多個樣本的條件速度以不同加權的形式組合起來,從而在僅引入較小偏差的代價下,將監督方差從依賴單樣本的水平降到與批次大小相關的更低水平。這種替換直接帶來更穩定的訓練監督和更低的一步推理誤差,是對核心瓶頸的直接修正。
2. 漸進時間采樣器(Gradual time sampler)
第二個改進方向是設計一個漸進式時間采樣器,讓訓練從更容易的目標開始,再平滑過渡到完整的 MeanFlow 式時間采樣。這個設計直接回應了前面對 “固定終止時間還是隨機化” 的分析:固定 r=0 的監督更像去噪任務,早期更容易優化,而覆蓋完整區間的隨機時間采樣雖然更難,卻更有助于模型學習全局 “捷徑模式”。基于這一點,本文采用了一個 漸進式課程策略:在訓練初期,以較高概率采樣 r=0,讓模型先獲得更穩定、更直接的監督;隨著訓練推進,再逐步衰減這一概率,最終完全過渡到 MeanFlow 的采樣方式。這樣做的好處是同時兼顧了前期收斂速度和后期全局建模能力,本質上是在易學性與最終性能之間建立了一條更平滑的訓練路徑。
3. 優化技巧 (Other techniques)
第三個方向則是吸收已有連續時間一步擴散方法中已經驗證有效的訓練技巧,并將其系統整合到新的框架中。包括一些已經在原始工作中驗證過的優化手段,例如 variational adaptive loss weighting,進一步提升優化過程的穩定性、緩解訓練早期的脆弱性,并改善不同監督項之間的平衡。此外,在 classifier-free guidance 的訓練場景下,本文額外引入類別一致的批次處理技術(class-consistent mini-batching):通過同類采樣避免 mini-batch 平均過程過度稀釋類別信號,以達到大規模訓練中真正 “跑得穩、訓得動、效果好” 的工程化方案。
綜合上述提出的技術,本文在SiT-B/2的網絡架構上,進行了實驗消融,說明幾種改進的結合可以穩定地提升模型的生成能力。
![]()
大規模實驗評估結果
我們在 ImageNet-256x256 的潛空間中使用 SiT-XL/2(約 676M 參數)評估 ESC。遵循 MeanFlow 在分類器自由引導下的訓練協議,ESC 從頭訓練 240 個 epoch(約 120 萬次迭代),ESC+ 延長至 480 個 epoch(約 240 萬次迭代)。
在 ImageNet-256x256 的 1-NFE 設置下,ESC 將 FID 從 3.43 提升至 2.85,ESC+ 進一步提升至 2.53。這一結果不僅大幅超越了此前所有從頭訓練的捷徑模型,甚至超過了 MeanFlow 的兩步結果(2-NFE,FID50k 2.93)。
此外,我們還可以從中觀察到:
1)類別一致的批次處理可以提高收斂速度。
2)即插即用的理想速度幾乎不增加計算開銷 (千分之九),同時提高穩定性。
3) 性能提升在更大的骨干網絡上更明顯,表明降低方差對于模型穩定性的提升,隨著模型容量增長,效果越來越明顯
![]()
![]()
ESC 在 ImageNet-256x256 上展現出卓越的一步生成能力(SiT-XL/2 架構),且能夠生成高質量、多樣化的圖像。(FID:2.85)
ExplicitShortCut (ESC)通過統一的設計空間分析,系統性地理解了一步擴散模型的關鍵設計選擇。通過改進目標構建和降低監督方差,ESC 在 ImageNet-256x256 上實現了 FID 2.53 的 SOTA 一步生成性能,也為之后設計高效的一步擴散生成提供了關鍵見解。
然而,目前的從頭訓練的一步生成模型仍然面臨著極大的挑戰,比如在連續情形下 jvp 函數的使用以及其導致的顯存開銷、對 CFG 的依賴導致設計不夠靈活、超參數空間龐大等,即便如此,一步生成的擴散模型仍然是未來的研究熱門,正如何凱明在 CVPR 上的演講所說:
“生成模型正處在一個類似「AlexNet 之前」的階段:從多步推理,走向端到端的一步生成。” 一步擴散不是終點,它是下一個 AIGC 時代的開始。
作者及團隊簡介
本文第一作者為林海濤,來自西湖大學,共同一作為胡佩炎,來自中科院數學所。其研究方向關注生成模型,一步擴散與蒸餾,離散擴散模型等,以第一作者在 ICML、ICLR、NeurIPS 等人工智能頂級會議或期刊上發表論文超過十篇;本文的通訊作者為西湖大學李子青、吳泰霖,以及深勢科技柯國霖。李子青實驗室關注人工智能及生命科學交叉領域。吳泰霖實驗室關注生成式 AI(包括擴散模型、多智能體等)及其在聚變、水下具身智能、虛擬細胞領域的應用,長期招募希望在以上領域做出有影響力工作的同學(博士生、博后、助理研究員、實習生)。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.