TiM團(tuán)隊 投稿
量子位 | 公眾號 QbitAI

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

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

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

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

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

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

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

什么是Transition Model？

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

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

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

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

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

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

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

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

因此，對于能夠帶來高保真度的局部監(jiān)督信號而言，時間區(qū)間，或者說單步步長理想情況下應(yīng)該是要能靈活改

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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