PixArt-α: 真实文本到图像合成的扩散Transformer的快速训练.
0. TL; DR
本文介绍了一种名为PixArt-α的文本到图像(T2I)扩散模型,它基于Transformer架构,能够在显著降低训练成本的同时,生成与现有最先进模型(如Imagen、SDXL和Midjourney)相媲美的高质量图像。
PixArt-α通过三个核心设计实现了这一目标:训练策略分解、高效的T2I Transformer和高信息量的数据。该模型仅需12%的Stable Diffusion v1.5训练时间(约753个A100 GPU天),节省了近30万美元的训练成本,并减少了90%的二氧化碳排放。此外,PixArt-α在图像质量、艺术性和语义控制方面表现出色,为AIGC社区和初创公司提供了构建高质量、低成本生成模型的新思路。
1. 背景介绍
近年来,文本到图像(T2I)生成模型取得了显著进展,如DALL·E 2、Imagen和Stable Diffusion等模型,它们能够生成逼真的图像,对图像编辑、视频生成和3D资产创建等下游应用产生了深远影响。然而,这些先进模型的训练需要大量的计算资源,例如Stable Diffusion v1.5需要6000个A100 GPU天,成本高达32万美元,且训练过程会产生大量的二氧化碳排放。这些高昂的成本对研究社区和创业者构成了重大障碍,限制了AIGC社区的发展。因此,开发一种高效、低成本的高质量图像生成器成为了一个重要的研究方向。
2. PixArt-α 模型
PIXART-α将复杂的T2I生成任务分解为三个子任务:学习自然图像的像素分布、学习文本与图像的对齐以及提升图像的审美质量。具体来说:
- 像素依赖学习:通过从预训练的ImageNet类条件模型初始化T2I模型,显著降低了学习成本。
- 文本-图像对齐学习:使用信息密度高的文本-图像对数据进行预训练,然后在具有更好审美质量的数据上进行微调,提高了训练效率。
- 高分辨率和审美图像生成:使用高质量的审美数据对模型进行微调,以生成高分辨率图像。由于前两个阶段建立了强大的先验知识,这一阶段的适应过程收敛速度更快。
(1)高效T2I Transformer
PIXART-α基于Diffusion Transformer(DiT)架构,并进行了以下改进:
- 交叉注意力模块:在DiT块中加入多头交叉注意力层,以便模型能够灵活地与从语言模型中提取的文本嵌入进行交互。为了保留预训练权重,交叉注意力层的输出投影层初始化为零,相当于恒等映射,保留了输入供后续层使用。
- AdaLN-Single:在DiT的自适应归一化层(adaLN)模块中,线性投影占据了大量参数(27%),但这些参数在T2I模型中并不必要。因此,PIXART-α提出了adaLN-Single,仅在第一个块中使用时间嵌入作为输入进行独立控制。具体来说,adaLN-Single在第一个块中计算全局的尺度和偏移参数,这些参数在所有块中共享。然后,通过一个求和函数和一个层特定的可训练嵌入来调整不同块中的尺度和偏移参数。
- 重参数化:为了利用预训练权重,所有层特定的嵌入都被初始化为0值,这些值在选定的时间步长(实验中使用$t = 500$)下产生与没有类别条件的DiT相同的尺度和偏移参数。这种设计有效地用全局MLP和层特定的可训练嵌入替换了层特定的MLP,同时保持了与预训练权重的兼容性。
(2)高信息量数据
现有的文本-图像对数据集(如LAION)存在一些问题,例如文本描述缺乏信息量、词汇使用频率低等,这些问题严重影响了T2I模型的训练效率。为了生成信息密度高的描述,PIXART-α利用最先进的视觉-语言模型LLaVA对SAM数据集进行自动标注。通过使用提示“非常详细地描述这张图片及其风格”,显著提高了描述的质量。此外,SAM数据集包含丰富多样的对象,是创建高信息密度文本-图像对的理想资源。
3. 实验分析
PIXART-α在COCO数据集上的零样本FID得分为7.32,仅用了Stable Diffusion v1.5的12%的训练时间和1.25%的训练样本量。与现有的最先进模型相比,PIXART-α在资源消耗方面显著减少,同时在FID性能上具有可比性。
在T2I-CompBench上,PIXART-α在属性绑定、对象关系和复杂组合等方面表现出色,表明该方法在组合生成能力上具有优势。
用户研究结果表明,PIXART-α在图像质量和文本对齐精度方面优于现有的顶级T2I模型,如DALL·E 2、Stable Diffusion v2和DeepFloyd等。