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

OpenSora/PixArt都能提速

原標題：無需訓練讓擴散模型提速2倍，上交大提出Token級緩存方案｜ICLR‘25
文章來源：量子位
內容字數：5376字

上海交大團隊提出ToCa：Token級緩存加速擴散Transformer模型

近日，上海交通大學等團隊提出的Toca（Token-wise Caching）模型，通過token粒度的緩存方法，實現了圖像和視頻生成模型上無需訓練的兩倍以上的加速，相關論文已被ICLR 2025接收。該研究針對擴散Transformer模型（如FLUX、Sora等）計算成本高昂的問題，提出了一種創新的解決方案。

1. 背景：擴散Transformer模型的挑戰

擴散模型在圖像和視頻生成領域表現出色，但基于擴散Transformer的模型計算成本巨大，推理速度緩慢，限制了其在實時應用中的發展。雖然已有特征緩存方法用于加速，但這些方法忽略了不同token對緩存的敏感性差異，某些token的緩存甚至可能導致生成質量大幅下降。

2. ToCa的核心貢獻

ToCa首次在擴散Transformer模型加速中引入token級的緩存復用策略，并從誤差積累與傳播的角度分析特征緩存方法。其主要貢獻包括：

1. 提出了基于token的特征緩存方法，自適應地選擇最適合緩存的token，并根據不同網絡層應用不同的緩存比率。
2. 設計了四種token選擇策略，分別基于Self-Attention Map、Cross-Attention Map、token連續緩存次數和空間分布加權，以更有效地選擇需要緩存的token。
3. 在PixArt-α、OpenSora、DiT和FLUX等多種模型上進行了廣泛實驗，證明了ToCa的有效性，在OpenSora和PixArt-α上分別實現了2.36倍和1.93倍的接近無損加速。

3. ToCa的工作機制

ToCa的緩存-復用流程包括三個步驟：

1. Cache初始化：推理一個完整的時間步，將各層的特征放入緩存。
2. 重要性得分計算：計算每個token的重要性得分，將重要性最低的token標記為緩存狀態，不參與后續計算。
3. Cache更新：從緩存中調出被緩存token的輸出，并將新計算的token輸出更新到緩存中。這個循環通常持續2-4個時間步。

ToCa設計了四種從不同角度計算token重要性得分的方法，并在實際應用中進行加權求和。

4. 實驗結果與分析

ToCa在圖像生成模型（PixArt-α、FLUX、DiT）和視頻生成模型（OpenSora）上的實驗結果表明，其加速效果顯著優于其他方法，并在生成質量上保持了較高的水平。例如，在OpenSora上實現了高達2.36倍的無損加速。實驗還顯示，ToCa在圖像文本對齊和細節還原方面也具有優勢。

5. 總結

ToCa作為首個基于token級緩存加速擴散Transformer模型的方法，具有良好的適配性和高效的加速效果，為擴散模型的實際應用提供了新的可能性。其無需訓練的特性以及在多種模型和任務上的優異表現，使其成為一種具有前景的模型加速方案。

聯系作者

文章來源：量子位
作者微信：
作者簡介：追蹤人工智能新趨勢，關注科技行業新突破