基于DiT，支持4K图像生成，华为诺亚0.6B文生图模型PixArt-Σ来了

　　众所周知，开发顶级的文生图(T2I)模型需要大量资源，因此资源有限的个人研究者基本都不可能承担得起，这也成为了 AIGC(人工智能内容生成)社区创新的一大阻碍。同时随着时间的推移，AIGC 社区又能获得持续更新的、更高质量的数据集和更先进的算法。

　　于是关键的问题来了:我们能以怎样的方式将这些新元素高效地整合进现有模型，依托有限的资源让模型变得更强大?

　　为了探索这个问题，华为诺亚方舟实验室等研究机构的一个研究团队提出一种新的训练方法:由弱到强式训练(weak-to-strong training)。

　　他们的研究基于他们去年十月提出的一种高效的文生图训练方法 PixArt-α，参阅机器之心报道《超低训练成本文生图模型 PixArt 来了，效果媲美 MJ，只需 SD10% 训练时间》。PixArt-α 是 DiT(扩散 Transformer)框架的一种早期尝试。而现在，随着 Sora 登上热搜以及 Stable Diffusion 层出不穷的应用，DiT 架构的有效性得到了研究社区越来越多工作的验证，例如 PixArt， Dit-3D， GenTron 等「1」。

　　该团队使用 PixArt-α 的预训练基础模型，通过整合高级元素以促进其持续提升，最终得到了一个更加强大的模型 PixArt-Σ。图1展示了一些生成结果示例。

　　PixArt-Σ 如何炼成?

　　具体来说，为了实现由弱到强式训练，造出 PixArt-Σ，该团队采用了以下改进措施。

　　更高质量的训练数据

　　该团队收集了一个高质量数据集 Internal-Σ，其主要关注两个方面:

　　(1) 高质量图像:该数据集包含3300万张来自互联网的高分辨率图像，全都超过1K 分辨率，包括230万张分辨率大约为4K 的图像。这些图像的主要特点是美观度高并且涵盖广泛的艺术风格。

　　(2) 密集且准确的描述:为了给上述图像提供更精准和详细的描述，该团队将 PixArt-α 中使用的 LLaVA 替换成了一种更强大的图像描述器 Share-Captioner。

　　不仅如此，为了提升模型对齐文本概念和视觉概念的能力，该团队将文本编码器(即 Flan-T5)的 token 长度扩展到了大约300词。他们观察到，这些改进可以有效消除模型产生幻觉的倾向，实现更高质量的文本 - 图像对齐。

　　下表1展示了不同数据集的统计数据。

　　高效的 token 压缩

　　为了增强 PixArt-α，该团队将其生成分辨率从1K 提升到了4K。为了生成超高分辨率(如2K/4K)的图像，token 数量会大幅增长，这就会导致计算需求大幅增长。

　　为了解决这一难题，他们引入了一种专门针对 DiT 框架调整过的自注意力模块，其中使用了键和值 token 压缩。具体来说，他们使用了步长为2的分组卷积来执行键和值的局部聚合，如下图7所示。

　　此外，该团队还采用了一种专门设计的权重初始化方案，可在不使用 KV(键 - 值)压缩的前提下从预训练模型实现平滑适应。这一设计可有效将高分辨率图像生成的训练和推理时间降低大约34%。

　　由弱到强式训练策略

　　该团队提出了多种微调技术，可快速高效地将弱模型调整为强模型。其中包括:

　　(1) 替换使用了一种更强大的变分自动编码器(VAE):将 PixArt-α 的 VAE 替换成了 SDXL 的 VAE。

　　(2) 从低分辨率到高分辨率扩展，这个过程为了应对性能下降的问题，他们使用了位置嵌入(PE)插值方法。

　　(3) 从不使用 KV 压缩的模型演进为使用 KV 压缩的模型。

　　实验结果验证了由弱到强式训练方法的可行性和有效性。

　　通过上述改进，PixArt-Σ 能以尽可能低的训练成本和尽可能少的模型参数生成高质量的4K 分辨率图像。

　　具体来说，通过从一个已经预训练的模型开始微调，该团队仅额外使用 PixArt-α 所需的9% 的 GPU 时间，就得到了能生成1K 高分辨率图像的模型。如此表现非常出色，因为其中还替换使用了新的训练数据和更强大的 VAE。

　　此外，PixArt-Σ 的参数量也只有0.6B，相较之下，SDXL 和 SD Cascade 的参数量分别为2.6B 和5.1B。

　　PixArt-Σ 生成的图像的美观程度足以比肩当前最顶级的文生图产品，比如 DALL・E3和 MJV6。此外，PixArt-Σ 还展现出了与文本 prompt 细粒度对齐的卓越能力。

　　图2展示了一张 PixArt-Σ 生成4K 高分辨率图像的结果，可以看到生成结果很好地遵从了复杂且信息密集的文本指令。

　　实验

　　实现细节

　　训练细节:对于执行条件特征提取的文本编码器，该团队按照 Imagen 和 PixArt-α 的做法使用了 T5的编码器(即 Flan-T5-XXL)。基础扩散模型就是 PixArt-α。不同于大多数研究提取固定的77个文本 token 的做法，这里将文本 token 的长度从 PixArt-α 的120提升到了300，因为 Internal-Σ 中整理的描述信息更加密集，可以提供高细粒度的细节。另外 VAE 使用了来自 SDXL 的已预训练的冻结版 VAE。其它实现细节与 PixArt-α 一样。