ControlNet:为文本到图像扩散模型添加条件控制
ControlNet 是一种神经网络架构,可以通过空间定位、特定于任务的图像条件来增强大型预训练文本到图像扩散模型。
ControlNet 将其他条件注入神经网络的块中,如下图。
神经块将特征图 x 作为输入并输出另一个特征图 y,如 (a) 所示。
为了将 ControlNet 添加到这样的块中,我们锁定原始块并创建一个可训练的副本,并使用零卷积层将它们连接在一起,即 1 × 1 卷积,权重和偏差都初始化为零。这里 c 是我们希望添加到网络中的条件向量,如 (b) 所示。
为了将 ControlNet 添加到这样一个预先训练的神经块中,我们锁定(冻结)原始块的参数 Θ,同时将块克隆到参数为 Θc 的可训练副本(图 2b)。可训练副本将外部条件向量 c 作为输入。当这种结构应用于 Stable Diffusion 等大型模型时,锁定的参数会保留使用数十亿张图像训练的生产就绪模型,而可训练副本则重用此类大规模预训练模型来建立深度、稳健和强大的主干来处理各种输入条件。
可训练副本连接到具有零卷积层的锁定模型,表示为 Z(·; ·)。具体来说,Z(·; ·) 是一个 1 × 1 卷积层,权重和偏置都初始化为零。为了构建 ControlNet,我们使用两个零卷积实例,参数分别为 Θz1 和 Θz2。然后,完整的 ControlNet 会计算
这里 yc 是 ControlNet 块的输出。在第一个训练步骤中,由于零卷积层的权重和偏置参数都初始化为零,因此方程 (2) 中的 Z(·; ·) 项的计算结果为零,并且
这样,当训练开始时,有害噪声就不会影响可训练副本中神经网络层的隐藏状态。此外,由于 Z(c; Θz1) = 0 并且可训练副本也接收输入图像 x,因此可训练副本功能齐全,并保留了大型预训练模型的功能,使其能够作为进一步学习的强大支柱。零卷积通过在初始训练步骤中消除作为梯度的随机噪声来保护此主干。我们详细介绍了补充材料中零卷积的梯度计算。
ControlNet for Text-to-Image Diffusion
我们以稳定扩散 [72] 为例,展示 ControlNet 如何向大型预训练扩散模型添加条件控制。稳定扩散本质上是一个 U-Net [73],带有一个编码器、一个中间块和一个跳跃连接的解码器。编码器和解码器都包含 12 个模块,完整模型包含 25 个模块,包括中间的模块。在这 25 个块中,8 个块是下采样或上采样卷积层,而其他 17 个块是主块,每个块包含 4 个 resnet 层和 2 个 Vision Transformer (ViT)。每个 ViT 都包含几种交叉注意力和自我注意机制。例如,在图 3a 中,“SD 编码器块 A”包含 4 个 resnet 层和 2 个 ViT,而“×3”表示该块重复 3 次。文本提示使用CLIP 文本编码器 [66] 和扩散时间步长使用位置编码的时间编码器进行编码。
图 3:Stable Diffusion 的 U-net 架构与编码器模块和中间模块上的 ControlNet 连接。锁定的灰色块显示了 Stable Diffusion V1.5(或 V2.1,因为它们使用相同的 U-net 架构)的结构。添加可训练的蓝色块和白色零卷积层以构建 ControlNet。
ControlNet 结构应用于 U-net 的每个编码器级别(图 3b)。具体而言,我们使用 ControlNet 创建 12 个编码块和 1 个 Stable Diffusion 中间块的可训练副本。12 个编码块有 4 种分辨率(64 × 64、32 × 32、16 × 16、8 × 8),每个编码块复制 3 次。输出被添加到 U-net 的 12 个跳过连接和 1 个中间块。由于 Stable Diffusion 是一种典型的 U-net 结构,因此这种 ControlNet 架构可能适用于其他模型。我们连接 ControlNet 的方式在计算上是高效的 — 由于锁定的复制参数是冻结的,因此不需要在最初锁定的编码器中进行梯度计算进行微调。这种方法可以加快训练速度并节省 GPU 内存。在单个 NVIDIA A100 PCIE 40GB 上测试后,与在没有 ControlNet 的情况下优化 Stable Diffusion 相比,使用 ControlNet 优化稳定扩散每次训练迭代只需要增加约 23% 的 GPU 内存和 34%d的时间。
图像扩散模型学习逐步对图像进行去噪,并从训练域生成样本。去噪过程可以发生在像素空间或由训练数据编码的潜在空间中。Stable Diffusion 使用潜在图像作为训练域,因为在这个空间中工作已被证明可以稳定训练过程 [72]。具体来说,Stable Diffusion 使用类似于 VQ-GAN [19] 的预处理方法将 512 × 512 像素空间的图像转换为更小的 64 × 64 个潜伏图像。为了将 ControlNet 添加到 Stable Diffusion 中,我们首先将每个输入调节图像(例如,边缘、姿态、深度等)从 512 × 512 的输入大小转换为与 Stable Diffusion 大小匹配的 64 × 64 特征空间向量。特别是,我们使用一个由四个卷积层组成的微小网络 E(·),具有 4 × 4 个内核和 2 × 2 个步幅(由 ReLU 激活,分别使用 16、32、64、128 个通道,用高斯权重初始化并与完整模型联合训练)将图像空间条件 ci 编码为特征空间条件向量 cf 为
条件向量 cf 被传递到 ControlNet 中。
Training
给定输入图像 z0,图像扩散算法逐渐向图像中添加噪声并产生噪声图像 zt,其中 t 表示添加噪声的次数。给定一组条件,包括时间步长 t、文本提示 ct 以及特定于任务的条件 cf,图像扩散算法学习网络 εθ 来预测添加到噪声图像 zt 中的噪声
其中 L 是整个扩散模型的总体学习目标。此学习目标直接用于使用 ControlNet 微调扩散模型。
在训练过程中,我们随机将 50% 的文本提示 ct 替换为空字符串。这种方法提高了 ControlNet 直接识别输入条件图像中的语义(例如,边缘、姿势、深度等)以替代提示的能力。
在训练过程中,由于零卷积不会给网络增加噪声,因此模型应该始终能够预测高质量的图像。我们观察到,该模型并没有逐渐学习控制条件,而是突然成功地遵循了输入条件图像;通常在不到 10K 的优化步骤中。如图 4 所示,我们称之为 “突然收敛现象”。
图 4:突然收敛现象。由于卷积为零,ControlNet 在整个训练过程中始终预测高质量的图像。在训练过程的某个步骤(例如,以粗体标记的 6133 个步骤)时,模型突然学会了遵循输入条件。
Inference
我们可以通过多种方式进一步控制 ControlNet 的额外条件如何影响去噪扩散过程。
无分类器指导分辨率加权。稳定扩散依赖于一种称为无分类器引导 (CFG) [29] 的技术来生成高质量的图像。CFG 公式化为 εprd = εuc + βcfg(εc − εuc),其中 εprd、εuc、εc、βcfg 分别是模型的最终输出、无条件输出、条件输出和用户指定的权重。通过 ControlNet 添加调节图像时,可以将其添加到 εuc 和 εc,也可以只添加到 εc。在具有挑战性的情况下,例如,当没有给出提示时,将其添加到 εuc 和 εc 将完全消除 CFG 指导(图 5b);仅使用 εc 将使导向非常强(图 5c)。我们的解决方案是首先将条件图像添加到 εc,然后根据每个块的分辨率 wi = 64/hi,将权重 wi 乘以 Stable Diffusion 和 ControlNet 之间的每个连接,其中 hi 是第 i 个块的大小,例如,h1 = 8,h2 = 16,…,h13 = 64。通过减小 CFG 引导强度 ,我们可以获得图 5d 所示的结果,我们称之为 CFG 分辨率加权。
编写多个 ControlNet。要将多个条件反射图像(例如,Canny 边缘和姿势)应用于 Stable Diffusion 的单个实例,我们可以直接将相应 ControlNet 的输出添加到 Stable Diffusion 模型中(图 6)。这种组合不需要额外的加权或线性插值。
Experiments Results
Qualitative Results
图 1 显示了几个提示设置中生成的图像。图 7 显示了我们在各种条件下的结果,没有提示,其中 ControlNet 稳健地解释各种输入条件图像中的内容语义。
图 1:使用学习条件控制稳定扩散。ControlNet 允许用户添加 Canny 边缘(顶部)、人体姿势(底部)等条件,以控制大型预训练扩散模型的图像生成。默认结果使用提示 “a high-quality, detailed, and professional image”。用户可以选择提供提示,例如 “chef in kitchen”。
图 7:在没有提示的情况下在各种条件下控制稳定扩散。第一行是输入条件,而所有其他行都是输出。我们使用空字符串作为输入提示。所有模型都使用通用域数据进行训练。模型必须识别输入条件 images 中的语义内容才能生成图像。
Ablative Study
我们通过以下方式研究 ControlNet 的替代结构:(1) 将零卷积替换为使用高斯权重初始化的标准卷积层,以及 (2) 将每个块的可训练副本替换为单个卷积层,我们称之为 ControlNet-lite。有关这些烧蚀结构的完整详细信息,另请参阅补充材料。我们提出了 4 个提示设置来测试真实用户可能的行为:(1) 无提示;(2) 提示不足,不能完全覆盖条件化图像中的对象,例如,本文的默认提示“高质量、详细和专业的图像”;(3) 改变条件反射图像语义的冲突提示;(4) 描述必要内容语义的完美提示,例如,“A nice house”。图 8a 显示 ControlNet 在所有 4 个设置表现优异。轻量级的 ControlNet-lite(图 8c)不够强大,无法解释条件反射图像,并且在不足和无提示的情况下会失败。当替换零卷积时,ControlNet 的性能下降到与 ControlNet-lite 大致相同,这表明可训练副本的预训练主干在微调过程中被破坏(图 8b)。
图 8:草图条件和不同提示设置下不同架构的烧蚀研究。对于每个设置,我们显示随机批次的 6 个样本,没有 cherry-picking。图像为 512 × 512,放大后查看效果最佳。左侧的绿色 “conv” 块是使用高斯权重初始化的标准卷积层。