ICLR2025|計算量僅DiT的一半!一個模型架構搞定T2X任務

點選下方卡片,關注“CVer”公眾號

AI/CV重磅乾貨,第一時間送達

新增微訊號:CVer2233,小助手會拉你進群!
掃描下方二維碼,加入CVer學術星球可以獲得最新頂會/頂刊上的論文ideaCV從入門到精通資料,及最前沿應用!發論文/搞科研/漲薪,強烈推薦!

Diffusion Transformer模型由於全域性self-attention,其計算複雜度與序列長度平方成正比,導致其在高解析度影像和長時間影片生成任務中面臨計算成本高的問題。
為此,來自中山大學和360 AI Research的研究人員基於Proxy token提出了一種高效的Diffusion Transformer 即PT-DiT, 能夠適用於文字生成影像、影片和Multi-View的等多種生成任務。作者基於PT-DiT進一步構建了包含Qihoo-T2I,Qihoo-T2V和Qihoo-T2MV等模型的Qihoo-T2X系列實現文字生成任意任務。
同參數規模下,PT-DiT的計算量僅為Lumina-Next的

,DiT的

。在降低計算量的同時,Qihoo-T2X展現出了與PixArt-

,Flux,Easyanimate和CogVideo相似的效能,能夠根據給定指令生成逼真的高質量影像和影片。

目前該研究已開源。
論文地址:https://arxiv.org/pdf/2409.04005
專案主頁:https://360cvgroup.github.io/Qihoo-T2X
程式碼倉庫:https://github.com/360CVGroup/Qihoo-T2X

研究動機

當前,基於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-

在解析度為

情況下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不足以表示區域性區域完成全域性資訊建模。
因此,研究團隊將壓縮比

設定為

分別在 256、512、1024 和 2048 解析度。

當輸入是影像時,

將被設定為1。

對於影片生成任務,研究團隊在不同解析度下均設定

= 4 以保持時間壓縮一致。由於幀、高度和寬度維度上的token壓縮,PT-DiT 可以訓練更長影片的生成器。

計算複雜度分析
PT-DiT 僅使用少量的代表性 token 注意力,就降低了原始全 token 自注意力的計算冗餘度。研究團隊進一步從理論上分析PT-DiT在計算複雜度方面的優勢。
自注意力的計算複雜度為

,計算如下:

表示潛在標記的長度,

表示特徵維度。

類似地,GIIM 和 TCM 的計算複雜度計算如下:   

較大和序列長度(

)較長的情況下。當

且影像解析度為 256

、512

、1024

和 2048

時,PT-DiT的計算複雜度僅為全域性自注意力的

。此外,PT-DiT 對於序列長度較長的影片生成任務提供了更大的好處。

實驗
作者在T2I T2V和T2MV任務上進行了定性和定量是實驗來評估Qihoo-T2X。
定性分析
Text-to-Image:如圖所示。Qihoo-T2I 表現出了極具競爭力的效能,能夠生成與提供的文字提示完美匹配的照片般逼真的影像。   
Text-to-Video:研究人員將 Qihoo-T2V 與最近釋出的開源文字轉影片模型(即 EasyAnimateV4 和 CogVideoX)在 512 解析度下進行了比較,取得了更好的效果,如圖所示。   
Text-to-MultiView:作者進一步探索了 PT-DiT 在文字到多檢視 (T2MV) 任務中的有效性。
經過訓練的 Qihoo-T2MV 能夠根據提供的文字指令從各個視點生成 512 \times 512 \times 24$ 影像,表現出很強的空間一致性,如圖所示。   
定量分析
研究團隊在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有潛力完成更長時間的影片生成任務。

消融實驗

何愷明在MIT授課的課件PPT下載


在CVer公眾號後臺回覆:何愷明,即可下載本課程的所有566頁課件PPT!趕緊學起來!

ECCV 2024 論文和程式碼下載


在CVer公眾號後臺回覆:ECCV2024,即可下載ECCV 2024論文和程式碼開源的論文合集

CVPR 2024 論文和程式碼下載


在CVer公眾號後臺回覆:CVPR2024,即可下載CVPR 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼,或者新增微訊號:CVer2233,即可新增CVer小助手微信,便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋:目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註:研究方向+地點+學校/公司+暱稱(如Mamba、多模態學習或者論文投稿+上海+上交+卡卡),根據格式備註,可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233,進交流群
CVer計算機視覺(知識星球)來了!想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料,歡迎掃描下方二維碼,加入CVer計算機視覺(知識星球),已彙集上萬人!

掃碼加入星球學習
▲點選上方卡片,關注CVer公眾號
整理不易,請贊和在看


相關文章