PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis
公众:EDPJ(添加 VX:CV_EDPJ 或直接进 Q 交流群:922230617 获取资料)
目录
0. 摘要
最先进的文本到图像(T2I)模型需要显著的训练成本(例如,数百万小时的 GPU),严重阻碍了AIGC 社区的基础创新,同时增加了二氧化碳排放。本文介绍了 PIXART-α,一种基于 Transformer 的 T2I 扩散模型,其图像生成质量与最先进的图像生成器(例如,Imagen、SDXL,甚至 Midjourney)相媲美,达到了接近商业应用标准的水平。此外,它支持高分辨率图像合成,最高可达 1024×1024 分辨率,训练成本低,如图 1 和 2 所示。为了实现这一目标,提出了三个核心设计:(1)训练策略分解:我们设计了三个不同的训练步骤,分别优化像素依赖性,文本-图像对齐和图像审美质量;(2)高效 T2I Transformer:我们在扩散 Transformer(Diffusion Transformer,DiT)中加入交叉注意力模块,以注入文本条件并简化计算密集型的类别条件分支;(3)高信息数据:我们强调文本-图像对的概念密度的重要性,并利用大型视觉语言模型自动标注密集的伪标题,以帮助文本-图像对齐学习。因此,PIXART-α 的训练速度明显超过现有的大规模 T2I 模型,例如,PIXART-α 仅占稳定扩散 v1.5 训练时间的 10.8%(约 675 对 6250 A100 GPU天),节省了近 30 万美元(26000 对 320000 美元),减少了 90% 的二氧化碳排放。此外,与更大的 SOTA 模型 RAPHAEL 相比,我们的训练成本仅为 1%。大量实验证明 PIXART-α 在图像质量、艺术性和语义控制方面表现出色。我们希望 PIXART-α 将为 AIGC 社区和初创公司提供新的见解,加速从零开始构建其自己高质量而低成本的生成模型。
2. 方法
2.1 动机
T2I 训练缓慢的原因在于两个方面:训练流程和数据。
T2I 生成任务可以分解为三个方面:捕捉像素依赖性:生成逼真的图像涉及理解图像内部复杂的像素级依赖关系并捕捉它们的分布;文本与图像之间的对齐:需要精确的对齐学习,以了解如何生成准确匹配文本描述的图像;高审美质量:除了忠实的文本描述之外,还要使生成的图像在审美上令人愉悦,这是生成图像的另一个重要属性。当前的方法将这三个问题纠缠在一起,并直接使用大量数据从头开始训练,导致训练效率低下。为了解决这个问题,我们将这些方面分解为三个阶段,如将在第 2.2 节中描述的那样。
另一个问题如图 3 所示,与当前数据集的标题质量有关。当前的文本-图像对经常受到文本-图像不对齐、描述不足、词汇使用不够丰富以及包含低质量数据的困扰。这些问题引入了训练的困难,导致不必要地进行数百万次迭代才能实现文本和图像之间的稳定对齐。为了解决这个挑战,我们引入了创新的自动标注流程,生成精确的图像标题,如将在第 2.4 节中描述的那样。
2.2 训练策略分解
通过将训练分为三个具有不同数据类型的阶段,可以逐渐优化模型的生成能力。
第一阶段:像素依赖性学习。当前的类别引导方法(Peebles&Xie,2023)在生成单个图像中语义连贯且合理的像素方面表现出色。为自然图像训练类别条件图像生成模型(Peebles&Xie,2023)相对容易且成本较低,详见附录 A.5 的解释。此外,我们发现适当的初始化可以显著提高训练效率。因此,我们从一个 ImageNet 预训练模型中提升我们的模型,并且我们的模型的架构被设计为与预训练权重兼容。
第二阶段:文本-图像对齐学习。从预训练的类别引导图像生成过渡到文本到图像生成的主要挑战在于如何实现显著增加的文本概念和图像之间的准确对齐。
这个对齐过程不仅耗时,而且本质上具有挑战性。为了有效促进这个过程,我们构建了一个由精确文本-图像对组成且具有高概念密度的数据集。数据创建流程将在第 2.4 节中描述。通过使用准确且信息丰富的数据,我们的训练过程可以在每次迭代中高效处理更多的名词,与以前的数据集相比遇到的模糊性明显较少。这种战略性的方法使我们的网络能够有效地将文本描述与图像对齐。
第三阶段:高分辨率和审美图像生成。在第三阶段,我们使用高质量审美数据对我们的模型进行微调,用于高分辨率图像生成。值得注意的是,我们观察到这个阶段的适应过程收敛速度显著更快,主要是因为在前面的阶段建立的强大先验知识。
将训练过程分解为不同阶段显著减轻了训练困难,实现了高效的训练。
2.3 高效 T2I Transformer
PIXART-α 采用扩散 Transformer(Diffusion Transformer,DiT)(Peebles&Xie,2023)作为基础架构,并创新地调整Transformer 块以处理 T2I 任务的独特挑战,如图 4 所示。提出了几个专门的设计,如下:
交叉注意力层。我们将一个多头的交叉注意力层整合到 DiT 块中。它位于自注意力层和前馈层之间,使模型可以灵活地与从语言模型中提取的文本嵌入进行交互。为了促进预训练权重的使用,我们将交叉注意力层中的输出投影层初始化为零,实际上起到身份映射(identity mapping)的作用,并保留输入以供后续层使用。
AdaLN-single。我们发现 DiT 的自适应归一化层(adaptive normalization layers,adaLN)模块中的线性投影占据了参数的相当大比例(27%)。由于类别条件没有用于我们的 T2I 模型,这么多的参数是没有用的。因此,我们提出 adaLN-single,它只在第一个块中使用时间嵌入作为输入进行独立控制(如图 4 右侧所示)。具体而言,在第 i 个块中,令
为 adaLN 中所有缩放和平移参数的元组。在 DiT 中,S^(i) 通过一个块特定的 MLP S^(i) = f^(i)(c+t) 获得,其中 c 和 t 分别表示类别条件和时间嵌入。然而,在 adaLN-single 中,一组全局的平移和缩放在第一个块中计算为 ~S = f(t),它在所有块之间共享。然后,S^(i) 通过 S^(i) = g(~S, E(i)) 获得,其中 g 是一个求和函数,E^(i) 是一个与 ~S 相同形状的层特定的可训练嵌入,它在不同块中自适应地调整缩放和平移参数。
重新参数化。为了利用前述的预训练权重,所有的 E^(i) 都初始化为一个值,该值获得与 DiT 在选定 t 且没有 c 的情况下相同的 S^(i) 值(在实证中,我们使用 t = 500)。这个设计有效地用全局MLP 和层特定的可训练嵌入替换了层特定的 MLPs,同时保持了与预训练权重的兼容性。
实验证明,整合全局 MLP 和逐层嵌入以获取时间步信息,以及使用交叉注意力层处理文本信息,可以保持模型的生成能力,同时有效地减小模型的尺寸。
2.4 数据集构建
图像文本对自动标注。LAION 数据集的标题存在各种问题,如文本图像不对齐、描述不足以及词汇使用不足,如图 3 所示。为了生成信息密度高的标题,我们利用了最先进的视觉语言模型 LLaVA(Liu等,2023)。通过使用提示 “以非常详细的方式描述这张图像及其风格”,我们显著提高了标题的质量,如图 3 所示。
然而,值得注意的是,LAION 数据集主要由来自购物网站的简单产品预览组成,这对于训练寻求对象组合多样性的文本到图像生成并不理想。因此,我们选择使用 SAM 数据集(Kirillov等,2023),该数据集最初用于分割任务,但其图像富含各种对象。通过将 LLaVA 应用于 SAM,我们成功地获得了以高概念密度为特征的高质量文本图像对,如附录中的图 10 和图 11 所示。
在第三阶段,我们通过整合 JourneyDB(Pan等,2023)和一个 1000 万的内部数据集来构建我们的训练数据集,以提高生成图像的审美质量,超越了真实照片。详细信息请参见附录 A.5。
因此,我们在表 1 中展示了词汇分析(NLTK,2023),并将出现在数据集中超过 10 次的词汇定义为有效的不同名词(distinct noun)。我们在 LAION 上应用 LLaVA 生成 LAION-LLaVA。LAION 数据集有 2.46 百万个不同的名词,但只有 8.5% 是有效的。这有效名词比例从 8.5% 显著增加到 13.3%。尽管LAION 的原始标题包含惊人的 21 万个不同名词,但其总名词数仅为 7200 万。然而,LAION-LLaVA 包含 234M 名词数,有 85K 个不同名词,每张图像的平均名词数从 6.4 增加到 21,表明原始 LAION 标题的不完整性。此外,SAM-LLaVA 在总名词数方面优于 LAION-LLaVA,为 328M,每张图像 30 个名词,表明 SAM 包含更丰富的目标和更高的信息密度。最后,内部数据还确保了足够的有效名词和用于微调的平均信息密度。LLaVA 标记的标题显著提高了有效比例和每张图像的平均名词数,提高了概念密度。
3. 实验
PIXART-α 在具有较少模型参数、较少训练数据的情况下,表现良好,且需要的训练资源少。
在图文对齐方面,PIXART-α 表现突出。
在人类评估中,人们认为 PIXART-α 生成的图像质量以及与文本对齐程度高。
消融研究的视觉结果清楚地表明,“adaLN” 和 “adaLN-single” 模型之间的 FID 分数差异相对较小,但前者视觉质量较差且参数量多。“无重新参数化(w/o re-param)” 模型在整个测试集上一贯显示出扭曲的目标图像,并且缺乏关键的细节。
附录
A.8 评估图像质量的 FID 度量的讨论
在我们的实验中,我们观察到 FID(Fréchet Inception Distance)分数可能不能准确反映生成图像的视觉质量。最近的研究,如 SDXL(Podell等人,2023)和 Pick-a-pic(Kirstain等人,2023),提供了证据表明 COCO 零样本 FID 与视觉美学呈负相关。
此外,Betzalel 等人(Betzalel等人,2022)指出,FID 中使用的特征提取网络在 ImageNet 数据集上进行了预训练,而 ImageNet 数据集与当前文本到图像生成数据的重叠有限。因此,FID 可能不是评估这些模型生成性能的合适指标,Betzalel 等人建议使用人类评估者进行更合适的评估。
A.11 限制和失败案例
在图 19 中,我们用红色文本和黄色圆圈突出了模型的失败案例。我们的分析揭示了模型在准确控制目标数量和处理特定细节(如人手的特征)方面的弱点。此外,由于我们的数据中字体和字母相关图像数量有限,模型的文本生成能力也相对较弱。我们的目标是在未来探索这些在生成领域尚未解决的问题,提升模型在文本生成、细节控制和数量控制方面的能力。
S. 总结
S.1 主要贡献
本文提出 PIXART-α,为了以低训练成本支持高分辨率图像合成,提出三个设计:
- 训练策略分解:设计三个不同的训练步骤,分别优化像素依赖性,文本-图像对齐和图像审美质量;
- 高效 T2I Transformer:在扩散 Transformer(Diffusion Transformer,DiT)中加入交叉注意力模块,以注入文本条件并简化计算密集型的类别条件分支;
- 高信息数据:强调文本-图像对的概念密度的重要性,并利用大型视觉语言模型自动标注密集的伪标题,以帮助文本-图像对齐学习。
S.2 方法
训练策略分解。
- 第一阶段:像素依赖性学习。使用类别引导方法生成像素,并从一个 ImageNet 预训练模型初始化来提升模型,并且模型的架构被设计为与预训练权重兼容。
- 第二阶段:文本-图像对齐学习。从预训练的类别引导图像生成过渡到文本到图像生成的主要挑战在于如何实现显著增加的文本概念和图像之间的准确对齐。为解决这个问题,构建了一个由精确文本-图像对组成且具有高概念密度的数据集。
- 第三阶段:高分辨率和审美图像生成。使用高质量审美数据对模型进行微调,用于高分辨率图像生成。
高效 T2I Transformer。 PIXART-α 使用的架构如图 4 所示。
- 交叉注意力层。将一个多头的交叉注意力层整合到 DiT 块中。为了促进预训练权重的使用,将交叉注意力层中的输出投影层初始化为零。
- AdaLN-single。在第一个块中只使用时间嵌入作为输入进行独立控制,在其他块使用可训练嵌入,在不同块中自适应地调整缩放和平移参数。
- 重新参数化。为了利用前述的预训练权重,所有的可训练嵌入都进行初始化。
数据集构建。利用了最先进的视觉语言模型 LLaVA,通过使用提示 “以非常详细的方式描述这张图像及其风格”,生成信息密度高的标题。