無需訓練讓擴散模型提速2倍，上交大提出Token級快取方案｜ICLR‘25

張林峰投稿

量子位 | 公眾號 QbitAI

Diffusion Transformer模型模型透過token粒度的快取方法，實現了影像和影片生成模型上無需訓練的兩倍以上的加速。

上海交通大學等團隊提出Toca（Token-wise Caching），相關論文已被ICLR 2025接收。

Diffusion Transformer在影像和影片生成中展現了顯著的效果，但代價是巨大的計算成本。

為了解決這一問題，特徵快取方法被引入，用於透過快取前幾個時間步的特徵並在後續時間步中複用它們來加速擴散Transformer。

然而，之前的快取方法忽略了不同的token對特徵快取表現出不同的敏感性，而對某些token的特徵快取可能導致生成質量整體上高達10倍的破壞，相較於其他token。

Toca團隊提出了基於token的特徵快取方法，允許自適應地選擇最適合進行快取的token，並進一步為不同型別和深度的神經網路層應用不同的快取比率。

透過在PixArt-α、OpenSora和DiT, 以及FLUX上的廣泛實驗，團隊證明了在影像和影片生成中無需訓練即可實現團隊方法的有效性。例如，在OpenSora和PixArt-α上分別實現了2.36倍和1.93倍的接近無損的生成加速。

背景 Backgrounds

擴散模型(Diffusion Models)在影像生成、影片生成等多種生成任務中展現了出色的效能。近年來，以FLUX, Sora, 可靈等模型為代表的 Diffusion Transformers 透過擴充套件引數量和計算規模進一步推動了視覺生成領域的發展。然而，Diffusion Transformers面臨的一個重大挑戰在於其高計算成本，這導致推理速度緩慢，從而阻礙了其在即時場景中的實際應用。為了解決這一問題，研究者們提出了一系列加速方法，主要集中在減少採樣步數和加速去噪網路模型。

近期，基於特徵快取來實現去噪模型加速的方法由於其優秀的無損加速效能，以及無需訓練的優良效能，受到工業界的廣泛關注。上海交通大學張林峰團隊進一步注意到一個自然而有趣的現象：不同計算層，以及同計算層的不同 Token 對於快取誤差的適應性不同，同樣的快取誤差在不同位置對模型影響最高可以達到數十，百倍的差異，因此有必要進一步將模型加速的粒度由特徵級進一步到 token 級，並考慮瞭如何衡量視覺生成模型中 token 的重要性，以實現重要 token 的篩選保留。

核心貢獻

ToCa首次在DiT加速中中引入 token 級的快取複用策略，並首次從誤差積累與傳播的角度分析特徵快取方法。
ToCa提出4種從不同角度出發，適用於不同情形的token selection策略：
基於Self-Attention Map來評估token對其它token的影響;
基於Cross-Attention Map評估文生圖/影片任務中 image token 對 text token的關注分佈，以加強控制能力;
基於該 token在先前去噪步中的被連續快取複用的次數設計增益策略，鼓勵token在時間步上被更均勻地計算，避免區域性誤差積累過大，破壞全域性影像;
將各個token的重要性得分基於空間分佈進行加權，鼓勵被計算的token在空間上分佈更均勻。
ToCa被應用於多種最新模型上開展實驗，證明了其相比現有方法更加優秀，包含文生圖模型PixArt-alpha，FLUX-dev和FLUX-schnell，文生影片模型 OpenSora，以及基於ImageNet類標籤生成影像的DiT模型。

研究動機

如圖1所示，不同token在相鄰兩步間進行特徵快取引入的誤差值的差異高達幾十上百倍；

圖2說明不同token上引入同樣大小的誤差，這最初幅度相同的誤差在模型推理過程經過積累和傳播，對模型的輸出的影響差異也極大。因此，有必要考慮token級別的特徵快取-複用策略，使得模型的計算更集中在關鍵被需要的token上。

方法

計算流程

ToCa的快取-複用流程如圖3(a)所示：

Cache 初始化首先推理一個完整的時間步，將各層的特徵放入cache中以便使用。

重要性得分計算在使用ToCa的時間步上，對於每一層：先計算各個token的重要性得分，將最低的部分token 標記為cache狀態（例如圖示中ID為1和3的token），不傳入網路層進行計算。

部分計算對於被傳入的token(2,4,5),執行正常的計算,得到它們的輸出。

Cache更新從cache中調出儲存的 token 1,3的輸出，並將計算得到的新的token 2,4,5輸出更到cache中。

通常這樣的一個迴圈長度為2~4個時間步，即1步充分計算後續搭配1至3個ToCa step。此外，ToCa還基於不同層的重要性，設計了隨著層深度上升而衰減的計算比例，詳情請參考論文。

重要性得分計算

如圖4所示，ToCa設計了基於4個不同方面考慮的重要性分數計算，在實際應用中它們以加權求和給出總的重要性得分，詳情請參考論文。

實驗結果

ToCa被應用於文字到影像生成模型 PixArt-alpha, FLUX, 類到影像生成模型 DiT, 以及文字到影片生成模型 OpenSora以驗證其方法有效性，充分的實驗結果證明，ToCa具有超越其他同類方法的加速效果。

影像生成模型: PixArt-alpha，FLUX， DiT

如上圖所示，ToCa相比另兩種加速方法和無加速的高質量原圖對齊效果更佳，且具有更佳的圖-文對齊能力（例如從左到右第四列的wooden dock）。

從FID-30k和 CLIP Score上衡量，ToCa也取得了遠超其他方法的表現。

如上圖所示，ToCa在FLUX 模型上的生成質量也極佳，可以看到和原圖基本沒有差異。但值得考慮的是在文字生成這類對細節要求極其高的任務上（例如左下角的地圖）仍有差異，這將作為團隊後續研究的出發點。

對於高階的模型，使用Image Reward通常能更好地對生成質量進行衡量，團隊分別在50step的FLUX-dev和4step的FLUX-schnell上開展了實驗，可以看到，ToCa在FLUX上1.5倍加速，相比未加速模型的數值指標基本不變，遠遠優於其他方法。

在基礎模型DiT上的結果也證明了ToCa的優越性。

影片生成模型：OpenSora

團隊製作了一個網頁來展示OpenSora上的加速效果。

https://toca2024.github.io/ToCa

此外，團隊將影片生成結果部分抽幀以供快速瀏覽：

在VBench上測試ToCa的加速效果，實驗結果表明，ToCa遠優於其他方法，取得了高達2.36倍的無損加速, 在加速效果和生成質量上都取得最優表現。

ToCa在VBench的大部分指標上都取得了和原模型幾乎相同的得分。

總結

ToCa作為首次被提出的從Token級來實現擴散模型加速的方法，相比以往加速方法具有更強的適配性，（儘管設計時作為專為DiT加速的方案，它的結構也可以被複用到U-Net結構的模型上），同時在多種任務上具有極佳的表現。近年來，包括ToCa在內的系列基於特徵快取的擴散模型加速方法兼具無需訓練的優越性和強勁的無損加速效果，取得了卓越的成效，是一種不同於蒸餾類方法的值得被進一步探索的加速方案。

論文：https://arxiv.org/abs/2410.05317
Github：https://github.com/Shenyi-Z/ToCa

— 完 —

學術投稿請於工作日發郵件到：

[email protected]

標題註明【投稿】，告訴我們：

你是誰，從哪來，投稿內容

附上論文/專案主頁連結，以及聯絡方式哦

我們會（儘量）及時回覆你