兼得快與好！訓練新範式TiM，原生支持FSDP+Flash Attention
發布時間：2025-09-14 來源：膽大心細網作者：D_a_i_d_i_a_o
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）任務上進行了驗證
在Genevals數據集上，分別比較了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{\mathrm{d} f_{\theta^{-}, t, r}}{\mathrm{d} 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以全局路徑視角嚐試解決“速度與質量難兩全”的根本矛盾，提供了更通用、可擴展且穩定的生成建模。
- 上一篇：{loop type="arclist" row=1 }{$vo.title}