ICLR 2025 | 计算量仅DiT的一半！一个模型架构搞定T2X任务

前言 Diffusion Transformer模型由于全局self-attention，其计算复杂度与序列长度平方成正比，导致其在高分辨率图像和长时间视频生成任务中面临计算成本高的问题。为此，来自中山大学和360 AI Research的研究人员基于Proxy token提出了一种高效的Diffusion Transformer 即PT-DiT, 能够适用于文本生成图像、视频和Multi-View的等多种生成任务。

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

本文转载自CVer

仅用于学术分享，若侵权请联系删除

CV方向的准研究生们，未来三年如何度过？

招聘高光谱图像、语义分割、diffusion等方向论文指导老师

目前该研究已开源。

论文地址：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的关注程度是不同的。

如上图所示，研究团队分析了在分辨率为 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不足以表示局部区域完成全局信息建模。

计算复杂度分析

PT-DiT 仅使用少量的代表性 token 注意力，就降低了原始全 token 自注意力的计算冗余度。研究团队进一步从理论上分析PT-DiT在计算复杂度方面的优势。

自注意力的计算复杂度计算如下：

类似地，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 能够根据提供的文本指令从各个视点生成 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有潜力完成更长时间的视频生成任务。