一、介绍
随着文本到图像的扩散模型出现,我们可以通过输入文本提示来创建视觉上令人惊叹的图像。然而,仅通过文本提示在对图像的空间组成提高的控制方面受到限制,例如无法精准表达复杂的空间布局、物体姿势、形状。
那么ControlNet的想法就是附加图像到输入条件中,来实现更细粒度的空间控制,例如通过输入边缘特征、骨骼图、分割图、深度图、涂鸦草图等来控制stable diffusion。
简单来说,ControlNet,是一种端到段神经网络架构,可学习大型预训练文本到图像扩散模型(stable diffusion)的条件控制。它通过锁定SD的参数并制作其编码层的可训练副本(trainable copy)来保留大型模型的质量和功能。该架构将大型预训练模型视为学习各种条件控制的强大支柱。trainable copy与原始的锁定模型进行零卷积层连接,权重初始化为零,以便它们在训练过程中逐渐增长。这种架构确保在训练开始时有害噪声不会添加到扩散模型的深层特征中,并保护可训练副本中的大规模预训练主干不被此类噪声损坏。
ControlNet的独特之处在于它解决了空间一致性问题。ControlNet允许对空间结构和几何方面进行明确的控制,同时保留了从文本标题中获得的语义控制。
空间一致性:是指在图像生成或处理中,生成的图像应该在空间上保持一致性,即相邻的像素或区域之间应该有合理的关系和连续性。例如,在生成一个人脸的图像时,眼睛、鼻子和嘴巴的相对位置应该是一致的,不能随机分布。
二、ControlNet
1.原理
ControlNet 将附加条件注入神经网络(neural network)的块中(如图 2所示)。
图2:neural network block通过x作为输入,输出特征图y,如(a)所示,这好比扩散模型当中UNet的编码层(下采样层)。为了将ControlNet添加到这样的块中,我们锁定原始块并创建可训练副本,并使用零卷积层(权重和偏差都初始化为0的1x1卷积)将它们连接在一起。其中c是添加到网络中的条件向量,如(b)所示。零卷积层连接在trainable copy network block的两端,然后跟locked copy以残差形式并列在一起作为输出。
在文章中,使用network block来指代一组神经层,这些神经层组合在一起形成神经网络的基础单元,例如resnet块、conv-bn-relu块、multi-head attention块、transformer块。
假设\(\textcolor{blue}{\mathcal{F}(\cdot;\Theta)}\)是一个预训练好的神经块、参数为\(\textcolor{blue}{\Theta}\)。它将输入特征图\(\textcolor{blue}{x}\)转换为另一个特征图\(\textcolor{blue}{y=\mathcal{F}(x;\Theta)}\),其中\(\textcolor{blue}{x \in R^{h \times w \times c}}\)。\(\textcolor{blue}{\{h,w,c\}}\)分别作为图中的高度、宽度和通道数。
为了将ControlNet添加到预训练的network block中,我们锁定(冻结)原始块的参数\(\textcolor{blue}{\Theta}\),同时将该块克隆到具有参数\(\textcolor{blue}{\Theta_c}\)的可训练副本。可训练副本采用外部条件向量c作为输入。这里将原始参数和新参数分别称为locked copy 和trainable copy。比起直接微调原模型参数 \(\textcolor{blue}{\Theta}\),微调\(\textcolor{blue}{\Theta_c}\)可以避免在数据集较小时出现过拟合,并且保持大模型的先验能力不变。
traniable copy与locked copy通过零卷积层进行连接,其表示为\(\textcolor{blue}{\mathcal{Z}(\cdot;\cdot)}\),这是一个1x1的卷积层,权重和偏差都初始化为0。为了构建ControlNet,使用两个零卷积层,参数分别为\(\textcolor{blue}{\Theta_{z1}}\)和\(\textcolor{blue}{\Theta_{z2}}\),ControlNet的输出\(\textcolor{blue}{y_c}\)可以通过下式进行计算:
在第一个训练步骤中,由于零卷积层的权重和偏置参数都初始化为零,因此等式(2)中的\(\textcolor{blue}{\mathcal{Z}(\cdot;\cdot)}\)项均计算为零,则有
因此,当训练开始时,有害噪声就无法影响可训练副本中神经网络层的隐藏状态。此外,由于\(\textcolor{blue}{\mathcal{Z}(c;\Theta_{z1})=0}\) 并且可训练副本还接收输入图像 \(\textcolor{blue}{x}\),因此可训练副本功能齐全,并保留大型预训练模型的功能,使其能够作为进一步学习的强大骨干。零卷积通过消除初始训练步骤中作为梯度的随机噪声来保护该主干网络。
在加入zero-convolution后,ControlNet的所有参数都具有非零梯度,因而是可学习的。假设zero-convolution的权重为\(\textcolor{blue}{W}\),偏置为\(\textcolor{blue}{B}\),在任意空间位置\(\textcolor{blue}{p}\)和channel索引\(\textcolor{blue}{i}\),给定输入特征为\(\textcolor{blue}{I \in R^{h \times w \times c}}\),则前向传播可以写成: \[ \huge \mathcal{Z}(I;\{W,B\})_{p,i}=B_i + \sum^c_jI_{p,i}W_{i,j} \] 由于zero-convolution优化前有(\(\textcolor{blue}{W=0}\), \(\textcolor{blue}{B=0}\)),对于\(\textcolor{blue}{I_{p,i}}\)非零的任何地方,梯度变为: \[ \huge\left\{ \begin{aligned} &\frac{\partial\mathcal{Z}(I;\{W,B\})_{p,i}}{\partial B_i}=1\\ &\frac{\partial\mathcal{Z}(I;\{W,B\})_{p,i}}{\partial I_{p,i}}=\sum^c_jW_{i,j}=0\\ &\frac{\partial\mathcal{Z}(I;\{W,B\})_{p,i}}{\partial W_{i,j}}=I_{p,i} \neq 0 \end{aligned} \right. \] 可以看到,虽然zero-convolution可以使特征输入\(\textcolor{blue}{I}\)的梯度变为零,但权重和偏置的梯度不受影响。只要特征输入\(\textcolor{blue}{I}\)不为零,权重\(\textcolor{blue}{W}\)就会在第一次梯度下降迭代中被优化为非零矩阵。
输入特征\(\textcolor{blue}{I}\)是输入数据或者条件向量,这自然确保了非零。考虑常规梯度下降算法,如果学习速率和前面链式梯度不为零,那么根据梯度下降算法可得: \[ \huge W^* = W - \beta_{lr} \cdot \frac{\partial \mathcal{L}}{\partial \mathcal{Z}(I;\{W,B\})} \odot \frac{\partial \mathcal{Z}(I;\{W,B\})}{\partial W} \neq 0 \] 这说明经过一个梯度下降步骤后权重\(\textcolor{blue}{W^*}\)将不为零,接着可以推导出: \[ \huge \frac{\partial\mathcal{Z}(I;\{W,B\})_{p,i}}{\partial I_{p,i}} = \sum_j^c W^*_{i,j} \neq 0 \] 这说明特征的梯度\(\textcolor{blue}{I}\)也非零,网络可以进一步学习了。
总的来说,zero-convolution是一种独的连接层,通过训练,它不会在初始阶段显著改变网络初始值,而是在迭代中渐进式拟合数据集。
2.将ControlNet用于文本到图像扩散
ControlNet 可以在大型预训练扩散模型添加条件控制。stable diffusion本质上是一个 U-Net,带有encoder、middle block和skip-connected decoder。并使用autoencoder,将512x512图像转换为64x64的Latent以进行快速和稳定的训练。
🌧Stable diffusion:如在图3a中所示,encoder和decoder分别包含 12 个块,完整模型包含 25 个块,包括middle block。在 25 个块中,8 个块是下采样或上采样卷积层,而其他 17 个块是主块,每个块包含 4 个 resnet 层和 2 个视觉transformer (ViT)。每个 ViT 都包含多个cross-attention 和self-attention。"SD Encoder Block A"包含4个resnet层和2个ViT,而“×3”表示该块重复3次。文本提示(Text Prompt \(\textcolor{blue}{c_t}\))使用 CLIP 文本编码器进行编码,扩散时间步长(Time \(\textcolor{blue}{t}\) )使用位置编码的时间编码器进行编码。
🌈ControlNet: 结构应用于U-Net的每个编码器级别(如图 3b)。ControlNet 共创建 12 个encoder block和 1 个middle block组成的可训练副本。 12个encoder block有4种分辨率(64×64、32×32、16×16、8×8),每一个复制3次。输出通过对应的zero-convolution连接到U-Net的 12 个skip-connected decoder block和 1 个middle block。
总的来说,ControlNet 在训练上算是高效的,因为原始网络权重被锁定,而且只对encoder block和middle block增加了trainable copy进行训练,相当于只进行原模型一半的梯度计算。
为了将 ControlNet 添加到stable diffusion中,首先将每个输入条件图像(例如边缘、姿态、深度等)从 512 × 512 输入大小转换为与stable diffusion大小匹配的 64 × 64 特征空间向量。特别地,使用一个由四个卷积层组成的微型网络 \(\textcolor{blue}{\mathcal{E}(·)}\),具有 4 × 4 内核(kernels)和 2 × 2 步幅(strides)(由 ReLU 激活,分别使用 16、32、64、128 个通道,用高斯权重初始化并联合训练与完整模型)将图像空间条件 \(\textcolor{blue}{c_i}\) 编码为特征空间条件向量\(\textcolor{blue}{c_f}\),条件向量\(\textcolor{blue}{c_f}\)被传递到 ControlNet:
3.训练
给定输入图像 \(\textcolor{blue}{z_0}\),图像扩散算法逐渐向图像添加噪声并生成噪声图像\(\textcolor{blue}{z_t}\),其中 t 表示添加噪声的次数。给定一组条件,包括时间步\(\textcolor{blue}{t}\)、文本提示 \(\textcolor{blue}{c_t}\) 以及特定于任务的条件\(\textcolor{blue}{c_f}\),图像扩散算法学习网络 \(\textcolor{blue}{ε_θ}\) 来预测添加到噪声图像\(\textcolor{blue}{z_t}\)的噪声:
其中\(\textcolor{blue}{\mathcal{L}}\)是整个扩散模型的总体学习目标。该学习目标直接用于通过 ControlNet 微调扩散模型。
在训练过程中,随机将50%的文本提示\(\textcolor{blue}{c_t}\)替换为空字符串。这种方法提高了 ControlNet 直接识别输入调节图像中的语义(例如边缘、姿势、深度等)的能力,以替代提示。由于零卷积不会给网络增加噪声,因此模型应该始终能够预测高质量的图像。可观察到该模型并不是逐渐学习控制条件,而是突然成功地遵循输入条件图像;通常优化步骤少于 10K。如图4所示,作者称之为“突然收敛现象”。
4.推理
作者表示我们可以通过多种方式进一步控制 ControlNet 的额外条件去影响去噪扩散过程。
无分类器引导分辨率加权(Classifier-free guidance resolution weighting):stable diffusion依赖于一种称为无分类器指导 (CFG) 的技术来生成高质量图像。 CFG 的公式为 \(\textcolor{blue}{ε_{prd} = ε_{uc} + β_{cfg}(ε_c − ε_{uc})}\),其中 \(\textcolor{blue}{ε_{prd}}\)、\(\textcolor{blue}{ε_{uc}}\)、\(\textcolor{blue}{ε_{c}}\)、\(\textcolor{blue}{β_{cfg}}\) 分别是模型的最终输出、无条件输出、条件输出和用户指定的权重。当通过 ControlNet 添加调节图像时,可以将其添加到\(\textcolor{blue}{ε_{uc}}\)和\(\textcolor{blue}{ε_{c}}\),或仅添加到\(\textcolor{blue}{ε_{c}}\)。在具有挑战性的情况下,例如,当没有给出提示时,将其添加到\(\textcolor{blue}{ε_{uc}}\)和\(\textcolor{blue}{ε_{c}}\)将完全删除 CFG 指导(图 5b);仅使用\(\textcolor{blue}{ε_{c}}\)将使指导非常强大(图 5c)。我们的解决方案是首先将调节图像添加到\(\textcolor{blue}{ε_{c}}\),然后根据每个块的分辨率\(\textcolor{blue}{w_i = 64/h_i}\),为Stable Diffusion和ControlNet之间的每个连接乘以一个权重\(\textcolor{blue}{w_i}\) ,其中\(\textcolor{blue}{h_i}\)是第i个块的大小,例如\(\textcolor{blue}{h_1 = 8, h_2 = 16, ..., h_{13} = 64}\)。通过降低CFG引导强度,我们可以得到如图5d所示的结果,我们称之为CFG分辨率加权。
组成多个 ControlNet:要将多个条件图像(例如 Canny 边缘和姿势)应用于稳定扩散的单个实例,我们可以直接将相应 ControlNet 的输出添加到稳定扩散模型中(图 6)。这种组合不需要额外的加权或线性插值。
三、实验
文章使用稳定扩散的 ControlNet 来测试各种条件,包括 Canny Edge、深度图、法线图、M-LSD 线、HED 软边缘、ADE20K 分割 、Openpose和用户草图。
1.定性结果
图 7 显示了在没有提示的情况下在各种条件下的结果,其中 ControlNet 稳健地解释了不同输入条件图像中的内容语义。
2.消融实验
作者通过以下方式研究 ControlNet 的替代结构:1️⃣用高斯权重初始化的标准卷积层替换零卷积(图 8b),以及2️⃣用一个卷积层替换每个块的可训练副本,我们称之为 ControlNet-lite(图 8c)。
作者提供了 4 种提示设置来测试现实世界用户可能的行为:
🔴无提示(No prompt);
🔵提示不足(Insufficient prompt),没有完全覆盖调节图像中的对象,例如本文默认的提示“高质量、详细、专业的图像”;
🟢 改变条件图像语义的冲突提示(Conflicting prompt);
🟡完善的提示(Perfect prompt),描述必要的内容语义。
例如“a beautiful house”。图 8a 显示 ControlNet 在所有 4 项设置中均成功。轻量级的 ControlNet-lite(图 8c)不够强大,无法解释调节图像,并且在不足和无提示的情况下会失败。当替换零卷积时,ControlNet 的性能下降到与 ControlNet-lite 大致相同,这表明可训练副本的预训练主干在微调过程中被破坏(图 8b)。
3.定量评价
对比实验:作者对 20 个看不见的手绘草图(sketches)进行采样,然后将每个草图分配给 5 种方法:1️⃣PITI 的草图模型、2️⃣具有默认边缘引导尺度(β = 1.6)的草图引导扩散(SGD)、3️⃣具有相对较高的边缘引导尺度(β=3.2)的SGD,以及4️⃣ControlNet-lite和5️⃣ControlNet。作者邀请 12 位用户根据“显示图像的质量”和“草图的保真度”分别对这 20 组 5 个结果进行排名。由此获得了 100 个结果质量排名和 100 个条件保真度排名。使用平均人类排名 (AHR) 作为偏好指标,用户按 1 到 5 的范围对每个结果进行排名(越低越差)。平均排名如表1所示:
条件重建和 FID 评分:作者使用ADE20K 的测试集来评估条件保真度(conditioning fidelity)。最先进的分割方法 OneFormer 在真实数据集上实现了 0.58 的交并(IoU)。作者使用不同的方法来生成具有 ADE20K 分割的图像,然后应用 One-Former 再次检测分割以计算重建的 IoU(表 2)。
此外,我们使用 Frechet Inception Distance (FID) 来测量使用不同分割条件方法随机生成的 512×512 图像集的分布距离,以及文本图像 CLIP 分数 和 CLIP 美学分数 在表 3 中。
4.与以前的方法进行比较
图 9 展示了baseline和(稳定扩散 + ControlNet)的视觉比较。具体来说,作者展示了 PITI 、Sketch-Guided Diffusion 和 Taming Transformers 的结果。 (请注意,PITI 的骨干是 OpenAI GLIDE ,它们具有不同的视觉质量和性能。)我们可以观察到 ControlNet 可以稳健地处理不同的条件图像并实现清晰的结果。
5.其他
训练数据集大小的影响:在图 10 中展示了 ControlNet 训练的鲁棒性。训练不会因有限的 1k 图像而崩溃,并且允许模型生成可识别的狮子。当提供更多数据时,学习是可扩展的。
解释内容的能力:图11展示了 ControlNet 从输入调节图像捕获语义的能力。
传输到社区模型:由于ControlNets不会改变预训练SD模型的网络拓扑,因此它可以直接应用于稳定扩散社区中的各种模型,例如Comic Diffusion 和Protogen 3.4,如图12所示。
四、总结
ControlNet 是一种神经网络结构,可学习大型预训练文本到图像扩散模型的条件控制。它重用源模型的大规模预训练层来构建深度且强大的编码器来学习特定条件。原始模型和可训练副本通过“零卷积”层连接,消除训练期间的有害噪声。大量的实验验证了ControlNet可以在单个或多个条件下、有或没有提示的情况下有效地控制稳定扩散。不同条件数据集的结果表明,ControlNet 结构可能适用于更广泛的条件,并促进相关应用。