ICLR 2025 | 計算量僅DiT一半!中山大學&360 AI研究院開源Qihoo-T2X:統一架構搞定T2X任務
論文地址：https://arxiv.org/pdf/2409.04005
項目主頁：https://360cvgroup.github.io/Qihoo-T2X
代碼倉庫：https://github.com/360CVGroup/Qihoo-T2X
?
作者信息：論文一作為來自中山大學的博士生王晶；論文共同一作和項目leader為來自360 AI Research視頻生成方向的負責人馬傲
亮點直擊
- 提出了Proxy-Tokenized擴散 transformer（Proxy-Tokenized Diffusion Transformer，PT-DiT）。
- 推出了Qihoo-T2X系列模型，包括文本到圖像（T2I）、文本到視頻（T2V）以及文本到多視圖（T2MV）生成模型。
- 實驗結果表明，該方法在保持競爭性能的同時顯著提高了效率。
- 在標準的3D VAE設置下（8倍空間下采樣率和4倍時間下采樣率），實驗表明，該方法可以在64GB Ascend 910B上訓練PT-DiT/XL（1.1B參數）模型，用于分辨率為2048×2048的圖像生成或分辨率為512×512×288的視頻生成。
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系列實現文本生成任意任務。
研究動機
當前，基于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的關注程度是不同的。
這表明計算所有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不足以表示局部區域完成全局信息建模。
計算復雜度分析
PT-DiT 僅使用少量的代表性 token 注意力，就降低了原始全 token 自注意力的計算冗余度。研究團隊進一步從理論上分析PT-DiT在計算復雜度方面的優勢。
其中N表示潛在標記的長度，D表示特征維度。
類似地，GIIM 和 TCM 的計算復雜度計算如下：
實驗
作者在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 能夠根據提供的文本指令從各個視點生成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，它利用代理標記化注意機制來緩解擴散Transformer中自注意力的計算冗余。基于時間和空間先驗計算出一系列代表性標記，并在它們之間進行全局交互。此外，引入了窗口注意和移位窗口注意來優化局部細節的建模。我們提出的代表性標記機制對于具有冗余信息的視頻任務特別有效，實現了3D時空建模，同時避免了計算復雜度的爆炸。實驗表明，PT-DiT在提供顯著效率的同時實現了具有競爭力的性能。我們進一步基于PT-DiT開發了Qihoo-T2X系列，包括T2I、T2V和T2MV等模型。希望PT-DiT和Qihoo-T2X能為擴散Transformer領域提供新的見解和參考。
本文轉自AI生成未來 ，作者：AI生成未來
