
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生成未来
