本文目录

1 扩散模型解读 (五)：PixArt 系列最新力作 PixArt-Σ：Weak-to-Strong 训练的 4K 高清文生图模型
(来自华为诺亚方舟实验室，大连理工大学，香港大学，香港科技大学)
1 PixArt-Σ 论文解读
1.1 生成 4K 高清大图的 PixArt-Σ：一个新的数据集，一个新的架构改进
1.2 高质量的训练数据
1.3 高效 DiT 架构设计
1.4 Weak-to-Strong 训练策略
1.5 实验结果

太长不看版

本文提出的 PixArt-Σ 是 PixArt 系列的最新力作，它是一个 Diffusion Transformer model (DiT)，可以直接生成 4K 分辨率的图像。PixArt-Σ 相比其前身 PixArt-α，有重大的进步：可以生成明显更高保真度的图像，而且提高了与文本提示的对齐质量，如图 1,2 所示。PixArt-Σ 的一个关键特性是其训练效率。从 PixArt-α 的基本预训练为起点，通过结合更高质量的数据从一个 "weaker" 的基线演化为一个 "stronger" 的模型，这个过程作者称为 "weak-to-strong training"。

PixArt-Σ 的贡献包含 2 个方面：

1. 高质量训练数据：PixArt-Σ 结合了高质量的图像数据，并与更精确和更详细的图像标题配对。
2. 高效的 token 压缩模块：作者提出了一种新颖的注意力模块，该模块同时压缩 Self-attention 的 Key 和 Value，显著地提高了效率并能有效促进模型生成超高分辨率的图像。

PixArt-Σ 比现有的文生图的扩散模型，例如 SDXL (2.6B 参数) 和 SD Cascade (5.1B 参数) 相比，其生成的图片质量卓越，而且拥有极佳的用户指令遵从性能，且模型参数明显更低 (0.6B 参数)。此外，PixArt-Σ 生成 4K 图像的能力支持创建高分辨率海报和墙纸，有效地支持了电影和游戏等这类需要高质量视觉内容的行业。

1 PixArt 系列最新力作 PixArt-Σ：Weak-to-Strong 训练的 4K 高清文生图模型

论文名称：PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation (ECCV 2024)

论文地址：

http://arxiv.org/pdf/2403.04692

项目主页：

http://pixart-alpha.github.io/PixArt-sigma-project/

1.1 生成 4K 高清大图的 PixArt-Σ：一个新的数据集，一个新的架构改进

高质量的文生图 (Text-to-Image, T2I) 模型在很大程度上影响了 AIGC 社区。私有模型包括 DALL·E 3 [1] 、Midjourney [2] 等，开源模型包括 Stable Diffusion [3] 和 PixArt-α [4] 等。但是尽管如此，开发一个顶级 T2I 模型需要相当多的资源：从头开始训练 SD1.5 需要大约 6000 A100 GPU days，这个问题对于资源有限的个人研究者构成了重大的障碍，也从一定程度上讲阻碍了 AIGC 社区的创新。随着时间的推移，AIGC 社区在持续不断地获得不断更新、更高质量的数据集和更高级的算法。那么一个关键的问题是： 如何有效地将这些新元素集成到现有的模型中，使我们得以在有限的资源约束下实现更强大的模型？

为了探索这个问题，本文聚焦于增强 PixArt-α，PixArt-α 是一种高效的 T2I 训练方法。PixArt-α 是使用 DiT [5] 架构的早期探索。为了最大化这种架构的潜力，作者在 PixArt-α 的预训练的基础之上，继续改进得到更强大的模型 PixArt-Σ。作者把这种从一个相对较弱的模型通过高效训练演化为相对较强的模型的过程称为 "weak-to-strong training"。那么为了实现 "weak-to-strong training"，作者使用了下面的改进：

1) 一个新的高质量的数据集

作者收集了一个优于 PixArt-α 中使用的高质量数据集。这个数据集有两个主要的特点： 1) 高质量图片： 该数据集包含来自互联网的 33M 高分辨率图像，均超过 1K 分辨率，也包括分辨率约为 4K 的 2.3M 张图像。这些图像的主要特征是它们的高美学，而且包含广泛的艺术风格。 2) 密集且精确的字幕： 为了为上述图像提供更精确和详细的字幕，作者将 PixArt-α 中使用的 LLaVA [6] 替换为更强大的图像字幕模型 Share-Captioner [7] 。此外，为了提高模型在文本和视觉概念之间的对齐能力，将文本编码器 (即 Flan-T5 [8] ) 的 token length 扩展到大约 300 个单词。可以观察到这些改进有效地消除了模型幻觉，从而带来更高质量的文本图像对齐。

2) 一个新的架构改进：Key Value token 压缩

为了增强 PixArt-α，将其生成分辨率从 1K 扩展到 4K。以超高分辨率 (例如 2K/4K) 生成图像会导致 token 的数量显着增加，导致计算需求量增加。为此，作者为 Self-attention 机制引入了 Key Value token 压缩方案。具体而言，利用 stride 为 2 的 Group Convolution 来在局部集成 Key 和 Value。而且，作者还采用了专门的权重初始化方案，允许从没有 KV 压缩的预训练模型来平滑微调。这种设计减少了约 34% 的高分辨率图像生成的训练和推理时间。

3) "weak-to-strong" 训练策略

作者提出了几种微调技术，以有效地从弱模型快速适应强模型：1) 用更强大的变分自动编码器 (VAE) 替换原来的 VAE。2) 从低分辨率扩展到高分辨率。3) 从没有键值 (KV) 压缩的模型演变为具有 KV 压缩的模型。这些结果证实了 "weak-to-strong" 训练策略的有效性。

1.2 高质量的训练数据

更高的美学和更高的分辨率

为了提高数据集的审美质量，作者将内部数据从 14M 扩展到 33M。为清楚起见，分别将两个数据集命名为 Inside-α 和 Inside-Σ。但是与一些开源模型 (比如 SD v1.5) 使用的巨量图片相比，这种数据集的扩展依然相对少。本文也证明了虽然数据集量有限，但是通过训练策略仍然可以获得强大的 T2I 模型。

Inside-Σ 内的图像高于 1K 分辨率。为了便于 4K 分辨率生成，作者还收集了 4K 分辨率的 8M 真实摄影图像数据集。为了确保审美质量，作者还采用了一个美学打分模型 (aesthetic scoring model, AES) 来过滤这些 4K 图片。这个过程最后会得到一个高度精细的 2M 超高分辨率和高质量图像数据集。

有趣的是作者观察到，随着图像的分辨率的增加，模型的保真度 (FID) 和语义分割 (CLIP Score) 性能有所提高。

更好的图文对齐

加强文本-图像描述的对齐对于提高模型能力至关重要。因此，为了进一步完善原始的图像描述，作者希望提高字幕的长度和准确性。值得注意的是，Internal-Σ 的描述在下面这几个方面比 Internal-α 更有优势：

1) 更好的字幕精度

如图 5 所示，PixArt-α 中使用的 LLaVA 具有一定的幻觉问题。PixArt-Σ 使用更强大的视觉语言模型，即 ShareCaptioner 来生成详细和正确的字幕。

2) 字幕长度增加

如下图 6 和 7 所示，平均字幕长度增加到 180 个单词，大大提高了字幕的描述性能力。而且，作者将文本编码器的 token 处理长度从 120 tokens 扩展到 300 tokens。本文呢的模型是在长字幕 (Share-Captioner) 和短字幕 (原始字幕) 上混合一起训练的，比例分别为 60% 和 40%。这种方法增强了文本描述的多样性。

图 7 是 Internal-α 和 Internal-Σ 的一个总结，作者通过各种指标来评估数据集的多样性，包括每个图像的名词多样性、总名词计数、平均标题长度和平均名词。

高质量的评测数据集

大多数 T2I 模型喜欢使用 MSCOCO [9] 作为评估集来评估 FID 和 CLIP 分数。但是本文作者观察到在 MSCOCO 数据集上进行的评估可能无法充分反映模型在美学和文本图像对齐方面的能力。因此，作者提出了一个包含 30,000 个高质量、美观的图文对的精选数据集，来帮助评估 T2I 模型。该数据集旨在提供对模型性能的更全面的评估，尤其是在捕捉美学吸引力的复杂性方面，以及文本描述和视觉内容之间对齐的保真度上面。

1.3 高效 DiT 架构设计

以超高分辨率生成图像时，计算需求显著增加，因此高效的 DiT 架构就是必不可少的。Self-Attention 机制在 DiT 中起着关键作用，但其二次计算需求显著限制了模型的可扩展性，尤其在更高的分辨率下，如 2K 和 4K 分辨率中。受 PVT v2 [10] 的启发，作者将 KV token 压缩技术合并到原始的 PixArt-α 中，以应对密集的计算量的挑战。这种设计仅增加了 0.018% 的总参数量，但通过 token 压缩技术有效降低了计算成本，同时仍然保留了空间和语义信息。

KV Token 压缩技术

本文提出了一个有趣的观察，即直接将 Key-Value (KV) Token 压缩应用于预训练的 PixArt-α，仍然可以生成合理的图像。这表明特征中存在冗余。

考虑到相邻的 Patch 内的高相似度, 作者假设窗口内的特征语义是多余的, 可以进行合理压缩。KV Token Compression 通过压缩相邻的 Patch 内压缩 token 的特征，以 表示, 如图 8 所示。

此外，为了减轻 Self-Attention 计算中 KV 压缩引起的潜在信息丢失，作者选择去保留了 Q 的所有 tokens。这种策略使作者能够有效地利用 KV Token 压缩，同时减轻丢失关键信息的风险。通过使用 KV Token 压缩，可以提高注意力机制计算的效率，减少计算复杂度从 到 ，从而更好地控制生成高分辨率图像的计算成本。

作者使用具有特定初始化的卷积算子 "Conv2×2" 压缩深层。作者设计了专门的卷积初始化 "Conv Avg Init"，它利用 Group Convolution 初始化权重 ，相当于平均算子。这种初始化策略可以加速微调过程，同时只引入了 0.018% 的额外参数。

关于 KV Token 压缩的消融实验

作者也做了一些关于 KV Token 压缩的消融实验。首先是压缩的位置，如图 9 所示，在深层使用 KV Token 压缩可以显著提升性能。作者推测这是因为浅层通常编码详细的纹理内容，而深层表征高级语义内容。由于压缩往往会影响图像质量而不是语义信息，所以压缩深层可以达到最小的信息损失。

作者也在图 10 中探索了不同压缩算子的影响。共计使用了 3 种技术，随机丢弃、平均池化和参数卷积，将 2×2 Token 压缩为单个 Token。"Conv 2×2" 方法优于其他方法，使用可学习卷积核比其他方法更有效地减少冗余。

作者在图 11 中研究了不同压缩比对不同分辨率的影响。发现 Token 压缩不影响文本和图像 (CLIP Score) 之间的对齐，但会影响不同分辨率的图像质量 (FID)。尽管随着压缩率的增加，图像质量略有下降，但本文策略使训练速度提高了 18% 到 35%。

作者进一步在图 12 中全面验证了推理速度的优势。在 4K 生成中，本文方法可以将训练和推理速度提高大约 35%。值得注意的是，可以观察到训练的加速随着分辨率的增加而增加：随着分辨率从 1K 增加到 4K，训练加速从 18% 增至 35%。这些结果也证明了本文方法对更高分辨率的图像生成任务的适用性。

1.4 Weak-to-Strong 训练策略

作者提出了几种有效的训练策略，以增强从 "weak" 模型到 "strong" 模型的转换。这些策略包括 VAE 快速适配、高分辨率微调和 KV Token 压缩。

将模型适配新的 VAE

随着 VAE 技术的持续演进，从头开始训练 T2I 模型需要巨量的计算资源。本文作者将 PixArt-α 的 VAE 替换为 SDXL 的 VAE，并继续微调扩散模型。作者观察到了快速收敛的现象，只需要微调 2K 步就可以迅速收敛，如图 13(a) 所示。对于处理那些 VAE 模型的迁移学习时，微调无疑是一种高效的选择。

适配更高的分辨率

当从低分辨率 (LR) 模型微调到高分辨率 (HR) 模型时，作者观察到了性能下降，如图 13(b) 所示，并将其归因于不同分辨率之间位置编码 (PE) 的差别。为了缓解这个问题，作者利用 "PE 插值" [11] 技术：通过插值 LR 模型的 PE 来初始化 HR 模型的 PE，显著地增强了 HR 模型的初始状态并加速微调过程。

为模型应用于 KV Token 压缩

当现在有一个没有 KV Token 压缩的 LR 预训练模型，我们要对其进行微调时，就可以直接使用 KV Token 压缩。如图 13(c) 所示，使用本文的 "Conv Avg Init" 策略时，PixArt-Σ 可以从更好的初始状态开始，更容易收敛。值得注意的是，即使是只使用 100 训练步骤，PixArt-Σ 也能达到令人满意的视觉效果。通过 KV Token 技术和压缩层设计，训练和推理时间可以减少约 34%。

1.5 实验结果

作者遵循 Imagen [12] 和 PixArt-α 使用 T5 的编码器 (即 Flan-T5 [8] ) 作为条件特征提取的文本编码器，并使用 PixArt-α 作为基本扩散模型。与大多数工作提取固定 77 个 text tokens 不同，作者把 text tokens 的数量从 PixArt-α 的 120 增加至 300，因为 Inside-Σ 数据集的标题要密集得多，可以提供更高细粒度的细节。为了捕捉输入图像的 latent representation，作者使用来自 SDXL [3] 的预训练 VAE，保持冻结。其他实现细节与 PixArt-α 相同。模型使用位置编码插值技术在 PixArt-α 的 256px 预训练 ckpt 上微调。最终模型，包括 1K 分辨率，是在 32 个 V100 GPU 上训练的。作者还使用 16 个 A100 GPU 来训练 2K 和 4K 图像生成模型。

作者使用了权重衰减为 0 的 CAME 优化器 [13] 和 2e-5 的恒定学习率，而不是常规的 AdamW [14] 优化器。这有助于降低优化器状态的维度，导致 GPU 内存较低，而不会导致性能下降。

评价指标方面，为了更好地说明美学和语义能力，作者收集了 30K 高质量的文本图像对来评测 T2I 模型。主要通过人类和人工智能偏好研究来评估 PixArt-Σ，因为 FID 指标可能不能充分反映生成能力。

图片质量评价

如图 1 所示，PixArt-Σ 可以生成高质量、逼真的图像，而且在不同的纵横比和风格上生成的图像也具有复杂的细节。如图 3 所示，作者将 PixArt-Σ 与开源模型 SDXL 和 PixArt-α 进行了比较，PixArt-Σ 增强了肖像的真实感，提高了语义分析的能力。与 SDXL 相比，PixArt-Σ 在遵守用户指令方面也有卓越的表现。

不仅优于开源模型，而且 PixArt-Σ 也与当前 T2I 闭源产品相比也是非常具有竞争力，如图 4 所示。PixArt-Σ 产生照片般逼真的结果，并密切遵守用户指令，与目前的一些商业产品的水平相当。

高分辨率图像生成

本文方法能够直接生成 4K 分辨率的图像，不需要任何后处理。此外，它擅长准确遵循用户提供的复杂、详细和长文本，如图 2 所示。因此，用户不需要提示工程就可以获得满意的结果。

Human/AI (GPT4V) 偏好研究

作者也进行了 Human/AI (GPT4V) 偏好研究，使用 300 个 captions，这 300 个 captions 是从 1.2 节提到的高质量评测集中抽到的。作者收集了由共 6 个开源模型生成的图像，包括 PixArt-α、PixArt-Σ、SD1.5 [15] 、Stable Turbo [16] 、Stable XL、Stable Cascade [17] 和 Playground-V2.0。作者开发了一个网站，并分发给评估者，要求他们评估图像，根据质量对它们进行排名以及它们与文本提示的匹配程度。结果表明 PixArt-Σ 相对于其他 6 个 T2I 生成器有显著的偏好。

此外，作者还使用多模态模型 GPT-4 Vision [18] 作为评估器。每个试验为 GPT-4 Vision 提供 2 张图片：一个来自 PixArt-Σ，另一个来自其他竞争对手 T2I 模型。作者制作了不同的 Prompts，引导 GPT-4 Vision 根据图像质量和图像和文本对齐进行投票。图 14 中的橙色和绿色条表示的结果显示了人类和 AI 偏好研究中一致的结果。PixArt-Σ 在超过了基线 PixArt-α。而且与 Stable Cascade 等模型相比，PixArt-Σ 在生成图像质量和指令跟随能力方面表现出极具竞争力的优越性能。

参考

1. ^ https://openai.com/index/dall-e-3/
2. ^ https://www.midjourney.com/home
3. ^ ab SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis
4. ^ PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis
5. ^ Scalable Diffusion Models with Transformers
6. ^ Visual Instruction Tuning
7. ^