ICLR 2025 | 計(jì)算量?jī)HDiT一半!中山大學(xué)&360 AI研究院開源Qihoo-T2X:統(tǒng)一架構(gòu)搞定T2X任務(wù)
論文地址:https://arxiv.org/pdf/2409.04005
項(xiàng)目主頁:https://360cvgroup.github.io/Qihoo-T2X
代碼倉(cāng)庫(kù):https://github.com/360CVGroup/Qihoo-T2X
?
作者信息:論文一作為來自中山大學(xué)的博士生王晶;論文共同一作和項(xiàng)目leader為來自360 AI Research視頻生成方向的負(fù)責(zé)人馬傲
亮點(diǎn)直擊
- 提出了Proxy-Tokenized擴(kuò)散 transformer(Proxy-Tokenized Diffusion Transformer,PT-DiT)。
- 推出了Qihoo-T2X系列模型,包括文本到圖像(T2I)、文本到視頻(T2V)以及文本到多視圖(T2MV)生成模型。
- 實(shí)驗(yàn)結(jié)果表明,該方法在保持競(jìng)爭(zhēng)性能的同時(shí)顯著提高了效率。
- 在標(biāo)準(zhǔn)的3D VAE設(shè)置下(8倍空間下采樣率和4倍時(shí)間下采樣率),實(shí)驗(yàn)表明,該方法可以在64GB Ascend 910B上訓(xùn)練PT-DiT/XL(1.1B參數(shù))模型,用于分辨率為2048×2048的圖像生成或分辨率為512×512×288的視頻生成。
Diffusion Transformer模型由于全局self-attention,其計(jì)算復(fù)雜度與序列長(zhǎng)度平方成正比,導(dǎo)致其在高分辨率圖像和長(zhǎng)時(shí)間視頻生成任務(wù)中面臨計(jì)算成本高的問題。
為此,來自中山大學(xué)和360 AI Research的研究人員基于Proxy token提出了一種高效的Diffusion Transformer 即PT-DiT, 能夠適用于文本生成圖像、視頻和Multi-View的等多種生成任務(wù)。作者基于PT-DiT進(jìn)一步構(gòu)建了包含Qihoo-T2I,Qihoo-T2V和Qihoo-T2MV等模型的Qihoo-T2X系列實(shí)現(xiàn)文本生成任意任務(wù)。
研究動(dòng)機(jī)
當(dāng)前,基于Diffusion Transformer 的模型(Sora , Vidu, Flux等)能夠生成高保真度圖像或視頻并與文本指令具有強(qiáng)一致性,極大的促進(jìn)了視覺生成的進(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è)問題也阻礙了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í)間視頻生成的探索。
一些視覺理解和識(shí)別領(lǐng)域的研究發(fā)現(xiàn),由于視覺信息的稀疏和重復(fù)性質(zhì),global self-attention具有一定的冗余性。研究團(tuán)隊(duì)通過可視化注意力圖,發(fā)現(xiàn)同一窗口內(nèi)的不同token對(duì)于空間上距離較遠(yuǎn)的token的關(guān)注程度是相似的,對(duì)于空間上相近的token的關(guān)注程度是不同的。
這表明計(jì)算所有token的注意力是冗余的,而對(duì)計(jì)算空間相鄰token的注意力至關(guān)重要。
所以,研究團(tuán)隊(duì)提出了一種基于proxy token的稀疏注意力策略,從每個(gè)窗口采樣有限的proxy token來執(zhí)行自注意力,從而減少冗余并降低復(fù)雜性。
Qihoo-T2X方法
如圖所示,研究團(tuán)隊(duì)提出的 PT-DiT引入了proxy token來減少計(jì)算global self-attention所涉及的token數(shù)量,高效地建立全局視覺信息的關(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,其通過window attention和shift window attention進(jìn)一步細(xì)化局部細(xì)節(jié)。
下面將詳細(xì)介紹這兩部分:
Global Information Interaction Module
給定一系列l(wèi)atent token,首先根據(jù)空間和時(shí)間先驗(yàn)(即位于同一個(gè)空間窗口)通過計(jì)算平均token得到一系列proxy tokens。
每個(gè)proxy token代表圖像或視頻內(nèi)的一個(gè)局部區(qū)域的信息,并與其他局部區(qū)域中的proxy token通過self-attention進(jìn)行交互以建立全局視覺關(guān)聯(lián)。
隨后,proxy tokens中蘊(yùn)含的信息被通過與latent token的cross-attention傳播到全部latent token中,從而實(shí)現(xiàn)高效的全局視覺信息交互。
Texture Complement Module
由于稀疏proxy tokens交互的特點(diǎn)并且缺乏空間鄰近token的相互關(guān)聯(lián),生成模型對(duì)于建模細(xì)節(jié)紋理的能力有限,難以滿足生成任務(wù)的高質(zhì)量需求。
為了解決這個(gè)問題,研究團(tuán)隊(duì)引入了局部window attention,補(bǔ)充模型的細(xì)節(jié)建模和平滑能力。
僅有window attention 會(huì)導(dǎo)致窗口間token缺乏聯(lián)系,導(dǎo)致生成圖像格子現(xiàn)象明顯,如圖所示。因此,TCM 中還引入了shift window attention,緩解局部window attention引起的問題。
由于計(jì)算window attention涉及的token數(shù)量較少,所以模型的計(jì)算復(fù)雜度并沒有大規(guī)模增加。
壓縮比例
對(duì)于圖像生成任務(wù),研究團(tuán)隊(duì)發(fā)現(xiàn)在不同分辨率下保持相同數(shù)量的窗口對(duì)于確保一致的語義層次結(jié)構(gòu)至關(guān)重要,這有助于從低分辨率到高分辨率的訓(xùn)練過程。 同時(shí),窗口應(yīng)該維持較多的數(shù)量以防止窗口內(nèi)的語義信息太過豐富導(dǎo)致單個(gè)token不足以表示局部區(qū)域完成全局信息建模。
計(jì)算復(fù)雜度分析
PT-DiT 僅使用少量的代表性 token 注意力,就降低了原始全 token 自注意力的計(jì)算冗余度。研究團(tuán)隊(duì)進(jìn)一步從理論上分析PT-DiT在計(jì)算復(fù)雜度方面的優(yōu)勢(shì)。
其中N表示潛在標(biāo)記的長(zhǎng)度,D表示特征維度。
類似地,GIIM 和 TCM 的計(jì)算復(fù)雜度計(jì)算如下:
實(shí)驗(yàn)
作者在T2I T2V和T2MV任務(wù)上進(jìn)行了定性和定量是實(shí)驗(yàn)來評(píng)估Qihoo-T2X。
定性分析
Text-to-Image:如圖所示。Qihoo-T2I 表現(xiàn)出了極具競(jìng)爭(zhēng)力的性能,能夠生成與提供的文本提示完美匹配的照片般逼真的圖像。
Text-to-Video:研究人員將 Qihoo-T2V 與最近發(fā)布的開源文本轉(zhuǎn)視頻模型(即 EasyAnimateV4 和 CogVideoX)在 512 分辨率下進(jìn)行了比較,取得了更好的效果,如圖所示。
Text-to-MultiView:作者進(jìn)一步探索了 PT-DiT 在文本到多視圖 (T2MV) 任務(wù)中的有效性。經(jīng)過訓(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ù),在相同參數(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)
結(jié)論
鑒于視覺信息的稀疏性和冗余性,本文提出了PT-DiT,它利用代理標(biāo)記化注意機(jī)制來緩解擴(kuò)散Transformer中自注意力的計(jì)算冗余?;跁r(shí)間和空間先驗(yàn)計(jì)算出一系列代表性標(biāo)記,并在它們之間進(jìn)行全局交互。此外,引入了窗口注意和移位窗口注意來優(yōu)化局部細(xì)節(jié)的建模。我們提出的代表性標(biāo)記機(jī)制對(duì)于具有冗余信息的視頻任務(wù)特別有效,實(shí)現(xiàn)了3D時(shí)空建模,同時(shí)避免了計(jì)算復(fù)雜度的爆炸。實(shí)驗(yàn)表明,PT-DiT在提供顯著效率的同時(shí)實(shí)現(xiàn)了具有競(jìng)爭(zhēng)力的性能。我們進(jìn)一步基于PT-DiT開發(fā)了Qihoo-T2X系列,包括T2I、T2V和T2MV等模型。希望PT-DiT和Qihoo-T2X能為擴(kuò)散Transformer領(lǐng)域提供新的見解和參考。
本文轉(zhuǎn)自AI生成未來 ,作者:AI生成未來
