StableDiffusion3：Scaling Rectified Flow Transformers for High-Resolution Image Synthesis——高分辨率图像合成

这篇文章的核心内容是关于如何改进和扩展校正流（Rectified Flow）模型，以实现更高效的高分辨率图像合成，特别是在文本到图像合成任务上。文章通过一系列创新的方法和实验，展示了校正流模型在生成质量和效率上的显著提升。以下是文章的主要研究内容总结：

研究背景与动机

• 扩散模型：扩散模型是一种强大的生成模型，通过从噪声中重建数据来生成新的图像和视频。尽管有效，但传统的扩散模型在高分辨率图像生成时存在效率问题，如采样速度慢和计算成本高。

• 校正流模型：校正流模型通过直线连接数据和噪声，具有更好的理论属性和概念上的简洁性，但在实际应用中尚未成为主流。文章的目标是改进校正流模型，使其在高分辨率图像合成中更具竞争力。

研究方法与创新

1. 改进的噪声采样技术：

   • 作者提出了一种新的噪声采样方法，通过偏向感知相关的尺度来改进校正流模型的训练。这种方法通过调整时间步的采样分布，使得模型在训练过程中更加关注中间步骤，从而提高模型的性能。

   • 通过大规模实验，作者比较了不同的扩散模型和校正流公式，发现特定的校正流公式（如rf/lognorm(0.00, 1.00)）在高分辨率文本到图像合成中表现优于传统的扩散模型。

2. 新型多模态变换器架构（MM-DiT）：

   • 作者设计了一个新的基于变换器的架构，该架构允许文本和图像标记之间的双向信息流动。这种架构使用两组独立的权重来处理文本和图像模态，从而提高了模型对文本的理解能力和图像生成质量。

   • 通过实验，作者证明了这种架构在文本到图像合成任务中的有效性，特别是在处理复杂的文本提示和高分辨率图像生成时。

3. 大规模训练与扩展：

   • 作者进行了大规模的训练实验，将模型扩展到8B参数，并在5×10^22训练FLOPs下进行训练。实验结果表明，随着模型大小和训练步数的增加，验证损失持续下降，且与现有的文本到图像评估指标和人类偏好评分高度相关。

   • 作者还提出了一些技术来提高训练的稳定性和效率，如QK归一化和分辨率依赖的时间步移动，这些技术有助于在高分辨率图像生成中保持模型的稳定性和性能。

实验结果与结论

• 性能提升：通过改进的噪声采样技术和新型架构，作者的模型在高分辨率图像合成任务上优于现有的最先进的模型，包括SDXL、Pixart-α和DALL-E 3等。

• 可扩展性：文章展示了模型性能随着模型大小和训练步数的增加而持续提升，且没有显示出饱和迹象，这表明有潜力进一步提高模型性能。

• 灵活性：新型架构允许在推理时灵活选择文本编码器，从而在性能和内存效率之间进行权衡。

广泛影响

文章指出，其研究目标是推进机器学习和图像合成领域的发展，并讨论了扩散模型可能带来的社会影响。作者强调，他们的工作旨在提高图像合成技术的效率和质量，同时确保其安全和负责任的使用。

总的来说，这篇文章通过一系列创新的方法和实验，展示了校正流模型在高分辨率图像合成中的潜力，特别是在文本到图像合成任务上。通过改进的噪声采样技术、新型架构和大规模训练，作者成功地提高了模型的性能和效率，为未来的研究和应用提供了新的方向。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

图 1. 我们的 8B 校正流模型的高分辨率样本，展示了其在排版、精确提示遵循、空间推理、关注细节以及各种风格的高质量图像生成方面的能力。

摘要

扩散模型通过反转数据向噪声的前向路径来从噪声中创建数据，并已成为一种强大的高维感知数据（如图像和视频）生成建模技术。校正流是一种最近的生成模型公式，它通过直线连接数据和噪声。尽管它具有更好的理论属性和概念上的简洁性，但尚未确立为标准实践。在这项工作中，我们通过将噪声采样技术偏向感知相关的尺度，改进了用于训练校正流模型的现有噪声采样技术。通过大规模研究，我们展示了这种方法与已建立的扩散公式相比，在高分辨率文本到图像合成方面的优越性能。此外，我们提出了一个新的基于变换器的文本到图像生成架构，该架构使用单独的权重用于两种模态，并允许图像和文本标记之间的双向信息流动，从而提高了文本理解、排版和人类偏好评分。我们证明了这种架构遵循可预测的扩展趋势，并且较低的验证损失与通过各种指标和人类评估衡量的改进的文本到图像合成相关。我们最大的模型优于最先进的模型，我们将公开实验数据、代码和模型权重。

1. 引言

扩散模型通过从噪声中创建数据（Song et al., 2020）。它们被训练以反转数据向随机噪声的前向路径，因此，结合神经网络的近似和泛化属性，可以用来生成训练数据中不存在但遵循训练数据分布的新数据点（Sohl-Dickstein et al., 2015; Song & Ermon, 2020）。这种生成建模技术已被证明对于建模高维感知数据（如图像）非常有效（Ho et al., 2020）。近年来，扩散模型已成为从自然语言输入生成高分辨率图像和视频的首选方法，具有令人印象深刻的泛化能力（Saharia et al., 2022b; Ramesh et al., 2022; Rombach et al., 2022; Podell et al., 2023; Dai et al., 2023; Esser et al., 2023; Blattmann et al., 2023b; Betker et al., 2023; Blattmann et al., 2023a; Singer et al., 2022）。由于它们的迭代性质以及相关的计算成本，以及推理过程中的长采样时间，关于这些模型的更高效训练和/或更快采样的公式的研究有所增加（Karras et al., 2023; Liu et al., 2022）。

虽然指定从数据到噪声的前向路径可以实现高效训练，但它也引发了选择哪条路径的问题。这种选择对采样可能有重要影响。例如，未能从数据中去除所有噪声的前向过程可能导致训练和测试分布之间的差异，并导致诸如灰色图像样本等伪影（Lin et al., 2024）。重要的是，前向过程的选择也影响所学习的后向过程，因此影响采样效率。虽然弯曲路径需要许多积分步骤来模拟该过程，但直线路径可以通过单步模拟，并且较少受到误差累积的影响。由于每个步骤对应于神经网络的评估，这直接影响采样速度。

前向路径的一个特定选择是所谓的校正流（Liu et al., 2022; Albergo & Vanden-Eijnden, 2022; Lipman et al., 2023），它通过直线连接数据和噪声。尽管这个模型类别具有更好的理论属性，但尚未在实践中确立。到目前为止，一些优势已在小型和中型实验中得到实证证明（Ma et al., 2024），但这些大多限于类别条件模型。在这项工作中，我们通过引入类似于噪声预测扩散模型（Ho et al., 2020）的校正流模型的噪声尺度重新加权来改变这一现状。通过大规模研究，我们将我们的新公式与现有的扩散公式进行了比较，并展示了其优势。

我们展示了广泛使用的文本到图像合成方法（将固定的文本表示直接输入模型，例如通过交叉注意力（Vaswani et al., 2017; Rombach et al., 2022））并非理想，并提出了一个新的架构，该架构引入了图像和文本标记的可学习流，从而实现了它们之间的双向信息流动。我们将这一架构与我们改进的校正流公式相结合，并研究了其可扩展性。我们展示了验证损失的可预测扩展趋势，并证明较低的验证损失与通过诸如 T2I-CompBench（Huang et al., 2023）、GenEval（Ghosh et al., 2023）和人类评分等指标评估的改进的文本到图像性能强烈相关。我们最大的模型优于最先进的开放模型，如 SDXL（Podell et al., 2023）、SDXL-Turbo（Sauer et al., 2023）、Pixart-α（Chen et al., 2023）以及闭源模型，如 DALL-E 3（Betker et al., 2023），无论是在提示理解的定量评估（Ghosh et al., 2023）还是人类偏好评分方面。

我们工作的核心贡献是：（i）我们进行了一个大规模的系统研究，以确定最佳设置，比较不同的扩散模型和校正流公式。为此，我们引入了新的噪声采样器，用于校正流模型，其性能优于以前已知的采样器。（ii）我们设计了一个新的可扩展的文本到图像合成架构，允许文本和图像标记流在网络内双向混合。我们展示了其相对于已建立的骨干网络（如 UViT（Hoogeboom et al., 2023）和 DiT（Peebles & Xie, 2023））的优势。最后，我们（iii）对我们的模型进行了扩展研究，并展示了它遵循可预测的扩展趋势。我们展示了较低的验证损失与通过诸如 T2I-CompBench（Huang et al., 2023）、GenEval（Ghosh et al., 2023）和人类评分等指标评估的改进的文本到图像性能强烈相关。我们公开提供结果、代码和模型权重。

2. 无需模拟的流训练

我们考虑生成模型，它们通过常微分方程（ODE）定义从噪声分布 p1 的样本 x1 到数据分布 p0 的样本 x0 之间的映射，

其中速度 v 由神经网络的权重 Θ 参数化。Chen et al.（2018）的先前工作建议通过可微分的 ODE 求解器直接求解方程（1）。然而，对于参数化 vΘ(yt, t) 的大型网络架构，这一过程计算成本高昂。一个更有效的替代方法是直接回归一个向量场 ut，它在 p0 和 p1 之间生成一个概率路径。为了构建这样一个 ut，我们定义一个前向过程，对应于 p0 和 p1 = N(0, 1) 之间的概率路径 pt，为

对于 a0 = 1, b0 = 0, a1 = 0 和 b1 = 1，边际分布，

与数据和噪声分布一致。

为了表达 zt、x0 和 ϵ 之间的关系，我们引入 ψt 和 ut 为

由于 zt 可以写成 ODE z′ t = ut(zt|ϵ) 的解，初始值为 z0 = x0，ut(·|ϵ) 生成 pt(·|ϵ)。值得注意的是，可以使用条件向量场 ut(·|ϵ) 构造一个边际向量场 ut，以生成边际概率路径 pt（Lipman et al., 2023）（见附录 B.1），使用条件向量场 ut(·|ϵ)：

虽然直接回归 ut 的流匹配目标

由于方程（6）中的边际化而难以处理，但条件流匹配（见附录 B.1），

与条件向量场 ut(z|ϵ) 提供了一个等效但可处理的目标。

为了将损失转换为显式形式，我们将 ψ′ t(x0|ϵ) = a′ tx0 + b′ tϵ 和 ψ−1 t (z|ϵ) = z−btϵ

at 代入（5）

现在，考虑信噪比 有了，我们可以将方程（9）重写为

接下来，我们使用方程（10）对方程（8）进行重参数化，得到噪声预测目标：

其中我们定义了

请注意，在引入时间依赖权重后，上述目标的最优值不会改变。因此，可以推导出各种加权损失函数，这些函数可以为期望的解提供信号，但可能会影响优化轨迹。为了对不同方法（包括经典的扩散公式）进行统一分析，我们可以将目标写成以下形式（遵循 Kingma & Gao（2023））：

其中 对应于 LCF M。

3. 流轨迹

在本工作中，我们简要描述了上述形式主义的不同变体。

校正流 校正流（RFs）（Liu et al., 2022; Albergo & Vanden-Eijnden, 2022; Lipman et al., 2023）定义前向过程为数据分布和标准正态分布之间的直线路径，即

并使用 LCF M，其对应于 wRF t = t

1−t。网络输出直接参数化速度 vΘ。

EDM EDM（Karras et al., 2022）使用形式为

的前向过程，其中（Kingma & Gao, 2023）bt = exp F−1 N (t|Pm, P2 s)，F−1 N 是均值为 Pm、方差为 P2 s 的正态分布的分位数函数。注意，这种选择导致

网络通过 F 预测（Kingma & Gao, 2023; Karras et al., 2022）进行参数化，损失可以写成 LwEDM t，其中

余弦（Nichol & Dhariwal, 2021）提出了形式为

当与 ϵ 参数化和损失结合时，这对应于权重 wt = sech(λt/2)。当与 v 预测损失（Kingma & Gao, 2023）结合时，权重为 wt = e−λt/2。

3.1. 为 RF 模型量身定制的 SNR 采样器

RF 损失在 [0, 1] 中的所有时间步上均匀地训练速度 vΘ。直观上，然而，对于 [0, 1] 中间的时间步，结果速度预测目标 ϵ − x0 更难，因为对于 t = 0，最优预测是 p1 的均值，对于 t = 1，最优预测是 p0 的均值。一般来说，将时间 t 上的分布从常用的均匀分布 U(t) 改为具有密度 π(t) 的分布，等同于加权损失 Lwπ t，其中

因此，我们旨在通过更频繁地采样它们来给予中间时间步更多的权重。接下来，我们描述我们用于训练模型的时间步密度 π(t)。

对数正态采样 对于一个在中间步骤上分配更多权重的分布选项，是对数正态分布（Atchison & Shen, 1980）。其密度，

其中 logit(t) = log t

1−t，具有位置参数 m 和尺度参数 s。位置参数使我们能够将训练时间步偏向数据 p0（负 m）或噪声 p1（正 m）。如图 11 所示，尺度参数控制分布的宽度。

在实践中，我们从正态分布 u ∼ N(u; m, s) 中采样随机变量 u，然后将其通过标准逻辑函数进行映射。

具有重尾的模式采样 对数正态密度总是在端点 0 和 1 处消失。为了研究这是否对性能产生不利影响，我们还使用了一个在 [0, 1] 上具有严格正密度的时间步采样分布。对于尺度参数 s，我们定义

对于 −1 ≤ s ≤ 2

π−2，这个函数是单调的，我们可以用它来采样来自隐含密度 πmode(t; s) = d

dtf −1 mode(t) 的样本。如图 11 所示，尺度参数控制在采样期间是中间点（正 s）还是端点（负 s）受到青睐的程度。这种公式还包括一个均匀加权 πmode(t; s = 0) = U(t) 用于 s = 0，这在以前关于校正流（Liu et al., 2022; Ma et al., 2024）的工作中被广泛使用。

余弦映射 最后，我们还考虑了余弦时间表（Nichol & Dhariwal, 2021）中的校正流设置。特别是，我们正在寻找一个映射 f : u → f(u) = t，u ∈ [0, 1]，使得对数信噪比与余弦时间表相匹配：2 log

4. 文本到图像架构

对于文本条件的图像采样，我们的模型必须同时考虑文本和图像两种模态。我们使用预训练模型来推导合适的表示，然后描述我们的扩散骨干架构。图 2 展示了概述。

我们的通用设置遵循 LDM（Rombach et al., 2022）用于在预训练自编码器的潜在空间中训练文本到图像模型。与将图像编码为潜在表示类似，我们还遵循先前的方法（Saharia et al., 2022b; Balaji et al., 2022）并使用预训练的、冻结的文本模型对文本条件 c 进行编码。详细信息可在附录 B.2 中找到。

多模态扩散骨干 我们的架构基于 DiT（Peebles & Xie, 2023）架构。DiT 只考虑类别条件图像生成，并使用调制机制使网络同时依赖于扩散过程的时间步和类别标签。同样，我们使用时间步 t 和 cvec 的嵌入作为调制机制的输入。然而，由于池化文本表示仅保留关于文本输入的粗粒度信息（Podell et al., 2023），网络还需要来自序列表示 cctxt 的信息。

我们构建一个由文本和图像输入的嵌入组成的序列。具体来说，我们添加位置编码，并将潜在像素表示 x ∈ Rh×w×c 的 2×2 块展平为长度为 1

2 · h · 1

2 · w 的补丁编码序列。在将补丁编码和文本编码嵌入到共同维度后，我们将这两个序列连接起来。然后，我们按照 DiT 的方式应用一系列调制注意力和 MLP。

由于文本和图像嵌入在概念上相当不同，我们为这两种模态使用两组单独的权重。如图 2b 所示，这相当于为每种模态拥有两个独立的变换器，但在注意力操作中将两种模态的序列连接起来，以便两种表示可以在自己的空间中工作，同时考虑对方。

对于我们的扩展实验，我们用模型的深度 d 来参数化模型的大小，即注意力块的数量，通过将隐藏大小设置为 64·d（在 MLP 块中扩展到 4·64·d 通道），并将注意力头的数量设置为 d。

5. 实验

5.1. 改进校正流

我们的目标是了解哪种用于无需模拟的归一化流训练的方法（如方程 1）最有效。为了使不同方法之间的比较成为可能，我们控制了优化算法、模型架构、数据集和采样器。此外，不同方法的损失是不可比的，也不一定与输出样本的质量相关；因此，我们需要评估指标，以便在不同方法之间进行比较。我们在 ImageNet（Russakovsky et al., 2014）和 CC12M（Changpinyo et al., 2021）上训练模型，并在训练期间使用验证损失、CLIP 分数（Radford et al., 2021; Hessel et al., 2021）和 FID（Heusel et al., 2017）在不同采样器设置（不同的引导尺度和采样步数）下评估模型的训练和 EMA 权重。我们在 COCO-2014 验证分割（Lin et al., 2014）上评估所有指标。我们在附录 B.3 中提供了训练和采样超参数的详细信息。

5.1.1. 结果

我们在两个数据集上训练了 61 种不同的公式。我们在第 3 节中包含了以下变体：

• ϵ 和 v 预测损失，具有线性（eps/linear，v/linear）和余弦（eps/cos，v/cos）时间表。

• RF 损失与 πmode(t; s)（rf/mode(s)）结合，s 在 −1 和 1.75 之间均匀选择 7 个值，以及对于 s = 1.0 和 s = 0，这对应于均匀的时间步采样（rf/mode）。

• RF 损失与 πln(t; m, s)（rf/lognorm(m, s)）结合，(m, s) 在网格中有 30 个值，m 在 −1 和 1 之间均匀分布，s 在 0.2 和 2.2 之间均匀分布。

• RF 损失与 πCosMap(t)（rf/cosmap）结合。

• EDM（edm(Pm, Ps)）与 15 个值结合，Pm 在 −1.2 和 1.2 之间均匀分布，Ps 在 0.6 和 1.8 之间均匀分布。注意，Pm, Ps = (−1.2, 1.2) 对应于（Karras et al., 2022）中的参数。

• 与 rf 的对数信噪比权重相匹配的 EDM（edm/rf）和与 v/cos 的对数信噪比权重相匹配的 EDM（edm/cos）。

对于每次运行，我们选择在 EMA 权重评估时具有最小验证损失的步骤，然后收集使用 6 种不同采样器设置（有和没有 EMA 权重）获得的 CLIP 分数和 FID。

对于所有 24 种采样器设置、EMA 权重和数据集选择的组合，我们使用非支配排序算法对不同公式进行排名。为此，我们反复计算根据 CLIP 和 FID 分数是帕累托最优的变体，为这些变体分配当前迭代索引，移除这些变体，然后继续处理剩余的变体，直到所有变体都被排名。最后，我们在 24 种不同控制设置上平均这些排名。

我们在表 1 中展示了结果，其中我们只展示了对于那些具有不同超参数评估的变体中表现最好的两个变体。我们还在限制采样器设置为 5 步和 50 步的平均排名上进行了展示。

我们观察到 rf/lognorm(0.00, 1.00) 一致地获得了较好的排名。它优于具有均匀时间步采样的校正流公式（rf），从而证实了我们的假设，即中间时间步更为重要。在所有变体中，只有经过修改的时间步采样的校正流公式优于先前使用的 LDM-Linear（Rombach et al., 2022）公式（eps/linear）。

我们还观察到一些变体在某些设置中表现良好，但在其他设置中表现较差，例如 rf/lognorm(0.50, 0.60) 是具有 50 个采样步的最佳表现变体，但在 5 个采样步时表现较差（平均排名为 8.5）。我们在表 2 中观察到关于两个指标的类似行为。第一组展示了代表性变体及其在两个数据集上具有 25 个采样步的指标。下一组展示了在 CLIP 和 FID 分数上表现最佳的变体。除了 rf/mode(1.75) 外，这些变体通常在一个指标上表现非常好，但在另一个指标上表现相对较差。相比之下，我们再次观察到 rf/lognorm(0.00, 1.00) 在不同指标和数据集上均表现出良好的性能，其中它在四次中有两次获得了第三好的分数，一次获得了第二好的性能。

最后，我们在图 3 中用不同颜色展示了不同公式组（edm、rf、eps 和 v）的定性行为。校正流公式通常表现良好，与其他公式相比，其性能在减少采样步数时下降较少。

5.2. 改进模态特定表示

在前一节中，我们找到了一种公式，它不仅使校正流模型能够与已建立的扩散公式（如 LDM-Linear（Rombach et al., 2022）或 EDM（Karras et al., 2022））竞争，甚至优于它们，我们现在转向将该公式应用于高分辨率文本到图像合成。相应地，我们算法的最终性能不仅取决于训练公式，还取决于通过神经网络进行参数化以及我们使用的图像和文本表示的质量。在接下来的部分中，我们描述了如何在第 5.3 节中扩展最终方法之前改进所有这些组件。

5.2.1. 改进的自编码器

潜在扩散模型通过在预训练自编码器的潜在空间中运行来实现高效率，该自编码器将输入 RGB 图像 X ∈ RH×W ×3 映射到低维空间 x = E(X) ∈ Rh×w×d。该自编码器的重建质量为可实现的图像质量提供了上限，经过潜在扩散训练后。与 Dai et al.（2023）类似，我们发现增加潜在通道的数量 d 可显著提高重建性能，如表 3 所示。直观上，预测具有更高 d 的潜在变量是一项更具挑战性的任务，因此具有更大容量的模型应该能够为更大的 d 表现更好，最终实现更高的图像质量。我们在图 10 中证实了这一假设，其中我们看到 d = 16 的自编码器在样本 FID 方面展现出更好的扩展性能。在本文的其余部分，我们选择 d = 16。

5.2.2. 改进的标题

Betker et al.（2023）证明，合成生成的标题可以极大地提高大规模训练的文本到图像模型的性能。这是由于大规模图像数据集附带的人类生成标题通常具有简单性，它们过于关注图像主题，通常省略描述背景或场景构成的细节，或者在适用的情况下，显示的文本（Betker et al., 2023）。我们遵循他们的方法，并使用现成的最先进的视觉语言模型 CogVLM（Wang et al., 2023）为我们的大规模图像数据集创建合成注释。由于合成标题可能会导致文本到图像模型忘记某些不在 VLM 知识库中的概念，我们使用 50% 原始和 50% 合成标题的比例。

为了评估在该标题混合上训练的效果，我们训练了两个 d = 15 的 MM-DiT 模型，训练步数为 250k 步，一个仅使用原始标题，另一个使用 50/50 混合。我们在表 4 中使用 GenEval 基准（Ghosh et al., 2023）评估训练后的模型。结果表明，使用合成标题添加训练的模型明显优于仅使用原始标题的模型。因此，我们在本文的其余部分使用 50/50 的合成/原始标题混合。

5.2.3. 改进的文本到图像骨干

在本节中，我们将现有基于变换器的扩散骨干的性能与我们在第 4 节中介绍的新型多模态变换器基础扩散骨干 MM-DiT 进行比较。MM-DiT 特别设计用于处理不同领域（此处为文本和图像标记）的标记，使用（两组）不同的可训练模型权重。更具体地说，我们遵循第 5.1 节中的实验设置，并在 CC12M 上比较 DiT、CrossDiT（DiT，但通过交叉关注文本标记而不是序列连接（Chen et al., 2023））和我们的 MM-DiT 的文本到图像性能。对于 MM-DiT，我们比较了具有两组权重和三组权重的模型，后者分别处理 CLIP（Radford et al., 2021）和 T5（Raffel et al., 2019）标记（见第 4 节）。注意，DiT（带有文本和图像标记的连接，如第 4 节所述）可以被视为具有单一共享权重集处理所有模态的 MM-DiT 的特例。最后，我们考虑 UViT（Hoogeboom et al., 2023）架构，它是广泛使用的 UNets 和变换器变体之间的混合体。

我们在图 4 中分析了这些架构的收敛行为：纯 DiT 的表现不如 UViT。CrossDiT 变体（带有交叉关注）的性能优于 UViT，尽管 UViT 在初始阶段的学习速度似乎更快。我们的 MM-DiT 变体明显优于交叉关注和纯变体。当使用三组权重而不是两组时，我们只观察到较小的增益（以增加参数数量和 VRAM 使用为代价），因此我们在本文的其余部分选择前者选项。

5.3. 大规模训练

在扩展之前，我们对数据进行过滤和预编码，以确保安全和高效的预训练。然后，所有先前关于扩散公式、架构和数据的考虑都汇集在最后一节中，我们将模型扩展到 8B 参数。

5.3.1. 数据预处理

预训练缓解措施 训练数据对生成模型的能力有显著影响。因此，数据过滤在限制模型的不期望能力方面是有效的（Nichol, 2022）。在大规模训练之前，我们对数据进行过滤，以去除以下类别：（i）性内容：我们使用 NSFW 检测模型来过滤明确内容。（ii）美学：我们移除我们评分系统预测分数较低的图像。（iii）重复：我们使用基于聚类的去重方法，从训练数据中移除感知和语义上的重复项；见附录 E.2。

预计算图像和文本嵌入 我们的模型使用多个预训练的、冻结的网络的输出作为输入（自编码器潜在变量和文本编码器表示）。由于这些输出在训练期间是恒定的，因此我们只需为整个数据集预计算一次。我们在附录 E.1 中详细讨论了我们的方法。

5.3.2. 在高分辨率上进行微调

QK 归一化 一般来说，我们在低分辨率图像（256×256 像素）上预训练所有模型。接下来，我们在更高分辨率（具有混合纵横比）上对模型进行微调（见下一段关于详细信息）。我们发现，当转向高分辨率时，混合精度训练可能会变得不稳定，导致损失发散。可以通过切换到全精度训练来解决这一问题，但这与混合精度训练相比会带来大约 2 倍的性能下降。更有效的替代方法是在（判别式）ViT 文献中报告的：Dehghani et al.（2023）观察到，大型视觉变换器模型的训练会因为注意力熵不受控制地增长而发散。为了避免这种情况，Dehghani et al.（2023）建议在注意力操作之前对 Q 和 K 进行归一化。我们遵循这种方法，并在我们模型的 MMDiT 架构的两个流中使用 RMSNorm（Zhang & Sennrich, 2019），带有可学习的尺度，如图 2 所示。如图 5 所示，额外的归一化可以防止注意力对数增长不稳定性，证实了 Dehghani et al.（2023）和 Wortsman et al.（2023）的发现，并在与 AdamW（Loshchilov & Hutter, 2017）优化器中的 ϵ = 10−15 结合使用时，启用高效的 bf16-mixed 精度训练。当在预训练期间未使用 qk 归一化的预训练模型上应用这种方法时，模型可以快速适应额外的归一化层，并更稳定地进行训练。最后，我们指出，尽管这种方法通常可以帮助稳定大型模型的训练，但它并非万能药，可能需要根据具体的训练设置进行调整。

用于变化纵横比的位置编码 在固定 256×256 分辨率上进行训练后，我们旨在（i）提高分辨率和（ii）启用具有灵活纵横比的推理。由于我们使用 2D 位置频率嵌入，因此我们需要根据分辨率调整它们。在多纵横比设置中，直接插值嵌入（如 Dosovitskiy et al., 2020）将无法正确反映边长。相反，我们使用扩展和插值位置网格的组合，然后对它们进行频率嵌入。

对于目标分辨率为 S2 像素的情况，我们使用分桶采样（NovelAI, 2022; Podell et al., 2023），使得每个批次包含具有均匀大小 H×W 的图像，其中 H·W ≈ S2。对于最大和最小训练纵横比，这导致了将遇到的最大宽度值 Wmax 和高度值 Hmax。设 hmax = Hmax/16, wmax = Wmax/16，s = S/16 为在潜在空间（下采样因子为 8）中对应的大小，经过打补丁（下采样因子为 2）后。基于这些值，我们构建一个垂直位置网格，其值为 ((p−hmax−s

2 )·256

S )hmax−1 p=0，水平位置也相应地构建。然后，我们从得到的 2D 位置网格中进行中心裁剪，再进行嵌入。

分辨率依赖的时间步移动 当分辨率更高时，由于像素更多，我们需要更多的噪声来破坏其信号。假设我们正在处理一个具有 n = H·W 像素的分辨率。现在，考虑一个“常数”图像，即每个像素的值都为 c。前向过程产生 zt = (1 − t)c1 + tϵ，其中 1 和 ϵ ∈ Rn。因此，zt 提供了随机变量 Y = (1 − t)c + tη 的 n 个观测值，其中 c 和 η ∈ R，η 遵循标准正态分布。因此，E(Y ) = (1 − t)c，σ(Y ) = t。我们可以通过 c = 1

1−tE(Y ) 来恢复 c，样本估计值 ˆc = 1

1−t ∑n i=1 zt,i 的误差的标准差为 σ(t, n) = t

1−t √

1

n（因为 Y 的均值的标准误差为 t

√n）。因此，如果已经知道图像 z0 在其像素上是常数，σ(t, n) 表示关于 z0 的不确定性程度。例如，我们立即可以看出，将宽度和高度加倍会导致在任何给定时间 0 < t < 1 时不确定性减半。但是，我们可以通过假设 σ(tn, n) = σ(tm, m) 来将分辨率 n 处的时间步 tn 映射到分辨率 m 处的时间步 tm，从而得到相同程度的不确定性。解出 tm 得到

我们在图 6 中可视化了这个移动函数。注意，假设常数图像并不现实。为了找到在推理期间使用的好值，我们将移动值 α :− m

n 应用于 1024×1024 分辨率的模型的采样步骤，并进行人类偏好研究。图 6 中的结果显示，对于移动值大于 1.5 但小于较高移动值的情况，样本具有强烈的偏好。在我们的后续实验中，我们在训练和采样期间都使用移动值 α = 3.0，分辨率为 1024×1024。在图 6 中可以找到使用和不使用这种移动的样本的定性比较。最后，注意方程（24）意味着对数信噪比移动为 log n

m，类似于（Hoogeboom et al., 2023）：

λtm = 2 log 1 − tn m

n tn（25）

= λtn − 2 log α = λtn − log m

n。（26）

5.3.3. 结果

在图 8 中，我们检查了大规模训练我们的 MM-DiT 的效果。对于图像，我们进行了大规模的扩展研究，并在 2562 像素分辨率上训练了不同参数数量的模型 500k 步，使用预编码的数据（见附录 E.1），批量大小为 4096。我们在 2×2 补丁（Peebles & Xie, 2023）上进行训练，并在 CoCo 数据集（Lin et al., 2014）上每 50k 步报告验证损失。特别是，为了减少验证损失信号中的噪声，我们在 t ∈ (0, 1) 中等距采样损失水平，并分别计算每个水平的验证损失。然后，我们取平均值，除了最后一个（t = 1）水平。

同样，我们在视频领域对我们的 MM-DiT 进行了初步的扩展研究。为此，我们从预训练的图像权重开始，并额外使用 2x 时间补丁。我们遵循 Blattmann et al.（2023b）的方法，通过将时间维度折叠到批次轴中，将数据输入预训练模型。在每个注意力层中，我们在空间注意力操作之后、最终前馈层之前重新排列视觉流的表示，并添加一个完整的时空标记上的注意力操作。我们的视频模型训练了 140k 步，批量大小为 512，在包含 16 帧、分辨率为 2562 像素的视频上进行训练。我们在 Kinetics 数据集（Carreira & Zisserman, 2018）上每 5k 步报告验证损失。注意，我们报告的视频训练的 FLOPs 仅来自视频训练，不包括图像预训练的 FLOPs。

对于图像和视频领域，我们观察到当增加模型大小和训练步数时，验证损失平稳下降。我们发现验证损失与全面评估指标（CompBench（Huang et al., 2023）、GenEval（Ghosh et al., 2023））和人类偏好高度相关。这些结果支持验证损失作为模型性能的简单而通用的衡量标准。我们的结果没有显示出图像或视频模型的饱和迹象。

图 12 描述了训练更大模型更长时间对样本质量的影响。表 5 展示了 GenEval 的完整结果。当应用第 5.3.2 节中介绍的方法并增加训练图像分辨率时，我们最大的模型在大多数类别中表现出色，并优于 DALLE 3（Betker et al., 2023），后者是当前在提示理解方面的最先进的模型。

我们的 d = 38 模型在人类偏好评估中优于当前的专有（Betker et al., 2023; ide, 2024）和开放（Sauer et al., 2023; pla, 2024; Chen et al., 2023; Pernias et al., 2023）的最先进的生成图像模型，在

表 5. GenEval 比较。我们最大的模型（深度 = 38）在 GenEval（Ghosh et al., 2023）上优于所有当前开放模型和 DALLE-3（Betker et al., 2023）。我们用粗体突出显示了最佳、第二好和第三好的条目。有关 DPO，请参见附录 C。

Parti-prompts 基准测试（Yu et al., 2022）中，这些类别包括视觉美学、提示遵循和排版生成，见图 7。为了评估这些类别中的人类偏好，评分者被展示了来自两个模型的成对输出，并被要求回答以下问题：提示遵循：哪幅图像看起来更符合上面显示的文本，并且更忠实地遵循它？视觉美学：鉴于提示，哪幅图像质量更高，且在审美上更令人愉悦？排版：哪幅图像更准确地显示/呈现了上面描述中的文本？更准确的拼写是首选！忽略其他方面。

最后，表 6 突出显示了一个有趣的结果：不仅更大的模型表现更好，而且它们需要更少的步骤就能达到最佳性能。

灵活的文本编码器 尽管使用多个文本编码器的主要动机是提高整体模型性能（Balaji et al., 2022），我们现在展示这种选择还增加了我们在推理期间的灵活性。如附录 B.3 所述，我们使用三个文本编码器训练模型，每个编码器的个体丢弃率为 46.3%。因此，在推理时，我们可以使用所有三个文本编码器的任意子集。这为在模型性能和改进内存效率之间进行权衡提供了手段，这在 4.7B 参数的 T5-XXL（Raffel et al., 2019）需要大量 VRAM 时特别相关。有趣的是，我们发现在使用所有三个文本编码器时，对于大多数提示，性能下降有限。我们提供了一个定性可视化的图 9。只有对于涉及场景高度详细描述或大量文本的复杂提示，我们才发现在使用所有三个文本编码器时有显著的性能提升。这些观察结果也在图 7（Ours w/o T5）中的人类偏好评估结果中得到了验证。移除 T5 对审美质量评分（50% 胜率）没有影响，对提示遵循（46% 胜率）只有很小的影响，而其对生成书面文本的能力贡献更为显著（38% 胜率）。

6. 结论

在这项工作中，我们提出了对文本到图像合成的校正流模型的扩展分析。我们提出了一个新的校正流训练的时间步采样，与以前的扩散训练公式相比，它改进了潜在扩散模型，并保留了校正流在少步采样制度中的有利属性。我们还展示了我们基于变换器的 MM-DiT 架构的优势，该架构考虑了文本到图像任务的多模态性质。最后，我们对这种组合进行了扩展研究，模型大小达到 8B 参数，训练 FLOPs 达到 5×1022。我们展示了验证损失的改进与现有的文本到图像基准以及人类偏好评估相关。这种改进的生成建模和可扩展的多模态架构的结合，实现了与最先进的专有模型相媲美的性能。扩展趋势没有显示出饱和迹象，这使我们对继续改进模型性能充满信心。