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

TiM團隊 投稿
量子位 | 公眾號 QbitAI

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

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

這便是訓練目標引發的“生成質量”與“生成速度”之間的矛盾根源

要么只監督無窮小局部動力學（PF-ODE），要么只學習有限區間的端點映射，兩者都各有內在限制。

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

它放棄了傳統擴散模型學習“瞬時速度場”或Few-step模型學習“端點映射”的做法，轉而直接建模任意兩個時間點之間的完整狀態轉移。

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

什么是Transition Model？

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

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

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

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

所以，對于生成模型，什么才是一個合適的訓練目標呢？

從擴散模型與Few-step模型的訓練目標的局限性出發，可以得到以下分析——

局部（無窮小）監督：PF-ODE/SDE類目標。

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

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

全局端點監督：few-step/一致性/蒸餾一類目標/mean-flow/short-cut。

這類訓練目標學習固定跨度的端點映射（或者平均速度場），核心是一步 “吃掉” 整段軌跡，因而少步很強；但因為 “把整條軌跡平均化”，細節動力學被抹掉，再加步也難以繼續提升——出現質量飽和。

因此，訓練目標應該要求沿軌跡保持一致，要存在中間步驟充當單個軌跡的細化，而不是偏離新的軌跡，這使得sampler對采樣規劃不敏感，并能夠通過更多步驟實現穩定的質量改進。

因此，一個能兼得快速生成（few-step）與高保真度生成（擴散模型）的訓練目標應該是：

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

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

這一工作基于此設計了Transition Model：

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

設計1：實現“靈活的單步尺寸”

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

設計2：實現“多段細化軌跡的生成路徑”

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

通過設計1和設計2，這篇文章提出的Transition Model將“在任意狀態下，任意時間間隔內，與前狀態之間的狀態轉移的動力學方程”作為訓練目標，它就滿足了兼得推理速度與高保真度質量的核心設計。

Transition Model的數學本質

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

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

不同于前兩者，Transition Model做的是任意時間區間的任意狀態間的狀態轉移，可以認為是任意速度場，自然而然地包含了瞬時速度和平均速度；

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

作者們主要在圖文生成（Text-to-Image）任務上進行了驗證

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

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

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

Transition Model的訓練穩定性與擴展性

讓Transition Model訓練具有可擴展性.

在Transition Model的訓練過程中，它的訓練目標的關鍵在于計算網絡關于時間的導數$\frac{\mathrmdjd33nj f_{\theta^{-}, t, r}}{\mathrmdjd33nj t}$

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

然而，JVP在可擴展性上構成了根本性瓶頸：

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

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

如表中所示，這篇文章所提出的DDE計算方式不僅比JVP約快2倍，更關鍵的是其僅依賴前向傳播，與FSDP天然兼容，從而將原本不可擴展的訓練流程變為可大規模并行計算的方案.

讓Transition Model訓練更加穩定.

除了可擴展性，基于任意時間間隔訓練的另一大挑戰是控制梯度方差

比如，當轉移跨越很大的時間間隔（$\Delta t \to t$）時，更容易出現損失突增。

為緩解這一問題，作者們引入一種損失加權策略，優先考慮短間隔轉移——這類轉移更為常見，也能提供更穩定的學習信號。

其中，$\tau(\cdot)$是對時間軸進行重新參數化的單調函數。

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

其中，$\sigma_{\text{data}}$表示干凈數據（clean data）的標準差，這一方法有效地提升了訓練的穩定性。

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

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

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

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

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