今天小編分享的科學經驗：無需訓練讓擴散模型提速2倍，上交大提出Token級緩存方案，歡迎閱讀。
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
— 完 —
學術投稿請于工作日發郵件到：
标題注明【投稿】，告訴我們：
你是誰，從哪來，投稿内容
附上論文 / 項目主頁鏈接，以及聯系方式哦
我們會（盡量）及時回復你
一鍵關注 點亮星标
科技前沿進展每日見
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法！
