兼得快與好！訓(xùn)練新范式TiM，原生支持FSDP+Flash Attention

TiM團(tuán)隊(duì) 投稿
量子位 | 公眾號(hào) QbitAI

生成式AI的快與好，終于能兼得了

從Stable Diffusion到DiT、FLUX系列，社區(qū)探索了很多技術(shù)方法用于加速生成速度和提高生成質(zhì)量，但是始終圍繞擴(kuò)散模型和Few-step模型兩條路線進(jìn)行開(kāi)發(fā)，不得不向一些固有的缺陷妥協(xié)。

這便是訓(xùn)練目標(biāo)引發(fā)的“生成質(zhì)量”與“生成速度”之間的矛盾根源

要么只監(jiān)督無(wú)窮小局部動(dòng)力學(xué)（PF-ODE），要么只學(xué)習(xí)有限區(qū)間的端點(diǎn)映射，兩者都各有內(nèi)在限制。

一項(xiàng)新研究提出了名為Transition Model（TiM）的新范式，試圖從根本上解決這一矛盾。

它放棄了傳統(tǒng)擴(kuò)散模型學(xué)習(xí)“瞬時(shí)速度場(chǎng)”或Few-step模型學(xué)習(xí)“端點(diǎn)映射”的做法，轉(zhuǎn)而直接建模任意兩個(gè)時(shí)間點(diǎn)之間的完整狀態(tài)轉(zhuǎn)移。

這意味著TiM在理論上支持任意步長(zhǎng)的采樣，并能將生成過(guò)程分解為多段粒度可靈活調(diào)整的細(xì)化軌跡

什么是Transition Model？

為什么說(shuō)“PF-ODE”與“概率分布匹配”對(duì)于生成模型都不是理想的訓(xùn)練目標(biāo)？

來(lái)看擴(kuò)散模型，它以迭代去噪獲得高保真，在于它學(xué)習(xí)的是PF-ODE的局部向量場(chǎng)，訓(xùn)練時(shí)只對(duì)無(wú)窮小時(shí)間步的瞬時(shí)動(dòng)力學(xué)做監(jiān)督，采樣時(shí)必須用很小步長(zhǎng)或高階多步求解器來(lái)壓離散誤差，導(dǎo)致NFEs居高不下。

又比如少步生成（如 Consistency/Shortcut/Distillation/Meanflow）雖快，但因?yàn)闆](méi)有刻畫(huà)中間動(dòng)力學(xué)，增步后收益很快飽和，常遭遇 “質(zhì)量天花板”，增加步數(shù)反而不再帶來(lái)收益，生成能力上限不及擴(kuò)散模型。

這些固有的缺陷來(lái)源于模型訓(xùn)練過(guò)程中監(jiān)督信號(hào)的引入方式，或是求解局部的PFE方程，或是匹配固定的概率分布；換句話說(shuō)，生成過(guò)程中，模型做出預(yù)測(cè)被clean data所監(jiān)督的粒度，直接決定了模型在推理過(guò)程中的離散誤差和生成質(zhì)量上限。

所以，對(duì)于生成模型，什么才是一個(gè)合適的訓(xùn)練目標(biāo)呢？

從擴(kuò)散模型與Few-step模型的訓(xùn)練目標(biāo)的局限性出發(fā)，可以得到以下分析——

局部（無(wú)窮小）監(jiān)督：PF-ODE/SDE類(lèi)目標(biāo)。

這類(lèi)目標(biāo)只在極小時(shí)間步上擬合瞬時(shí)動(dòng)力學(xué)（Δt→0），要想維持連續(xù)時(shí)間解的精度，采樣時(shí)就必須用很小步長(zhǎng)/很多步，于是NFEs很高；一旦把步數(shù)壓到很少，質(zhì)量就會(huì)明顯掉隊(duì)。

因此，對(duì)于能夠帶來(lái)高保真度的局部監(jiān)督信號(hào)而言，時(shí)間區(qū)間，或者說(shuō)單步步長(zhǎng)理想情況下應(yīng)該是要能靈活改

全局端點(diǎn)監(jiān)督：few-step/一致性/蒸餾一類(lèi)目標(biāo)/mean-flow/short-cut。

這類(lèi)訓(xùn)練目標(biāo)學(xué)習(xí)固定跨度的端點(diǎn)映射（或者平均速度場(chǎng)），核心是一步 “吃掉” 整段軌跡，因而少步很強(qiáng)；但因?yàn)?“把整條軌跡平均化”，細(xì)節(jié)動(dòng)力學(xué)被抹掉，再加步也難以繼續(xù)提升——出現(xiàn)質(zhì)量飽和。

因此，訓(xùn)練目標(biāo)應(yīng)該要求沿軌跡保持一致，要存在中間步驟充當(dāng)單個(gè)軌跡的細(xì)化，而不是偏離新的軌跡，這使得sampler對(duì)采樣規(guī)劃不敏感，并能夠通過(guò)更多步驟實(shí)現(xiàn)穩(wěn)定的質(zhì)量改進(jìn)。

因此，一個(gè)能兼得快速生成（few-step）與高保真度生成（擴(kuò)散模型）的訓(xùn)練目標(biāo)應(yīng)該是：

在“多段細(xì)化軌跡”里實(shí)現(xiàn)“靈活的單步尺寸”（任意步長(zhǎng)），這便是Transition Model。

想要兼得推理速度與高保真度質(zhì)量，需要一個(gè)核心設(shè)計(jì)，“在多段細(xì)化的軌跡”里面實(shí)現(xiàn)“靈活的單步尺寸”。

這一工作基于此設(shè)計(jì)了Transition Model：

將模型的訓(xùn)練從單一時(shí)刻t，拓展到建模任意兩個(gè)時(shí)刻t與r的狀態(tài)x_t, x_r.

設(shè)計(jì)1：實(shí)現(xiàn)“靈活的單步尺寸”

對(duì)于給定的兩個(gè)時(shí)刻t與r之間的狀態(tài)轉(zhuǎn)移，通過(guò)化簡(jiǎn)其微分方程得到了“通用狀態(tài)轉(zhuǎn)移恒等式”（State Transition Identity）；基于通用狀態(tài)轉(zhuǎn)移恒等式，得以描述任意的一個(gè)時(shí)間間隔內(nèi)的具體狀態(tài)轉(zhuǎn)移，而不是作為數(shù)值擬合求解。

設(shè)計(jì)2：實(shí)現(xiàn)“多段細(xì)化軌跡的生成路徑”

在設(shè)計(jì)1中，已經(jīng)實(shí)現(xiàn)了任意步長(zhǎng)（任意時(shí)間間隔）， 因此對(duì)于多段細(xì)化軌跡的生成路徑，這個(gè)方法就可以直接的描述任意時(shí)刻t下對(duì)于此前任意時(shí)刻r之間的狀態(tài)轉(zhuǎn)移，那么“多段細(xì)化的生成路徑”就變成了“任意狀態(tài)與前狀態(tài)之間的狀態(tài)轉(zhuǎn)移動(dòng)態(tài)（state transition dynamics）”，這樣就能在保持快速生成的同時(shí)保證高保真度的生成質(zhì)量。

通過(guò)設(shè)計(jì)1和設(shè)計(jì)2，這篇文章提出的Transition Model將“在任意狀態(tài)下，任意時(shí)間間隔內(nèi)，與前狀態(tài)之間的狀態(tài)轉(zhuǎn)移的動(dòng)力學(xué)方程”作為訓(xùn)練目標(biāo)，它就滿足了兼得推理速度與高保真度質(zhì)量的核心設(shè)計(jì)。

Transition Model的數(shù)學(xué)本質(zhì)

Diffusion model是建模瞬時(shí)速度場(chǎng)，局限性是瞬時(shí)速度需要時(shí)間區(qū)間趨近于0；

Meanflow核心是建模平均速度場(chǎng)，局限性是平均速度丟了局部?jī)?yōu)化的dynamics細(xì)節(jié)，生成質(zhì)量早早收斂，過(guò)了few-step后近乎為定值；

不同于前兩者，Transition Model做的是任意時(shí)間區(qū)間的任意狀態(tài)間的狀態(tài)轉(zhuǎn)移，可以認(rèn)為是任意速度場(chǎng)，自然而然地包含了瞬時(shí)速度和平均速度；

從解的形式上講 Diffusion是局部PF-ODE的數(shù)值解，meanflow是局部平均速度場(chǎng)中的解集，transition model求的是全局生成路徑上的解的流型，special case情況下可以退化為平均速度場(chǎng)，解的流型退化為局部解集。

作者們主要在圖文生成（Text-to-Image）任務(wù)上進(jìn)行了驗(yàn)證

在Geneval數(shù)據(jù)集上，分別比較了Transition Model在不同推理步數(shù)（NFE）， 不同分辨率，不同橫縱比下的生成能力：

這篇文章發(fā)現(xiàn)865M參數(shù)大小的Transition Model（TiM）可以在明確地超過(guò)FLUX.1-Schnell（12B參數(shù)）這一蒸餾模型；與此同時(shí)，在生成能力上限上也可以超過(guò)FLUX.1-Dev（12B參數(shù)）

并且由于TiM結(jié)合了Native-Resolution預(yù)訓(xùn)練的訓(xùn)練策略（詳見(jiàn)Native-Resolution Image Synthesis），這篇文章所提出的模型在分辨率和橫縱比上也更加靈活。

Transition Model的訓(xùn)練穩(wěn)定性與擴(kuò)展性

讓Transition Model訓(xùn)練具有可擴(kuò)展性.

在Transition Model的訓(xùn)練過(guò)程中，它的訓(xùn)練目標(biāo)的關(guān)鍵在于計(jì)算網(wǎng)絡(luò)關(guān)于時(shí)間的導(dǎo)數(shù)$\frac{\mathrmdjd33nj f_{\theta^{-}, t, r}}{\mathrmdjd33nj t}$

以MeanFlow和Short-cut Model為代表的既有方法通常依賴(lài)雅可比—向量乘積（JVP）來(lái)完成這一計(jì)算。

然而，JVP在可擴(kuò)展性上構(gòu)成了根本性瓶頸：

不僅計(jì)算開(kāi)銷(xiāo)高，更麻煩的是它依賴(lài)Backward自動(dòng)微分，這與諸如FlashAttention和分布式框架Fully Sharded Data Parallel（FSDP）等關(guān)鍵訓(xùn)練優(yōu)化并不兼容，致使基于JVP的方法難以實(shí)際用于十億參數(shù)級(jí)的基礎(chǔ)模型訓(xùn)練。

為此，他們提出差分推導(dǎo)方程（DDE），用一種有原則且高效的有限差分近似來(lái)突破該限制：

如表中所示，這篇文章所提出的DDE計(jì)算方式不僅比JVP約快2倍，更關(guān)鍵的是其僅依賴(lài)前向傳播，與FSDP天然兼容，從而將原本不可擴(kuò)展的訓(xùn)練流程變?yōu)榭纱笠?guī)模并行計(jì)算的方案.

讓Transition Model訓(xùn)練更加穩(wěn)定.

除了可擴(kuò)展性，基于任意時(shí)間間隔訓(xùn)練的另一大挑戰(zhàn)是控制梯度方差

比如，當(dāng)轉(zhuǎn)移跨越很大的時(shí)間間隔（$\Delta t \to t$）時(shí)，更容易出現(xiàn)損失突增。

為緩解這一問(wèn)題，作者們引入一種損失加權(quán)策略，優(yōu)先考慮短間隔轉(zhuǎn)移——這類(lèi)轉(zhuǎn)移更為常見(jiàn)，也能提供更穩(wěn)定的學(xué)習(xí)信號(hào)。

其中，$\tau(\cdot)$是對(duì)時(shí)間軸進(jìn)行重新參數(shù)化的單調(diào)函數(shù)。

在這篇文章最終模型中，他們采用正切空間變換（tangent space transformation來(lái)有效拉伸時(shí)間域，從而得到具體的加權(quán)形式：

其中，$\sigma_{\text{data}}$表示干凈數(shù)據(jù)（clean data）的標(biāo)準(zhǔn)差，這一方法有效地提升了訓(xùn)練的穩(wěn)定性。

研究團(tuán)隊(duì)提出了Transition Model（TiM）作為生成模型的新的范式：

不再只學(xué)習(xí)瞬時(shí)向量場(chǎng)或固定跨度的端點(diǎn)映射，而是直接建模任意兩時(shí)刻間的狀態(tài)轉(zhuǎn)移，用“通用狀態(tài)轉(zhuǎn)移恒等式”支撐任意步長(zhǎng)與多段細(xì)化軌跡，從而兼顧少步速度與高保真質(zhì)量。

在理論上，從學(xué)習(xí)生成路徑上特定的解拓展到學(xué)習(xí)全局生成路徑的解的流形；在實(shí)踐上，通過(guò)DDE的前向有限差分替代JVP，原生兼容 FSDP/FlashAttention、訓(xùn)練更快更可擴(kuò)展；同時(shí)用時(shí)間重參化+核函數(shù)的損失加權(quán)優(yōu)先短間隔，降低梯度方差、提升穩(wěn)定性。

實(shí)驗(yàn)表明，TiM-865M在多分辨率與多橫縱比設(shè)置下，少步即可超越 FLUX.1-Schnell/Dev（12B）的速度-質(zhì)量權(quán)衡。

總體而言，TiM以全局路徑視角嘗試解決“速度與質(zhì)量難兩全”的根本矛盾，提供了更通用、可擴(kuò)展且穩(wěn)定的生成建模。