性能不輸SOTA模型，計(jì)算開銷卻更低了——

中山大學(xué)和360 AI Research聯(lián)合推出PT-DiT，同參數(shù)規(guī)模下，計(jì)算量?jī)H為DiT的51.4%，Lumina-Next的17.5%。

具體來(lái)說(shuō)，PT-DiT基于Proxy token機(jī)制，能用于文生圖（Qihoo-T2I）、文生視頻（Qihoo-T2V）和文生多視圖（Qihoo-T2MV）等多種任務(wù)。

（Qihoo-T2X指文本到任意視覺(jué)任務(wù)）

圖片

話不多說(shuō)，我們直接看幾個(gè)最終生成效果，文生圖be like：

圖片

接下來(lái)是今年火熱的視頻生成，prompt如下：

Sunset cityscape with spires, buildings, clouds, warm glow, and trees.(夕陽(yáng)下的城市景觀，有尖塔、建筑物、云朵、溫暖的光芒和樹木。)

最后是多視圖生成，寶劍、小黃鴨等任意素材均可實(shí)現(xiàn)轉(zhuǎn)3D效果。

圖片

目前該研究已經(jīng)開放了論文、項(xiàng)目主頁(yè)和代碼倉(cāng)庫(kù)，即將開源。

研究動(dòng)機(jī)

當(dāng)前，基于Diffusion Transformer的模型（Sora , Vidu, Flux等）能夠生成高保真圖像或視頻，并與文本指令具有強(qiáng)一致性，極大促進(jìn)了視覺(jué)生成的進(jìn)步。

然而，global self-attention關(guān)于序列長(zhǎng)度的二次復(fù)雜度增加了Diffusion Transformer的計(jì)算開銷，導(dǎo)致了實(shí)際應(yīng)用時(shí)更長(zhǎng)的生成時(shí)間和更高的訓(xùn)練成本。

這個(gè)問(wèn)題也阻礙了Diffusion Transformer在高質(zhì)量和長(zhǎng)時(shí)間視頻生成中的應(yīng)用。

例如，優(yōu)于2D spatial attention+ 1D temporal attention的3D full attention卻由于計(jì)算開銷的限制而難以進(jìn)行更高分辨率和更長(zhǎng)時(shí)間視頻生成的探索。

一些視覺(jué)理解和識(shí)別領(lǐng)域的研究發(fā)現(xiàn)，由于視覺(jué)信息的稀疏和重復(fù)性質(zhì)，global self-attention具有一定的冗余性。

研究團(tuán)隊(duì)通過(guò)可視化注意力圖發(fā)現(xiàn)：

同一窗口內(nèi)的不同token對(duì)于空間上距離較遠(yuǎn)的token的關(guān)注程度是相似的，對(duì)于空間上相近的token的關(guān)注程度是不同的。

如上圖所示，研究團(tuán)隊(duì)分析了PixArt-??在分辨率為512x512情況下self-attention中的注意力圖。

然后將位于同一個(gè)空間窗口token的注意力圖組合到一起，如圖右側(cè)所示，其中垂直軸表示窗口中的不同token，水平軸表示窗口內(nèi)token與所有token的關(guān)聯(lián)程度。

很明顯，同一窗口內(nèi)不同token的注意力對(duì)于空間上距離較遠(yuǎn)的token幾乎是一致的，即在相同的水平位置，垂直值幾乎相同；而空間相鄰的token表現(xiàn)出不同的關(guān)注。

這表明計(jì)算所有token的注意力是冗余的，而對(duì)計(jì)算空間相鄰token的注意力至關(guān)重要。

所以，研究團(tuán)隊(duì)提出了一種基于proxy token的稀疏注意力策略，從每個(gè)窗口采樣有限的proxy token來(lái)執(zhí)行自注意力，從而減少冗余并降低復(fù)雜性。

Qihoo-T2X方法

如下圖所示，研究團(tuán)隊(duì)提出的PT-DiT引入了proxy token來(lái)減少計(jì)算global self-attention所涉及的token數(shù)量，高效地建立全局視覺(jué)信息的關(guān)聯(lián)。

PT-DiT包含的兩個(gè)核心模塊是：

• Global Information Interaction Module (GIIM)
• Texture Complement Module (TCM)

圖片

其中，GIIM使用稀疏proxy token機(jī)制促進(jìn)所有潛在代碼之間的高效交互，而空間相鄰token的關(guān)聯(lián)是不可忽略的，特別是對(duì)于細(xì)節(jié)紋理要求高的圖像生成任務(wù)。

為此研究團(tuán)隊(duì)設(shè)計(jì)了TCM，其通過(guò)window attention和shift window attention進(jìn)一步細(xì)化局部細(xì)節(jié)。

下面將詳細(xì)介紹這兩個(gè)部分：

Global Information Interaction Module

給定一系列l(wèi)atent token，首先根據(jù)空間和時(shí)間先驗(yàn)（即位于同一個(gè)空間窗口）通過(guò)計(jì)算平均token得到一系列proxy tokens。

每個(gè)proxy token代表圖像或視頻內(nèi)的一個(gè)局部區(qū)域的信息，并與其他局部區(qū)域中的proxy token通過(guò)self-attention進(jìn)行交互以建立全局視覺(jué)關(guān)聯(lián)。

隨后，proxy tokens中蘊(yùn)含的信息被通過(guò)與latent token的cross-attention傳播到全部latent token中，從而實(shí)現(xiàn)高效的全局視覺(jué)信息交互。

Texture Complement Module

由于稀疏proxy tokens交互的特點(diǎn)并且缺乏空間鄰近token的相互關(guān)聯(lián)，生成模型對(duì)于建模細(xì)節(jié)紋理的能力有限，難以滿足生成任務(wù)的高質(zhì)量需求。

為了解決這個(gè)問(wèn)題，研究團(tuán)隊(duì)引入了局部window attention，補(bǔ)充模型的細(xì)節(jié)建模和平滑能力。

不過(guò)僅有window attention會(huì)導(dǎo)致窗口間token缺乏聯(lián)系，導(dǎo)致生成圖像格子現(xiàn)象明顯。

因此，TCM中還引入了shift window attention，緩解局部window attention引起的問(wèn)題。

圖片

由于計(jì)算window attention涉及的token數(shù)量較少，所以模型的計(jì)算復(fù)雜度并沒(méi)有大規(guī)模增加。

壓縮比例

對(duì)于圖像生成任務(wù)，研究團(tuán)隊(duì)發(fā)現(xiàn)在不同分辨率下保持相同數(shù)量的窗口對(duì)于確保一致的語(yǔ)義層次結(jié)構(gòu)至關(guān)重要，這有助于從低分辨率到高分辨率的訓(xùn)練過(guò)程。

同時(shí)，窗口應(yīng)該維持較多的數(shù)量以防止窗口內(nèi)的語(yǔ)義信息太過(guò)豐富，導(dǎo)致單個(gè)token不足以表示局部區(qū)域完成全局信息建模。

因此，研究團(tuán)隊(duì)將壓縮比（????，???，????）設(shè)置為(1,2,2)、(1,4,4)、(1,8,8)和(1,16,16)分別在256、512、1024和2048分辨率。

當(dāng)輸入是圖像時(shí)，??和????將被設(shè)置為1。

對(duì)于視頻生成任務(wù)，研究團(tuán)隊(duì)在不同分辨率下均設(shè)置????=4以保持時(shí)間壓縮一致。

由于幀、高度和寬度維度上的token壓縮，PT-DiT可以訓(xùn)練更長(zhǎng)視頻的生成器。

計(jì)算復(fù)雜度分析

PT-DiT僅使用少量的代表性token注意力，就降低了原始全token自注意力的計(jì)算冗余度。

研究團(tuán)隊(duì)進(jìn)一步從理論上分析PT-DiT在計(jì)算復(fù)雜度方面的優(yōu)勢(shì)。

自注意力的計(jì)算復(fù)雜度為2N2D，計(jì)算如下：

圖片

其中N表示潛在標(biāo)記的長(zhǎng)度，D表示特征維度。

類似地，GIIM和TCM的計(jì)算復(fù)雜度計(jì)算如下：

顯然，由于代理標(biāo)記化策略，PT-DiT具有顯著的優(yōu)勢(shì)，尤其是在壓縮比（????，???，????）較大和序列長(zhǎng)度（N）較長(zhǎng)的情況下。

當(dāng)（????，???，????）為(1,2,2)、(1,4,4)、(1,8,8)和(1,16,16)且圖像分辨率為256（N=256）、512（N=1024） 、1024（N=4096）和2048（N=16384）時(shí)，PT-DiT的計(jì)算復(fù)雜度僅為全局自注意力的34.3%、9.7%、4.7%、2.3%。

此外，PT-DiT對(duì)于序列長(zhǎng)度較長(zhǎng)的視頻生成任務(wù)提供了更大的好處。

實(shí)驗(yàn)

作者在T2I、T2V和T2MV任務(wù)上進(jìn)行了定性和定量實(shí)驗(yàn)來(lái)評(píng)估Qihoo-T2X。

定性分析

對(duì)于Text-to-Image，如圖所示，Qihoo-T2I能夠生成與提供的文本提示非常匹配的逼真圖像。

圖片

對(duì)于Text-to-Video，研究人員將Qihoo-T2V與最近發(fā)布的開源文本轉(zhuǎn)視頻模型（即EasyAnimateV4和CogVideoX）在512分辨率下進(jìn)行了比較，如圖取得了更好的效果。

圖片

最后，作者進(jìn)一步探索了PT-DiT在文本到多視圖 (T2MV) 任務(wù)中的有效性。

經(jīng)過(guò)訓(xùn)練的Qihoo-T2MV能夠根據(jù)提供的文本指令從各個(gè)視點(diǎn)生成512x512x24圖像，表現(xiàn)出強(qiáng)空間一致性。

圖片

定量分析

研究團(tuán)隊(duì)在MS-COCO FID-30K, UCF-101和MSR-VTT等benchmark上定量評(píng)估Qihoo-T2I和Qihoo-T2V。

結(jié)果顯示，Qihoo-T2I和Qihoo-T2V均能實(shí)現(xiàn)有競(jìng)爭(zhēng)力的性能，證明了PT-DiT的有效性。

圖片

計(jì)算復(fù)雜度分析

如圖所示，無(wú)論是圖像或視頻生成任務(wù)，在相同參數(shù)規(guī)模下，PT-DiT相比現(xiàn)有Diffusion Transformer方法，均有大幅度的計(jì)算復(fù)雜度優(yōu)勢(shì)。

同時(shí)對(duì)比3D full attention建模的EasyanimateV4，其訓(xùn)練顯存隨著幀數(shù)的增加而爆炸增長(zhǎng)，而PT-DiT的顯存僅有微弱增長(zhǎng)，表明PT-DiT有潛力完成更長(zhǎng)時(shí)間的視頻生成任務(wù)。

圖片

圖片

最后的消融實(shí)驗(yàn)也驗(yàn)證了PT-DiT中關(guān)鍵設(shè)計(jì)的合理性。

圖片

更多細(xì)節(jié)歡迎查閱原論文。

論文地址：
https://arxiv.org/pdf/2409.04005項(xiàng)目主頁(yè)：
https://360cvgroup.github.io/Qihoo-T2X
代碼倉(cāng)庫(kù)：
https://github.com/360CVGroup/Qihoo-T2X