Adding Conditional Control to Text-to-Image Diffusion Models

date

Mar 10, 2023

Last edited time

Mar 27, 2023 08:27 AM

status

Published

slug

Adding_Conditional_Control_to_Text-to-Image_Diffusion_Models

摘要

我们提出了一个神经网络结构 controlnet 网络来控制预训练的大扩散模型以支持额外的输入条件。controlnet 网络以端到端的方式学习任务特定条件，即使在训练数据集较小 (< 50k) 的情况下，学习也是鲁棒的。此外，训练一个 controlnet 网络和微调一个扩散模型一样快，并且模型可以在个人设备上训练。或者，如果有强大的计算集群，模型可以扩展到大量的 (百万到十亿) 数据。我们报告了像 Stable Diffusion 这样的大型扩散模型可以用 ControlNets 来增强，以启用边缘图、分割图、关键点等条件输入。这可能会丰富控制大型扩散模型的方法，进而方便相关应用。

介绍

提出的问题

基于文本提示的控制来生成图像是否满足我们的需求?

在图像处理中，考虑到许多长期存在的具有明确问题表述的任务，能否应用这些大模型来方便这些特定的任务?

我们应该构建什么样的框架来处理广泛的问题条件和用户控制?

在特定的任务中，大型模型可以保留从数十亿张图像中获得的优势和能力?

文章的发现

任务特定领域的可用数据规模并不总是像一般图文领域那样大。像姿势理解这些领域的数据集就很小。这就要求在针对特定问题训练大型模型时，采用稳健的神经网络训练方法以避免过拟合并保持泛化能力。

当使用数据驱动的解决方案处理图像处理任务时，大型计算集群并不总是可用的，因为耗时耗内存。所以需要利用预训练的权重，以及微调策略或迁移学习。

各种图像处理问题具有不同形式的问题定义、用户控制或图像标注。考虑到一些特定的任务，如 depth-to-image、pose-to-person 等，这些问题本质上需要将原始输入解释为对象级或场景级的理解，这使得 hand-crafted 的过程性方法变得不可行。为了在许多任务中实现学习到的解决方案，端到端的学习是必不可少的。

根据以上的发现，文章提出了控制大型图像扩散模型 (像 Stable Diffusion) 去学习任务特定输入条件的端到端神经网络架构 ControlNet。

模型介绍

ControlNet 将大型扩散模型的权重克隆为 “trainale copy” 和 “locked copy”：locked copy 保留了从数十亿张图像中学习到的网络能力，而 trainale copy 则在任务特定的数据集上训练学习条件控制。

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

令人十分欣喜的是，在深度到图像等任务中，在个人计算机 (1 个 Nvidia RTX 3090TI) 上训练 ControlNets 可以取得与在拥有 TB 级 GPU 内存和数千 GPU 小时的大型计算集群上训练的商业模型相当的结果。

网络概览

解释：

整体还是 U-Net 框架

注意整个网络是一起训练的，右边绿色部分的初始参数是从左边 copy 过来的，左边部分的参数不能训练，右边部分的可以训练。

使用 ControlNet 创建稳定扩散的 12 个编码块和 1 个中间块的 trainable copy，对应右边绿色部分的。

将条件图像变成 64 × 64 的特征空间方法：使用一个 4 × 4 kernel，2 × 2 步长的 4 个卷积层，卷积层通过 ReLU 激活，通道数分别为 16、32、64、128，以高斯权重初始化，与整个模型联合训练) 的网络，公式描述如下：，是转换后的特征图。该网络将 512 × 512 图像条件转换为 64 × 64 特征图。

方法

所作的工作

初始网络块

如图 a 所示：没有使用 ControlNet 之前的一个网络块，用公式描述就是：。

其中是一个特征图，代表一个神经网络块，这个网络块将转换0成另一个特征图。

使用了 ControlNet 的网络块

如图 b 所示：使用了 ControlNet 之后的一个网络块，它所作的修改如下：

它从原来的网络中复制了一份参数，这份参数是可以训练的（论文中叫 trainable copy，它的参数为），而原来网络的参数是不能训练的（论文中叫 locked copy），也就是图中加锁的部分。这样做的目的是为了防止数据量小的时候过拟合并且保留之前模型的生成质量。

增加了 zero-convolution，它的特点是权重和偏差都初始化为 0。

经过修改后，新的网络块用公式描述就是：

其中 Z 表示 zero convolution 操作，和是两个 zero convolution 的参数。

zero-convolution 的训练介绍

在第一个训练阶段：

因此我们可以推出：

因为优化之前卷积部分的参数都是 0，根据卷积操作的运算过程可知卷积的结果都是 0。

由此可以看到：在第一个训练步骤中，神经网络块的 trainable copy 和 locked copy 的所有输入和输出都与 ControlNet 网络不存在时一致。也就是说在开始优化参数之前，ControlNet 不会对深度神经特征造成任何影响。任何进一步的优化都将变得像微调 (相比于从头训练这些层) 一样快。

梯度下降过程：

一个卷积层的公式描述如下：

是输入的特征图，是偏差，是权重。

它的梯度为：（注意它的权重和偏差初始值都是）

从这个公式可以看出，zero-convolution 可以导致特征项上的梯度变为零，但是权重和偏置的梯度不受影响。在本文中，为输入数据或从数据集中采样的条件向量，肯定不为，则在第一次梯度下降迭代中将权重优化为非零矩阵。梯度下降过程如下：

然后特征上的梯度变为：

此时特征上的梯度不为，由此神经网络开始学习。