Adding Conditional Control to Text-to-Image Diffusion Models

10 Feb 2023

Lvmin Zhang, Maneesh Agrawala

[Stanford University]

https://arxiv.org/abs/2302.05543

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

1简介

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

为了回答这些问题，我们调查了各种图像处理应用程序，并得出了三个发现。首先，特定任务域中的可用数据规模并不总是像一般图像文本域中的那么大。许多特定问题（例如，物体形状/法线、姿势理解等）的最大数据集大小通常小于100k，即比LAION-5B小5×104倍。这将需要鲁棒的神经网络训练方法，以避免过拟合，并在针对特定问题训练大型模型时保持泛化能力。

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

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

本文介绍了ControlNet，这是一种端到端的神经网络架构，它控制大型图像扩散模型（如稳定扩散）来学习特定任务的输入条件。ControlNet将大型扩散模型的权重克隆为“可训练副本”和“锁定副本”：锁定副本保留了从数十亿图像中学习的网络能力，而可训练副本在特定任务的数据集上进行训练，以学习条件控制。可训练和锁定的神经网络块与一种称为“零卷积”的独特类型的卷积层相连，其中卷积权重以学习的方式从零逐渐增长到优化的参数。由于保留了生产就绪权重，因此在不同规模的数据集上训练是稳健的。由于零卷积不会向深层特征添加新的噪声，因此与从头开始训练新层相比，训练的速度与微调扩散模型一样快。

我们用不同条件的各种数据集训练了几个ControlNets，例如Canny边、Hough线、用户涂鸦、人类关键点、分割图、形状法线、深度等。我们还用小数据集（样本小于50k甚至1k）和大数据集（数百万样本）对ControlNets进行了实验。我们还表明，在深度到图像等一些任务中，在个人计算机（一台Nvidia RTX 3090TI）上训练ControlNets可以获得与在具有TB GPU内存和数千GPU小时的大型计算集群上训练的商业模型相比具有竞争力的结果

2相关工作

2.1超网络和神经网络结构

HyperNetwork源于一种神经语言处理方法[14]，用于训练一个小的递归神经网络来影响一个大网络的权重。HyperNetwork在使用生成对抗性网络[1，10]和其他机器学习任务[51]的图像生成中也取得了成功。受这些想法的启发，[15]提供了一种将较小的神经网络连接到稳定扩散[44]的方法，以改变其输出图像的艺术风格。在[28]提供了几个超级网络的预训练权重后，这种方法变得更加流行。ControlNet和HyperNetwork在影响神经网络行为的方式上有相似之处。

ControlNet使用一种特殊类型的卷积层，称为“零卷积”。早期的神经网络研究[31，47，32]广泛讨论了网络权重的初始化，包括用高斯分布初始化权重的合理性，以及用零初始化权重可能带来的风险。最近，[37]讨论了一种在扩散模型中缩放几个卷积层的初始权重以改进训练的方法，该方法与零卷积的思想相似（并且他们的代码包含一个称为“zero_module”的函数）。ProGAN[21]和StyleGAN[22]以及Noise2Noise[33]和[65]中也讨论了对初始卷积权重的操作。稳定性的模型卡[55]也提到了在神经层中使用零权重。

2.2扩散概率模型

[52]中提出了扩散概率模型。图像生成的成功结果首先在小规模上报道[25]，然后在相对大的规模上报道[9]。该架构通过重要的训练和采样方法进行了改进，如去噪扩散概率模型（DDPM）[17]、去噪扩散隐式模型（DDIM）[53]和基于分数的扩散[54]。图像扩散方法可以直接使用像素颜色作为训练数据，在这种情况下，研究人员经常考虑在处理高分辨率图像时节省计算能力的策略[53，50，26]，或者直接使用基于金字塔的方法或多阶段方法[18，43]。这些方法本质上使用U-net[45]作为其神经网络架构。为了降低训练扩散模型所需的计算能力，基于潜像[11]的思想，提出了潜在扩散模型（LDM）[44]方法，并将其进一步扩展到稳定扩散。

2.3文本到图像的扩散

扩散模型可以应用于文本到图像的生成任务，以实现最先进的图像生成结果。这通常是通过使用像CLIP[41]这样的预训练语言模型将文本输入编码为潜在向量来实现的。例如，Glide[38]是一个文本引导的扩散模型，支持图像生成和编辑。Disco Diffusion是[9]的一个剪辑引导实现，用于处理文本提示。稳定扩散是潜在扩散[44]的大规模实现，以实现文本到图像的生成。Imagen[49]是一种文本到图像的结构，它不使用潜像，而是使用金字塔结构直接扩散像素。

2.4预训练扩散模型的个性化、定制和控制

由于最先进的图像扩散模型主要由文本到图像的方法主导，因此增强对扩散模型控制的最直接的方法通常是文本引导[38，24，2，3，23，43，16]。这种类型的控制也可以通过操纵CLIP功能来实现[43]。图像扩散过程本身可以提供一些功能来实现颜色级别的细节变化[35]（稳定扩散社区称之为img2img）。图像扩散算法自然支持将修复作为控制结果的重要方法[43，2]。文本反转[12]和DreamBooth[46]被提议使用具有相同主题或对象的一小组图像来定制（或个性化）生成的结果中的内容。

2.5图像到图像的翻译

我们想指出的是，尽管ControlNet和图像到图像翻译可能有几个重叠的应用程序，但它们的动机本质上是不同的。图像到图像的翻译旨在学习不同领域中图像之间的映射，而ControlNet旨在控制具有特定任务条件的扩散模型。

Pix2Pix[20]提出了图像到图像翻译的概念，早期的方法以条件生成神经网络为主[20，69，60，39，8，63，68]。在Transformer和视觉Transformer（ViTs）流行起来之后，使用自回归方法已经报道了成功的结果[42，11，7]。一些研究还表明，多模型方法可以从各种翻译任务中学习鲁棒生成器[64，29，19，40]。

我们讨论了当前图像到图像翻译中最强大的方法。驯服Transformer[11]是一种视觉Transformer，具有生成图像和执行图像到图像转换的能力。Palette[48]是一个统一的基于扩散的图像到图像翻译框架。PITI[59]是一种基于扩散的图像到图像转换方法，它利用大规模预训练来提高生成结果的质量。在草图引导扩散等特定领域，[58]是一种基于优化的方法，可以操纵扩散过程。这些方法在实验中得到了验证。

3方法

ControlNet是一种神经网络架构，可以在特定任务条件下增强预训练的图像扩散模型。我们在第3.1节中介绍了ControlNet的基本结构和各部分的动机。我们使用第3.2节中的稳定扩散示例详细介绍了将ControlNets应用于图像扩散模型的方法。我们在第3.3节中详细阐述了学习目标和一般训练方法，然后在第3.4节中描述了在极端情况下改进训练的几种方法，例如使用一台笔记本电脑或使用大型计算集群进行训练。最后，我们在第3.5节中介绍了几种具有不同输入条件的ControlNet实现的细节。

3.1ControlNet

ControlNet操纵神经网络块的输入条件，以便进一步控制整个神经网络的整体行为。这里，“网络块”是指一组神经层，它们被放在一起作为构建神经网络的常用单元，例如“resnet”块、“conv-bn-relu”块、多头注意力块、transformer块等。

以2D特征为例，给定特征映射x∈R h×w×c，其中{h，w，c}是高度、宽度和通道数，具有一组参数的神经网络块F（·；Θ）将x变换为另一个具有

该过程如图2-（a）所示。

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

神经网络块由一种称为“零卷积”的独特类型的卷积层连接，即1×1卷积层，权重和偏差都用零初始化。我们将零卷积运算表示为Z（·；·），并使用参数{θz1，θz2}的两个实例来组成ControlNet结构

其中yc成为该神经网络块的输出，如图2-（b）所示。因为零卷积层的权重和偏差都被初始化为零，所以在第一个训练步骤中，我们

这可以转换为

和Eq-（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和学习率βlr 6=0的经典梯度下降，如果“外部”梯度ŞL/õZ（I；{W，B}）不为零，我们将得到

其中W*是一个梯度下降步骤后的重量；

是哈达玛产品。在这一步之后，我们将

其中获得非零梯度并且神经网络开始学习。通过这种方式，零卷积成为一种独特类型的连接层，以学习的方式从零逐渐增长到优化的参数。

3.2图像扩散模型中的ControlNet

我们以稳定扩散[44]为例，介绍了使用ControlNet控制具有特定任务条件的大型扩散模型的方法。

稳定扩散是在数十亿张图像上训练的大型文本到图像扩散模型。该模型本质上是一个带有编码器、中间块和跳过连接解码器的U-net。编码器和解码器都有12个块，整个模型有25个块（包括中间块）。在这些块中，8个块是下采样或上采样卷积层，17个块是主块，每个块包含四个resnet层和两个视觉Transformer（ViT）。每个Vit都包含几个交叉注意和/或自我注意机制。文本由OpenAI CLIP编码，扩散时间步长由位置编码编码。

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

其中cf是转换后的特征图。该网络将512×512的图像条件转换为64×64的特征图。

如图3所示，我们使用ControlNet来控制U-net的每个级别。请注意，我们连接ControlNet的方式在计算上是高效的：由于原始权重被锁定，因此不需要在原始编码器上进行梯度计算来进行训练。这可以加快训练并节省GPU内存，因为可以避免原始模型上一半的梯度计算。使用ControlNet训练稳定的扩散模型在每次训练迭代中只需要增加约23%的GPU内存和34%的时间（在单个Nvidia A100 PCIE 40G上测试）。

具体来说，我们使用ControlNet来创建稳定扩散的12个编码块和1个中间块的可训练副本。这12个块有4个分辨率（64×64，32×32，16×16，8×8），每个块有3个块。输出被添加到U-net的12个跳过连接和1个中间块。由于SD是一种典型的U-net结构，因此这种ControlNet架构可能可用于其他扩散模型。

3.3训练

图像扩散模型学习逐步对图像进行去噪以生成样本。去噪可以发生在像素空间或根据训练数据编码的“潜在”空间中。稳定扩散使用潜像作为训练域。在这种情况下，术语“图像”、“像素”和“去噪”都指的是“感知潜在空间”中的相应概念[44]。

给定图像z0，扩散算法逐渐将噪声添加到图像中，并产生有噪声的图像zt，其中t是噪声添加的次数。当t足够大时，图像近似于纯噪声。给定一组条件，包括时间步长t、文本提示ct以及任务特定条件cf，图像扩散算法学习网络cθ来预测添加到噪声图像zt的噪声

其中L是整个扩散模型的总体学习目标。该学习目标可以直接用于微调扩散模型。

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

3.4改进训练

我们讨论了几种改进ControlNets训练的策略，特别是在计算设备非常有限（例如在笔记本电脑上）或非常强大（例如在具有大规模GPU的计算集群上）的极端情况下。在我们的实验中，如果使用这些策略中的任何一种，我们都会在实验环境中提到。

小规模训练

当计算设备有限时，我们发现部分断开ControlNet和稳定扩散之间的连接可以加速收敛。默认情况下，我们将ControlNet连接到“SD中间块”和“SD解码器块1,2,3,4”，如图3所示。我们发现，断开与解码器1、2、3、4的链接，只连接中间块，可以将训练速度提高约1.6倍（在RTX 3070TI笔记本电脑GPU上测试）。当模型显示出结果和条件之间的合理关联时，可以在持续的训练中再次连接那些断开的链接，以便于精确控制。

大规模训练

这里，大规模训练是指强大的计算集群（至少8个Nvidia A100 80G或等效）和大型数据集（至少100万个训练图像对）都可用的情况。这通常适用于数据容易获得的任务，例如Canny检测到的边缘图。在这种情况下，由于过度拟合的风险相对较低，我们可以首先训练ControlNets进行足够多的迭代（通常超过50k步），然后解锁稳定扩散的所有权重，并将整个模型作为一个整体进行联合训练。这将导致一个更具体的问题模型。

3.5实施

我们介绍了具有不同基于图像的条件的ControlNets的几种实现，以各种方式控制大型扩散模型。

Canny边缘

我们使用Canny边缘检测器[5]（具有随机阈值）从互联网上获得3M边缘图像字幕对。该模型使用Nvidia A100 80G进行了600 GPU小时的训练。基本模型是稳定扩散1.5。（另请参见图4。）

Canny边缘（Alter）

我们对上述Canny边缘数据集的图像分辨率进行了排序，并对几个具有1k、10k、50k、500k样本的子集进行了采样。我们使用相同的实验设置来测试数据集规模的效果。（另请参见图22。）

霍夫线

我们使用基于学习的深度霍夫变换[13]来检测Places2[66]中的直线，然后使用BLIP[34]来生成字幕。我们获得了600k个边缘图像字幕对。我们使用上面的Canny模型作为起始检查点，并使用Nvidia A100 80G对模型进行150 GPU小时的训练。（另请参见图5。）

HED边界

我们使用HED边界检测[62]从互联网上获得3M边缘图像字幕对。该模型使用Nvidia A100 80G进行了300 GPU小时的训练。基本模型是稳定扩散1.5。（另请参见图7。）

用户绘制

我们使用HED边界检测[62]和一组强数据增强（随机阈值、随机屏蔽涂鸦的随机百分比、随机形态变换和随机非最大值抑制）的组合，从图像中合成人类涂鸦。我们从互联网上获得了50万对涂鸦图片说明。我们使用上面的Canny模型作为起始检查点，并使用Nvidia A100 80G对模型进行150 GPU小时的训练。请注意，我们还尝试了一种更“类似人类”的合成方法[57]，但该方法比简单的HED慢得多，而且我们没有注意到明显的改进。（另请参见图6。）

人体姿势（Openpifpaf）

我们使用基于学习的姿势估计方法[27]，使用一个简单的规则从互联网上“找到”人类：一张有人类的图像必须检测到至少30%的全身关键点。我们获得了80k个姿势的图片字幕对。注意，我们直接使用具有人体骨骼的可视化姿势图像作为训练条件。该模型在Nvidia RTX 3090TI上使用400 GPU小时进行训练。基本模型是稳定扩散2.1。（另请参见图8。）

人体姿势（开放式）

我们使用基于学习的姿态估计方法[6]，在上述Openpifpaf设置中使用相同的规则从互联网上找到人类。我们获得了20万个姿势图片字幕对。注意，我们直接使用具有人体骨骼的可视化姿势图像作为训练条件。该模型使用Nvidia A100 80G进行了300 GPU小时的训练。其他设置与上述Openpifpaf相同。（另请参见图9。）

语义分割（COCO）

由BLIP[34]标题的COCO Stuff数据集[4]。我们获得了164K个分割图像字幕对。该模型在Nvidia RTX 3090TI上使用400 GPU小时进行训练。基本模型是稳定扩散1.5。（另请参见图12。）

语义分割（ADE20K）

由BLIP[34]标题的ADE20K数据集[67]。我们获得了164K个分割图像字幕对。该模型在Nvidia A100 80G上使用200 GPU小时进行训练。基本模型是稳定扩散1.5。（另请参见图11。）

深度（大型）

我们使用Midas[30]从互联网上获得3M深度的图像字幕对。该模型使用Nvidia A100 80G进行了500 GPU小时的训练。基本模型是稳定扩散1.5。（另见图23、24、25。）

深度（小规模）

我们对上述深度数据集的图像分辨率进行排序，以对200k对的子集进行采样。该集合用于实验训练模型所需的最小数据集大小。（另请参见图14。）

法线贴图

DIODE数据集[56]由BLIP[34]加上标题。我们获得了25452个正常的图片字幕对。该模型在Nvidia A100 80G上使用100 GPU小时进行训练。基本模型是稳定扩散1.5。（另请参见图13。）

法线贴图（扩展）

我们使用Midas[30]来计算深度图，然后从距离执行法线，以实现“粗略”法线图。我们使用上面的Normal模型作为起始检查点，并使用Nvidia A100 80G对模型进行200 GPU小时的训练。（另见图23、24、25。）

卡通线条画

我们使用卡通线条图提取方法[61]从互联网上的卡通插图中提取线条图。通过对受欢迎的卡通图像进行排序，我们获得了排名前1M的线性卡通字幕对。该模型使用Nvidia A100 80G进行了300 GPU小时的训练。基本模型是外复扩散（一个由稳定扩散发展而来的有趣的社区变化模型[36]）。（另请参见图15。）

4实验

4.1实验设置

本文的所有结果都是在CFG等级为9.0的情况下获得的。采样器为DDIM。默认情况下，我们使用20个步骤。我们使用三种类型的提示来测试模型：

（1） 无提示：我们使用空字符串“”作为提示。

（2） 默认提示：由于稳定扩散本质上是用提示训练的，空字符串可能是模型的意外输入，如果不提供提示，SD倾向于生成随机纹理贴图。一个更好的设置是使用无意义的提示，如“一张图片”、“一张漂亮的图片”和“一张专业的图片”等。在我们的设置中，我们使用“专业、详细、高质量的图片”作为默认提示。

（3） 自动提示：为了测试最先进的全自动管道的最大化质量，我们还尝试使用自动图像字幕方法（例如，BLIP[34]），使用“默认提示”模式获得的结果生成提示。我们再次使用生成的提示进行扩散。

（4） 用户提示：用户给出提示。