OpenSora/PixArt都能提速

原標(biāo)題：無(wú)需訓(xùn)練讓擴(kuò)散模型提速2倍，上交大提出Token級(jí)緩存方案｜ICLR‘25
文章來(lái)源：量子位
內(nèi)容字?jǐn)?shù)：5376字

上海交大團(tuán)隊(duì)提出ToCa：Token級(jí)緩存加速擴(kuò)散Transformer模型

近日，上海交通大學(xué)等團(tuán)隊(duì)提出的Toca（Token-wise Caching）模型，通過(guò)token粒度的緩存方法，實(shí)現(xiàn)了圖像和視頻生成模型上無(wú)需訓(xùn)練的兩倍以上的加速，相關(guān)論文已被ICLR 2025接收。該研究針對(duì)擴(kuò)散Transformer模型（如FLUX、Sora等）計(jì)算成本高昂的問(wèn)題，提出了一種創(chuàng)新的解決方案。

1. 背景：擴(kuò)散Transformer模型的挑戰(zhàn)

擴(kuò)散模型在圖像和視頻生成領(lǐng)域表現(xiàn)出色，但基于擴(kuò)散Transformer的模型計(jì)算成本巨大，推理速度緩慢，限制了其在實(shí)時(shí)應(yīng)用中的發(fā)展。雖然已有特征緩存方法用于加速，但這些方法忽略了不同token對(duì)緩存的敏感性差異，某些token的緩存甚至可能導(dǎo)致生成質(zhì)量大幅下降。

2. ToCa的核心貢獻(xiàn)

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

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

3. ToCa的工作機(jī)制

ToCa的緩存-復(fù)用流程包括三個(gè)步驟：

1. Cache初始化：推理一個(gè)完整的時(shí)間步，將各層的特征放入緩存。
2. 重要性得分計(jì)算：計(jì)算每個(gè)token的重要性得分，將重要性最低的token標(biāo)記為緩存狀態(tài)，不參與后續(xù)計(jì)算。
3. Cache更新：從緩存中調(diào)出被緩存token的輸出，并將新計(jì)算的token輸出更新到緩存中。這個(gè)循環(huán)通常持續(xù)2-4個(gè)時(shí)間步。

ToCa設(shè)計(jì)了四種從不同角度計(jì)算token重要性得分的方法，并在實(shí)際應(yīng)用中進(jìn)行加權(quán)求和。

4. 實(shí)驗(yàn)結(jié)果與分析

ToCa在圖像生成模型（PixArt-α、FLUX、DiT）和視頻生成模型（OpenSora）上的實(shí)驗(yàn)結(jié)果表明，其加速效果顯著優(yōu)于其他方法，并在生成質(zhì)量上保持了較高的水平。例如，在OpenSora上實(shí)現(xiàn)了高達(dá)2.36倍的無(wú)損加速。實(shí)驗(yàn)還顯示，ToCa在圖像文本對(duì)齊和細(xì)節(jié)還原方面也具有優(yōu)勢(shì)。

5. 總結(jié)

ToCa作為首個(gè)基于token級(jí)緩存加速擴(kuò)散Transformer模型的方法，具有良好的適配性和高效的加速效果，為擴(kuò)散模型的實(shí)際應(yīng)用提供了新的可能性。其無(wú)需訓(xùn)練的特性以及在多種模型和任務(wù)上的優(yōu)異表現(xiàn)，使其成為一種具有前景的模型加速方案。

聯(lián)系作者

文章來(lái)源：量子位
作者微信：
作者簡(jiǎn)介：追蹤人工智能新趨勢(shì)，關(guān)注科技行業(yè)新突破