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

中山大學和360 AI Research聯合推出PT-DiT，同參數規模下，計算量僅為DiT的51.4%，Lumina-Next的17.5%。

具體來說，PT-DiT基于Proxy token機制，能用于文生圖（Qihoo-T2I）、文生視頻（Qihoo-T2V）和文生多視圖（Qihoo-T2MV）等多種任務。

（Qihoo-T2X指文本到任意視覺任務）

圖片

話不多說，我們直接看幾個最終生成效果，文生圖be like：

圖片

接下來是今年火熱的視頻生成，prompt如下：

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

最后是多視圖生成，寶劍、小黃鴨等任意素材均可實現轉3D效果。

圖片

目前該研究已經開放了論文、項目主頁和代碼倉庫，即將開源。

研究動機

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

然而，global self-attention關于序列長度的二次復雜度增加了Diffusion Transformer的計算開銷，導致了實際應用時更長的生成時間和更高的訓練成本。

這個問題也阻礙了Diffusion Transformer在高質量和長時間視頻生成中的應用。

例如，優于2D spatial attention+ 1D temporal attention的3D full attention卻由于計算開銷的限制而難以進行更高分辨率和更長時間視頻生成的探索。

一些視覺理解和識別領域的研究發現，由于視覺信息的稀疏和重復性質，global self-attention具有一定的冗余性。

研究團隊通過可視化注意力圖發現：

同一窗口內的不同token對于空間上距離較遠的token的關注程度是相似的，對于空間上相近的token的關注程度是不同的。

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

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

很明顯，同一窗口內不同token的注意力對于空間上距離較遠的token幾乎是一致的，即在相同的水平位置，垂直值幾乎相同；而空間相鄰的token表現出不同的關注。

這表明計算所有token的注意力是冗余的，而對計算空間相鄰token的注意力至關重要。

所以，研究團隊提出了一種基于proxy token的稀疏注意力策略，從每個窗口采樣有限的proxy token來執行自注意力，從而減少冗余并降低復雜性。

Qihoo-T2X方法

如下圖所示，研究團隊提出的PT-DiT引入了proxy token來減少計算global self-attention所涉及的token數量，高效地建立全局視覺信息的關聯。

PT-DiT包含的兩個核心模塊是：

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

圖片

其中，GIIM使用稀疏proxy token機制促進所有潛在代碼之間的高效交互，而空間相鄰token的關聯是不可忽略的，特別是對于細節紋理要求高的圖像生成任務。

為此研究團隊設計了TCM，其通過window attention和shift window attention進一步細化局部細節。

下面將詳細介紹這兩個部分：

Global Information Interaction Module

給定一系列latent token，首先根據空間和時間先驗（即位于同一個空間窗口）通過計算平均token得到一系列proxy tokens。

每個proxy token代表圖像或視頻內的一個局部區域的信息，并與其他局部區域中的proxy token通過self-attention進行交互以建立全局視覺關聯。

隨后，proxy tokens中蘊含的信息被通過與latent token的cross-attention傳播到全部latent token中，從而實現高效的全局視覺信息交互。

Texture Complement Module

由于稀疏proxy tokens交互的特點并且缺乏空間鄰近token的相互關聯，生成模型對于建模細節紋理的能力有限，難以滿足生成任務的高質量需求。

為了解決這個問題，研究團隊引入了局部window attention，補充模型的細節建模和平滑能力。

不過僅有window attention會導致窗口間token缺乏聯系，導致生成圖像格子現象明顯。

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

圖片

由于計算window attention涉及的token數量較少，所以模型的計算復雜度并沒有大規模增加。

壓縮比例

對于圖像生成任務，研究團隊發現在不同分辨率下保持相同數量的窗口對于確保一致的語義層次結構至關重要，這有助于從低分辨率到高分辨率的訓練過程。

同時，窗口應該維持較多的數量以防止窗口內的語義信息太過豐富，導致單個token不足以表示局部區域完成全局信息建模。

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

當輸入是圖像時，??和????將被設置為1。

對于視頻生成任務，研究團隊在不同分辨率下均設置????=4以保持時間壓縮一致。

由于幀、高度和寬度維度上的token壓縮，PT-DiT可以訓練更長視頻的生成器。

計算復雜度分析

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

研究團隊進一步從理論上分析PT-DiT在計算復雜度方面的優勢。

自注意力的計算復雜度為2N2D，計算如下：

圖片

其中N表示潛在標記的長度，D表示特征維度。

類似地，GIIM和TCM的計算復雜度計算如下：

顯然，由于代理標記化策略，PT-DiT具有顯著的優勢，尤其是在壓縮比（????，???，????）較大和序列長度（N）較長的情況下。

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

此外，PT-DiT對于序列長度較長的視頻生成任務提供了更大的好處。

實驗

作者在T2I、T2V和T2MV任務上進行了定性和定量實驗來評估Qihoo-T2X。

定性分析

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

圖片

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

圖片

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

經過訓練的Qihoo-T2MV能夠根據提供的文本指令從各個視點生成512x512x24圖像，表現出強空間一致性。

圖片

定量分析

研究團隊在MS-COCO FID-30K, UCF-101和MSR-VTT等benchmark上定量評估Qihoo-T2I和Qihoo-T2V。

結果顯示，Qihoo-T2I和Qihoo-T2V均能實現有競爭力的性能，證明了PT-DiT的有效性。

圖片

計算復雜度分析

如圖所示，無論是圖像或視頻生成任務，在相同參數規模下，PT-DiT相比現有Diffusion Transformer方法，均有大幅度的計算復雜度優勢。

同時對比3D full attention建模的EasyanimateV4，其訓練顯存隨著幀數的增加而爆炸增長，而PT-DiT的顯存僅有微弱增長，表明PT-DiT有潛力完成更長時間的視頻生成任務。

圖片

圖片

最后的消融實驗也驗證了PT-DiT中關鍵設計的合理性。

圖片

更多細節歡迎查閱原論文。

論文地址：
https://arxiv.org/pdf/2409.04005項目主頁：
https://360cvgroup.github.io/Qihoo-T2X
代碼倉庫：
https://github.com/360CVGroup/Qihoo-T2X