论文题目:Adding Conditional Control to Text-to-Image Diffusion Models论文链接:https://arxiv.org/abs/2302.05543
本文提出一种神经网络结构 ControlNet,用来控制预训练的大型扩散模型,并使其支持额外的输入条件。它以端到端的方式学习特定任务条件,即使在训练数据集很小时,学习也是鲁棒的。此外,训练 ControlNet 与微调扩散模型一样快,并且模型可以在个人设备上进行训练。如果有强大的计算集群,该模型可以扩展到大量数据。此外,Stable diffusion 等大型扩散模型可以用 ControlNet 增强,以实现边缘图、分割图、关键点等条件输入。
随着大型文本到图像模型的出现,生成一个视觉上具有吸引力的图像也许需要仅仅一个由用户输入的短的描述性prompt。在输入一些文本并获取图像后,我们自然会想到一些问题:① 这个基于prompt的控制方式是否满足我们的需求?② 例如在图像处理中,考虑许多长期存在的具有明确问题定义的任务,这些大型模型是否可以应用于促进这些特定任务?③ 我们应该构建什么样的框架来处理广泛的问题条件和用户控制?④ 在特定的任务中,大模型能否保留从数十亿图像中获得的优势和功能?
为了回答这些问题,我们调查了各种图像处理应用程序并得出了三个发现。
第一,在特定任务领域中可用的数据规模并不总是像图像-文本领域中那样大。许多特定问题(例如对象形状/法线、姿态理解等)的最大数据集大小通常不到100k,即比LAION5B小5×10^4倍。这就需要鲁棒的神经网络训练方法以避免过拟合,并在训练大型模型处理特定问题时保持泛化能力。 第二,在使用数据驱动的解决方案处理图像处理任务时,并非总是可用大型计算集群。因此,快速的训练方法对于在可接受的时间和内存空间内优化大型模型以执行特定任务非常重要(例如,在个人设备上)。这还需要利用预训练权重,以及微调策略或迁移学习。
第三,各种图像处理问题有不同形式的问题定义、用户控制或图像注释。在解决这些问题时,虽然图像扩散算法可以以“程序化”的方式进行调节,例如约束去噪过程、编辑多头注意激活等,但这些手工制作的规则的行为本质上是由人类指令规定的。考虑到一些特定的任务,如深度到图像、姿态到人等,这些问题本质上需要从原始输入到对象级或场景级理解的解释,这使得手工制作的程序化方法不那么可行。为了在许多任务中实现学习到的解决方案,端到端学习是必不可少的。
图1:用Canny边缘映射控制Stable Diffusion。Canny边缘映射是输入,当我们生成右边的图像时,不使用源图像。输出是通过一个默认提示“高质量、详细、专业的图像”来实现的。此提示在本文中用作默认提示,不提及任何关于图像内容和对象名称的内容。本文中的大部分图像都是高分辨率的图像,在放大时观看效果最好。 本文提出ControlNet,一种端到端神经网络架构,它控制大型图像扩散模型(如Stable Diffusion)来学习特定任务的输入条件。ControlNet将一个大型扩散模型的权重克隆为一个“可训练副本”和一个“锁定副本”:锁定副本保留了从数十亿张图像中学习到的网络能力,而可训练副本在特定任务的数据集上进行训练,以学习条件控制。可训练的和锁定的神经网络块通过一种独特类型的卷积层相连接,称为“零卷积”,其中卷积权值以学习的方式从零逐渐增长到优化参数。由于保留了进行生产所要准备好的权重,因此在不同规模的数据集上进行的训练是鲁棒的。由于零卷积不会给深度特征添加新的噪声,因此与从头开始训练新的层相比,该训练与微调一个扩散模型(的速度)一样快。
我们用不同条件下的各种数据集来训练几个ControlNets,例如,Canny边缘、霍夫线、用户涂鸦、人体关键点、分割图、形状法线、深度等。我们还用小数据集(样本小于50k甚至1k)和大数据集(数百万个样本)进行了实验。我们还表明,在一些任务中,如从深度到图像,在个人电脑(一台Nvidia RTX 3090TI)上训练ControlNets,可以实现与在具有兆兆字节的GPU内存和数千个GPU小时的大型计算集群上训练的商业模型(相比)具有竞争力的结果。
ControlNet 操纵神经网络块的输入条件,从而进一步控制整个神经网络的整体行为。以 2D 的特征图为例,对于给定的特征图 ,其中 分别为高度、宽度和通道数。神经网络模块 通过一组参数 将 转化为特征图 。
将 中的所有参数锁定,并克隆一个可训练的副本 ,用外部条件 训练副本 。将原始参数和新参数分别称为"锁定副本"和"可训练副本"。这样可以在数据集较小时避免过拟合,并且保持从数十亿张图像中学习到的大型模型的质量。神经网络块由一种称为“零卷积”的独特类型的卷积层连接,即权重和偏置都以零初始化的 卷积层,零卷积操作表示为 。由此,ControlNet 的结构可以表示为
由于零卷积层的权重和偏置都被初始化为零,所以在第一个训练步骤中
在第一个训练步骤中,神经网络块的可训练副本和锁定副本的所有输入和输出都与 ControlNet 不存在时一致。换句话说,当 ControlNet 应用于某些神经网络块时,在进行优化之前,不会对特征造成任何影响。考虑权重为 偏置为 的 卷积层,在空间位置 和通道索引 处,对于给定的输入 ,前向过程可以写作
由于零卷积 且 ,对于任意 不等于零,梯度
尽管零卷积可以导致特征项 上的梯度为零,但权重和偏差的梯度不受影响。只要特征 是非零的,权重 就会在第一次梯度下降迭代中被更新为非零矩阵。例如,考虑具有整体损失函数为 ,学习率为 的梯度下降,如果“外部”梯度 不为零,有
是一次梯度下降更新后的权重。
由此,得到非零梯度,神经网络开始学习。通过这种方式,零卷积逐步从零开始学习到优化参数。
以 Stable Diffusion 为例,介绍了利用 ControlNet 对具有任务特定条件的大型扩散模型进行控制的方法。Stable Diffusion 是一个在数十亿张图像上训练的大型文本到图像扩散模型。该模型本质上是一个 U-net,具有编码器、中间块和解码器。它使用了类似于 VQ-GAN 的预处理方法,将整个 图像数据集转换为更小的 的“潜在图像”进行稳定训练。这需要 ControlNet 将基于图像的条件 转换为 的特征图 ,以匹配卷积大小。
由于原始权重是锁定的,因此在训练时不需要对原始编码器进行梯度计算,可以加快训练速度并节省 GPU 内存。具体来说,使用 ControlNet 创建 12 个编码块和 1 个中间块的可训练副本。输出被连接到 U-net 的 12 个解码块和 1 个中间块中。 设置中,采样器为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的子集,用于训练模型所需的最小数据集大小。