什么是扩散 Transformer(Diffusion Transformer)?
扩散 Transformer(Diffusion Transformer,DiT)是 arXiv 论文《Scalable Diffusion Models with Transformers》提出的一类扩散模型,用 Transformer 替换扩散模型中常用的 U-Net 主干,在潜空间图块上运算。该论文以 Gflops 衡量前向计算复杂度,发现 Gflops 更高的 DiT 在 FID 上表现更好,其最大的 DiT-XL/2 在 ImageNet 512x512 与 256x256 类别条件基准上取得当时最优结果。
扩散 Transformer(Diffusion Transformer,缩写 DiT)是一类基于 Transformer 架构的扩散模型。按照 arXiv 论文《Scalable Diffusion Models with Transformers》的描述,它训练图像的潜扩散模型(latent diffusion model),把扩散模型中常用的 U-Net 主干替换为一个在潜空间图块(latent patch)上运算的 Transformer。它要解决的问题是:在扩散模型里,主干网络的选择长期以 U-Net 为主,而 Transformer 在其它生成任务中已经展现出良好的可扩展性,因此该论文尝试把 Transformer 直接用作扩散模型的主干,并考察其规模扩展行为。
为什么重要
在 DiT 出现之前,扩散模型的图像生成主干普遍采用 U-Net。U-Net 通过编码器—解码器结构与跳跃连接(skip connection)在多个分辨率上处理特征,这种设计对图像这类具有空间结构的信号很自然,但它并不是为「按算力规模放大」而设计的架构。当研究者希望用更多算力、更大参数量来提升生成质量时,U-Net 的扩展路径并不清晰:加深、加宽或改变分辨率层级都会同时牵动结构中的多个部分,难以把「计算量」与「生成质量」之间的关系单独拎出来研究。
与此同时,Transformer 在自然语言处理以及视觉识别等任务上已经表现出较为平滑的规模扩展特性:增加深度、宽度或输入 token 数量,通常能带来可预期的性能变化。该论文正是把这一思路迁移到扩散模型上,用 Transformer 作为扩散主干,并明确以「前向计算复杂度」作为衡量规模的标尺,从而让扩散模型的扩展性问题变得可以度量。
该论文报告的结果也说明了这一方向的意义:在类别条件(class-conditional)的 ImageNet 512x512 与 256x256 基准上,其最大的 DiT-XL/2 模型超过了此前所有的扩散模型,在 256x256 上取得了 2.27 的 FID。这一结果由该论文给出,属于该论文自身的实验结论,而非整个领域的统一共识。
工作机制
该论文描述的 DiT 做法可以拆成几个要点。
- 在潜空间而非像素空间工作。 DiT 训练的是图像的潜扩散模型,扩散过程发生在由自编码器压缩得到的潜表示上,而不是直接在原始像素上做去噪。这与潜扩散模型的通用思路一致,其好处是把建模对象从高维像素转移到维度更低的潜空间。
- 把潜表示切成图块。 Transformer 的输入不是整张特征图,而是被划分成的一系列潜空间图块(latent patch)。图块序列化之后,Transformer 就像处理 token 序列一样处理这些图块,注意力机制在全部图块之间建立联系。这一点与视觉 Transformer 把图像切成 patch 的做法在形式上相通。
- 用 Transformer 替换 U-Net 主干。 扩散模型需要一个在每一步去噪中预测噪声(或等价目标)的网络,DiT 把这个位置交给 Transformer,而不是 U-Net。这是该论文最核心的结构改动。
- 以 Gflops 度量规模。 该论文通过前向传播复杂度来考察可扩展性,具体指标是 Gflops。规模可以通过三条途径提高:增加 Transformer 的深度、增加 Transformer 的宽度,或者增加输入 token 的数量。
- 规模与质量的关系。 该论文发现,Gflops 更高的 DiT 一致地具有更低的 FID。也就是说,在上述三种放大方式下,计算量更大的模型在生成质量指标上表现更好,这构成了该论文所称的良好可扩展性。
需要强调的是,上述机制与结论都出自这一篇论文的实验设定,包括其使用的潜扩散框架、图块化方式与 Gflops 度量口径。换用不同的自编码器、不同的图块划分或不同的训练配置,结论是否同样成立,论文本身并未给出跨设定的普遍性断言。
典型例子
该论文中最具代表性的系统是 DiT-XL/2。按照论文摘要的表述,它是该工作中最大的 DiT 模型,在类别条件的 ImageNet 512x512 与 256x256 两个基准上超过了此前所有扩散模型,并在 256x256 上取得 FID 2.27 的结果。这里的「XL」与「/2」属于该论文给出的模型命名,用于区分不同规模的配置。
论文中另一类被反复引用的例子是规模对比本身:通过增加 Transformer 深度或宽度、或增加输入 token 数量来提高 Gflops,并观察对应的 FID 变化。这种「以 Gflops 为横轴、以 FID 为纵轴」的呈现方式,是该论文论证可扩展性的主要证据形式。
该论文的提交信息也构成其出处的一部分:它于 2022 年 12 月 19 日提交第一版,2023 年 3 月 2 日修订为第二版,作者为 William Peebles 与 Saining Xie,发表在 arXiv 的计算机视觉与模式识别(cs.CV)分类下,同时涉及机器学习(cs.LG)分类。论文还提供了代码、项目页面与视频资源。
边界与常见误解
第一,DiT 并不等于「扩散模型」这一整类方法。它特指以 Transformer 作为扩散主干、并在潜空间图块上运算的那一类模型。扩散模型本身有更广的范围,U-Net 主干在 DiT 之外依然是常见选择。
第二,DiT 也不等于「把图像切成 patch 的视觉 Transformer」。虽然两者在输入序列化方式上有相通之处,但 DiT 的用途是在扩散过程中执行去噪,其训练目标、条件注入方式与推理流程都服从扩散模型的框架,而不是分类或表征学习。
第三,该论文关于「Gflops 更高则 FID 更低」的结论,是在其自身的实验范围内得出的。它说明的是该论文所测试的配置下规模与质量之间的经验关系,不能被外推为「任何扩散模型只要加大算力就一定会变好」,也不能被当作对整个领域的普适规律。
第四,FID 与 Gflops 都是特定口径下的度量。FID 依赖用于计算特征分布的参考模型与数据集划分,Gflops 依赖具体的算子实现与输入尺寸。论文报告的 2.27 这一数值,应当连同其基准设定(类别条件 ImageNet 256x256)一起理解,脱离设定单独引用这个数字容易造成误读。
第五,DiT 的代价在于计算与工程复杂度。把主干换成 Transformer 意味着注意力在全部图块上两两交互,输入 token 数量增加会显著抬高计算量;论文正是用 Gflops 来刻画这种代价,并把「增加 token 数量」列为提升规模的途径之一。换言之,可扩展性并不等于低成本,而是指在投入更多计算时能获得可预期的质量回报。
最后需要说明的是,本词条所述内容来自该论文的摘要与著录信息,未包含对该方法的独立复现或实测。论文中出现的模型名称、指标数值与结论,均应视为该论文作者的主张。