ControlNet-有条件图文生成论文阅读

论文： 《Adding Conditional Control to Text-to-Image Diffusion Models》
github： https://github.com/lllyasviel/ControlNet

摘要

ControlNet控制大的预训练扩散模型支持额外输入条件，ControlNet端到端的方式学习特定任务条件，即使训练集小（<50k），学习也比较鲁棒。作者基于Stable Diffusion训练ControlNets，可支持边缘map、分割map、关键点为条件输入；这丰富了控制扩散模型方法，为相关应用提供便利。

算法：

ControlNet

ControlNet通过控制神经网络模块输入条件进一步控制整个神经网络输出。如式1，神经网络模块F通过参数Θ将特征图x变换为另一个特征图y，即为图2a过程；

ControlNet应用于任意神经网络模块过程如图2b所示，$c$为额外条件向量，$\Theta$为locked copy参数，${\Theta }_c$为$\Theta$的clone，且可训练，没有直接训练是为了防止过拟合，${\Theta }_{z1},{\Theta }_{z2}$为ControlNet中两个零卷积$Z$参数，零卷积表示初始化weight及bias为0的1*1卷积，该过程如式2，

由于两个零卷积$Z$参数初始化为0，因此训练第一步时$y_c=y$，如式3,


零卷积参数以一种可学习方式从零开始优化；

ControlNet in Image Diffusion Model

Stable Diffusion使用数十亿图片训练的文本图像生成模型，本质上为包含编码器、中间层、skip-connected解码器的U-net，整个模型有25个block，编码器和解码器各自有12个block，所有block中，8个为上采样或下采样卷积层，17个为主要block，每个包括4个resnet层、2个ViT。文本由CLIP进行编码，扩散时间步长使用位置编码。
Stable Diffusion类似VQ-GAN为了稳定训练过程，将512 * 512图片映射到64 * 64隐空间，因此需要ControlNet将基于图片的条件转换到64 * 64特征空间，该过程通过4个kernel=4，steide=2卷积实现，如式9.

如图3，ControlNet控制U-net每个层级，因为原始权重固定，因此计算高效；ControlNet使用与SD（Stable Diffusion）相同的12个编码block及1个middle block，其中12个block有4个分辨率(64 × 64, 32 × 32, 16 × 16, 8 × 8)，每个分辨率3个block；输出部分增加12个skip-connections及1个middle block至U-net

Training

扩散模型学习图像逐渐去噪，生成样本；
对于图片$z_0$，通过扩散算法逐渐增加噪声生成噪声图$z_t$，其中$t$为噪声添加次数，给出step t、text prompt $c_t$及特定任务条件$c_f$，扩散算法通过网络${ϵ}_{\theta }$预测增加到噪声图$z_t$上的噪声，如式10，L为整体损失函数：

训练过程中，作者随机替换50% test prompt $c_t$为空，使得SD模型编码器可以从输入的控制图$c_f$学习更多语义信息。

Improved Training

Small-Scale Training：当资源有限时，作者发现断开ControlNet与SD Decoder Block 1,2,3,4训练速度提升1.6倍，当模型输出结果与条件相关联时，可重新连接Decoder Block 1,2,3,4进行训练；

Large-Scale Training：当训练资源充足，数据量大时，模型过拟合风险低，可先充分训练ControlNet，而后解锁SD模型权重，与ControlNet作为一个整体模型进行联合训练；

实验

实验设置，作者使用三种prompt：
1、No prompt；“”
2、Default prompt：无意义prompt
3、Automatic prompt：BLIP生成
4、User prompt：用户输入

Canny edges

Hough lines

Human scribbles

HED boundary map

Openpifpaf pose

Openpose

ADE20K segmentation map

COCO-Stuff segmentation map

DIODE normal map

Depth-to-Image

cartoon line drawings

限制

如图28，当输入分割容易引起歧义时，难以生成合理内容。

结论

ControlNet呈现效果比较惊艳，在SD模型基础上更进一步，支持多种条件控制文本至图像生成。