为文本到图像的 diffusion 模型添加条件控制

来源
：arxiv

作者
：Lvmin Zhang, Maneesh Agrawala

论文题目
：Adding Conditional Control to Text-to-Image Diffusion Models

论文链接
：https://arxiv.org/abs/2302.05543

内容整理
：汪奕文

本文提出一种神经网络结构 ControlNet，用来控制预训练的大型扩散模型，并使其支持额外的输入条件。它以端到端的方式学习特定任务条件，即使在训练数据集很小时，学习也是鲁棒的。此外，训练 ControlNet 与微调扩散模型一样快，并且模型可以在个人设备上进行训练。如果有强大的计算集群，该模型可以扩展到大量数据。此外，Stable diffusion 等大型扩散模型可以用 ControlNet 增强，以实现边缘图、分割图、关键点等条件输入。

引言
方法

ControlNet
ControlNet 在图像扩散模型中的应用

训练

训练策略

实验

实验设置
实验结果
消融性实验

引言

随着大型文本到图像模型的出现，生成一张吸引人的图像可能只需要用户输入一个简短的描述性提示。在研究了各种图像处理应用方法后，有三个发现。首先，对于特定任务的可用数据集规模并不总是像一般图像-文本对数据集那样大。许多具体问题的最大数据集大小通常在 100k 以下。这需要鲁棒的神经网络训练方法，以避免过度拟合，并在大模型针对特定问题进行训练时保持泛化能力。其次，当使用数据驱动解决图像处理任务时，需要用到大量的计算资源。这使得快速训练方法对于在可接受的时间和内存空间内优化大型模型以执行特定任务非常重要。这将进一步需要利用预训练的权重，以及微调策略或迁移学习。第三，各种图像处理任务具有不同形式的问题定义、用户控制或图像标注。在解决这些问题时，虽然图像扩散算法可以通过人为的方式进行调节，如约束去噪过程、编辑多头注意力激活等。但是考虑到一些特定的任务，如深度-图像转换、姿态-人体转换等，这些问题本质上需要将原始输入解释为对象级或场景级的理解，这使得人为调整过程的方法不太可行。要在许多任务中实现解决方案，端到端学习是必不可少的。

方法

ControlNet

ControlNet 操纵神经网络块的输入条件，从而进一步控制整个神经网络的整体行为。以 2D 的特征图为例，对于给定的特征图，其中分别为高度、宽度和通道数。神经网络模块通过一组参数将转化为特征图。

将中的所有参数锁定，并克隆一个可训练的副本，用外部条件训练副本。将原始参数和新参数分别称为"锁定副本"和"可训练副本"。这样可以在数据集较小时避免过拟合，并且保持从数十亿张图像中学习到的大型模型的质量。神经网络块由一种称为“零卷积”的独特类型的卷积层连接，即权重和偏置都以零初始化的卷积层，零卷积操作表示为。由此，ControlNet 的结构可以表示为

由于零卷积层的权重和偏置都被初始化为零，所以在第一个训练步骤中

在第一个训练步骤中，神经网络块的可训练副本和锁定副本的所有输入和输出都与 ControlNet 不存在时一致。换句话说，当 ControlNet 应用于某些神经网络块时，在进行优化之前，不会对特征造成任何影响。考虑权重为偏置为的卷积层，在空间位置和通道索引处，对于给定的输入，前向过程可以写作

由于零卷积且，对于任意不等于零，梯度

尽管零卷积可以导致特征项上的梯度为零，但权重和偏差的梯度不受影响。只要特征是非零的，权重就会在第一次梯度下降迭代中被更新为非零矩阵。例如，考虑具有整体损失函数为，学习率为的梯度下降，如果“外部”梯度不为零，有

是一次梯度下降更新后的权重。

由此，得到非零梯度，神经网络开始学习。通过这种方式，零卷积逐步从零开始学习到优化参数。

ControlNet 在图像扩散模型中的应用

以 Stable Diffusion 为例，介绍了利用 ControlNet 对具有任务特定条件的大型扩散模型进行控制的方法。Stable Diffusion 是一个在数十亿张图像上训练的大型文本到图像扩散模型。该模型本质上是一个 U-net，具有编码器、中间块和解码器。它使用了类似于 VQ-GAN 的预处理方法，将整个图像数据集转换为更小的的“潜在图像”进行稳定训练。这需要 ControlNet 将基于图像的条件转换为的特征图，以匹配卷积大小。

由于原始权重是锁定的，因此在训练时不需要对原始编码器进行梯度计算，可以加快训练速度并节省 GPU 内存。具体来说，使用 ControlNet 创建 12 个编码块和 1 个中间块的可训练副本。输出被连接到 U-net 的 12 个解码块和 1 个中间块中。

训练

对于给定的特征图，扩散算法逐步向图像添加噪声，并产生噪声图像，其中是噪声添加的次数。当足够大时，图像近似于纯高斯噪声。给定一组条件，包括时间步长，文本提示，以及特定任务条件，图像扩散算法学习一个网络来预测添加到噪声图像的噪声。

在训练期间，50%的文本提示被随机替换为空文本，这有助于提升 ControlNet 从输入条件图中识别语义内容的能力。这主要是由于当文本提示对于 SD 模型不可见时，编码器倾向于从条件输入中学习更多语义。

训练策略

小规模训练当计算设备有限时，断开部分 ControlNet 与 Stable Diffusion 之间的联系可以加快收敛速度。断开 ControlNet 与解码器1、2、3、4的链接，仅连接中间块，可以将训练速度提高约1.6倍。当模型显示出结果与条件之间的关联合理时，可以在继续训练中重新连接，以便准确控制。
大规模训练大规模训练通常适用于数据很容易获得的任务，例如Canny检测到的边缘图。在这种情况下，由于过拟合的风险相对较低，可以首先训练控制网络进行足够多的迭代次数，然后解锁 Stable Diffusion 的所有权重，并将整个模型作为一个整体进行训练。