超越α！PixArt家族新秀PixArt-Σ: 由弱到强训练的文本生成4K图像DiT（华为诺亚）

文章链接： https://arxiv.org/pdf/2403.04692
开源地址：https://pixart-alpha.github.io/PixArt-sigma-project/

PixArt-Σ，这是一个能够直接生成4K分辨率图像的Diffusion Transformer（DiT）。PixArt-Σ相比其前身PixArt-α具有显著的进步，提供了质量更高、与文本提示更好对齐的图像。

PixArt-Σ的一个关键特性是其训练效率。利用PixArt-α的基础预训练，它通过整合更高质量的数据从“较弱”的基线发展到“较强”的模型，这一过程称之为“弱到强训练”。

PixArt-Σ的进步有两方面：

高质量训练数据：PixArt-Σ整合了质量更高的图像数据，并配以更精确和详细的图像描述。
高效的Token压缩：在DiT框架内提出了一个新颖的注意力模块，可以压缩键和值，显著提高效率，促进超高分辨率图像生成。

由于这些改进，PixArt-Σ实现了优越的图像质量和用户提示遵循能力，并且具有显著更小的模型大小（0.6B参数），远远小于现有的文本到图像扩散模型，如SDXL（2.6B参数）和SD Cascade（5.1B参数）。此外，PixArt-Σ生成4K图像的能力支持高分辨率海报和壁纸的创建，有效地提升了电影和游戏等行业的高质量视觉内容的生产。

介绍

最近,高质量文本到图像（T2I）模型的出现深刻影响了人工智能生成内容（AIGC）社区。这包括专有模型，如DALL·E 3 、Midjourney ，以及开源模型，如Stable Diffusion 和PixArt-α。然而，开发顶级的T2I模型需要相当大的资源；例如，从头开始训练SD1.5需要约6000个A100 GPU 天，这给资源有限的个人研究人员造成了重大障碍，并阻碍了AIGC社区的创新。随着时间的推移，AIGC社区将获得持续更新、质量更高的数据集和更先进的算法。一个关键问题是：如何在有限资源的约束下有效地将这些新元素整合到现有模型中，实现更强大的版本？

为了探讨这个问题，本研究重点是增强PixArt-α，这是一种高效的T2I训练方法。PixArt-α代表了DiT框架内的一个早期尝试，这是一个具有巨大潜力的模型结构，正如Sora和Stable Diffusion 3 等工作所证明的那样。为了最大化这一潜力，我们建立在PixArt-α的预训练基础上，整合先进的元素来促进其持续改进，结果是一个更强大的模型，PixArt-Σ。将这个通过高效训练从相对较弱的基线发展到更强大模型的过程称为“弱到强训练”。具体来说，为了实现“弱到强训练”，我们引入了以下增强措施：

高质量训练数据

收集了一份高质量的数据集，比PixArt-α中使用的数据更好，重点关注两个关键方面：

高质量图像：该数据集包括来自互联网的3300万高分辨率图像，全部超过1K分辨率，其中包括约230万张分辨率约为4K的图像。这些图像主要以其高美感特性为特征，并涵盖了各种艺术风格。
密集且准确的描述：为了为上述图像提供更精确和详细的描述，用一个更强大的图像描述生成器Share-Captioner替换了PixArt-α中使用的LLaVA。

此外，为了提高模型在文本和视觉概念之间的对齐能力，将文本编码器（即Flan-T5）的Token长度扩展到约300个单词。我们观察到，这些改进有效地消除了模型产生幻觉的倾向，导致了更高质量的文本图像对齐。

高效Token压缩

为了增强PixArt-α，我们将其生成分辨率从1K扩展到4K。在超高分辨率（例如2K/4K）生成图像会导致Token数量显著增加，进而导致计算需求大幅上升。为了解决这一挑战，在DiT框架中引入了一个自注意力模块，用于定制键和值的Token压缩。具体来说，使用步幅为2的分组卷积来对键和值进行局部聚合。此外，采用了一种专门的权重初始化方案，允许在没有KV压缩的情况下从预训练模型平稳过渡。这一设计有效地将高分辨率图像生成的训练和推理时间减少了约34%。

弱到强训练策略

我们提出了几种微调技术，以快速有效地从弱模型适应为强模型。这包括：

用更强大的变分自编码器（VAE）替换
从低分辨率到高分辨率进行缩放
从没有键值（KV）压缩的模型演化到有KV压缩的模型。

这些结果证实了“弱到强训练”方法的有效性和合理性。

通过提出的改进，PixArt-Σ在最小的训练成本和模型参数下实现了高质量的4K分辨率图像生成。具体来说，从一个预训练模型微调，仅额外利用了PixArt-α所需GPU天数的9%，就实现了一个强大的1K高分辨率图像生成模型，这令人印象深刻，考虑到这仅仅使用了新的训练数据和更强大的VAE。

此外，我们仅使用了0.6B参数，而SDXL和SD Cascade分别使用了2.6B和5.1B参数。PixArt-Σ生成的图像具有与当前顶级T2I产品（如DALL·E 3 和MJV6）相媲美的美感质量（如下图4所示）。

此外，PixArt-Σ还展现了与文本提示的细粒度对齐的出色能力（如下图2和3所示）。

框架

数据分析

更具美感和更高的分辨率。 为了增强我们数据集的美感质量，将内部数据从14M扩展到33M。为了清晰起见，将两个数据集分别命名为Internal-α和Internal-Σ。需要注意的是，与当前可用的开源模型如SD v1.5使用的20亿数据相比，这种扩展仍然不足。我们证明了即使数据量有限，有效的训练策略仍然可以获得强大的T2I模型。

Internal-Σ中的图像分辨率均超过1K。为了促进4K分辨率的生成，另外收集了一组8M的真实摄影图像，分辨率为4K。为了确保美感质量，使用美感评分模型（AES）对这些4K图像进行筛选。这个过程产生了一个高度精炼的、包含2M张超高分辨率和高质量图像的数据集。

有趣的是，我们观察到随着图像分辨率的增加，模型的保真度（Fréchet Inception Distance (FID)）和语义对齐（CLIP评分）都有所提升，突显了生成高分辨率图像能力的重要性。

更好的文本-图像对齐。最近的研究，如PixArt-α和DALL-E 3，强调了文本-图像描述对齐的重要性。加强这种对齐对于提升模型能力至关重要。为了进一步改进收集到的“原始”描述，专注于提高我们描述的长度和准确性。值得注意的是，我们的描述（Internal-Σ）在以下方面与PixArt-α（Internal-α）中使用的描述相比具有几个优势。

提高了描述的准确性：如下图5所示，PixArt-α中使用的LLaVa存在一定的幻觉问题。我们利用更强大的视觉语言模型，即Share-Captioner，生成详细和正确的描述，增强了收集的原始提示。

增加了描述的长度：如下表1和下图6所示，平均描述长度显著增加至180个单词，极大地增强了描述的描述能力。此外，将文本编码器的Token处理长度从120个Token（如在Internal-α中）扩展到300个Token。我们的模型在长（Share-Captioner）和短（原始）描述的混合数据集上训练，比例分别为60%和40%。这种方法增强了文本描述的多样性，并减轻了仅依赖生成式描述可能产生的潜在偏见。

上面表1总结了Internal-α和-Σ的情况，通过各种指标评估了数据集的多样性，包括名词种类、总名词数、平均描述长度和每个图像的平均名词数。

高质量评估数据集。大多数SOTA的 T2I模型选择MSCOCO作为评估集来评估FID和CLIP Scores。然而，我们观察到在MSCOCO数据集上进行的评估可能无法充分反映模型在美感和文本-图像对齐方面的能力。因此，我们提出了一个精心策划的数据集，包括30,000个高质量、美感愉悦的文本-图像对，以便进行评估。数据集的选定样本见附录。此数据集旨在更全面地评估模型的性能，特别是在捕捉美感吸引力的复杂性和文本描述与视觉内容之间对齐的保真度方面。除非另有说明，本文中的评估实验均在收集的高质量评估数据集上进行。

高效的DiT设计

高效的DiT网络至关重要，因为在生成超高分辨率图像时，计算需求会显著增加。注意力机制在Diffusion Transformers的有效性中起着关键作用，然而，其二次计算需求显著限制了模型的可扩展性，特别是在更高的分辨率下，例如2K和4K。受PVT v2的启发，在原始PixArt-α框架中引入KV压缩，以解决计算挑战。这种设计仅增加了总参数的0.018%，但通过Token压缩实现了计算成本的有效减少，同时仍然保留了空间和语义信息。

键-值（KV）Token压缩。 动机源于一个有趣的观察，即将键-值（KV）Token压缩直接应用于预训练的PixArt-α仍然可以生成合理的图像。这表明特征中存在冗余。考虑到相邻的R×R块之间的高相似性，假设窗口内的特征语义是冗余的，并且可以合理地压缩。我们提出了KV Token压缩，用fc(·)表示，通过压缩运算符在一个R×R窗口内压缩Token特征，如图7所示。

此外，为了减轻KV压缩在自注意力计算中可能导致的信息丢失，我们选择保留所有query（Q）的Token。这种策略使能够有效地利用KV压缩，同时减轻丢失关键信息的风险。通过采用KV压缩，增强了注意力计算的效率，并将计算复杂度从降低到，从而使直接生成高分辨率图像的计算成本可管理。

使用卷积运算符“Conv2×2”和特定初始化来压缩深层。其他设计变体的详细实验在第5节中讨论。具体来说，设计了一个专门的卷积核初始化方法“Conv Avg Init”，它利用了分组卷积，并将权重w初始化为w=，相当于一个平均运算符。这种初始化策略可以初始产生粗糙的结果，加速微调过程，同时只引入了0.018%的额外参数。

弱到强训练策略

本文提出了几种有效的训练策略，以增强从“弱”模型到“强”模型的过渡。这些策略涵盖了VAE快速适应、高分辨率微调和KV Token压缩。

适应新的VAE模型 随着VAE的不断发展，从头开始训练T2I模型需要大量资源。我们用SDXL的VAE替换了PixArt-α的VAE，并继续微调扩散模型。如下图8(a)所示，观察到微调很快在2K训练步骤内收敛。在处理VAE模型转移时，微调更加高效，避免了从头开始训练的必要性。

适应更高分辨率 当从低分辨率（LR）模型微调到高分辨率（HR）模型时，观察到性能下降，如上图8(b)所示，我们将其归因于不同分辨率之间位置embedding（PE）的差异。为了减轻这个问题，我们利用了“PE插值”技巧：通过插值LR模型的PE来初始化HR模型的PE，显著增强HR模型的初始状态并加快微调过程。即使在仅100次训练迭代内，也可以获得视觉上令人满意的图像。此外在下表2中定量评估了模型性能的变化。微调在1K步骤内迅速收敛，并进一步训练稍微提高了性能。这说明使用“PE插值”技巧使更高分辨率的生成快速收敛，无需从头开始训练来生成更高分辨率的图像。

适应KV压缩模型 可以在从未进行KV压缩的LR预训练模型进行微调时直接使用KV压缩。如上图8(c)所示，通过“Conv Avg Init.”策略，PixArt-Σ从更好的初始状态开始，使收敛更加容易和快速。值得注意的是，即使在100个训练步骤内，PixArt-Σ也能产生满意的视觉结果。最后，通过前面的KV压缩运算符和压缩层设计，我们可以减少约34%的训练和推理时间。

实验

实现细节

训练细节。遵循Imagen和PixArt-α，采用T5的编码器（即Flan-T5-XXL）作为条件特征提取的文本编码器，并使用PixArt-α作为我们的基础扩散模型。与大多数提取固定的77个文本Token的工作不同，我们将文本Token的长度从PixArt-α的120调整为300，因为在Internal-Σ中策划的描述更密集，以提供高度细粒度的细节。

为了捕捉输入图像的潜在特征，我们采用了来自SDXL的预训练且冻结的VAE。其他实现细节与PixArt-α相同。模型在PixArt-α的256px预训练检查点上进行微调，使用位置embedding插值技巧。我们的最终模型，包括1K分辨率，是在32个V100 GPU上训练的。还使用16个A100 GPU来训练2K和4K图像生成模型。有关更多信息，请参阅附录。

请注意，我们使用CAME优化器，权重衰减为0，恒定学习率为，而不是常规的AdamW优化器。这有助于减小优化器状态的维度，从而降低GPU内存使用量而不降低性能。

评估指标。 为了更好地说明美感和语义能力，我们收集了30K个高质量的文本-图像对，以评估最强大的T2I模型。主要通过人类和AI偏好研究来评估PixArt-Σ，因为FID指标可能无法充分反映生成质量。然而，仍然在附录中提供了在收集的数据集上的FID结果。

性能比较

图像质量评估。对我们的方法进行了定性评估，与闭源的文本到图像（T2I）产品和开源模型进行了比较。如前面图1所示，我们的模型能够以复杂的细节生成高质量的照片逼真的图像，覆盖了各种不同的宽高比和风格。这种能力凸显了我们的方法在从文本描述中生成视觉上引人注目的内容方面的卓越性能。如下面图所示，将PixArt-Σ与开源模型SDXL和PixArt-α进行了比较，我们的方法增强了肖像的逼真度，并提高了语义分析的能力。与SDXL相比，我们的方法在遵循用户指令方面表现出更高的熟练度。

我们的方法不仅优于开源模型，在与当前T2I闭源产品的比较中，也表现得非常有竞争力，如下图所示。PixArt-Σ生成了逼真的照片，并且与用户指令紧密一致，与当代商业产品不相上下。

高分辨率生成。我们的方法能够直接生成4K分辨率的图像，无需任何后处理。此外，它在准确地遵循用户提供的复杂、详细和长文本方面表现出色，如前面图2所示。因此，用户无需进行提示工程即可获得满意的结果。我们的方法实现了直接的4K图像生成。与此同时，研究[10, 15]引入了无调整后处理技术，旨在从LR模型生成HR图像或使用超分辨率模型[49]生成HR图像。

然而，由于两个主要原因，它们的相应结果通常会出现伪影：

级联流水线可能导致累积误差。
这些方法不捕捉4K图像的真实分布，也不学习文本与4K图像之间的对齐关系。

我们认为，我们的方法可能是生成高分辨率图像的更有前景的方法。

我们的方法产生了优越的结果，下面包含了更多的视觉比较。

人类/人工智能（GPT4V）偏好研究。 使用前面提到的高质量评估数据集中随机收集的300个描述的子集，在人类和人工智能偏好研究中评估经过充分训练的模型。收集了六个开源模型生成的图像，包括PixArt-α、PixArt-Σ、SD1.5、Stable Turbo、Stable XL、Stable Cascade 和 Playground-V2.0。我们为人类偏好研究开发了一个网站，以显示提示及其相应的图像。

该网站分发给经过训练的评估者，他们被要求评估这些图像，并根据质量和与文本提示的匹配程度对它们进行排名。结果显示在下图9中的蓝色条形图中，表明与其他六个T2I生成器相比，PixArt-Σ受到了明显的青睐。PixArt-Σ生成的高质量图像优于其他生成器，且与用户提示密切匹配，而其参数规模较小（0.6B参数），相比现有的T2I扩散模型（如SDXL（2.6B参数）和SD Cascade（5.1B参数））更为高效。

此外，在我们的人工智能偏好研究中，采用了先进的多模态模型GPT-4 Vision作为评估器。对于每个试验，向GPT-4 Vision提供了两个图像：一个来自PixArt-Σ，另一个来自竞争的T2I模型。设计了不同的提示，引导GPT-4 Vision根据图像质量和图像与文本对齐情况进行投票。结果显示在上图9中的橙色和绿色条形图中，表明在人类和人工智能偏好研究中均得到了一致的结果。具体而言，PixArt-Σ在效果上超越了基线PixArt-α。与诸如Stable Cascaded等当代先进模型相比，PixArt-Σ在图像质量和指示跟随能力方面表现出具有竞争力或更优越的性能。

消融研究

对各种KV压缩设计的生成性能进行了消融研究。除非另有说明，否则实验均在512px生成上进行。每个消融实验的详细设置都包含在附录中。

实验设置

使用前面描述的测试集进行评估。采用FID来计算收集和生成数据之间的分布差异，作为比较指标。此外，利用CLIP-Score来评估提示与生成图像之间的对齐情况。

压缩设计

压缩位置。 我们在Transformer结构的不同深度实现了KV压缩：浅层（1∼14）、中间层（7∼20）和深层（14∼27）。如下表3a所示，对深层进行KV压缩明显实现了更优越的性能。我们推测这是因为浅层通常编码了详细的纹理内容，而深层抽象了高级语义内容。由于压缩往往影响图像质量而不是语义信息，对深层进行压缩可以实现最少的信息丢失，这使得它成为加速训练但不损害生成质量的实用选择。

压缩操作符。 探讨了不同压缩操作符的影响。我们采用了三种技术，即随机丢弃、平均池化和参数化卷积，将2×2的token压缩成单个token。如表3b所示，“Conv 2×2”方法优于其他方法，强调了使用可学习的核来更有效地减少冗余特征的优势，而不是简单的丢弃方法。

在不同分辨率下的压缩比率。 研究了不同分辨率下不同压缩比率对结果的影响。如表3c所示，令人惊讶的是，我们发现Token压缩并不影响文本和生成图像之间的对齐（CLIP得分），但会影响不同分辨率下的图像质量（FID）。虽然随着压缩比率的增加，图像质量略有下降，但我们的策略将训练速度提高了18%至35%。这表明我们提出的KV压缩方法既有效又高效，可实现高分辨率的T2I生成。

不同分辨率下的速度比较。 在表3d中进一步全面验证了训练和推断的速度加速情况。我们的方法可以将4K生成的训练和推断速度加快约35%。值得注意的是，我们观察到随着分辨率的提高，训练加速度也会增加。例如，随着分辨率从1K提高到4K，训练逐渐加速了18%至35%。这表明我们的方法在分辨率提高时的有效性，展示了其对更高分辨率图像生成任务的潜在适用性。

结论

本文介绍了PixArt-Σ，一个能够直接生成4K分辨率高质量图像的文本到图像（T2I）扩散模型。基于PixArt-α的预训练基础，PixArt-Σ通过一种新颖的“从弱到强的训练”方法实现了高效的训练。这种方法的特点是整合了更高质量的数据和高效的Token压缩。PixArt-Σ在生成高保真度图像的同时，能够紧密遵循文本提示，超越了其前身PixArt-α的高标准。相信PixArt-Σ所呈现的创新将不仅有助于AIGC社区的进步，还将为个体提供更高效、更高质量的生成模型铺平道路。