Adding Conditional Control to Text-to-Image Diffusion Models

摘要

• ControlNet，控制预训练大型扩散模型，以支持额外的输入条件。ControlNet以端到端方式学习特定任务的条件输入，即使训练数据集很小(< 50k)，效果也很健壮。
• 此外，训练ControlNet的速度与微调扩散模型一样快，而且该模型可以在个人设备上训练。或者，如果强大的计算集群可用，该模型可以扩展到大量(数百万到数十亿)的数据。
• 论文表明，像Stable Diffusion这样的大型扩散模型可以用ControlNet来增强，以支持像边缘map、分割map、关键点等条件输入。这将丰富大型扩散模型的控制方法，并进一步促进相关应用：
• https://github.com/lllyasviel/ControlNet

背景

• 随大型文本到图像模型的出现，生成视觉效果生动的图像只需要用户输入一个简短的描述性提示文本（prompt）。对此，可能有几个问题:
• 这种基于文本提示的生成，是否满足了我们的需求?
• 在图像处理中，许多长期存在的任务都有明确的问题解决公式，这些大型模型能否应用于促进这些特定的任务？应该构建什么样的框架可以统一来处理各种各样的问题条件和用户控制?
• 在特定的任务中，大型模型能否保持从数十亿张图像中获得的优势和能力?
• 为回答这些问题，本文研究了各种图像处理应用，并有三个发现：
• 首先，特定任务域的可用数据规模并不总是像一般图像-文本域那样大。许多具体任务(如对象形状、姿态理解等)的最大数据集大小往往在100k以下，即比LAION5B 小50000倍。这将需要鲁棒的神经网络训练方法，以避免过拟合，并在大模型针对特定问题进行训练时保持泛化能力。
• 其次，大型计算集群并不总是可用的。这使得快速训练方法对于在可接受的时间和内存空间内(例如在个人设备上)优化大型模型以执行特定任务非常重要。这将进一步需要利用预先训练的权重，以及微调策略或迁移学习。
• 第三，各种图像处理问题具有不同形式的问题定义、用户控制或图像标注。在解决这些问题时，虽然图像扩散算法可以手工做一些调整，例如，约束去噪过程、编辑多头注意力激活等，但考虑到一些特定的任务，如深度图转图像、姿势关键点到人像等，需要将原始输入解释为对象级或场景级的理解，要在许多任务中实现通用的学习解决方案，端到端学习是必不可少的。
• 端到端神经网络ControlNet，控制大型图像扩散模型来学习特定于任务的输入条件。ControlNet将大型扩散模型的权重克隆为一个“可训练副本”和一个“锁定副本”:锁定副本保留了从数十亿张图像中学习到的网络能力，而可训练副本则在特定任务的数据集上进行训练，以学习条件控制。
• 可训练和锁定神经网络块与一种称为“零卷积”的独特类型的卷积层连接，其中卷积权值以学习的方式从零逐步转变到优化后的参数。由于保留了权值，这样一来在不同规模的数据集上也具有鲁棒性。由于零卷积不会给深度特征添加新的噪声，因此与从头开始训练新层相比，训练速度与微调扩散模型一样快。
• 用不同条件的各种数据集训练几个ControlNet，例如，Canny边缘、Hough线、用户涂鸦、人体关键点、分割图、深度图等。还用小数据集(样本小于50k甚至1k)和大数据集(数百万样本)对ControlNets进行了实验。表明在一些任务中，如深度到图像，在个人计算机(一台Nvidia RTX 3090TI)上训练ControlNets，与在大型计算集群上训练的商业模型具有竞争力的结果。

HyperNetwork和神经网络结构

• HyperNetwork起源于一种语言处理方法，用于训练一个小的循环神经网络来影响一个大的神经网络的权重。HyperNetwork也被应用于生成对抗网络等图像生成任务。ControlNet和HyperNetwork在影响神经网络行为的方式上有相似之处。ControlNet使用一种特殊类型的卷积层，称为“零卷积”。早期的神经网络研究广泛讨论了网络权值的初始化问题，包括高斯分布初始化权值的合理性以及用零初始化权值可能带来的风险。在ProGAN和StyleGAN以及Noise2Noise等工作中也讨论了初始卷积权值的操作。

扩散概率模型

• 近来，扩散概率模型取得了巨大的研究进展，例如去噪扩散概率模型(DDPM)、去噪扩散隐式模型(DDIM)和基于分数的扩散模型等。研究人员在处理高分辨率图像时往往会考虑节省计算能力的策略，或直接使用基于金字塔或多阶段的方法，常见比如使用U-net作为神经网络结构。为了降低训练扩散模型所需的计算能力，基于图像潜在特征的思想，提出了潜在扩散模型(latent diffusion model, LDM)方法，并进一步扩展到Stable Diffusion。

文本到图像扩散模型

• 扩散模型可以应用于文本到图像生成任务，以达到最先进的图像生成效果。这通常是通过使用CLIP等预训练好的语言模型将文本输入编码为潜在向量来实现的。例如，Glide是一个文本引导的扩散模型，同时支持图像生成和编辑。Disco Diffusion是clip引导实现、用于处理文本提示。Stable Diffusion实现潜在扩散。Imagen是一个文本到图像的结构，它不使用潜在图像，而是使用金字塔结构直接扩散像素。

预训练扩散模型的控制生成

• 最先进的图像扩散模型，是由文本到图像的方法主导的，所以增强对扩散模型控制的最直接的方法，通常是文本引导。这种类型的控制也可以通过操纵CLIP特性来实现。图像扩散过程本身可以提供一些功能来实现颜色级别的细节变化(社区称之为img2img)。图像扩散算法当然支持inpainting作为控制结果的重要方式。而文本inversion和DreamBooth方法，基于具有相同主题或目标对象的小样本图像数据集，去自定义控制生成结果中的内容。

图像转换

• Pix2Pix提出了图像到图像转换的概念，早期的方法以条件生成神经网络为主。在Vision Transformers(ViTs)得到普及后，使用自回归方法也获得了不错的结果。目前最强大的图像到图像转换方法里，Taming Transformer属于vision transformer类，具有生成图像和执行图像到图像转换的功能。Palette是一个统一的基于扩散的图像到图像转换框架。PITI是一种基于扩散的图像到图像转换方法，它利用大规模的预训练来提高生成结果的质量。在草图引导扩散等特定领域，像论文《Sketch-guided text-to-image diffusion models》则属于基于优化的方法，用于控制扩散过程。

方法

• ControlNet通过对神经网络模块的输入条件进行操作，从而进一步控制整个神经网络的整体行为。其中，“神经网络模块”是指将一组神经层作为一个常用单元组合在一起构建神经网络，如“resnet”块、“convn-bn-relu”块、多头注意力块、transformer模块等。
• 以二维特征为例，给定特征x∈h×w×c，{h, w, c}为高度、宽度和通道，神经网络模块F(·;Θ)和一组参数Θ将x转换为另一个特征y：

• 如果将所有参数锁定在Θ中，然后将其克隆为可训练的副本Θc。复制的Θc使用外部条件向量c进行训练。在本文中，称原始参数和新参数为“锁定副本”和“可训练副本”。制作这样的副本而不是直接训练原始权重的动机是：避免数据集较小时的过拟合，并保持从数十亿张图像中学习到的大型模型的能力。
• 神经网络模块由一种称为“零卷积”的独特类型的卷积层连接，即1×1卷积层，权重和偏差都用零初始化。将零卷积运算表示为Z(·;·)，使用参数{Θz1，Θz2}的两个实例组成ControlNet结构：

• 其中yc成为该神经网络模块的输出。因为零卷积层的权值和偏差都初始化为零，所以在第一个训练步骤中，有：

• 这可以转换为yc = y
• 表明，在第一个训练步骤中，神经网络块的可训练副本和锁定副本的所有输入和输出都与它们的状态一致，就像ControlNet不存在一样。换句话说，当一个ControlNet应用于一些神经网络块时，在进行任何优化之前，它不会对深层神经特征造成任何影响。任何神经网络块的能力、功能和结果质量都得到了完美的保留，任何进一步的优化都将变得像微调一样快(与从零开始训练这些层相比)。
• 下面简单地推导零卷积层的梯度计算。考虑权值W和偏差B的1 × 1卷积层，在任意空间位置p和通道索引i处，给定输入特征i∈h×w×c，正向通过可写成

• 零卷积有W = 0和B = 0(在优化之前)，对于Ip,i非零的任何地方，梯度变为：

• 可以看到，尽管零卷积可以导致特征项i的梯度变为零，但权值和偏差的梯度不受影响。在第一次梯度下降迭代中，只要特征I是非零，权重W就会被优化为非零矩阵。值得注意的是，在例子中，特征项是从数据集中采样的输入数据或条件向量，这自然地确保了I不为零。例如，考虑一个具有总体损失函数L和学习率βlr的经典梯度下降，如果“外部”梯度∂L/∂Z(I;{W, B})不为零，有：

• 其中W∗是一阶梯度下降后的权值;是Hadamard乘积。在这一步之后，可得到：

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

扩散模型中的ControlNet

• 以Stable Diffusion为例，介绍利用ControlNet对具有任务特定条件的大型扩散模型进行控制的方法。
• Stable Diffusion是一个大型的文本到图像扩散模型，训练了数十亿张图像。该模型本质上是一个带有编码器、中间块和跳接解码器的U-net。编码器和解码器都有12个块，完整模型有25个块(包括中间块)。其中8个块是下采样或上采样卷积层，17个块是主块，每个主块包含4个resnet层和2个Vision transformer (vit)。每个Vit都包含若干交叉注意和/或自注意机制。文本采用OpenAI CLIP编码，扩散时间步长采用位置编码。
• Stable Diffusion使用类似于VQ-GAN的预处理方法，将整个512 × 512图像数据集转换为较小的64 × 64“潜在图像”进行训练。这需要ControlNets将基于图像的条件转换为64 × 64的特征空间，以匹配卷积大小。使用一个由4个4核和2 × 2步长卷积层组成的微型网络E(·)(由ReLU激活，通道为16,32,64,128，初始化为高斯权值，与完整模型联合训练)将图像空间条件ci编码为特征映射：

• 其中cf为转换后的特征图。该网络将512 × 512的图像条件转换为64 × 64的特征图。如图所示，使用ControlNet来控制U-net的各个层。

• ControlNet的方式在计算上是高效的:由于原始权重是锁定的，因此训练不需要在原始编码器上进行梯度计算。这样可以加快训练速度，节省GPU内存，因为可以避免原来模型一半的梯度计算。
• 使用ControlNet训练一个扩散模型只需要在每次训练迭代中增加约23%的GPU内存和34%的时间(在单个Nvidia A100 PCIE 40G上测试)。具体来说，使用ControlNet创建了Stable Diffusion的12个编码块和1个中间块的可训练副本。12个块有4种分辨率(64 × 64，32 × 32，16 × 16，8 × 8)，每种3个块。输出被添加到U-net的12个跳过连接和1个中间块。SD是一个典型的U-net结构，这种ControlNet结构也可用于其他扩散模型。

训练

• 图像扩散模型学习逐步去噪来生成样本。去噪可以发生在像素空间或从训练数据编码的“潜在”空间。Stable Diffusion算法使用潜在空间作为训练域。
• 给定一个图像z0，扩散算法逐步向图像中添加噪声，并产生一个噪声图像zt，其中t是噪声添加的次数。当t足够大时，图像近似于纯噪声。给定一组条件，包括时间步长t，文本提示ct，以及特定任务条件cf，图像扩散算法学习一个网络θ来预测噪声加入到噪声图像zt中

• 其中L为整个扩散模型的总体学习目标。该学习目标可直接用于精细调整扩散模型。在训练过程中，随机将50%的文本提示替换为空字符串。这有助于提升ControlNet从输入条件（例如Canny边缘图或人类涂鸦等）中识别语义内容的能力。这主要是因为当prompt对于SD模型不可见时，编码器倾向于从条件输入中学习更多语义，以替代prompt。
• 接下来讨论几种改进ControlNet训练的策略，比如计算设备非常有限(例如，在笔记本电脑上）时。
• 小规模训练。当计算设备有限时，切断部分ControlNet和Stable Diffusion之间的连接可以加速收敛。默认情况下，将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步)，然后解锁Stable Diffusion的所有权重，共同训练整个模型。

实验

设置中，采样器为DDIM。默认使用20个步骤，三种类型的prompt来测试模型:

(1)No prompt:使用空字符串“”作为prompt。

(2)Default prompt:由于Stable diffusion本质上是用prompt训练的，空字符串可能是模型的一个意外输入，如果没有提供prompt，SD倾向于生成随机纹理。更好的设置是使用无意义的prompt，如“一张图片”、“一张漂亮的图片”、“一张专业的图片”等。在设置中，使用“专业、详细、高质量的图像”作为默认prompt。

(3)Automatic prompt:为了测试全自动流程最好的效果，还使用自动图像caption方法(如BLIP)。使用“Default prompt”模式获得的结果生成prompt，再次使用生成的prompt进行扩散生成。

(4)User prompt:用户给出prompt。

实现

• 提出几种基于不同图像条件输入的方式，以控制生成。
• Canny边缘图。使用Canny边缘检测，用随机阈值从互联网上获得3M的边缘-图像-caption数据对。该模型使用Nvidia A100 80G进行600个gpu小时的训练。使用的基础模型是Stable Diffusion 1.5。此外，对上述Canny边缘数据集按图像分辨率进行排序，并采样1k、10k、50k、500k样本的子集。使用相同的实验设置来测试数据集规模的影响。

• Hough线图。使用基于深度学习的Hough变换方法对Places2数据集检测线图，并且使用BLIP生成caption。由此，得到了600k对的边缘-图像-caption。使用前面的Canny模型作为初始化checkpoint，并使用Nvidia A100 80G用150个gpu小时训练。

• HED边界检测。从互联网上获取3M的边缘-图像-caption的数据对，使用Nvidia A100 80G进行300个gpu小时训练。基础模型是Stable Diffusion 1.5。

• 用户草图。结合HED边界检测和一组强数据增强(随机阈值、随机屏蔽涂鸦、随机形态变换、随机非极大抑制)，从图像中合成人类涂鸦。从互联网上获得了50万对的涂鸦图像-caption数据对。使用前面的Canny模型作为初始化checkpoint，并使用Nvidia A100 80G用150个gpu小时训练。

• 人体姿势(Openpifpaf)。使用基于学习的姿势估计方法来从互联网上“找到”人类，使用一个简单的规则:含有人类的图像必须有至少30%的全身关键点被检测到。获得了80k个姿势-图像-caption数据对。请注意，直接使用人体骨骼的可视化姿态图像作为训练条件。在Nvidia RTX 3090TI上使用400个gpu小时进行训练。基础模型是Stable Diffusion 1.5。

• 人体姿势(Openpose)。使用基于学习的姿态估计方法在上面的Openpifpaf设置中使用相同的规则从互联网上找到人类。获得了200k个姿势-图像-caption数据对。直接使用人体骨骼的可视化姿态图像作为训练条件。使用Nvidia A100 80G进行300个gpu小时的训练，其他设置与上面的Openpifpaf相同。

• 语义分割(COCO)。由BLIP对COCO-stuff数据集生成caption，获得了164K分割-图像-caption数据对。该模型在Nvidia RTX 3090TI上使用400个gpu小时进行训练。基础模型是Stable Diffusion 1.5。

• 语义分割(ADE20K)。由BLIP对ADE20K数据集生成caption，获得了164K分割-图像-caption数据对。该模型在Nvidia A100 80G上使用200个gpu小时进行训练。基础模型是Stable Diffusion 1.5。

• 深度(大尺度)。使用Midas从互联网上获取3M的深度-图像-caption数据对。使用Nvidia A100 80G进行500 gpu小时的训练。基础模型是Stable Diffusion 1.5。
• 深度(小尺度)。对上述深度数据集的图像分辨率进行排序，采样200k的子集，用于训练模型所需的最小数据集大小。

• 卡通线图提取方法。从网络上的卡通插图中提取线描。通过对受欢迎程度的卡通图像进行排序，得到了前1M的素描-卡通-caption数据对，使用Nvidia A100 80G进行300个gpu小时的训练。基础模型是Waifu Diffusion。