PIXART-Σ是由华为诺亚方舟实验室、大连理工和香港大学的研究团队共同开发的扩散变换模型DiT，专门用于4K分辨率的文本到图像生成。模型于24年3月发布，在基础模型PixArt-a上进行了显著改进，提供了更高保真图像和更好的文本提示对齐。PixArt-Σ生成4K图像的能力有效地支持了电影和游戏等行业高质量视觉内容的产生。
总体说来，PixArt-Σ相对基础模型PixArt-α的主要改进如下：
- 具有高美感特性和艺术气质的数据集
- 更高效的自动化标注工具，更好的文本-图像对齐。采用性能更好的Share-Captioner作为标注工具，而非PixArt-α中幻觉问题严重的SAM-LLaVA。同时把文本编码器（即 Flan-T5 ）的token长度扩展到大约 300 个单词。为图像提供了更密集、更精确的描述。
- 更具美感和更高分辨率的训练数据集。构建了一个4K分辨率的图片数据集，并通过一个美学评分模型AES对数据进行进一步筛选。最终形成分辨率为1K的 33M 高分辨率图像，和分辨率约为 4K 的 2.3M 图像。
- 高质量评估数据集。提出了一个精心策划的数据集，包括30,000个高质量、美感愉悦的文本-图像对，以便进行评估。
- 从弱到强的训练Weak-to-Strong Training
- 将PixArt-α的VAE替换成SDXL的VAE。
- 以低分辨率模型做初始化，在高分辨率模型上Fine-tune。
- 高效的token压缩，在DiT框架内提出了一个新颖的注意力模块，可以压缩K和V，显著提高效率，促进超高分辨率图像生成。从没有KV压缩的模型演化到有KV压缩的模型。
- AI 偏好研究中，使用多模态模型GPT-4 Vision作为评估器，进行打分。
本文由深圳季连科技有限公司AIgraphX自动驾驶大模型团队编辑。如有错误，欢迎在评论区指正。
论文名称：PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation
论文地址：https://arxiv.org/abs/2403.04692
项目地址： https://pixart-alpha.github.io/PixArt-sigma-project/
Abstract
PixArt-Σ，是一种扩散 Transformer 模型 DiT，能够直接生成 4K 分辨率的图像。PixArt-Σ 代表了其前身 PixArt-α 的重大进步，提供了明显更高保真度的图像，更高的与文本提示对齐。PixArt-Σ 的一个关键特性是它的训练效率。利用PixArt-α基本预训练，通过结合更高质量的数据从"较弱"基线演变为"更强"模型，这个过程我们称之为"弱到强训练"。
PixArt-Σ 的进步是双重的：
- **高质量训练数据：**PixArt-Σ 结合了高质量的图像数据，并与更精确和更详细的图像描述配对。
- **高效的token压缩：**我们在 DiT 框架内提出了一种新颖的注意力模块，该模块同时压缩K和V值，显著提高了效率并促进超高分辨率图像生成。
由于这些改进，PixArt-Σ 比现有的文本到图像扩散模型，如 SDXL(2.6B 参数)和 SD Cascade(5.1B 参数)，实现了卓越的图像质量、更好依从用户提示；模型只有0.6B 参数，明显更小。此外，PixArt-Σ 生成 4K 图像的能力支持创建高分辨率海报和墙纸，有效地支持了电影和游戏等行业高质量视觉内容的产生。
1 Introduction
最近出现的高质量的文本到图像 (T2I) 模型深刻影响了 AI 生成内容 (AIGC) 社区。这包括DALL·E 3、Midjourney等专有模型，以及Stable Diffusion和PixArt-α等开源模型。尽管如此，开发一个顶级 T2I 模型涉及相当大的资源；例如，从头开始训练 SD1.5 需要大约 6000 A100 GPU 天，对研究人员和AIGC社区创新构成了重大障碍。随着时间的推移，AIGC社区将获得不断更新、更高质量数据集和更高级算法的访问。一个关键问题是：我们如何有效地将这些新元素集成到现有模型中，在有限的资源约束下实现更强大的版本？
为了探索这个问题，我们的研究侧重于增强 PixArt-α，这是一种高效的 T2I 训练方法。PixArt-α 代码在 DiT 框架内的早期探索，这是一个有显著潜力的模型结构，如GenTron、Sora和Stable Diffusion 3等著作所证明的那样。为了最大化这种潜力，我们在 PixArt-α 的预训练基础之上，集成高级元素以促进其持续改进，从而产生更强大的模型 PixArt-Σ。我们称这种为"由弱到强的训练"。具体来说，我们做了以下增强：
Higher-Quality Training Data
- **High-quality images：**该数据集包含来自互联网分辨率为1K的 33M 高分辨率图像，和分辨率约为 4K 的 2.3M 图像。这些图像主要特征是它们具有高度的审美，并涵盖了广泛的艺术风格。
- **Dense and accurate captions：**为了为上述图像提供更精确和详细的描述，我们将 PixArt-α 中使用的 LLaVA 替换为更强大的图像描述器 Share-Captioner。此外，为了提高模型在文本和视觉概念之间的对齐能力，我们将文本编码器（即 Flan-T5 ）的token长度扩展到大约 300 个单词。我们观察到这些改进有效地消除了模型幻觉倾向，从而导致更高质量的文本图像对齐。
Efficient Token Compression
为了增强PixArt-α，我们将其生成分辨率从1K扩展到4K。生成超高分辨率（例如 2K/4K）的图像，token数量会显著增加，导致计算需求显著增加。为了应对这一挑战，我们引入了一个自注意力模块，该模块具有针对 DiT 框架量身定制的key and value token compression。具体来说，我们利用步长为 2 的组卷积来局部聚合K和V。此外，我们采用了专门的权重初始化方案，允许从没有 KV 压缩的预训练模型平滑适应。这种设计有效地减少了约34%的高分辨率图像生成训练和推理时间。
Weak-to-Strong Training Strategy
我们提出了几种微调技术，以有效地从弱模型快速适应强模型。这包括
- 用更强大的变分自动编码器 (VAE) 替换。
- 从低分辨率扩展到高分辨率。
- 从没有键值 (KV) 压缩的模型演变为具有 KV 压缩的模型。
这些结果证实了"弱到强的训练"方法的有效性。
通过这些改进，PixArt-∑以最小的训练成本和模型参数实现了高质量的4K分辨率图像生成。具体来说，对预训练的模型进行微调，我们仅利用PixArt-α所需9%的GPU天数实现了强大的1K高分辨率图像生成，考虑到我们用新的训练数据和更强大的VAE，这令人印象深刻。PixArt-Σ生成的图像具有与当前顶级T2I产品相当的美学质量，以及与文本提示进行细粒度对齐的特殊能力。
2 Related Work
Diffusion Transformers
Transformer 架构在各个领域取得了显着的成功，如语言建模、计算机视觉和其他领域。在扩散模型领域，DiT和UViT开创了Transformer架构使用的先河。随后的工作，包括DiffiT、SiT和FiT，在DiT的体系结构上进行了改进，而MDT通过掩蔽建模技术提高了训练效率。对于文本到图像 T2I 合成，PixArt-α 探索了有效的 T2I 训练方案，实现了第一个能够生成1024px的高质量图像。GenTron探索了扩散transformer在图像和视频生成方面的灵活性和可扩展性。最近强大的视频生成模型Sora的出现进一步强调了扩散transformer的潜力。在这项工作中，我们首次探索了使用 Transformer 架构直接生成 4K 超高分辨率图像，解决了涉及长序列token带来的计算复杂性挑战。High Resolution Image Generation 高分辨率图像生成大大提高了视觉质量，在电影和游戏等各种行业都很重要。然而，由于计算需求的大幅增加，增加图像分辨率会带来挑战。在这个方向上已经探索了许多方法。例如，Imagen、GigaGAN和Stable Diffusion引入了一个额外的超分辨率网络，而Stable Cascade使用多个扩散网络逐步提高分辨率。然而，这些组合模型解决方案可能会引入累积误差。另一方面，SDXL、DALL·E 2、Playground和PixArt-α等工作试图直接使用扩散模型生成高分辨率图像。然而由于计算复杂性，这些努力仅限于生成分辨率高达 1024px 的图像。在本文中，我们将此边界推到 4K 分辨率，显着提高了生成内容的视觉质量。
Efficient Transformer architecture
Transformer 中自注意力机制在token数量方面存在二次计算复杂度，这阻碍了令牌数量的扩展。许多工作都在在这方面寻求改进：
- 稀疏注意力，它通过选择性地处理token的子集来减少整体计算负载。例如，PVT v2 采用卷积核来压缩K和V的空间，从而降低计算注意力所涉及的复杂性。
- 局部注意，将注意力集中在附近的区域内；值得注意的是，Swin Transformer利用基于窗口的注意将计算限制在指定的窗口大小内。
- 低秩/线性注意，Linformer 通过低秩近似降低了自注意力机制的计算复杂度。
在本文中，受 PVT v2 的启发，我们采用了一种基于KV压缩的自注意力机制来减轻处理 4K 图像的高度复杂性。
3 Framework
3.1 Data Analysis
Higher Aesthetic and higher Resolution
为了提高数据集的审美质量，我们将内部数据从 14M 扩展到 33M。为清楚起见，我们分别将两个数据集命名为 Inside-α 和 Inside-Σ。请注意，与当前可用开源模型（如 SD v1.5）使用 2B 数据相比，这种扩展仍然很小。我们证明了在数据量有限的有效训练策略下，仍然可以获得强大的 T2I 模型。
Inside-Σ 内的图像高于 1K 分辨率。为了便于 4K 分辨率生成，我们还收集了 4K 分辨率的 8M 真实摄影图像数据集。为了确保审美质量，我们采用了美学评分模型AES来过滤这4K图像。这个过程产生了一个高度精细的 2M 超高分辨率和高质量图像数据集。
有趣的是，我们观察到，随着图像的分辨率的增加，模型的保真度FID和语义分割CLIP Score有所提高，这强调了生成高分辨率图像的能力的重要性。
Better Text-Image Alignment
最近的研究如PixArt-α和DALL-E 3强调了文本-图像描述对齐的重要性。加强这种对齐对于提高模型能力至关重要。为了进一步完善我们收集的"原始"描述，我们专注于提高描述的长度和准确性。值得注意的是，我们的描述（内部 Σ）在以下方面显示出比 PixArt-α（内部 α）中使用的几个优势：
-
Enhanced caption accuracy: 如图5所示，PixArt-α中使用的LLaVa具有一定的幻觉问题。我们利用更强大的视觉语言模型，即 Share-Captioner，来生成详细和正确的描述，增强收集到的原始提示。
-
**Increased caption length:**如表1 和图 6所示，平均描述长度增加到 180 个单词，大大提高了描述性能力。此外，我们将文本编码器的token处理长度从 120 个token（如在内部 α 中）扩展到 300 个token。我们的模型是在长（Share-Captioner）和短（原始）描述的混合上训练的，比率分别为 60% 和 40%。这种方法增强了文本描述的多样性，并减轻了可能仅由依赖生成描述引起的潜在偏差。
备注：PixArt-Σ构建了一个4K分辨率的图片数据集，并通过一个美学评分模型AES对数据进行进一步筛选。新的内部数据集，平均caption长度更长，image caption信息量更大。
High-Quality Evaluation Dataset
大多数 SoTA T2I 模型选择 MSCOCO 作为评估集来评估 FID 和 CLIP 分数。然而，我们观察到在 MSCOCO 数据集上进行的评估可能无法充分反映模型在美学和文本图像对齐方面的能力。因此，我们提出了一个包含 30,000 个高质量、美观的文本图像对的精选集，以促进评估。数据集的选定样本在附录中给出。该数据集旨在提供对模型性能的更全面的评估，特别是在捕捉美学吸引力的复杂性以及文本描述和视觉内容之间对齐的保真度方面。除非另有说明，本文中的评估实验是在收集的高质量评估数据集上进行的。
3.2 Efficient DiT Design
一个有效的DiT网络是必不可少的，因为在以超高分辨率生成图像时，计算需求显著增加。注意机制在DiT的有效性中起着关键作用，但其二次计算需求显著限制了模型的可扩展性，特别是在更高的分辨率下，如2K和4K。受PVT v2的启发，我们将KV压缩合并到原始的PixArt-α框架中，以解决计算挑战。这种设计对总参数仅增加了 0.018%，但通过token压缩实现了计算成本的有效降低，同时仍然保留了空间和语义信息。
Key-Value (KV) Token Compression
我们的动机源于一个有趣的观察，即直接将 KV的token压缩应用于预训练的 PixArt-α 仍然可以生成合理的图像。这表明特征中存在冗余。考虑到相邻R × R patch内的高相似度，我们假设窗口内的特征语义是多余的，可以合理压缩。我们提出了KV token compression，表示为 fc(·)，通过压缩算子在 R × R 窗口内压缩token特征，如图 7 所示。
此外，为了减轻自注意力计算中 KV 压缩引起的潜在信息丢失，我们选择保留查询 Q 的所有token。这种策略使我们能够有效地利用 KV 压缩，同时减轻丢失关键信息的风险。通过使用KV压缩，我们提高了注意力计算的效率，减少了计算的复杂性，从而使直接生成高分辨率图像的计算成本可管理。
我们使用具有特定初始化的卷积算子"Conv2×2"压缩深层。其他设计变体的详细实验在第 5 节中讨论。具体来说，我们设计了一个专门的卷积核初始化"Conv Avg Init"，它利用组卷积初始化权重 w = 1/R² ，相当于平均算子。这种初始化策略最初可以产生粗略结果，加速了微调过程，同时只引入了 0.018% 的附加参数。
备注：PixArt-Σ像潜在空间模型LDM一样，在压缩上做文章。本文提出了一种更为激进的压缩策略------KV Compression，即在attention上做压缩。KV Compression [PVT v2: Improved Baselines with Pyramid Vision Transformer]提出的动机在于解决attention计算中，key和value存在的特征冗余。模型用压缩后的key和value进行运算，并不会对PixArt-Σ生成结果带来太大性能损失。
KV Compression流程简洁，如上图所示，attention运算之前对key和value压缩成更小的特征。即，给定key和value的形状为R×R，PixArt-Σ采用步长为2的Group Conv，并将其权重初始化为1/R²。key和value卷积下采样后分辨率为原来1/2倍，训练初期精度上的损失，可以通过模型训练来弥补。
3.3 Weak-to-Strong Training Strategy
我们提出了几种有效的训练策略，以增强从"弱"模型到"强"模型的转换。这些策略包括 VAE 快速适应、高分辨率微调和 KV token压缩。
Adapting model to new VAEs
随着 VAE 继续开发，从头开始训练 T2I 模型是资源密集型的。我们将PixArt-α的VAE替换为SDXL的VAE，并继续微调扩散模型。我们观察到快速收敛现象，微调在 2K 训练步骤中快速收敛，如图 8 (a) 所示。在处理 VAE 模型迁移和否定从头开始训练的必要性时，微调更有效。
Adapting to Higher-Resolution
当我们从低分辨率 (LR) 模型微调到高分辨率 (HR) 模型时，我们观察到性能下降，如图 8 (b) 所示，我们将其归因于不同分辨率之间位置嵌入 (PE) 的差异。为了缓解这个问题，我们利用"PE插值"技巧：通过插值 LR 模型的 PE 来初始化 HR 模型的 PE，显着增强了 HR 模型的初始状态并加快了微调过程。即使在只有 100 次训练迭代内，我们也可以获得视觉上令人愉悦的图像。此外，我们定量评估模型的性能变化，如表 2 所示，微调在 1K 步快速收敛，进一步的训练略微提高了性能。这表明，使用"PE Interpolation"技巧可以实现更高分辨率的生成的快速收敛，无需从头开始训练以更高的分辨率生成。
Adapting model to KV compression
当在没有 KV 压缩的 LR 预训练模型进行微调时，我们可以直接使用 KV 压缩。如图 8 (c) 所示，使用我们的"Conv Avg Init."策略，PixArt-Σ 从更好的初始状态开始，使收敛更容易和更快。值得注意的是，即使在 100 个训练步骤内，PixArt-Σ 也能执行令人满意的视觉效果。最后，通过第 3.2 节中的 KV 压缩算子和压缩层设计，我们可以减少训练和推理时间的34%。
备注：Weak-to-Strong Training 包含三方面，即更换VAE、高分辨率微调和 KV token压缩。
- 将PixArt-α模型中的VAE替换成SDXL的VAE。SDXL VAE用了更大Training Batch Size（9→256），并额外引入EMA用于模型权重更新。
- 和PixArt-α一样，第一阶段以低分辨率图片做预训练。PixArt-Σ模型后续在高分辨率图片上fine-tune。值得注意的是，一旦分辨率发生变化，原有Positional Encoding（PE）无法继续适用，这里PixArt-Σ用到了DiffFit中的一个PE Interpolation技巧。即将原来的PE除以分辨率增大的倍数。
- 从没有KV Compression的模型到有KV Compression的模型Fine-tune。
4 Experiment
4.1 Implementation Details
Training Details
我们遵循Imagen和PixArt-α使用T5的编码器(即Flan-T5-XXL)作为条件特征提取的文本编码器，并使用PixArt-α作为我们的基本扩散模型。与大多数提取固定 77 个文本标记的作品不同，我们从 PixArt-α 的 120 调整到 300 文本标记的长度，因为在内部-Σ中组织的描述要密集得多，以提供高度精细的细节。
为了捕捉输入图像的潜在特征，我们使用来自 SDXL 的预训练和冻结的 VAE。其他实现细节与 PixArt-α 相同。
模型使用位置嵌入插值技巧，在 PixArt-α 的 256px 预训练检查点上进行微调。我们的最终模型，包括 1K 分辨率，是在 32 个 V100 GPU 上训练的。我们还使用 16 个 A100 GPU 来训练 2K 和 4K 图像生成模型。使用权重衰减为 0 的 CAME 优化器和 2e-5 的恒定学习率。这有助于我们降低优化器状态的维度，导致 GPU 内存较低，而不会导致性能下降。
Evaluation Metrics
为了更好地说明美学和语义能力，我们收集了 30K 高质量的文本图像对，来对最强大的 T2I 模型进行基准测试。我们主要通过人类和人工智能偏好研究来评估PixArt-Σ，因为FID指标可能不能充分反映生成质量。但是，我们仍然在附录中收集的数据集上提供 FID 结果。
4.2 Performance Comparisons
Image Quality Assessment
我们根据闭源文本到图像 (T2I) 产品和开源模型定性地评估了我们的方法。我们的模型可以生成高质量、逼真的图像，在不同的长宽比和风格上具有复杂的细节。这种能力强调了我们的方法在从文本描述中生成视觉上引人注目的内容方面的卓越性能。
如图 3 所示，我们将PixArt-Σ与开源模型SDXL和PixArt-α进行了比较，我们的方法增强了肖像的真实感，提高了语义分析的能力。与 SDXL 相比，我们的方法在遵守用户指令方面表现出卓越的熟练程度。
不仅优于开源模型，而且我们的方法也与当前T2I封闭源产品非常具有竞争力，如图 4 所示。PixArt-Σ产生照片般逼真的结果，并密切遵守用户指令，这与当代商业产品相当。
High-resolution Generation
我们的方法能够直接生成 4K 分辨率的图像，而不需要任何后处理。此外，它擅长准确遵循用户提供的复杂、详细和长文本，如图 2 所示。因此，用户不需要提示工程才能获得满意的结果。我们的方法实现了直接的 4K 图像生成。同时，研究引入了无调优后处理技术，旨在从LR模型生成HR图像，或使用超分辨率模型生成HR图像。然而，由于两个主要原因，它们对应的结果往往表现出伪影：（1）由于级联管道可能会出现累积错误。(2)这些方法没有捕捉到4K图像的真实分布，也没有学习文本和图像之间的对齐。我们认为我们的方法可能是生成高分辨率图像的一种更有前途的方法。我们的方法产生了更好的结果，更多的视觉比较包含在附录中。
图 2，4K image generation with complex dense instructions. PixArt-Σ 可以直接生成 4K 分辨率的图像，无需后处理，展现了与文本提示细粒度对齐的出色能力。
Human/AI (GPT4V) Preference Study
我们使用从上述高质量评估数据集中随机收集的300个描述子集，在人类和人工智能偏好研究中评估训练有素的模型。我们收集由总共六个开源模型生成的图像，包括 PixArt-α、PixArt-Σ、SD1.5 、Stable Turbo、Stable XL 、Stable Cascade and Playground-V2.0。我们为人类偏好研究开发了一个网站来显示提示及其相应的图像。该网站分发给训练有素的评估者，他们被要求评估图像，根据质量对它们进行排名以及它们与文本提示的匹配程度。如图 9 中的蓝色条所示，结果表明 PixArt-Σ 相对于其他六个 T2I 生成器的显着偏好。与现有的 SDXL（2.6B 参数）和 SD Cascade（5.1B 参数）等 T2I 扩散模型相比，PixArt-Σ 生成了更接近用户提示的高质量图像，使用更小的尺寸（0.6B 参数）。
此外，在我们的 AI 偏好研究中，我们使用高级多模态模型 GPT-4 Vision 作为评估器。对于每个试验，我们提供具有两个图像的 GPT-4 vision：一个来自 PixArt-Σ，另一个来自竞争 T2I 模型。我们制作了不同的提示，引导GPT-4V ision根据图像质量和图像和文本对齐进行投票。图 9 中的橙色和绿色条表示的结果显示了人类和 AI 偏好研究中一致的结果。具体来说，PixArt-Σ 在有效性上超过了基线 PixArt-α。与稳定级联等当代高级模型相比，PixArt-Σ 在图像质量和指令跟随能力方面表现出具有竞争力或优越的性能。
5 Ablation Studies
我们对各种 KV 压缩设计的生成性能进行了消融研究。除非另有说明，实验是在 512px 上进行。
5.1 Experimental settings
我们使用第3.1节中描述的测试集进行评估。我们使用FID来计算收集数据和生成数据之间的分布差异，用于比较指标。此外，我们利用CLIP-Score来评估提示和生成图像之间的对齐。
5.2 Compression Designs
Compression positions
我们在 Transformer 结构内的不同深度的 KV 压缩：在浅层 (1∼14)、中间层 (7∼20) 和深层 (14∼27) 中。如表3a所示，在深层使用 KV 压缩可以显著提高性能。我们推测这是因为浅层通常编码详细的纹理内容，而深层抽象高级语义内容。由于压缩往往会影响图像质量而不是语义信息，压缩深层可以达到信息的最小损失，使其成为加速训练但不损害生成质量的实用选择。
Compression operators
我们探索了不同压缩算子的影响。我们使用了三种技术，随机丢弃、平均池化和参数卷积，将 2×2 令牌压缩为单个令牌。如表 3b 所示，"Conv 2×2"方法优于其他方法，强调了使用可学习内核比简单丢弃方法更有效地减少冗余特征的优势。
Compression ratios on different resolutions
我们研究了不同压缩比对不同分辨率的影响。如表3c中所示，值得注意的是，我们发现令牌压缩不影响文本和图像 (CLIP Score) 之间的对齐，但会影响不同分辨率的图像质量 (FID)。尽管随着压缩比的增加，图像质量略有下降，但我们的策略使训练速度提高了 18% 到 35%。这表明我们提出的 KV 压缩对于实现高分辨率 T2I 生成既有效且高效。
Speed comparisons on different resolutions
我们在表3d中进一步全面验证了训练和推理中的速度加速。我们的方法在 4K 生成中可以将训练和推理速度提高了大约 35%。值得注意的是，我们观察到训练加速度随着分辨率的增加而增加。例如，随着分辨率从 1K 增加到 4K，训练逐渐从 18% 加速到 35%。这表明我们的方法在分辨率不断提高方面的有效性，证明了它对更高分辨率的图像生成任务的潜在适用性。
6 conclusion
我们介绍了 PixArt-Σ，这是一种文本到图像 T2I 扩散模型，能够直接生成 4K 分辨率的高质量图像。基于PixArt-α预训练，PixArt-Σ通过一种新的从弱到强训练方法实现了有效的训练。这种方法的特点是结合了更高质量的数据和高效的token压缩的集成。PixArt-Σ 擅长生成高保真图像，同时紧跟文本提示，超过了其前身 PixArt-α 设置的高标准。
