Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

放大原图 ↗

1. 引言

Diffusion model 从噪声中生成数据（Song et al., 2020）。它们通过训练来逆转数据走向随机噪声的前向路径；因此，结合神经网络的逼近与泛化能力，它们能够生成训练数据中不存在、但仍服从训练数据分布的新数据点（Sohl-Dickstein et al., 2015; Song & Ermon, 2020）。事实证明，这种生成建模技术对于图像等高维感知数据非常有效（Ho et al., 2020）。近年来，Diffusion model 已成为根据自然语言输入生成 high-resolution 图像和视频的事实标准方法，并展现出令人印象深刻的泛化能力（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）。更重要的是，前向过程的选择也会影响所学习的反向过程，从而影响采样效率。弯曲路径需要许多积分步才能模拟这一过程，而直线路径可以只用一步模拟，并且更不容易累积误差。由于每一步都对应一次神经网络求值，所以这会直接影响采样速度。

前向路径的一种特定选择是所谓的 Rectified Flow（Liu et al., 2022; Albergo & Vanden-Eijnden, 2022; Lipman et al., 2023），它用直线连接数据与噪声。虽然这类模型具有更好的理论性质，但尚未在实践中得到明确确立。到目前为止，一些小型和中型实验已从经验上展示了它的若干优势（Ma et al., 2024），不过这些实验大多局限于 class-conditional 模型。在本工作中，我们通过对 Rectified Flow 模型中的噪声尺度进行重新加权来改变这一局面，这与 noise-predictive diffusion model（Ho et al., 2020）的做法相似。通过一项大规模研究，我们将新表述与现有 diffusion 表述进行比较，并证明了它的优势。

我们表明，Text-to-Image 中广泛使用的做法——把固定的文本表示直接输入模型，例如通过 cross-attention（Vaswani et al., 2017; Rombach et al., 2022）——并不理想；为此，我们提出一种新架构，其中图像 token 与文本 token 都拥有可学习的 stream，从而允许信息在二者之间双向流动。我们将该架构与改进后的 Rectified Flow 表述结合，并研究其 scalability。我们证明，validation loss 呈现可预测的 scaling 趋势，而且更低的 validation loss 与更好的自动评测和人类评测结果高度相关。

我们最大的模型在 prompt 理解的定量评测（Ghosh et al., 2023）和人类偏好评分上，都超过了 SDXL（Podell et al., 2023）、SDXL-Turbo（Sauer et al., 2023）、Pixart-$\alpha$（Chen et al., 2023）等 state-of-the-art 开放模型，以及 DALL-E 3（Betker et al., 2023）等闭源模型。

本工作的核心贡献如下：

我们将公开实验结果、代码和模型权重。

2. 无需模拟的 Flow 训练

我们考虑这样一类 Generative Model：它通过一个 ordinary differential equation（ODE）来定义从噪声分布 $p_1$ 的样本 $x_1$ 到数据分布 $p_0$ 的样本 $x_0$ 的映射：

其中，速度 $v$ 由神经网络的权重 $\Theta$ 参数化。Chen et al.（2018）的先前工作提出使用可微 ODE solver 直接求解 Equation (1)。然而，这一过程的计算代价很高，尤其是在用大型网络架构参数化 $v_\Theta(y_t,t)$ 时。更高效的替代方案，是直接回归一个 Vector Field $u_t$，使其生成一条连接 $p_0$ 与 $p_1$ 的 Probability Path。为了构造这样的 $u_t$，我们把 $p_0$ 与 $p_1=\mathcal{N}(0,1)$ 之间的一条 Probability Path $p_t$ 所对应的前向过程定义为：

当 $a_0=1$、$b_0=0$、$a_1=0$ 且 $b_1=1$ 时，边缘分布为：

这些边缘分布与数据分布和噪声分布一致。

为表达 $z_t$、$x_0$ 与 $\epsilon$ 之间的关系，我们引入 $\psi_t$ 和 $u_t$：

由于 $z_t$ 可以写成 ODE $z_t'=u_t(z_t\mid\epsilon)$ 在初值 $z_0=x_0$ 下的解，因此 $u_t(\cdot\mid\epsilon)$ 会生成 $p_t(\cdot\mid\epsilon)$。

值得注意的是，我们可以利用 Conditional Vector Field $u_t(\cdot\mid\epsilon)$ 构造一个生成 Marginal Probability Path $p_t$ 的 Marginal Vector Field $u_t$（Lipman et al., 2023；见 Appendix B.1）：

如果直接使用 Flow Matching objective 回归 $u_t$：

由于 Equation (6) 中存在边缘化，上述目标无法直接计算；而 Conditional Flow Matching（见 Appendix B.1）为：

使用 Conditional Vector Field $u_t(z\mid\epsilon)$ 后，这就成为一个与原目标等价、但可以实际计算的目标。

为了把 loss 化为显式形式，我们将
$\psi_t'(x_0\mid\epsilon)=a_t'x_0+b_t'\epsilon$
与
$\psi_t^{-1}(z\mid\epsilon)=\frac{z-b_t\epsilon}{a_t}$
代入 Equation (5)：

现在考虑 signal-to-noise ratio（SNR）
$\lambda_t\coloneqq\log\frac{a_t^2}{b_t^2}$。
利用
$\lambda_t'=2\left(\frac{a_t'}{a_t}-\frac{b_t'}{b_t}\right)$，
我们可以把 Equation (9) 改写为：

接下来，我们使用 Equation (10)，把 Equation (8) 重新参数化为一个 noise-prediction objective：

其中我们定义
$\epsilon_\Theta\coloneqq\frac{-2}{\lambda_t'b_t}\left(v_\Theta-\frac{a_t'}{a_t}z\right)$。

需要注意的是，引入随时间变化的权重并不会改变上述目标的最优解。因此，我们可以推导出多种加权 loss：它们都能提供指向目标解的信号，却可能影响优化轨迹。为了在统一框架下分析包括经典 diffusion 表述在内的不同方法，我们可以把目标写成以下形式（沿用 Kingma & Gao, 2023）：

其中，$w_t=-\frac{1}{2}\lambda_t'b_t^2$ 对应于 $\mathcal{L}_{\mathrm{CFM}}$。

3. Flow Trajectories

在本工作中，我们考虑上述形式体系的多种变体，下面对它们作简要介绍。

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

RF 使用 $\mathcal{L}_{\mathrm{CFM}}$；此时它对应的权重为
$w_t^{\mathrm{RF}}=\frac{t}{1-t}$。
网络输出直接参数化速度 $v_\Theta$。

EDM. EDM（Karras et al., 2022）采用如下形式的前向过程：

其中（Kingma & Gao, 2023），
$b_t=\exp F_{\mathcal{N}}^{-1}(t\mid P_m,P_s^2)$；
$F_{\mathcal{N}}^{-1}$ 是均值为 $P_m$、方差为 $P_s^2$ 的正态分布的分位数函数。需要注意的是，这一选择会得到：

网络通过 F-prediction（Kingma & Gao, 2023; Karras et al., 2022）进行参数化，其 loss 可以写为使用下列权重的 $\mathcal{L}_{w_t^{\mathrm{EDM}}}$：

Cosine. Nichol & Dhariwal（2021）提出如下形式的前向过程：

当它与 $\epsilon$-parameterization 和相应 loss 结合时，对应的权重为
$w_t=\operatorname{sech}(\lambda_t/2)$。
当它与 v-prediction loss（Kingma & Gao, 2023）结合时，权重为
$w_t=e^{-\lambda_t/2}$。

(LDM-)Linear. LDM（Rombach et al., 2022）使用 DDPM schedule（Ho et al., 2020）的一种修改形式。二者都是 variance-preserving schedule，即
$b_t=\sqrt{1-a_t^2}$，
并根据 diffusion coefficient $\beta_t$，为离散 timestep $t=0,\ldots,T-1$ 定义
$a_t=\left(\prod_{s=0}^{t}(1-\beta_s)\right)^{1/2}$。
给定边界值 $\beta_0$ 与 $\beta_{T-1}$ 后，DDPM 使用
$\beta_t=\beta_0+\frac{t}{T-1}(\beta_{T-1}-\beta_0)$，
而 LDM 使用
$\beta_t=\left(\sqrt{\beta_0}+\frac{t}{T-1}(\sqrt{\beta_{T-1}}-\sqrt{\beta_0})\right)^2$。

3.1. 面向 RF 模型定制的 SNR sampler

RF loss 在 $[0,1]$ 的所有 timestep 上均匀训练速度 $v_\Theta$。然而，直观来看，由此得到的速度预测目标 $\epsilon-x_0$ 在 $[0,1]$ 的中部更难预测：因为当 $t=0$ 时，最优预测是 $p_1$ 的均值；而当 $t=1$ 时，最优预测是 $p_0$ 的均值。

一般而言，把 $t$ 上常用的均匀分布 $\mathcal{U}(t)$ 改为密度为 $\pi(t)$ 的分布，等价于使用权重如下的加权 loss $\mathcal{L}_{w_t^\pi}$：

因此，我们希望通过更频繁地采样中间 timestep，为它们赋予更大的权重。下面介绍我们用于训练模型的 timestep 密度 $\pi(t)$。

Logit-Normal 采样。 一种能给中间 timestep 更高权重的分布，是 logit-normal 分布（Atchison & Shen, 1980）。它的密度为：

其中
$\operatorname{logit}(t)=\log\frac{t}{1-t}$；
$m$ 是位置参数，$s$ 是尺度参数。位置参数使我们能够让训练 timestep 偏向数据 $p_0$（负的 $m$）或噪声 $p_1$（正的 $m$）。如 Figure 11 所示，尺度参数控制分布的宽度。

在实践中，我们从正态分布 $u\sim\mathcal{N}(u;m,s)$ 中采样随机变量 $u$，再通过标准 logistic 函数进行映射。

具有重尾的 Mode Sampling。 Logit-normal 密度在端点 $0$ 和 $1$ 处总是趋于零。为了研究这是否会对性能造成不利影响，我们还采用一种在 $[0,1]$ 上密度严格为正的 timestep 采样分布。对于尺度参数 $s$，定义：

当
$-1\le s\le\frac{2}{\pi-2}$
时，该函数单调，因此我们可以用它从隐含密度
$\pi_{\mathrm{mode}}(t;s)=\left|\frac{d}{dt}f_{\mathrm{mode}}^{-1}(t)\right|$
中采样。如 Figure 11 所示，尺度参数控制采样偏向中点（正的 $s$）还是端点（负的 $s$）的程度。

该表述也包含 uniform weighting：当 $s=0$ 时，
$\pi_{\mathrm{mode}}(t;s=0)=\mathcal{U}(t)$。
这一设置已被此前关于 Rectified Flow 的工作广泛采用（Liu et al., 2022; Ma et al., 2024）。

CosMap. 最后，我们还在 RF 设置下考虑 Section 3 中的 cosine schedule（Nichol & Dhariwal, 2021）。具体而言，我们希望找到一个映射
$f:u\mapsto f(u)=t$、$u\in[0,1]$，
使其 log-SNR 与 cosine schedule 相匹配：
$2\log\frac{\cos(\frac{\pi}{2}u)}{\sin(\frac{\pi}{2}u)} =2\log\frac{1-f(u)}{f(u)}$。
对 $f$ 求解后，当 $u\sim\mathcal{U}(u)$ 时可得：

由此得到密度：

4. Text-to-Image Architecture

为了根据文本条件采样图像，我们的模型必须同时考虑文本和图像两种 modality。我们使用 pretrained model 获得合适的表示，随后描述 diffusion backbone 的架构。Figure 2 给出了整体概览。

我们的总体设置沿用 LDM（Rombach et al., 2022），在 pretrained autoencoder 的 latent space 中训练 Text-to-Image 模型。与把图像编码成 latent representation 类似，我们也沿用先前方法（Saharia et al., 2022b; Balaji et al., 2022），使用预训练且冻结的文本模型对文本条件 $c$ 进行编码。更多细节见 Appendix B.2。

multimodal Diffusion Backbone。 我们的架构建立在 DiT（Peebles & Xie, 2023）架构之上。DiT 只考虑 class-conditional 图像生成，并使用 modulation 机制，以 diffusion process 的 timestep 和类别标签共同作为网络条件。类似地，我们把 timestep $t$ 的 embedding 和 $c_{\mathrm{vec}}$ 作为 modulation 机制的输入。然而，pooled text representation 只保留了文本输入的粗粒度信息（Podell et al., 2023），因此网络还需要 sequence representation $c_{\mathrm{ctxt}}$ 中的信息。

我们构造一个由文本输入和图像输入的 embedding 组成的序列。具体而言，我们加入 positional encoding，并把 latent pixel representation
$x\in\mathbb{R}^{h\times w\times c}$
的 $2\times2$ patch 展平为长度
$\frac{1}{2}h\cdot\frac{1}{2}w$
的 patch encoding sequence。随后，我们把该 patch encoding 与文本编码 $c_{\mathrm{ctxt}}$ embedding 到相同维度。

放大原图 ↗

接着，我们拼接这两个序列。随后沿用 DiT，依次应用经 modulation 的 attention 和 MLP。

由于文本 embedding 与图像 embedding 在概念上差异很大，我们为两种 modality 分别使用两套独立权重。如 Figure 2b 所示，这等价于为每种 modality 各自设置一个独立的 Transformer，但在 attention 操作时把两种 modality 的序列合并起来；这样，两种表示既可以在各自的空间中工作，又能把另一种表示纳入考虑。

在 scaling 实验中，我们用模型深度 $d$（即 attention block 的数量）来参数化模型规模：将 hidden size 设为 $64d$，在 MLP block 中扩展为 $4\cdot64d$ 个 channel，并把 attention head 的数量设为 $d$。

5. 实验

5.1 改进 Rectified Flow

我们的目标是弄清楚：在式 (1) 所描述的 normalizing flow 无模拟训练方法中，哪一种最高效。为了能够比较不同方法，我们控制 optimization algorithm、模型架构、dataset 和 sampler 等变量。此外，不同方法的 loss 彼此不可直接比较，而且 loss 也未必与输出样本质量相关；因此，我们需要能够在方法之间进行比较的评价指标。

我们在 ImageNet (Russakovsky et al., 2014) 和 CC12M (Changpinyo et al., 2021) 上训练模型，并在训练过程中，使用 validation loss、CLIP score (Radford et al., 2021; Hessel et al., 2021) 和 FID (Heusel et al., 2017)，在不同 sampler 设置（不同 guidance scale 与 sampling step 数）下评估模型的训练权重和 EMA 权重。我们按照 Sauer et al. (2021) 的方法在 CLIP feature 上计算 FID。所有指标均在 COCO-2014 validation split (Lin et al., 2014) 上评估。训练与采样 hyperparameter 的完整细节见 Appendix B.3。

5.1.1 结果

我们在两个 dataset 上分别训练了 61 种不同的 formulation。我们纳入了 Section 3 中的以下 variant：

对每次运行，我们先依据 EMA 权重下的 evaluation 选择 validation loss 最小的 step，然后收集在 6 种不同 sampler 设置下、分别使用和不使用 EMA 权重时得到的 CLIP score 与 FID。

针对 sampler 设置、EMA 权重与 dataset 选择构成的全部 24 种组合，我们使用 non-dominated sorting algorithm 对不同 formulation 进行排名。具体而言，我们反复找出依据 CLIP 和 FID score 达到 Pareto optimal 的 variant，将当前迭代索引赋给这些 variant，移除它们，再对剩余 variant 继续这一过程，直到所有 variant 都获得排名。最后，我们在 24 种不同的控制设置上对这些 rank 取平均。

结果见 Table 1。对于那些使用不同 hyperparameter 评估过的 variant，我们只展示其中表现最好的两个。我们还给出了只在 5-step sampler 设置和 50-step sampler 设置上取平均时的 rank。

我们观察到，rf/lognorm(0.00, 1.00) 始终取得较好的 rank。它优于采用 uniform timestep sampling 的 Rectified Flow formulation（rf），从而验证了我们的假设：中间 timestep 更重要。在全部 variant 中，只有采用修改后 timestep sampling 的 Rectified Flow formulation 优于此前使用的 LDM-Linear (Rombach et al., 2022) formulation（eps/linear）。

我们还观察到，一些 variant 在某些设置下表现很好，但在另一些设置下更差。例如，rf/lognorm(0.50, 0.60) 在 50 个 sampling step 时表现最佳，但在 5 个 sampling step 时要差得多（平均 rank 为 8.5）。对于 Table 2 中的两个指标，我们也观察到类似现象。第一组展示若干代表性 variant 在两个 dataset 上、采用 25 个 sampling step 时的指标；下一组展示取得最佳 CLIP 和 FID score 的 variant。除 rf/mode(1.75) 外，这些 variant 通常在某一个指标上非常好，但在另一个指标上相对较差。相比之下，我们再次观察到 rf/lognorm(0.00, 1.00) 在不同指标和 dataset 上都表现良好：四项 score 中有两项排名第三，一项排名第二。

最后，我们在 Figure 3 中展示不同 formulation 的定性行为，并用不同颜色区分 formulation 分组：edm、rf、eps 和 v。Rectified Flow formulation 通常表现良好；与其他 formulation 相比，当 sampling step 数减少时，其性能下降得更少。

5.2 改进 modality-specific representation

上一节找到的 formulation 使 Rectified Flow 模型不仅能与 LDM-Linear (Rombach et al., 2022) 或 EDM (Karras et al., 2022) 等成熟 diffusion formulation 竞争，甚至还能胜过它们。接下来，我们将这一 formulation 应用于 high-resolution text-to-image synthesis。因此，我们算法的最终性能不仅取决于训练 formulation，还取决于 neural network 的 parameterization，以及所使用的图像与文本表示的质量。在以下各节中，我们会说明如何改进所有这些组件，然后在 Section 5.3 中扩展最终方法的规模。

放大原图 ↗

5.2.1 改进 Autoencoder

Latent diffusion model 通过在 pretrained autoencoder 的 latent space 中运行而获得较高效率 (Rombach et al., 2022)。该 autoencoder 将输入 RGB 图像 $X\in\mathbb{R}^{H\times W\times3}$ 映射到更低维的空间 $x=E(X)\in\mathbb{R}^{h\times w\times d}$。Autoencoder 的重建质量，为 latent diffusion 训练完成后能够达到的图像质量设定了上限。与 Dai et al. (2023) 类似，我们发现增加 latent channel 数 $d$ 能显著提升重建性能，见 Table 3。直观地说，预测具有更大 $d$ 的 latent 是一项更困难的任务，因此 capacity 更大的模型应能在更大的 $d$ 下表现得更好，最终获得更高的图像质量。我们在 Figure 10 中验证了这一假设：在 sample FID 方面，$d=16$ 的 autoencoder 呈现出更好的 scaling 性能。因此，本文其余部分均选择 $d=16$。

5.2.2 改进 Caption

Betker et al. (2023) 证明，合成生成的 caption 能显著改进 Training at Scale 的 text-to-image model。这是因为，大规模图像 dataset 附带的人工 caption 往往过于简单：它们过度关注图像主体，通常会省略背景、场景构图，以及在适用时对图中所显示文字的描述 (Betker et al., 2023)。我们沿用其方法，使用现成的 state-of-the-art vision-language model CogVLM (Wang et al., 2023) 为大规模图像 dataset 创建合成标注。由于合成 caption 可能使 text-to-image model 遗忘 VLM 知识语料中不存在的某些概念，我们采用 50% 原始 caption 与 50% 合成 caption 的比例。

为了评估使用这种 caption mixture 训练的效果，我们训练了两个 $d=15$ 的 MM-DiT 模型，共训练 250k step：一个只使用原始 caption，另一个使用 50/50 mixture。我们在 Table 4 中使用 GenEval benchmark (Ghosh et al., 2023) 评估训练后的模型。结果表明，加入合成 caption 进行训练的模型明显优于只使用原始 caption 的模型。因此，本文其余部分均使用 50/50 的合成／原始 caption mixture。

5.2.3 改进 Text-to-Image Backbone

本节比较现有的 Transformer-based diffusion backbone 与我们在 Section 4 中引入的新型 multimodal Transformer-based diffusion backbone——MM-DiT。MM-DiT 专门用不同的（两套）可训练模型权重处理不同 domain，这里即文本 token 与图像 token。更具体地说，我们沿用 Section 5.1 的实验设置，在 CC12M 上比较 DiT、CrossDiT（即不采用 sequence-wise concatenation，而是让 DiT 对文本 token 做 cross-attention (Chen et al., 2023)）与 MM-DiT 的 text-to-image 性能。对于 MM-DiT，我们比较使用两套权重和三套权重的模型；后者分别处理 CLIP (Radford et al., 2021) 与 T5 (Raffel et al., 2019) token（参见 Section 4）。注意，DiT（如 Section 4 所述，将文本 token 与图像 token 连接起来）可以看作 MM-DiT 的一个特例：所有 modality 共享同一套权重。最后，我们还考虑 UViT (Hoogeboom et al., 2023) 架构，它是广泛使用的 UNet 与 Transformer variant 之间的一种 hybrid。

放大原图 ↗

我们在 Figure 4 中分析这些架构的收敛行为：Vanilla DiT 的表现不如 UViT。Cross-attention DiT variant CrossDiT 优于 UViT，不过 UViT 在初始阶段似乎学得快得多。我们的 MM-DiT variant 显著优于 cross-attention variant 与 vanilla variant。使用三套参数而不是两套参数只带来很小的增益，却会增加参数量与 VRAM 使用量，因此本文其余部分选择前一种方案，也就是两套参数。

放大原图 ↗

5.3 Training at Scale

在扩大规模之前，我们先过滤并预编码数据，以保证预训练的安全性与效率。至此，前文对 diffusion formulation、架构与数据的全部考量汇聚到最后一节：我们把模型扩展到 8B 参数。

5.3.1 Data Preprocessing

预训练缓解措施。 训练数据会显著影响 Generative Model 的能力。因此，数据过滤是限制不良能力的有效手段 (Nichol, 2022)。在开展 Training at Scale 之前，我们针对以下类别过滤数据：

1. 性内容： 使用 NSFW-detection model 过滤露骨内容。
2. 美学质量： 移除评分系统预测为低分的图像。
3. 复现训练样本（regurgitation）： 使用基于 cluster 的 deduplication 方法，移除训练数据中在感知或语义层面重复的样本，见 Appendix E.2。

预计算图像与文本 Embedding。 我们的模型将多个预训练且冻结的 network 的输出作为输入，包括 autoencoder latent 与 text encoder representation。由于这些输出在训练期间保持不变，我们会针对整个 dataset 一次性预计算它们。Appendix E.1 对这一做法给出了详细讨论。

5.3.2 在 high-resolution 上 Finetune

放大原图 ↗

QK-Normalization。 一般而言，我们先在尺寸为 $256^2$ pixel 的低分辨率图像上预训练所有模型，随后再用具有混合宽高比的更 high-resolution 数据 finetune（细节见下一段）。我们发现，转向 high-resolution 时，mixed-precision training 可能变得不稳定并导致 loss 发散。切换到 full-precision training 可以解决这一问题，但与 mixed-precision training 相比会带来约 $2\times$ 的性能下降。Discriminative ViT 文献给出了一种更高效的替代方案：Dehghani et al. (2023) 观察到，大型 vision Transformer 的训练会因 attention entropy 不受控制地增长而发散。为避免这一点，他们建议在 attention operation 之前归一化 Q 与 K。我们采用该方法，在 MMDiT 架构的两个 stream 中都使用带可学习 scale 的 RMSNorm (Zhang & Sennrich, 2019)，见 Figure 2。Figure 5 表明，额外的归一化能防止 attention logit 增长不稳定，这与 Dehghani et al. (2023) 和 Wortsman et al. (2023) 的发现一致；与 AdamW optimizer (Loshchilov & Hutter, 2017) 中的 $\epsilon=10^{-15}$ 结合后，它还能实现高效的 bf16-mixed precision training (Chen et al., 2019)。该技术也可以用于预训练阶段没有采用 qk-normalization 的模型：模型能很快适应额外的 normalization layer，并以更稳定的方式训练。最后需要指出，尽管此方法通常有助于稳定大型模型训练，但它并不是一种普适方案，仍可能需要依据具体训练设置作调整。

适用于不同宽高比的 Positional Encoding。 在固定 $256\times256$ 分辨率上训练后，我们希望：(i) 提 high-resolution；(ii) 支持具有灵活宽高比的 inference。由于我们使用二维 positional frequency embedding，因此必须根据分辨率调整这些 embedding。在多宽高比设置中，若像 Dosovitskiy et al. (2020) 那样直接插值 embedding，就不能正确反映边长。为此，我们组合使用扩展后的 position grid 与插值后的 position grid，再对其进行 frequency embedding。

放大原图 ↗

放大原图 ↗

Figure 6 展示了这一 shifting function。需要注意，“常数图像”假设并不现实。为寻找 inference 时合适的 shift 值 $\alpha\coloneq\sqrt{m/n}$，我们把不同 shift 应用于一个在 $1024\times1024$ 分辨率上训练的模型的 sampling step，并开展人类偏好研究。Figure 6 的结果表明，样本在 shift 大于 $1.5$ 时受到明显偏好，但更大的各个 shift 值之间差异较小。因此，在后续实验中，无论训练还是在 $1024\times1024$ 分辨率采样，我们都采用 $\alpha=3.0$。Figure 6 还给出了训练 8k step 后、有无这种 shift 的样本定性对比。最后，Equation (23) 意味着一个与 Hoogeboom et al. (2023) 类似的 $\log\frac{n}{m}$ log-SNR shift：

完成分辨率 $1024\times1024$ 下的 shifted training 后，我们再按照 Appendix C 的方法，使用 Direct Preference Optimization（DPO）对模型进行对齐。

5.3.3 结果

在 Figure 8 中，我们考察 Training at Scale MM-DiT 的效果。对于图像，我们开展一项大规模 scaling study：使用预编码数据（参见 Appendix E.1），以 batch size 4096，在 $256^2$ pixel 分辨率上训练不同参数规模的模型 500k step。我们使用 $2\times2$ patch (Peebles & Xie, 2023) 训练，并每 50k step 在 COCO dataset (Lin et al., 2014) 上报告 validation loss。具体而言，为降低 validation loss 信号中的噪声，我们在 $t\in(0,1)$ 上等距采样 loss level，分别计算每个 level 的 validation loss，然后对除最后一个 level（$t=1$）之外的所有 level 取平均。

类似地，我们对 MM-DiT 在视频上的 scaling 做了一项初步研究。为此，我们从预训练图像权重出发，并额外采用 2 倍 temporal patching。我们沿用 Blattmann et al. (2023b) 的方法：把 temporal axis 折叠到 batch axis，将数据输入 pretrained model。在每个 attention layer 中，我们重新排列视觉 stream 中的表示，并在 spatial attention operation 之后、最终 feedforward layer 之前，增加一个覆盖全部 spatio-temporal token 的 full attention。视频模型训练 140k step，batch size 为 512，训练视频包含 16 帧、每帧 $256^2$ pixel。我们每 5k step 在 Kinetics dataset (Carreira & Zisserman, 2018) 上报告 validation loss。需要注意，Figure 8 中报告的视频训练 FLOPs 只包括视频训练本身，不包括图像预训练的 FLOPs。

无论图像还是视频 domain，随着模型规模与 training step 增加，我们都观察到 validation loss 平滑下降。我们发现，validation loss 与综合 evaluation metric（CompBench (Huang et al., 2023)、GenEval (Ghosh et al., 2023)）及人类偏好高度相关。这些结果支持把 validation loss 作为一种简单且通用的模型性能度量。无论图像模型还是视频模型，我们的结果都没有显示出饱和迹象。

Figure 12 展示了延长大型模型训练时间如何影响样本质量。Table 5 完整给出 GenEval 结果。应用 Section 5.3.2 中的方法并提高训练图像分辨率后，我们最大的模型在大多数类别上表现出色，并在总分上超过当前 prompt comprehension state of the art——DALL-E 3 (Betker et al., 2023)。

我们深度为 $d=38$ 的模型，在 PartiPrompts benchmark (Yu et al., 2022) 上针对视觉美学、prompt following 和文字排版生成三类任务进行的人类偏好评估中，优于当前闭源 (Betker et al., 2023; Ideogram, 2024) 与开源 (Sauer et al., 2023; Playground, 2024; Chen et al., 2023; Pernias et al., 2023) 的 SOTA 生成图像模型，见 Figure 7。评估这些类别的人类偏好时，评分者会看到两个模型的成对输出，并回答以下问题：

- Prompt following： 哪幅图像更能体现上方所示文本，并且更忠实地遵循该文本？
- 视觉美学： 在给定 prompt 的情况下，哪幅图像质量更高、在美学上更悦目？
- 文字排版： 哪幅图像能更准确地显示／呈现上方描述中指定的文字？拼写越准确越好！忽略其他方面。

最后，Table 6 突出显示了一项很有意思的结果：更大的模型不仅表现更好，而且达到峰值性能所需的 step 也更少。

灵活的 Text Encoder。 使用多个 text encoder 的主要动机，是提升模型整体性能 (Balaji et al., 2022)；这里我们进一步说明，这一选择还能提升基于 MM-DiT 的 Rectified Flow 在 inference 时的灵活性。如 Appendix B.3 所述，我们使用三个 text encoder 训练模型，每个 encoder 的独立 dropout rate 为 46.3%。因此，在 inference 时，我们可以任意选择三个 text encoder 的一个子集。这使得我们能够用模型性能换取更高的 memory efficiency；对于拥有 4.7B 参数、需要大量 VRAM 的 T5-XXL (Raffel et al., 2019)，这一点尤其重要。

放大原图 ↗

放大原图 ↗

有意思的是，只使用两个基于 CLIP 的 text encoder 处理 text prompt，并将 T5 embedding 替换为零时，性能下降有限。Figure 9 给出了定性可视化。只有在复杂 prompt 涉及极为详细的场景描述或大量待写文字时，我们才发现同时使用三个 text encoder 会带来显著的性能增益。Figure 7 的人类偏好评估结果（Ours w/o T5）也验证了这些观察。移除 T5 不影响美学质量评分（win rate 为 50%），对 prompt adherence 只有很小影响（win rate 为 46%），但它对生成书面文字能力的贡献更显著（win rate 为 38%）。

6. 结论

本文对用于 text-to-image synthesis 的 Rectified Flow 模型进行了 scaling analysis。我们为 Rectified Flow 训练提出了一种新的 timestep sampling 方法；对于 latent diffusion model，它优于以往的 diffusion training formulation，同时保留了 Rectified Flow 在 few-step sampling 场景中的有利性质。我们还展示了基于 Transformer 的 MM-DiT 架构的优势，该架构显式考虑了 text-to-image 任务的 multimodal 特性。最后，我们把这一组合扩展到 8B 参数和 $5\times10^{22}$ training FLOPs，并证明 validation loss 的改善既与现有 text-to-image benchmark 相关，也与人类偏好评估相关。凭借这些 generative modeling 改进和可扩展的 multimodal 架构，我们取得了能够与 state-of-the-art 闭源模型竞争的性能。Scaling 趋势没有显示出饱和迹象，这使我们有理由乐观地认为，未来仍能继续提升模型性能。

更广泛影响

本文工作的目标是推动整个 machine learning 领域，尤其是 image synthesis 领域的发展。我们的工作可能产生许多社会影响，但我们认为其中没有哪一项必须在此特别强调。关于 diffusion model 一般性影响的广泛讨论，我们建议感兴趣的读者参阅 (Po et al., 2023)。

补充材料

A. 背景

Diffusion Models。 Diffusion Models（Sohl-Dickstein et al., 2015；Song et al., 2020；Ho et al., 2020）通过逼近某个随机前向过程的逆 ODE 来生成数据；该前向过程会把数据逐渐变成噪声。它们已经成为图像（Dhariwal & Nichol, 2021；Ramesh et al., 2022；Saharia et al., 2022b；Rombach et al., 2022；Balaji et al., 2022）与视频（Singer et al., 2022；Ho et al., 2022；Esser et al., 2023；Blattmann et al., 2023b；Gupta et al., 2023）生成建模的标准方法。由于这些模型既可以从负对数似然的变分下界（Sohl-Dickstein et al., 2015）推导，也可以从 score matching（Hyvärinen, 2005；Vincent, 2011；Song & Ermon, 2020）推导，因此，前向与逆向过程的不同表述（Song et al., 2020；Dockhorn et al., 2021）、不同模型参数化（Ho et al., 2020；Ho & Salimans, 2022；Karras et al., 2022）、不同 loss weighting（Ho et al., 2020；Karras et al., 2022）以及不同 ODE solver（Song et al., 2022；Lu et al., 2023；Dockhorn et al., 2022），共同产生了大量不同的训练目标与采样流程。近期，Kingma & Gao（2023）和 Karras et al.（2022）的奠基性工作提出了统一表述，并为训练（Karras et al., 2022；Kingma & Gao, 2023）与推理（Karras et al., 2022）引入了新的理论与实践洞见。然而，即使经过这些改进，常见 ODE 的轨迹仍可能包含相当显著的曲率（Karras et al., 2022；Liu et al., 2022），这要求更多 solver steps，因而使快速推理变得困难。为解决这一问题，我们采用 Rectified Flow 模型，其表述允许学习笔直的 ODE 轨迹。

Rectified Flow Models。 Liu et al.（2022）、Albergo & Vanden-Eijnden（2022）以及 Lipman et al.（2023）通过 ordinary differential equation（ODE）在两个分布之间构造 transport map，以此实现生成建模。该方法与 continuous normalizing flow（CNF；Chen et al., 2018）以及 Diffusion Models 都有紧密联系。相较于 CNF，Rectified Flow 和 Stochastic Interpolant 的优势在于训练期间无须模拟 ODE。相较于 Diffusion Models，它们所得 ODE 的模拟速度可以快于与 Diffusion Models 相关的 probability flow ODE（Song et al., 2020）。不过，它们并不会直接得到 optimal transport 解，已有多项工作试图进一步降低轨迹曲率（Lee et al., 2023；Tong et al., 2023；Pooladian et al., 2023）。Dao et al.（2023）与 Ma et al.（2024）证明了 Rectified Flow 表述用于 class-conditional 图像合成的可行性；Fischer et al.（2023）将其用于 latent-space 上采样；Liu et al.（2023）则把 Liu et al.（2022）的 reflow 流程用于蒸馏预训练 text-to-image 模型（Rombach et al., 2022）。本文关注的是：将 Rectified Flow 作为一种只需较少采样步数的 text-to-image 合成基础。我们广泛比较了不同表述与 loss weighting，并提出一种性能更好的新 timestep schedule 来训练 Rectified Flow。

Scaling Diffusion Models。 Transformer 架构（Vaswani et al., 2017）在 NLP（Kaplan et al., 2020）和计算机视觉任务（Dosovitskiy et al., 2020；Zhai et al., 2022）中的 scaling 性质早已广为人知。对于 Diffusion Models，U-Net 架构（Ronneberger et al., 2015）一直占据主导地位（Ho et al., 2020；Rombach et al., 2022；Balaji et al., 2022）。尽管近期已有一些工作探索 diffusion Transformer backbone（Peebles & Xie, 2023；Chen et al., 2023；Ma et al., 2024），text-to-image Diffusion Models 的 scaling laws 仍未得到研究。

放大原图 ↗

放大原图 ↗

B. 关于 Flow Matching

B.1 Flow 的免仿真训练细节

沿用 Lipman et al.（2023）的论证。为了说明 $u_t(z)$ 会生成 $p_t$，我们注意到 continuity equation 给出了如下充要条件（Villani, 2008）：

因此，只需证明：

在从式 (28) 到式 (29) 的过程中，我们对 $u_t(z\mid\epsilon)$ 使用了式 (26) 的 continuity equation，因为 $u_t(z\mid\epsilon)$ 会生成 $p_t(z\mid\epsilon)$；在式 (27) 中，我们使用了式 (6) 的定义。

目标函数 $\mathcal{L}_{\mathrm{FM}}\leftrightharpoons\mathcal{L}_{\mathrm{CFM}}$ 的等价性（Lipman et al., 2023）可由下式得到：

其中 $c$ 与 $c'$ 都不依赖于 $\Theta$；从式 (31) 到式 (32) 的关系来自：

在式 (35) 中，我们乘上了 $p_t(z)/p_t(z)$；从式 (35) 到式 (36) 则使用了式 (6) 的定义。

B.2 图像与文本表示的细节

Latent 图像表示。 我们沿用 LDM（Rombach et al., 2022），使用 pretrained autoencoder，把 RGB 图像 $X\in\mathbb{R}^{H\times W\times3}$ 表示到更小的 latent space 中：$x=E(X)\in\mathbb{R}^{h\times w\times d}$。我们使用 $8$ 倍空间下采样，因此 $h=H/8$、$w=W/8$，并在第 5.2.1 节实验不同的 $d$。我们始终在 latent space 中应用式 (2) 的前向过程；当通过式 (1) 采样得到表示 $x$ 时，再使用 decoder $D$ 将其解码回 pixel space：$X=D(x)$。按照 Rombach et al.（2022），我们用 latent 的均值和标准差对其归一化；这些统计量在训练数据的一个子集上全局计算得到。Figure 10 展示了不同 $d$ 下的 Generative Model 训练如何随模型容量变化，相关讨论见第 5.2.1 节。

文本表示。 类似于把图像编码成 latent 表示，我们也沿用既有方法（Saharia et al., 2022b；Balaji et al., 2022），使用预训练且冻结的文本模型编码文本条件 $c$。具体来说，所有实验都使用 CLIP 模型（Radford et al., 2021）与一个 encoder-decoder 文本模型的组合。我们分别用 Radford et al.（2021）的 CLIP L/14 和 Cherti et al.（2023）的 OpenCLIP bigG/14 的 text encoder 编码 $c$。把二者维度分别为 $768$ 与 $1280$ 的 pooled outputs 拼接起来，得到 vector conditioning $c_{\text{vec}}\in\mathbb{R}^{2048}$。我们还沿 channel 维拼接倒数第二层的 hidden representations，得到 CLIP context conditioning $c_{\text{ctxt}}^{\text{CLIP}}\in\mathbb{R}^{77\times2048}$。接着，我们还用 T5-v1.1-XXL 模型（Raffel et al., 2019）的 encoder 编码 $c$，取其最终 hidden representation：$c_{\text{ctxt}}^{\text{T5}}\in\mathbb{R}^{77\times4096}$。最后，沿 channel 轴把 $c_{\text{ctxt}}^{\text{CLIP}}$ 用零填充到 $4096$ 维，以匹配 T5 表示，再沿 sequence 轴与 $c_{\text{ctxt}}^{\text{T5}}$ 拼接，得到最终 context representation $c_{\text{ctxt}}\in\mathbb{R}^{154\times4096}$。这两种 caption 表示 $c_{\text{vec}}$ 与 $c_{\text{ctxt}}$ 的使用方式不同，详见第 4 节。

放大原图 ↗

B.3 第 5.1 节实验的预备设置

Datasets。 由于缺少标准的 text-to-image benchmark，我们使用两个 datasets。ImageNet（Russakovsky et al., 2014）被广泛使用；我们为其中图像添加形如 “a photo of a <class name>” 的 caption，把它转成适用于 text-to-image 模型的数据集，其中 <class name> 从该图像 class label 所提供的名称中随机选择。作为更贴近真实 text-to-image 场景的数据集，我们使用 CC12M（Changpinyo et al., 2021）训练。

Optimization。 在该实验中，所有模型均使用 global batch size $1024$，采用 AdamW optimizer（Loshchilov & Hutter, 2017），learning rate 为 $10^{-4}$，并进行 $1000$ 个 linear warmup steps。我们使用 mixed-precision training，并保留一份模型权重副本；每经过 $100$ 个 training batches，就以 decay factor $0.99$ 的 exponential moving average（EMA）更新这份副本。对于 unconditional diffusion guidance（Ho & Salimans, 2022），我们分别以 $46.4\%$ 的概率把三个 text encoders 中每一个的输出置零，使所有训练 steps 中约有 $10\%$ 在训练 unconditional model。

Evaluation。 如第 5.1 节所述，我们使用 CLIP scores、FID 与 validation losses，在 COCO-2014 validation split（Lin et al., 2014）上定期评估模型。由于不同 timesteps 的 loss 值在量级与方差上差异很大，我们在时间区间 $[0,1]$ 上取八个等间距的值进行分层评估。

为了分析不同方法在不同 sampler settings 下的表现，我们为每个 sampler 生成 $1000$ 个样本；各 sampler 的 guidance scale 与 sampling steps 数量不同。我们使用 CLIP L/14（Radford et al., 2021）计算这些样本的 CLIP scores，并计算这些样本与 validation set 图像之间、基于 CLIP L/14 image features 的 FID。采样时，我们始终使用式 (1) 的 Euler discretization（Euler, 1768），共使用六种设置：$50$ steps，classifier-free-guidance scales 分别为 $1.0$、$2.5$、$5.0$；以及 $5$、$10$、$25$ steps，classifier-free-guidance scale 固定为 $5.0$。

B.4 改进 Rectified Flow 模型的 SNR Samplers

如第 2 节所述，我们为训练 Rectified Flow 模型所用的 timesteps 引入新的密度 $\pi(t)$。Figure 11 可视化了第 3.1 节提出的 logit-normal sampler 与 mode sampler 的分布。值得注意的是，如第 5.1 节所示，logit-normal sampler 优于经典的均匀 Rectified Flow 表述（Liu et al., 2022），也优于 EDM（Karras et al., 2022）与 LDM-Linear（Rombach et al., 2022）等已有 diffusion baselines。

放大原图 ↗

图中四个 PartiPrompts，按从左到右的顺序：

1. “一只穿着正装、戴礼帽、手持拐杖的浣熊。浣熊拿着一个垃圾袋。抽象立体主义风格的油画。”
2. “一碗看起来像由橡皮泥做成的怪物的汤。”
3. “两杯咖啡，其中一杯的 latte art 是心形；另一杯的 latte art 是星星。”
4. “一只微笑的树懒穿着皮夹克、戴牛仔帽、穿苏格兰短裙并系蝴蝶结。树懒拿着一根 quarterstaff 和一本大书。它站在草地上，前方几英尺处是一辆画满鲜花、闪闪发亮的大众面包车。使用从下方向上拍摄的广角镜头。”

放大原图 ↗

C. Direct Preference Optimization

列标题 prompt 保留原文，并可附中文：

- “a peaceful lakeside landscape with migrating herd of sauropods”（宁静的湖畔景观，一群迁徙的蜥脚类恐龙）
- “a book with the words ‘Don't Panic!’, written on it”（一本封面写着 “Don't Panic!” 的书）

行标签：2B base、2B w/ DPO、8B base、8B w/ DPO。

放大原图 ↗

Direct Preference Optimization（DPO；Rafailov et al., 2023）是一种使用偏好数据对 LLM 进行 finetune 的技术。最近，该方法已被改造用于 text-to-image Diffusion Models 的偏好 finetuning（Wallace et al., 2023）。本节验证我们的模型同样适用于 preference optimization。具体来说，我们把 Wallace et al.（2023）提出的方法应用到参数量为 2B 和 8B 的 base model。按照常见做法，我们并不 finetune 整个模型，而是为所有 linear layers 引入可学习的 Low-Rank Adaptation（LoRA）矩阵，其 rank 为 $128$。对于 2B 和 8B base model，我们分别对这些新增参数 finetune $4\text{k}$ 和 $2\text{k}$ iterations。随后，我们从 PartiPrompts 集合（Yu et al., 2022）中选取 $128$ 条 captions 的子集进行 human preference study，每个 prompt 与 comparison 大约由三名投票者评估。Figure 14 表明，base models 可以有效地针对 human preference 调整；Figure 13 则给出了各 base model 与 DPO-finetuned model 的样本。

D. 针对 instruction-based image editing 的 Finetuning

训练 instruction-based image editing 与一般 image-to-image Diffusion Models 的常见做法是：在将输入送入 U-Net 之前，沿 channel 维把输入图像的 latents 与 diffusion target 的 noised latents 拼接起来（Brooks et al., 2023；Sheynin et al., 2023；Saharia et al., 2022a,c）。我们采用同样的方法，在 patching 之前沿 channels 拼接输入与 target，并证明该方法同样适用于本文提出的架构。我们在一个由多种 image-to-image editing tasks 构成的数据集上 finetune 2B base model：其中既包含分布类似 InstructPix2Pix dataset（Brooks et al., 2023）的任务，也包含类似 Emu Edit 与 Palette（Sheynin et al., 2023；Saharia et al., 2022a）的 inpainting、segmentation、colorization、deblurring 和 ControlNet tasks。如 Figure 15 所示，所得 2B Edit model 能够操纵给定图像中的文字，尽管训练数据并未包含任何文字操纵任务。我们在相同数据上训练基于 SDXL（Podell et al., 2023）的 editing model 时，未能复现类似结果。

放大原图 ↗

E. 大规模 Text-to-Image 训练的 Data Preprocessing

E.1 预计算图像与文本 Embeddings

我们的模型把多个预训练且冻结的网络输出作为输入，包括 autoencoder latents 与 text encoder representations。由于这些输出在训练过程中保持不变，我们对整个 dataset 只预计算一次。这样做有两个主要优点：其一，训练时无须把 encoders 放在 GPU 上，从而降低显存需求；其二，训练时可以跳过前向编码过程，因此从第一个 epoch 之后开始能够节省时间与总计算量，见 Table 7。

这种方法也有两个缺点。第一，无法在每个 epoch 中为每个样本执行随机 augmentation；因此，在预计算 image latents 时，我们使用正方形 center crop。为了在更 high-resolution 上 finetune 模型，我们预先指定若干 aspect-ratio buckets，先把图像 resize 并 crop 到最接近的 bucket，再按该 aspect ratio 预计算。第二，text encoders 的 dense output 尤其庞大，会带来额外存储成本与更长的训练加载时间（参见 Table 7）。我们以 half precision 保存语言模型的 embeddings；实践中并未观察到性能下降。

E.2 防止 image memorization

在生成图像模型的语境下，对训练样本的 memorization 会引发多种问题（Somepalli et al., 2023a；Carlini et al., 2023；Somepalli et al., 2023b）。为避免训练出的模型逐字逐像地复制图像，我们仔细扫描训练 dataset 中的重复样本并将其移除。

列标签：Input、Output 1、Output 2。四条 editing instruction 依次为：

1. “改为写 ‘go small go home’。”
2. “黑板上写着 ‘GO BIG OR GO UNET’。”
3. “把这个词改成 ‘UNOT’。”
4. “让标牌写成 ‘MMDIT rules’。”

放大原图 ↗

deduplication 细节。 按照 Carlini et al.（2023）与 Somepalli et al.（2023a）所述方法，我们选择 SSCD（Pizzi et al., 2022）作为 deduplication 流程的 backbone。SSCD 是用于大规模检测近重复图像的 state-of-the-art 技术，它能生成高质量 image embeddings，可用于 clustering 及其他 downstream tasks。我们还按照 Nichol（2022）的方法确定 cluster 数量 $N$；实验中使用 $N=16{,}000$。

我们使用 autofaiss（2023）进行 clustering。autofaiss 是一个简化 Faiss（Facebook AI Similarity Search）大规模 clustering 使用流程的库。具体而言，我们利用 FAISS index factory 的功能，按预定义的 centroid 数量训练自定义 index。该方法能够高效、准确地聚类 image embeddings 这类高维数据。

Algorithm 1 详述了我们的 deduplication 方法。我们还通过实验考察不同 SSCD threshold 会移除多少数据，结果见 Figure 16b。根据这些结果，我们为最终运行选择了四个 thresholds，见 Figure 16a。

E.3 评估 deduplication 工作的有效性

Carlini et al.（2023）设计了一种两阶段数据 extraction attack：先用标准方法生成图像，再把 membership-inference score 超过某些判定标准的图像标记出来。Carlini et al.（2023）把搜索偏向训练数据中的重复样本，因为这类样本被记忆的概率比非重复样本高出多个数量级（Somepalli et al., 2023a,a；Lee et al., 2021）。

为了评估基于 SSCD 的 deduplication 效果，我们沿用 Carlini et al.（2023），从专门为该目的训练的小模型中提取被记忆的样本，并比较 deduplication 前后的情况。该流程主要包含两个步骤：1）以标准采样方式和已知 prompts，用 Diffusion Model 生成大量样本；2）执行 membership inference，把模型的新生成结果与模型记忆的训练样本分开。Algorithm 2 给出了依据 Carlini et al.（2023）寻找 memorized samples 的步骤。需要注意的是，我们把这套技术运行两次：一次用于只移除完全重复项的 SD-2.1 model，作为 baseline；另一次用于具有 SD-2.1 架构、但在训练前同时使用 SSCD（Pizzi et al., 2022）移除了完全重复与近重复项的模型。

我们基于 threshold $0.5$ 的 SSCD（Pizzi et al., 2022），从训练 dataset 中选出重复程度最高的 $350{,}000$ 个样本，并针对每个文本 prompt 生成 $500$ 张候选图像，以提高发现 memorization 的概率。其直觉是：对于 Diffusion Models，给定两个不同的随机初始 seeds $r_1,r_2$，$\operatorname{Gen}(p;r_1)\approx_d\operatorname{Gen}(p;r_2)$ 以高概率成立。反过来，如果在某种距离度量 $d$ 下 $\operatorname{Gen}(p;r_1)\approx_d\operatorname{Gen}(p;r_2)$，这些生成样本很可能就是被记忆的样本。为了计算两张图像之间的距离 $d$，我们采用修改后的 Euclidean $\ell_2$ distance。具体来说，我们发现很多生成结果会被 $\ell_2$ distance 错误地判为相似，例如它们都具有灰色背景。因此，我们改为把每张图像划分成 $16$ 个互不重叠的 $128\times128$ tiles，再在两张图像的任意 tile pair 之间计算 $\ell_2$ distance，并取其中的最大值。

Figure 17 比较了在使用 threshold $0.5$ 的 SSCD 移除近重复样本之前与之后，memorized samples 的数量。Carlini et al.（2023）把 clique size 为 $10$ 的图像标记为 memorized samples；本文还探索了不同的 clique sizes。对于所有 clique thresholds，SSCD 都能显著减少 memorized samples 的数量。具体来说，当 clique size 为 $10$ 时，在按 SSCD $=0.5$ 截断后的 deduplication 训练样本上训练 SD models，可把潜在 memorized examples 减少到原来的 $1/5$。

放大原图 ↗

放大原图 ↗