[ControlNet]Adding Conditional Control to Text-to-Image Diffusion Models

文章链接：https://arxiv.org/pdf/2302.05543v1.pdf

代码：GitHub - lllyasviel/ControlNet: Let us control diffusion models!

ControlNet的核心在于给stable diffusion施加了精准控图，stable diffusion生成的图本质上是通过prompt进行引导从噪声中生成的，text是唯一的输入，controlnet输入了text之外额外的一个维度，目前论文支持canny边缘图，hough线条，hed边缘图，简笔图，pose人体姿态，分割图，深度图等。在流程上，依然是输入图片和prompt，controlnet会先从图片中提取线框图，再输入prompt和线框图即可生成新图。从技术上，controlnet对stable diffusion进行微调，核心通过两个零卷积对权重部分进行映射调整，具体来说，对于预训练好的模型（比如作者使用sd1.5-unet里encoder和middlelayer的resenet和transformer层）里的一层结构，作者固定其参数，并将该层的输入额外添加一个全连接映射后的条件c，输入到一个和该层结构一致的复制网络里，再映射一次后重新添加回原结构里的输出。在训练上，需要成对的text-线框和原图这种数据对。此外作者也提出了小规模训练和大规模训练，力争将stable diffusion的微调平民化。

摘要：

我们提出了一个神经网络结构，ControlNet，来控制预训练的大型扩散模型，以支持额外的输入条件。ControlNet以端到端方式学习特定于任务的条件，即使训练数据集很小(< 50k)，学习也是鲁棒的。稳定扩散可以通过ControlNets增强，以支持边缘图、分割图、关键点等条件输入。

介绍：

三个问题：

1. 首先，任务特定领域的可用数据规模并不总是像一般图像-文本领域那样大。这将需要鲁棒的神经网络训练方法来避免过拟合，并在针对特定问题训练大型模型时保持泛化能力。
2. 其次，当使用数据驱动的解决方案处理图像处理任务时，大型计算集群并不总是可用的。这使得快速训练方法对于在可接受的时间和内存空间(例如，在个人设备上)内优化大型模型以完成特定任务非常重要。这将进一步需要使用预训练的权重，以及微调策略或迁移学习。
3. 第三，各种图像处理问题具有不同形式的问题定义、用户控件或图像注释。在解决这些问题时，尽管图像扩散算法可以以“程序化”的方式进行调节，例如，约束去噪过程，编辑多头注意激活等，但这些手工制作规则的行为基本上是由人类指令规定的。考虑到一些特定的任务，如深度到图像、姿势到人等，这些问题本质上需要将原始输入解释为对象级或场景级理解，这使得手工制作的程序方法不太可行。为了在许多任务中获得学习解决方案，端到端学习是必不可少的。

本文介绍了ControlNet，一个端到端神经网络架构，它控制大型图像扩散模型(如Stable diffusion)来学习特定于任务的输入条件。ControlNet将大型扩散模型的权重克隆为“trainable copy（可训练副本）”和“locked copy（锁定副本）”: 锁定副本保留了从数十亿张图像中学习到的网络能力，而可训练副本则在特定任务数据集上进行训练，以学习条件控制。

可训练和锁定的神经网络块与一种称为“zero convolution（零卷积）”的独特类型的卷积层相连，其中卷积权以学习的方式从零逐步增长到优化参数。由于保留了可生产的权值，因此训练在不同规模的数据集上具有鲁棒性。由于零卷积不会向深度特征中添加新的噪声，与从头开始训练新层相比，训练与微调扩散模型一样快。

相关工作：

HyperNetwork起源于一种神经语言处理方法[14]，目的是训练一个小的递归神经网络来影响一个大的递归神经网络的权重。即将较小的神经网络附加到Stable Diffusion上，从而改变其输出图像的艺术风格。ControlNet和HyperNetwork在影响神经网络行为的方式上有相似之处。

ControlNet使用一种特殊类型的卷积层，称为“零卷积”。早期的神经网络研究[31,47,32]对网络权值的初始化进行了广泛的讨论，包括用高斯分布初始化权值的合理性以及用零初始化权值可能带来的风险。最近，[37]讨论了一种在扩散模型中缩放多个卷积层的初始权值以改进训练的方法，该方法与零卷积的思想有相似之处(并且它们的代码包含一个称为“zero_module”的函数)。

方法：

ControlNet：

以2D特征为例，给定特征映射x∈R h×w×c，其中{h, w, c}为高度，宽度和通道数，神经网络块F(·;Θ)和一组参数Θ将x转换为另一个特征映射y：

我们锁定Θ中的所有参数，然后将其克隆到一个可训练的副本Θc中。复制后的Θc使用外部条件向量c进行训练。本文将原始参数和新参数称为“锁定副本”和“可训练副本”。 制作这样的副本，而不是直接训练原始权重，是为了避免在数据集较小时过度拟合，并保持从数十亿图像中学习的大型模型的生产准备质量。

神经网络块通过一种称为“零卷积”的独特类型的卷积层连接。即1×1权值和偏置都初始化为零的卷积层。我们将零卷积运算记为Z(·;·)，并使用两个参数实例{Θz1， Θz2}来组成ControlNet结构：

因为零卷积层的权值和偏置都被初始化为零，所以在第一个训练步骤中，我们有：

这个可以转化成：

公式(2,3,4)表明，在第一个训练步骤中，神经网络块的可训练副本和锁定副本的所有输入和输出都与ControlNet不存在时的情况一致。换句话说，当ControlNet应用于某些神经网络块时，在进行任何优化之前，它不会对深度神经特征产生任何影响。任何神经网络块的能力、功能和结果质量都得到了完美的保留，任何进一步的优化都将变得和微调一样快(与从头开始训练这些层相比)。

我们简要地推导了零卷积层的梯度计算。考虑一个权重为W，偏置为B的1 × 1卷积层，

在任意空间位置p和通道级索引i，给定输入映射I∈R h×w×c，前向传播可以被写成：

        

由于零卷积W=0， B=0(优化前)，对于任意Ip,i不为零的地方，梯度变成：

我们可以看到，虽然零卷积可以使特征项I上的梯度变为零，但权重和偏置梯度不受影响。只要特征 I 不为零，权值 W  将在第一次梯度下降迭代中被优化为非零矩阵。值得注意的是，在我们的例子中，特征项是从数据集中采样的输入数据或条件向量，这自然确保了 I 不为零。

例如，考虑一个带有整体损失函数L和学习率βIr≠0的经典梯度下降：如果是“外部”梯度∂L/∂Z(I;{W,B})不是0，我们就有：

式中W*为梯度下降一步后的权值; ⊙是Hadamard乘积。经过这一步，我们将有：

其中得到非零梯度，神经网络开始学习。这样，零卷积就成为一种独特的连接层，以一种学习的方式从零逐步增长到优化参数。

图像扩散模式下的ControlNet：

我们以Stable Diffusion为例，介绍了使用ControlNet控制具有特定任务条件的大型扩散模型的方法。Stable Diffusion是一个大型文本到图像的扩散模型，训练了数十亿张图像。该模型本质上是一个带有编码器、中间块和skip-connected解码器的U-net。编码器和解码器都有12个块，整个模型有25个块(包括中间块)。在这些块中，8个块是下采样或上采样卷积层，17个块是主块，每个块包含4个resnet层和2个Vision Transformers (ViTs)。每个Vit包含几个交叉注意和/或自注意机制。文本采用OpenAI CLIP编码，扩散时间步长采用位置编码。

Stable Diffusion使用类似于VQ-GAN[11]的预处理方法，将整个512 × 512图像的数据集转换成更小的64 × 64“潜图像”进行稳定训练。这需要ControlNets将基于图像的条件转换为64 × 64的特征空间，以匹配卷积大小。我们使用一个由4个卷积层组成的小型网络E(.)，具有4 × 4核和2 × 2步长（由ReLU激活，通道分别为16、32、64、128，用高斯权值初始化，与完整模型联合训练）将图像空间条件ci编码为为特征映射:

其中cf为转换后的特征映射。该网络将512 × 512的图像条件转换为64 × 64的特征映射。如图3所示，我们使用ControlNet来控制U-net的每个级别。请注意，我们连接ControlNet的方式是计算效率高的:由于原始权重是锁定的，因此不需要在原始编码器上进行梯度计算来进行训练。具体来说，我们使用ControlNet来创建Stable Diffusion的12个编码块和1个中间块的可训练副本。12个块有4种分辨率(64 × 64、32 × 32、16 × 16、8 × 8)，每个块有3个块。将输出添加到U-net的12个skip-connections和1个middle block中。由于SD是典型的U-net结构，因此这种控制网架构可能可用于其他扩散模型

训练：

图像扩散模型学习逐步去噪图像以生成样本。去噪可以在像素空间或从训练数据编码的“潜在”空间中进行。Stable Diffusion使用潜在图像作为训练域。在这种情况下，术语“图像”、“像素”、“去噪”都是指“感知潜在空间”中的相应概念。给定图像z0，扩散算法逐步向图像中添加噪声并产生噪声图像zt，其中t为噪声添加的次数。当t足够大时，图像近似于纯噪声。给定一组条件，包括时间步长t、文本提示ct，以及特定于任务的条件cf，图像扩散算法学习一个网络ϵθ来预测添加到zt里的噪声：

其中L为整个扩散模型的总体学习目标。

这个学习目标可以直接用于微调扩散模型。在训练过程中，我们随机将50%的文本提示ct替换为空字符串。这有助于ControlNet从输入条件图中识别语义内容的能力，例如，Canny边缘图或人类涂鸦等。这主要是因为当Stable diffusion模型中提示符不可见时，编码器倾向于从输入控制映射中学习更多的语义，作为提示符的替代品。

实验：