华为诺亚发布 PixArt-Σ：基于 DiT 架构的 0.6B 4K 文生图模型
引言
众所周知，开发顶级的文生图（Text-to-Image, T2I）模型需要消耗巨大的计算资源和数据积累，这导致资源有限的个人研究者或小型团队难以承担，成为了 AIGC（人工智能内容生成）社区创新的一大阻碍。随着时间推移，AIGC 社区虽然能获得持续更新的、更高质量的数据集和更先进的算法，但如何将这些新元素高效地整合进现有模型，依托有限的资源让模型变得更强大，一直是关键问题。
为了解决这一挑战，华为诺亚方舟实验室等研究机构的研究团队提出了一种新的训练方法：由弱到强式训练（Weak-to-Strong Training）。基于此方法，他们发布了名为 PixArt-Σ 的文生图模型。该模型与 Sora 一样采用了 DiT（Diffusion Transformer）框架，支持 4K 图像生成，且参数量仅为 0.6B。
论文标题：PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation
PixArt-α 基础与 DiT 架构演进
该团队的研究基于其去年十月提出的一种高效的文生图训练方法 PixArt-α。PixArt-α 是 DiT（扩散 Transformer）框架的一种早期尝试。DiT 架构将扩散模型的 U-Net 结构替换为 Transformer 结构，利用自注意力机制处理全局依赖关系。随着 Sora 登上热搜以及 Stable Diffusion 层出不穷的应用，DiT 架构的有效性得到了研究社区越来越多工作的验证，例如 PixArt, Dit-3D, GenTron 等。
该团队使用 PixArt-α 的预训练基础模型，通过整合高级元素以促进其持续提升，最终得到了一个更加强大的模型 PixArt-Σ。图 1 展示了一些生成结果示例。
PixArt-Σ 核心改进措施
为了实现由弱到强式训练，造出 PixArt-Σ，该团队采用了以下主要改进措施。
1. 更高质量的训练数据
数据质量直接决定了生成模型的上限。该团队收集了一个高质量数据集 Internal-Σ，其主要关注两个方面：
(1) 高质量图像：该数据集包含 3300 万张来自互联网的高分辨率图像，全都超过 1K 分辨率，包括 230 万张分辨率大约为 4K 的图像。这些图像的主要特点是美观度高并且涵盖广泛的艺术风格，确保了模型能学习到多样化的视觉特征。
(2) 密集且准确的描述：为了给上述图像提供更精准和详细的描述，该团队将 PixArt-α 中使用的 LLaVA 替换成了一种更强大的图像描述器 Share-Captioner。此外，为了提升模型对齐文本概念和视觉概念的能力，该团队将文本编码器（即 Flan-T5）的 token 长度扩展到了大约 300 词。他们观察到，这些改进可以有效消除模型产生幻觉的倾向，实现更高质量的文本 - 图像对齐。
下表展示了不同数据集的统计数据对比。
2. 高效的 Token 压缩
为了增强 PixArt-α，该团队将其生成分辨率从 1K 提升到了 4K。为了生成超高分辨率（如 2K/4K）的图像，token 数量会大幅增长，这就会导致计算需求呈指数级增长，显存占用也会急剧上升。
为了解决这一难题，他们引入了一种专门针对 DiT 框架调整过的自注意力模块，其中使用了键和值 token 压缩（KV Compression）。具体来说，他们使用了步长为 2 的分组卷积来执行键和值的局部聚合。这种设计能够在保留关键空间信息的同时，显著减少序列长度。
