PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation

Motivation

1. 高分辩率图像生成： Transformer架构中序列越长，计算复杂度是O(n^2)，越长的分辨率，对于计算推理时间和训练成本来说就越高。如何实现更好更快的生成是一大难点。

2. 高质量的图文对数据：爬虫图文对在图片质量和文本质量上都存在问题，不够美观，图文相关性弱。利用MLLM进行recaption通常会出现幻觉问题，提升MLLM的精度对于图文一致性非常重要。

3. 从弱到强的训练策略：对于低分辨率训练模型、vae模型切换，从已经训练好的base模型，继承之前的训练权重，如何更好的迁移到新模型非常重要，节约训练成本。

主要贡献

• 高分辩率图像生成

  • 根据self-attention的计算原理，KV矩阵的长度可以比原序列更短。注意力维度由NxN变成Nx(N/(RxR)): $QK^{T}$的维度变换(NC) (CxN/(RxR))-> N x (N/(RxR))。这样可以实现计算的压缩，并且相邻token存在语义的相似性，这样相当于引入了空间的局部先验。这里压缩的函数$f_{compress}$可以是global average pooling或者是stride为R的卷积层（可以用avg的kernel初始化加速训练）。

• 高质量的图文对数据

  • PixArt-Σ采用更好的ShareCaptioner替代原始的LLava模型，幻觉率更低，训练时采用60%概率选择，让模型能够适用caption文本和更多样范式的其他文本。收集了8百万4K分辨率的真实摄影图片。

• 从弱到强的训练策略:

  • VAE: 从SD1.5的VAE替换到SDXL的VAE，2k训练steps
  • 512分辨率提升到1024分辨率：结合位置编码插值(PE Interp)，可以实现更快的尺度适应
  • KV压缩并采用avg的kernel权重初始化可以加速训练
  • 继承原有的权重训练，PixArt-Σ具有非常高的训练效率

一些思考

• 局部窗口进行kv的压缩对于用Transformer架构的生成模型来说都是适用的，也可以用于自回归范式的图像生成模型