Adding Conditional Control to Text-to-Image Diffusion Models (Paper reading)

Lvmin Zhang and Maneesh Agrawala, Stanford University, arXiv, Cited:113, Code, Paper

1. 前言

我们提出了一种名为ControlNet的神经网络结构，用于控制预训练的大规模扩散模型，以支持额外的输入条件。ControlNet以端到端的方式学习任务特定的条件，并且即使训练数据集较小（< 50k），学习效果也很稳健。此外，训练ControlNet的速度与微调扩散模型的速度一样快，而且该模型可以在个人设备上进行训练。或者，如果有强大的计算集群可用，该模型可以扩展到大量（百万到十亿级）的数据。我们报告了使用ControlNet扩展稳定扩散等大规模扩散模型的结果，以实现诸如边缘映射、分割映射、关键点等条件输入。这可以丰富控制大规模扩散模型的方法，并进一步促进相关应用。

2. 问题声明

随着大型文本到图像模型的出现，仅凭用户输入的简短描述性提示即可生成具有视觉吸引力的图像。在输入一些文本并获得图像后，我们可能会自然而然地提出几个问题：这种基于提示的控制是否满足我们的需求？例如，在图像处理中，考虑到许多既定的、具有明确问题表述的任务，这些大型模型是否可以应用于促进这些特定任务？我们应该构建什么样的框架来处理广泛的问题条件和用户控制？在特定任务中，大型模型是否能保留从数十亿张图像中获得的优势和能力？

为了回答这些问题，我们研究了各种图像处理应用，并得出了三个发现。首先，任务特定领域中可用的数据规模并不总是与图像文本领域中的数据规模一样大。许多特定问题（例如对象形状/法线、姿势理解等）的最大数据集大小通常不超过100k，即比LAION5B小5×104倍。这将要求具有稳健神经网络训练方法，以避免过拟合并在大型模型针对特定问题进行训练时保持泛化能力。

其次，在使用数据驱动的解决方案处理图像处理任务时，并不总是可以使用大型计算集群。这使得快速训练方法对于在可接受的时间和内存空间内优化大型模型以适应特定任务非常重要（例如在个人设备上）。这还需要利用预训练权重以及微调策略或迁移学习。

第三，各种图像处理问题具有不同形式的问题定义、用户控制或图像注释。在解决这些问题时，尽管可以以“过程化”的方式调节图像扩散算法，例如限制去噪过程、编辑多头注意力激活等，但这些手工制定的规则行为基本上是由人类指令规定的。考虑到一些特定任务，如深度到图像、姿势到人类等，这些问题基本上需要将原始输入解释为对象级或场景级的理解，使得手工制定的过程化方法不太可行。在许多任务中实现学习的解决方案是不可或缺的，需要端到端学习。

2. 整体思想

大模型微调的一种范式，用另一个path来辅助大模型，主要的思想是固定大模型，copy大模型中去噪器的一些结构，然后对于不同任务训练这些结构，condition使用的是edge等，可以控制生成一些内容。

4. 方法

ControlNet是一种神经网络架构，可以通过任务特定条件增强预训练的图像扩散模型。我们在第4.1节介绍了ControlNet的基本结构，并解释了每个部分的动机。在第4.2节中，我们详细说明了将ControlNet应用于图像扩散模型的方法，以Stable Diffusion为例。在第4.3节中，我们详细描述了学习目标和通用训练方法，然后在第4.4节中描述了在极端情况下改善训练的几种方法，例如在一台笔记本电脑上进行训练或使用大规模计算集群。最后，在第4.5节中，我们提供了几个具有不同输入条件的ControlNet实现的详细信息。

4.1 ControlNet

ControlNet通过操作神经网络块的输入条件来进一步控制整个神经网络的行为。在这里，“网络块”指的是一组作为经常使用的单元组合在一起构建神经网络的神经层，例如“resnet”块、“conv-bn-relu”块、多头注意力块、变换器块等。以2D特征为例，给定一个特征图 $x ∈ R^{h×w×c}$ ，其中 ${h, w, c}$ 表示高度、宽度和通道数，一个具有一组参数 $Θ$ 的神经网络块 $F (\cdot; Θ)$ 将 $x$ 转换为另一个特征图 $y$ ，其形状为 $y = F (x; Θ)$ 。如图：
在这里插入图片描述

零卷积是权重初始化为0

我们将 $Θ$ 中的所有参数锁定，然后将其克隆为一个可训练的副本 $Θ_c$ 。复制后的 $Θ_c$ 将与外部条件向量 $c$ 一起进行训练。在本文中，我们将原始参数和新参数称为“锁定副本”和“可训练副本”。采用这种复制的动机是为了避免在数据集较小时出现过拟合，并保持从数十亿张图像中学习到的大型模型的产品可用质量。神经网络块通过一种称为“零卷积”的独特类型的卷积层连接，即使用零进行初始化的 $1 \times 1$ 卷积层，包括权重和偏置。我们用 $Z (\cdot;\cdot)$ 表示零卷积操作，并使用两个参数实例 $Θ_{z1}，Θ_{z2}$ 组成ControlNet结构。应用0卷积时，我们训练的第一步是完全使用了锁定副本的信息。在第一次训练步骤中，神经网络块的可训练副本和锁定副本的所有输入和输出与不存在ControlNet时的输入和输出一致。换句话说，在对某些神经网络块应用ControlNet之前，在任何优化之前，它不会对深层神经特征产生任何影响。任何神经网络块的能力、功能和结果质量都完全保留，任何进一步的优化将变得与微调一样快速（与从头开始训练这些层相比）。

4.2 ControlNet in Image Diffusion Model

我们以Stable Diffusion [44]作为例子，介绍了如何使用ControlNet来控制具有任务特定条件的大型扩散模型。Stable Diffusion是一个在数十亿张图像上训练的大型文本到图像扩散模型。该模型本质上是一个U-Net，包括一个编码器、一个中间块和一个跳连解码器。编码器和解码器都有12个块，完整的模型包含25个块（包括中间块）。在这些块中，有8个块是下采样或上采样的卷积层，17个块是主要块，每个块包含四个ResNet层和两个Vision Transformer（ViT）。每个ViT包含几个交叉注意力和/或自注意力机制。文本使用OpenAI CLIP进行编码，扩散时间步使用位置编码进行编码。
在这里插入图片描述

Stable Diffusion使用类似于VQ-GAN的预处理方法，将整个512×512的图像数据集转换为较小的64×64的“潜在图像”，以进行稳定训练。这需要使用ControlNet将基于图像的条件转换为64×64的特征空间，以匹配卷积的大小。我们使用一个小型网络 $E (\cdot)$ ，包含四个卷积层，内核大小为 $4 \times 4$ ，步长为 $2 \times 2$ （通过ReLU激活，通道数为16、32、64、128，使用高斯权重初始化，与完整模型一起进行联合训练），将图像空间的条件 $c_i$ 编码为特征图，大小为 $64 \times 64$ 。具体而言，我们使用ControlNet创建了Stable Diffusion的12个编码块和1个中间块的可训练副本。这12个块分布在4个分辨率（64×64、32×32、16×16、8×8）中，每个分辨率有3个块。输出被添加到U-net的12个跳跃连接和1个中间块中。由于SD是一个典型的U-net结构，这种ControlNet架构很可能在其他扩散模型中也可以使用。

4.3 训练

给定一个图像 $z_0$ ，扩散算法逐渐向图像添加噪声，产生一个带有噪声的图像 $z_t$ ，其中 $t$ 表示添加噪声的次数。当 $t$ 足够大时，图像逼近纯噪声。给定一组条件，包括时间步长 $t$ 、文本提示 $c_t$ 以及任务特定条件 $c_f$ ，图像扩散算法学习一个网络 $θ$ 来预测添加到噪声图像 $z_t$ 中的噪声:
$\mathcal{L}=\mathbb{E}_{z_0,t,c_t,c_f,\epsilon \sim N(0,1)}[||\epsilon-\epsilon_{\theta}(z_t,t,c_t,c_f)||^2_2]$
在训练过程中，我们随机将50%的文本提示ct替换为空字符串。这有助于ControlNet从输入条件图中识别语义内容，例如Canny边缘图或人工涂鸦等。这主要是因为当SD模型看不到提示时，编码器倾向于从输入控制图中学习更多的语义作为提示的替代。

4.4 改进训练

我们讨论了几种改进ControlNet训练的策略，特别是在计算设备非常有限的极端情况下（例如，在笔记本电脑上）或计算能力非常强大的情况下（例如，在具有大规模GPU的计算集群上）。在我们的实验中，如果使用了这些策略，我们将在实验设置中进行提及。

小规模训练：当计算设备有限时，我们发现部分断开ControlNet和Stable Diffusion之间的连接可以加快收敛速度。默认情况下，我们将ControlNet连接到“SD Middle Block”和“SD Decoder Block 1,2,3,4”，如图3所示。我们发现断开与解码器1,2,3,4的连接，仅连接中间块可以将训练速度提高约1.6倍（在RTX 3070TI笔记本GPU上测试）。当模型在结果和条件之间显示合理的关联时，可以重新连接这些断开的连接，进行进一步的训练以便于精确控制。

大规模训练：在这里，大规模训练是指同时具有强大的计算集群（至少8个Nvidia A100 80G或等效的GPU）和大规模数据集（至少100万个训练图像对）的情况。这通常适用于数据易于获取的任务，例如Canny检测到的边缘图。在这种情况下，由于过拟合的风险相对较低，我们可以先对ControlNet进行足够多的迭代训练（通常超过5万步），然后解锁Stable Diffusion的所有权重，联合训练整个模型。这将导致更加针对特定问题的模型。

4.5 实现

我们提供了几种使用不同基于图像的条件来控制大规模扩散模型的ControlNet实现。

Canny边缘：我们使用Canny边缘检测器（使用随机阈值）从互联网上获取了300万个边缘-图像标题对。该模型在Nvidia A100 80G上进行了600个GPU小时的训练。基础模型是稳定扩散1.5。
Canny边缘（改变）：我们对上述Canny边缘数据集的图像分辨率进行了排序，并采样了包含1,000、10,000、50,000、500,000个样本的子集。我们使用相同的实验设置来测试数据集规模的影响。显然训练数据越大效果越好啊！
Hough线：我们使用基于学习的深度霍夫变换从Places2中检测直线，然后使用BLIP生成标题。我们获得了60万个边缘-图像-标题对。我们使用上述Canny模型作为起始检查点，并在Nvidia A100 80G上进行了150个GPU小时的训练。
HED边界:我们使用HED边界检测从互联网上获取了300万个边缘-图像-标题对。该模型在Nvidia A100 80G上进行了300个GPU小时的训练。基础模型是稳定扩散1.5。
用户草图:我们使用HED边界检测和一组强大的数据增强方法（随机阈值、随机屏蔽一定百分比的草图、随机形态变换和随机非极大值抑制）从图像中合成人类涂鸦。我们从互联网上获取了50万个涂鸦-图像-标题对。我们使用上述Canny模型作为起始检查点，并在Nvidia A100 80G上进行了150个GPU小时的训练。注意，我们还尝试了一种更“类似人类”的合成方法，但该方法比简单的HED方法慢得多，我们没有注意到明显的改进。
人体姿势:我们使用基于学习的姿势估计方法从互联网上“找到”人类，使用一个简单的规则：具有人体的图像必须至少检测到整个身体关键点的30%。我们获得了8万个姿势-图像-标题对。注意，我们直接使用可视化的带有人体骨架的姿势图像作为训练条件。该模型在Nvidia RTX 3090TI上进行了400个GPU小时的训练。基础模型是稳定扩散2.1。
人体姿势: 我们使用基于学习的姿势估计方法在上述Openpifpaf设置中找到互联网上的人类，使用相同的规则。我们获得了20万个姿势-图像-标题对。注意，我们直接使用可视化的带有人体骨架的姿势图像作为训练条件。该模型在Nvidia A100 80G上进行了300个GPU小时的训练。其他设置与上述Openpifpaf相同。