PhoenixPeng's blog.

扩散模型——15.PIXART-α

2024/10/28

1、引言

  作者表示目前高级AIGC模型训练需要大量的计算资源。例如训练SD v1.5 需要 6k A100 GPU天,大约花费320,000美元,训练RAPHAEL甚至需要60k A100 GPU天,大约花费3080,000美元,如图2所示

图2:T2I 生成器的二氧化碳排放量和培训成本比较。 PIXART-α 的培训成本极低,仅为 28,400 美元。与RAPHAEL相比,PIXART-α的二氧化碳排放量和培训成本分别仅为1.2%和0.91%。

  如此巨大的成本给研究界和企业家带来了巨大的障碍。访问这些模型,对 AIGC 社区的关键进步造成重大阻碍。考虑到这些挑战,作者引出了一个关键问题:我们能否开发出一种资源消耗可承受的高质量图像生成器?

  在本文中,作者引入了 PIXART-α,它显着降低了训练的计算需求,同时保持了与当前最先进的图像生成器相比具有竞争力的图像生成质量,如图 1 所示。

图1:PIXART-α 生成的样本具有卓越的品质,其特点是在遵循所提供的文本描述方面具有极高的保真度和精确度。

  为了实现这一目标,作者提出了三个核心设计:

  1️⃣训练策略分解:作者将复杂的文本到图像生成任务分解为三个简化的子任务:(1)学习自然图像的像素分布,(2)学习文本图像对齐,以及(3)增强图像的美学质量。对于第一个子任务,作者建议使用低成本的类条件模型初始化 T2I 模型,从而显着降低学习成本。对于第二个和第三个子任务,我们制定了由预训练和微调组成的训练范式:对信息密度丰富的文本图像对数据进行预训练,然后对具有优越美学质量的数据进行微调,提高训练效率。

  2️⃣高效 T2I Transformer:基于扩散 Transformer(DiT),作者结合了交叉注意模块来注入文本条件并简化计算密集型类条件分支以提高效率。此外,引入了一种重新参数化技术,允许调整后的文本到图像模型直接加载原始类条件模型的参数。因此,可以利用从 ImageNet中学到的有关自然图像分布的先验知识,为 T2I Transformer 提供合理的初始化并加速其训练。

  3️⃣数百万次迭代来学习稳定的文本图像对齐:为了解决这些问题,我们提出了一种自动标记管道,利用最先进的视觉语言模型 (LLaVA ) 在 SAM 上生成字幕 。参考2.4节,SAM数据集因其丰富多样的对象集合而具有优势,使其成为创建高信息密度文本图像对的理想资源,更适合文本图像对齐学习。

二、方法

2.1 动机

  T2I训练缓慢的原因有两个方面:训练流程数据

  T2I生成任务可以分解为三个方面: 1️⃣捕获像素依赖性:生成真实图像涉及理解图像内复杂的像素级依赖性并捕获其分布;2️⃣文本和图像之间的对齐:需要精确的对齐学习来理解如何生成与文本描述精确匹配的图像;3️⃣高审美品质:除了忠实的文字描述之外,美观是生成图像的另一个重要属性。目前的方法将这三个问题纠缠在一起,直接使用大量数据从头开始训练,导致训练效率低下。为了解决这个问题,作者将这些方面分为三个阶段,如 2.2 节中所述。

图3:LAION 原始字幕与 LLaVA 精炼字幕。 LLaVA 提供高信息密度的字幕,帮助模型在每次迭代中掌握更多概念并提高文本图像对齐效率。

  如图 3 所示,另一个问题是当前数据集的标题质量。当前的文本-图像对经常遭受文本-图像错位、描述不足、词汇使用不频繁以及包含低质量数据的问题。这些问题给训练带来了困难,导致不必要的数百万次迭代才能实现文本和图像之间的稳定对齐。为了应对这一挑战,我们引入了一种创新的自动标记管道来生成精确的图像标题,如第 2.4 节中所述。

2.2 训练策略分解

  通过将训练分为不同数据类型的三个阶段,可以逐步优化模型的生成能力。

  第一阶段:像素依赖性学习:当前的类引导方法在单个图像中生成语义连贯且合理的像素方面表现出了示范性的性能。为自然图像训练类条件图像生成模型相对容易且成本低廉,如附录 A.5 中所述。此外,我们发现合适的初始化可以显着提高训练效率。因此,我们从 ImageNet 预训练模型中增强我们的模型,并且我们模型的架构被设计为与预训练权重兼容。  

  第二阶段:文本图像对齐学习:从预训练的类引导图像生成过渡到文本到图像生成的主要挑战是如何在显着增加的文本概念和图像之间实现准确对齐。

  这种对齐过程不仅耗时,而且本身就具有挑战性。为了有效地促进这一过程,我们构建了一个由具有高概念密度的精确文本图像对组成的数据集。数据创建管道将在 2.4 节中描述。通过使用准确且信息丰富的数据,我们的训练过程可以在每次迭代中有效地处理大量名词,同时与以前的数据集相比,遇到的歧义要少得多。这种战略方法使我们的网络能够有效地将文本描述与图像结合起来。

  第三阶段:高分辨率且美观的图像生成:在第三阶段,我们使用高质量的美学数据微调我们的模型以生成高分辨率图像。值得注意的是,我们观察到该阶段的适应过程收敛速度明显更快,这主要是由于前面阶段建立的强大先验知识。

  将训练过程解耦为不同阶段,显着缓解训练难度,实现高效训练。

2.3 高效 T2I Transfomer

  PIXART-α 采用 Diffusion Transformer (DiT) 作为基础架构,并创新地定制 Transformer 模块来应对 T2I 任务的独特挑战,如图 4 所示。提出了几种专用设计如下:

图4:LAION 原始字幕与 LLaVA 精炼字幕。 LLaVA 提供高信息密度的字幕,帮助模型在每次迭代中掌握更多概念并提高文本图像对齐效率。

  • 交叉注意力层。我们将多头交叉注意力层合并到 DiT 块中。它位于自注意力层和前馈层之间,使得模型可以灵活地与从语言模型中提取的文本嵌入进行交互。为了促进预训练权重,我们将交叉注意层中的输出投影层初始化为零,有效地充当恒等映射并保留后续层的输入。
  •   AdaLN-single。我们发现 DiT 的自适应归一化层 (adaLN) 模块中的线性投影占参数的很大比例 (27%)。如此大量的参数没有用,因为我们的 T2I 模型没有采用类条件。因此,我们提出了 adaLN-single,它仅使用时间嵌入作为第一个块中的输入来进行独立控制(如图 4 右侧所示)。具体来说,在第 i 个块中,令 S(i) = [β(i) 1 , β(i) 2 , γ(i) 1 , γ(i) 2 , α(i) 1 , α(i) 2 ]是 adaLN 中所有尺度和平移参数的元组。在DiT中,S(i)是通过特定于块的MLP S(i)=f(i)(c+t)获得的,其中c和t分别表示类别条件和时间嵌入。然而,在 adaLNsingle 中,一组全局移位和比例仅在第一个块处计算为 S = f (t),该块在所有块之间共享。然后,获得 S(i) 为 S(i) = g(S, E(i)),其中 g 是求和函数,E(i) 是与 S 形状相同的特定于层的可训练嵌入,它自适应地调整不同块中的缩放和移位参数。

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

    

  

  

  

  

  

CATALOG
  1. 1. 1、引言
  2. 2. 二、方法
    1. 2.1. 2.1 动机
    2. 2.2. 2.2 训练策略分解
    3. 2.3. 2.3 高效 T2I Transfomer