将扩散模型蒸馏为条件 GANs
Paper Title:Distilling Diffusion Models into Conditional GANs
Paper地址
Code地址
Abstract.
我们提出了一种方法,将复杂的多步扩散模型蒸馏为单步条件 GAN 学生模型,从而显著加快推理速度,同时保持图像质量。我们的方法将扩散蒸馏解释为成对的图像到图像转换任务,使用扩散模型 ODE 轨迹的噪声到图像对。为了高效计算回归损失,我们提出了 E-LatentLPIPS,这是一种直接在扩散模型的潜在空间中运行的感知损失,利用一组增强。此外,我们调整了扩散模型来构建具有文本对齐损失的多尺度鉴别器,以构建有效的基于条件 GAN 的公式。E-LatentLPIPS 比许多现有的蒸馏方法收敛得更有效,甚至考虑到数据集构建成本。我们证明我们的一步生成器在零样本 COCO 基准上优于尖端的一步扩散提炼模型 - DMD、SDXL-Turbo 和 SDXL-Lightning。
1 Introduction
扩散模型 [24, 90, 95] 在 LAION [88] 等具有挑战性的数据集上展示了前所未有的图像合成质量。然而,产生高质量的结果需要数十或数百个采样步骤。因此,大多数现有的基于扩散的图像生成模型,如 DALL·E 2 [73]、Imagen [81] 和 Stable Diffusion [77],都存在高延迟的问题,通常超过 10 秒,阻碍了实时交互。如果我们的模型只需要一个推理步骤,它不仅可以改善文本到图像合成的用户体验,还可以扩展其在 3D 和视频应用中的潜力 [23, 72]。但是我们如何建立一个一步到位的文本到图像模型呢?
一种简单的解决方案是从头开始训练一步式模型。例如,我们可以训练 GAN [16],它是简单领域领先的一步式模型 [36]。不幸的是,尽管最近取得了一些进展,但在大规模、多样化数据集上训练文本到图像的 GAN 仍然具有挑战性 [33, 85]。挑战在于 GAN 需要在没有任何监督的情况下同时解决两个困难的任务:(1)找到噪声和自然图像之间的对应关系,(2)有效优化生成器模型以执行从噪声到图像的映射。正如 CycleGAN [111] 中提到的,与成对学习相比,这种“非成对”学习通常被认为更不适当,在成对学习中,条件 GAN [29] 可以在给定基本事实对应关系的情况下学习将输入映射到输出。
我们的主要想法是逐一解决上述任务。我们首先通过使用预先训练的扩散模型模拟 ODE 求解器来找到噪声和图像之间的对应关系。给定已建立的对应对,然后我们要求条件 GAN 在成对的图像到图像转换框架中将噪声映射到图像 [29,67]。这种解耦方法使我们能够利用两种类型的生成模型来完成不同的任务,从而实现两者的优势:使用扩散模型找到高质量的对应关系,同时使用条件 GAN 实现快速映射。
在这项工作中,我们从预训练的扩散模型中收集了大量的噪声-图像对,并将该任务视为成对的图像-图像转换问题 [29],这使我们能够利用感知损失 [12, 30, 106] 和条件 GAN [16,29,62] 等工具。在此过程中,我们有了一个意想不到的发现。如果回归损失设计得当,收集大量噪声-图像对数据库并使用回归损失(没有 GAN 损失)进行训练,已经可以在计算预算显著降低的情况下实现与最近的蒸馏方法 [60,93] 相当的结果。
首先,在回归任务中,与基于点的损失(如 L2)相比,使用感知损失(如 LPIPS [106])可以更好地保留感知上重要的细节,然而,感知损失与潜在扩散模型 [77] 根本不兼容,因为它们需要从潜在空间到像素空间的昂贵解码。为了克服这个问题,我们提出了 LatentLPIPS,表明感知损失可以直接在潜在空间中起作用。与在像素空间中计算 LPIPS 相比,这使批量大小增加了四倍。不幸的是,我们观察到基于潜在的感知损失比其像素对应物有更多的盲点。虽然之前的研究发现集成对基于像素的 LPIPS [38] 很有帮助,但我们发现它对基于潜在的版本至关重要。使用我们的 Ensembled-LatentLPIPS 在潜在空间中工作,我们仅使用回归损失就展示了强大的性能,堪比引导式渐进式蒸馏 [60]。此外,我们在训练中使用了鉴别器来进一步提高性能。我们开发了一个多尺度条件扩散鉴别器,利用预先训练的权重并使用我们新的单样本 R1 损失和混合搭配增强。我们将蒸馏模型命名为 Diffusion2GAN。
使用提出的 Diffusion2GAN 框架,我们将 Stable Diffusion 1.5 [77] 蒸馏为单步条件 GAN 模型。与其他蒸馏方法相比,我们的 Diffusion2GAN 可以更好地学习目标扩散模型中固有的噪声与图像对应关系。在零样本单步 COCO2014 [49] 基准上,它的表现也优于最近提出的蒸馏模型 UFOGen [102] 和 DMD [103]。此外,我们进行了广泛的消融研究,并强调了 E-LatentLPIPS 和多尺度扩散鉴别器的关键作用。
除了 Stable Diffusion 1.5 的蒸馏之外,我们还证明了 Diffusion2GAN 在蒸馏更大的 SDXL [71] 方面的有效性,与单步 SDXL-Turbo [86] 和 SDXL-Lightning [48] 相比,其 FID [21] 和 CLIP-score [20] 更胜一筹。
图 1:与 SDXL 教师模型 [79](无分类器指导尺度 [25] 为 7)和选定的蒸馏学生模型(包括 SDXL-Turbo [86]、SDXLLightning [48] 和我们的 SDXL-Diffusion2GAN)进行视觉比较。给定行中的所有图像都是使用相同的噪声输入生成的,但 SDXL-Turbo 除外,它需要 4 × 64 × 64 4 \times 64 \times 64 4×64×64的不同噪声大小。与其他蒸馏模型相比,我们的 SDXL-Diffusion2GAN 更接近原始 ODE 轨迹。
图 2:与 Stable Diffusion 1.5 教师模型 [79](无分类器指导尺度 [25] 为 8)和选定的蒸馏学生模型(包括 InstaFlow0.9B [53]、LCM-LoRA [58] 和我们的 Diffusion2GAN)进行视觉比较。使用相同的噪声输入来生成同一行中的图像。我们的方法 Diffusion2GAN 比 2 步 LCM-LoRA 和 InstaFlow-0.9B 实现了更高的真实感。
图 3:使用我们的一步式 Diffusion2GAN 框架生成的高质量图像。
我们的模型可以在 A100 GPU 上以 0.09/0.16 秒的交互速度合成 512px/1024px 图像,而教师模型 Stable Diffusion 1.5 [79]/SDXL [71] 使用 DDIM [91] 的 50 个步骤在 2.59/5.60 秒内生成图像。请访问我们的网站获取更多视觉效果。
2 Related Work
扩散模型。扩散模型 (DM) [24, 90, 95] 是一类生成模型,由扩散过程和去噪过程组成。扩散过程逐步将高维数据从数据分布扩散到易于采样的高斯分布,而去噪过程则旨在使用在分数匹配目标上训练的深度神经网络来逆转该过程 [94–96]。经过训练后,这些模型可以使用数值积分器从随机高斯噪声中生成数据 [1, 2, 34]。扩散模型已经实现了众多视觉和图形应用,例如图像编辑 [6,19,59]、可控图像合成 [64, 105]、个性化生成 [14, 43, 80]、视频合成 [4,18,23] 和 3D 内容创建 [47,72]。然而,采样通常需要数十个采样步骤,导致图像生成速度比其他生成模型(如 GAN [16] 和 VAE [41])慢。在这项工作中,我们的目标是在保持图像质量的同时加速模型的推理。
扩散蒸馏。加快扩散模型的采样速度对于增强实际应用以及降低推理的能量成本至关重要。一些研究提出使用快速 ODE 求解器减少采样步骤数 [34,54,55] 或减少每步的计算时间 [8, 45, 46]。另一种有效的加速方法是知识蒸馏 [3, 17, 22, 52, 53, 56, 60, 83, 93, 100, 110]。在这种方法中,教师扩散模型的多个步骤被蒸馏为步骤更少的学生模型。 Luhman 等人 [56] 提出在高斯噪声 x T \textbf{x}_{T} xT的学生输出与其对应的 ODE 解 x 0 \textbf{x}_{0} x0之间进行 L p L_{p} Lp回归训练。尽管简单,但这种直接回归会产生模糊的输出,并且与其他生成模型所表现出的图像合成能力不符。为了提高图像质量,InstaFlow [53] 通过 ReFlow [52] 拉直高曲率 ODE 轨迹,并将线性化的 ODE 轨迹蒸馏到学生模型。一致性蒸馏 (CD) [57,93] 训练学生模型以预测任何噪声样本 x t + 1 \textbf{x}_{t+1} xt+1及其单步去噪 x t \textbf{x}_{t} xt的一致输出。最近,一些研究提出使用 GAN 鉴别器来增强蒸馏性能。例如,CTM [39] 和 SDXL-Turbo [86] 利用改进的 StyleGAN [85,87] 鉴别器来训练单步图像生成器。此外,UFOGen [102]、SDXL-Lightning [48] 和 ADD [84] 采用预训练扩散模型作为强鉴别器,展示了它们在一步式文本到图像合成中的能力。虽然这些工作是同时进行的,但我们将把我们的方法与 SDXL-Turbo 和 SDXL-Lightning 进行比较,这两者都在一步式图像合成的扩散蒸馏中展示了最先进的性能。
条件生成对抗网络 [29,62] 一直是条件图像合成的常用框架。条件可以是图像 [9,29,51,66,76,111]、类标签 [5,31,37,63,65] 和文本 [33,75,85,101,104]。特别是,当借助回归损失来稳定训练时,cGAN 表现出了令人印象深刻的性能,就像在图像翻译中一样 [29,67,68,97,107,111]。同样,我们通过使用图像条件 GAN 以及感知回归损失 [106] 来实现扩散模型蒸馏。早期的工作将 GAN 与前向扩散过程相结合 [99, 100],但并不旨在将预先训练的扩散模型蒸馏成 GAN。
3 Method
我们的目标是将一个预训练的文本到图像扩散模型蒸馏为一步生成器。也就是说,我们希望学习一个映射 x = G ( z , c ) \mathbf{x}=G(\mathbf{z}, \mathbf{c}) x=G(z,c),其中一步生成网络 G G G 将输入噪声 z \mathbf{z} z 和文本 c \mathbf{c} c 映射到扩散模型的输出 x \mathbf{x} x。我们假设学生和教师共享相同的架构,这样我们可以使用教师模型的权重来初始化学生模型 G G G。在我们的方法部分中,我们假设使用潜在扩散模型 [77],其中 x , z ∈ R 4 × 64 × 64 \mathbf{x}, \mathbf{z} \in \mathbb{R}^{4 \times 64 \times 64} x,z∈R4×64×64。随后,我们也将我们的方法应用于 SDXL 模型 [71]。
在本节的其余部分,我们将详细介绍我们框架的设计和训练原则。我们首先在 3.1 节将蒸馏描述为配对的图像到图像翻译问题。然后,我们在 3.2 节介绍我们的集成潜在 LPIPS 回归损失 (E-LatentLPIPS)。仅仅使用这个回归损失就提高了训练效率,并显著改善了潜在扩散模型的蒸馏性能。最后,我们提出了一种改进的判别器设计,该设计重新使用了一个预训练的扩散模型(3.3 节)。值得注意的是,我们的研究结果不仅限于特定类型的潜在空间扩散模型 [ 77 , 78 ] [77,78] [77,78],还适用于像素空间模型 [ 34 ] [34] [34]。
3.1 Paired Noise-to-Image Translation for One-step Generation
随着扩散概率模型的出现 [24,95],Luhman 等人 [56] 提出可以通过最小化以下蒸馏目标,将预训练扩散模型的多步去噪过程简化为单步:
L
distill
ODE
=
E
{
z
,
c
,
x
}
[
d
(
G
(
z
,
c
)
,
x
)
]
,
(
1
)
\mathcal{L}_{\text {distill }}^{\text {ODE }}=\mathbb{E}_{\{\mathbf{z}, \mathbf{c}, \mathbf{x}\}}[d(G(\mathbf{z}, \mathbf{c}), \mathbf{x})],\quad(1)
Ldistill ODE =E{z,c,x}[d(G(z,c),x)],(1)
其中
z
\mathbf{z}
z 是高斯噪声的样本,
c
\mathbf{c}
c 是文本提示,
G
G
G 表示具有可训练权重的 UNet 生成器,
x
\mathbf{x}
x 是使用 DDIM 采样器 [91] 模拟常微分方程(ODE)轨迹的扩散模型输出,
d
(
⋅
,
⋅
)
d(\cdot, \cdot)
d(⋅,⋅) 是距离度量。由于每次迭代获得
x
\mathbf{x}
x 的计算成本,该方法在训练开始前使用预先计算的噪声和对应 ODE 解对。在训练过程中,随机采样噪声-图像对并最小化 ODE 蒸馏损失(方程 1)。虽然所提出的方法看起来简单直接,但与最近的蒸馏方法 [53,60, 83, 93] 相比,直接蒸馏方法产生的图像合成效果较差。
在我们的工作中,我们假设直接蒸馏的全部潜力尚未被充分发挥。在 CIFAR10 的实验中,我们观察到通过(1)扩大 ODE 对数据集的规模和(2)使用感知损失 [106](而非 Luhman 等人中使用的像素空间 L2 损失),我们可以显著提高蒸馏的质量。在表 6 中,我们展示了 CIFAR10 数据集上的训练进展,并将其性能与一致性模型 [93] 进行比较。令人惊讶的是,使用 LPIPS 损失的直接蒸馏在较少的计算量下可以实现比一致性模型更低的 FID,即使考虑到收集 ODE 对所需的额外计算。
图 4:用于潜在空间蒸馏的 E-LatentLPIPS。在 NVIDIA A100 上使用 LPIPS [106] 训练单次迭代需要 117 毫秒和 15.0GB 额外内存,而我们的 E-LatentLPIPS 在同一设备上需要 12.1 毫秒和 0.6GB。因此,与 LPIPS 相比,ElatentLPIPS 将感知损失计算时间加快了 9.7 倍,同时减少了内存消耗。
3.2 Ensembled-LatentLPIPS for Latent Space Distillation
原始的 LPIPS [106] 观察到预训练分类器的特征可以校准得足够好以匹配人类的感知响应。此外,LPIPS 在许多图像翻译应用中作为有效的回归损失 [67,98]。然而,LPIPS 是为像素空间设计的,在潜在扩散模型 [77] 中使用起来很笨重。如图 4 所示,潜在编码在用特征提取器
F
F
F 和距离度量
ℓ
\ell
ℓ 计算 LPIPS 之前必须被解码到像素空间(例如,
64
→
512
64 \rightarrow 512
64→512 分辨率)。
d
LPIPS
(
x
0
,
x
1
)
=
ℓ
(
F
(
Decode
8
×
(
x
0
)
)
,
F
(
Decode
8
×
(
x
1
)
)
)
d_{\text {LPIPS }}\left(\mathbf{x}_0, \mathbf{x}_1\right)=\ell\left(F\left(\operatorname{Decode}^{8 \times}\left(\mathbf{x}_0\right)\right), F\left(\operatorname{Decode}^{8 \times}\left(\mathbf{x}_1\right)\right)\right)
dLPIPS (x0,x1)=ℓ(F(Decode8×(x0)),F(Decode8×(x1)))
这违背了 LDMs 在更高效的潜在空间中操作的主要动机。因此,我们能否绕过解码到像素的需要,直接在潜在空间中计算感知距离?
学习 LatentLPIPS。我们假设 LPIPS 的相同感知属性可以用于直接在潜在空间上计算的函数。根据 Zhang 等人 [106] 的程序,我们首先在 Stable Diffusion 的潜在空间中训练一个 VGG 网络 [89]。我们通过移除 3 个最大池化层稍微修改了架构,因为潜在空间已经是 8 × 8 \times 8× 下采样的,并将输入更改为 4 个通道。然后我们使用 BAPPS 数据集 [106] 线性校准中间特征。这成功地生成了一个在潜在空间中操作的函数: d LatentLPIPS ( x 0 , x 1 ) = d_{\text {LatentLPIPS }}\left(\mathbf{x}_0, \mathbf{x}_1\right)= dLatentLPIPS (x0,x1)= ℓ ( F ( x 0 ) , F ( x 1 ) ) \ell\left(F\left(\mathbf{x}_0\right), F\left(\mathbf{x}_1\right)\right) ℓ(F(x0),F(x1))
图 5:单幅图像重建。为了深入了解回归损失的损失情况,我们通过直接优化具有不同损失函数的单个潜在函数来进行图像重建实验。使用 LPIPS 进行重建可以粗略地再现目标图像,但需要解码成像素。单独的 LatentLPIPS 无法精确重建图像。但是,我们的集成增强 E-LatentLPIPS 可以在潜在空间中直接操作时更精确地重建目标。
有趣的是,我们观察到,虽然在潜在空间中 ImageNet 分类精度比在像素上略低,但感知一致性得以保持。这表明,尽管压缩到潜在空间破坏了一些有助于分类的低级信息 [28],但它保留了图像的感知相关细节,我们可以利用这些细节。附录 B 中有更多详细信息。
集成。我们观察到,直接应用 LatentLPIPS 作为蒸馏的新损失函数会导致产生波纹、斑块状伪影。我们在一个简单的优化设置中对此进行了进一步调查,如图 5 所示,通过优化随机采样的潜在编码以接近单个目标图像。这里我们旨在使用不同的损失函数恢复目标潜在变量: arg min x ^ d ( x ^ , x ) \arg \min _{\hat{\mathbf{x}}} d(\hat{\mathbf{x}}, \mathbf{x}) argminx^d(x^,x),其中 x \mathbf{x} x 是目标潜在变量, x ^ \hat{\mathbf{x}} x^ 是重建的潜在变量,而 d d d 是原始 LPIPS 或 LatentLPIPS。我们观察到,单个图像重建在 LatentLPIPS 下不收敛(图 5(c))。我们假设这种限制是由于 VGG 网络的潜在版本形成的次优损失景观。
受到 E-LPIPS [38] 的启发,我们将随机可微增强 [35, 108]、通用几何变换 [35] 和 cutout [11] 应用于生成和目标潜在变量。在每次迭代中,随机增强被应用于生成和目标潜在变量。当应用于单图像优化时,集成策略几乎完美地重建了目标图像,如图 4(d)所示。新的损失被命名为 Ensembled-LatentLPIPS,简称 E-LatentLPIPS。
d
E-LatentLPIPS
(
x
0
,
x
1
)
=
E
T
[
ℓ
(
F
(
T
(
x
0
)
)
,
F
(
T
(
x
1
)
)
)
]
,
d_{\text {E-LatentLPIPS }}\left(\mathbf{x}_0, \mathbf{x}_1\right)=\mathbb{E}_{\mathcal{T}}\left[\ell\left(F\left(\mathcal{T}\left(\mathbf{x}_0\right)\right), F\left(\mathcal{T}\left(\mathbf{x}_1\right)\right)\right)\right],
dE-LatentLPIPS (x0,x1)=ET[ℓ(F(T(x0)),F(T(x1)))],
其中
T
\mathcal{T}
T 是随机采样的增强。将损失函数应用于 ODE 蒸馏:
L
E-LatentLPIPS
(
G
,
z
,
c
,
x
)
=
d
E-LatentLPIPS
(
G
(
z
,
c
)
,
x
)
\mathcal{L}_{\text {E-LatentLPIPS }}(G, \mathbf{z}, \mathbf{c}, \mathbf{x})=d_{\text {E-LatentLPIPS }}(G(\mathbf{z}, \mathbf{c}), \mathbf{x})
LE-LatentLPIPS (G,z,c,x)=dE-LatentLPIPS (G(z,c),x)
其中
z
\mathbf{z}
z 表示高斯噪声,
x
\mathbf{x}
x 表示其目标潜在变量。如图 4(右)所示,与其 LPIPS 对应项相比,计算时间显著降低,因为(1)不需要解码到像素(在 A100 上为一张图像节省 79 毫秒),以及(2)(Latent)LPIPS 本身在比像素空间更低分辨率的潜在代码上操作(38
→
\rightarrow
→ 8 毫秒)。虽然增强需要一些时间(4 毫秒),但总的来说,使用我们的 E-LatentLPIPS,感知损失计算几乎便宜
10
×
10 \times
10×(117
→
\rightarrow
→ 12 毫秒)。此外,内存消耗大大减少(15
→
\rightarrow
→ 0.6 GB)。
表 1 的实验结果表明,使用 E-LatentLPIPS 学习 ODE 映射可获得更好的收敛性,表现出比其他度量(如 MSE、Pseudo Huber 损失 [27,92] 和原始 LPIPS 损失)更低的 FID。有关玩具重建实验和可微增强的更多详细信息,请参阅附录 B。
表 1:SD-CFG-3 数据集上的消融研究。我们使用 ODE 蒸馏损失(公式 1)将稳定扩散 1.5 [79] 蒸馏为一步生成器。除 LPIPS 模型和较大批量大小的模型外,所有模型均使用 8 个 A100-80GB GPU 以 256 的批量大小进行 20k 次迭代训练。对于 LPIPS 模型,我们使用 64 的批量大小并累积梯度四次。这种调整是必要的,因为当应用于 64 的批量大小时,LPIPS 模型每 A100-80GB 消耗 62GB 的 GPU 内存,而其他模型需要近 68GB 的 GPU 内存进行 256 批量训练。我们的 E-LatentLPIPS 实现了比传统 LPIPS 更强大的性能,而无需解码为像素。
图 6:我们的 Diffusion2GAN 用于一步式图像合成。首先,我们收集扩散模型输出潜伏信息以及输入噪声和提示。其次,使用 E-LatentLPIPS 回归损失(公式 4)和 GAN 损失(公式 6)训练生成器将噪声和提示映射到目标潜在信息。虽然生成器的输出可以由 SD 潜在信息解码器解码为 RGB 像素,但这是一项计算密集型操作,在训练期间从不执行。
3.3 Conditional Diffusion Discriminator
在第 3.1 和 3.2 节中,我们已经阐明可以通过将扩散蒸馏表述为配对噪声到潜在翻译任务来实现它。受配对图像到图像翻译的条件 GAN 有效性的启发 [29],我们采用了条件鉴别器。此鉴别器的条件不仅包括文本描述 c,还包括提供给生成器的高斯噪声 z。我们新的鉴别器在利用预训练的扩散权重的同时,结合了上述条件。形式上,我们为生成器 G G G 和鉴别器 D D D 优化以下极小极大目标:
min G max D E c , z , x [ log ( D ( c , z , x ) ) ] + E c , z [ log ( 1 − D ( c , z , G ( z , c ) ) ) ] \min _G \max _D \mathbb{E}_{\mathbf{c}, \mathbf{z}, \mathbf{x}}[\log (D(\mathbf{c}, \mathbf{z}, \mathbf{x}))]+\mathbb{E}_{\mathbf{c}, \mathbf{z}}[\log (1-D(\mathbf{c}, \mathbf{z}, G(\mathbf{z}, \mathbf{c})))] GminDmaxEc,z,x[log(D(c,z,x))]+Ec,z[log(1−D(c,z,G(z,c)))]
对于生成器,我们最小化以下非饱和 GAN 损失 [15]。
L G A N ( G , z , c , x ) = − E c , z [ log ( D ( c , z , G ( z , c ) ) ) ] . \mathcal{L}_{\mathrm{GAN}}(G, \mathbf{z}, \mathbf{c}, \mathbf{x})=-\mathbb{E}_{\mathbf{c}, \mathbf{z}}[\log (D(\mathbf{c}, \mathbf{z}, G(\mathbf{z}, \mathbf{c})))] . LGAN(G,z,c,x)=−Ec,z[log(D(c,z,G(z,c)))].
生成器的最终损失为 L G = L E − LatentLPIPS + λ G A N L G A N \mathcal{L}_G=\mathcal{L}_{\mathrm{E}-\text { LatentLPIPS }}+\lambda_{\mathrm{GAN}} \mathcal{L}_{\mathrm{GAN}} LG=LE− LatentLPIPS +λGANLGAN。我们提供关于鉴别器和损失函数的更多细节。
从预训练扩散模型初始化。我们证明了用预训练扩散模型初始化鉴别器权重对于扩散蒸馏是有效的。与 GigaGAN 鉴别器的实现 [33] 相比,使用预训练的 Stable Diffusion 1.5 U-Net [79] 并微调模型作为潜在空间中的鉴别器在表 2 中表现出更优越的 FID。对抗性损失在 U-Net 鉴别器输出的每个位置独立计算。注意,原始 U-Net 架构依赖于文本条件,但不依赖于输入噪声图 z。我们通过简单地添加通过单个卷积层(在通道维度上零初始化)处理的 z \mathbf{z} z 输入,进一步修改了鉴别器架构以支持 z \mathbf{z} z 条件。注意,扩散鉴别器的文本条件自然是由 SD U-Net 中的内置交叉注意层执行的。我们观察到所有指标均有适度改进。(正则)
单样本 R1 正则化。虽然来自预训练扩散权重的条件 U-Net 鉴别器在零样本 COCO2014 [49] 基准测试中已经取得了有竞争力的结果,但我们注意到在不同运行中训练方差显著,这可能是由于缺乏正则化和来自鉴别器的无界梯度造成的。为了缓解这一问题,我们在每个小批量上为训练扩散鉴别器引入 R1 正则化 [61]。然而,引入 R1 正则化会增加 GPU 内存消耗,尤其是在鉴别器是高容量 U-Net 时,这带来了实际挑战。为了最大限度地减少内存消耗并加速训练,我们不仅采用间隔为 16 的惰性正则化 [37],还仅对每个小批量的单个样本应用 R1 正则化。除了提高稳定性外,我们还观察到单样本 R1 正则化导致更好的收敛,如表 2 所示。
多尺度进出 U-Net 鉴别器。GigaGAN [33] 观察到 GAN 鉴别器倾向于关注特定频带,经常忽略高级结构,并引入了多尺度鉴别器来解决此问题。同样,我们提出了一种新的 U-Net 鉴别器设计,如图 7 所示,它在 U-Net 的各个部分强制执行独立的真/假预测。具体来说,我们修改了 U-Net 编码器以在每个下采样层接收调整大小的输入,并在 U-Net 解码器的每个尺度上附加三个读出层,以从 U-Net 跳过连接特征、U-Net 瓶颈的上采样特征和组合特征进行独立的真/假预测。在高层次上,新设计强制所有 U-Net 层参与最终预测,从浅层跳过连接到深层中间块。这种设计增强了低频结构一致性并显着增加了 FID,如表 2 所示。(正则)
图 7:我们的多尺度条件判别器设计。我们重用教师模型 U-Net 中的预训练权重,并使用多尺度输入和输出分支对其进行扩充。具体来说,我们将调整大小的输入潜在值版本馈送到编码器的每个下采样块。对于解码器部分,我们强制判别器在每个尺度的三个位置进行真实/虚假预测:跳过连接之前、跳过连接时和跳过连接之后。这种多尺度对抗训练进一步提高了图像质量。
混合搭配增强。为了进一步鼓励判别器专注于文本对齐和噪声调节,我们在判别器训练中引入了混合搭配增强,类似于 GigaGAN [33] 和早期的文本转图像 GAN 工作 [74, 104]。在判别器训练期间,我们用来自目标数据集的随机、不相关的潜在数据替换一部分生成的潜在数据,同时保持其他条件不变。这会将替换的潜在数据归类为假的,因为潜在数据与其配对的噪声和文本之间的对齐不正确,从而促进对齐的改进。此外,我们对文本和噪声进行了替换,有助于整体增强条件扩散判别器。(正则)