【论文阅读】ControlNet、文章作者 github 上的 discussions

文章目录

Introduction
Method
- ControlNet
- ControlNet for Text-to-Image Diffusion
- Training
- Inference
Experiments
- 消融实验
- 定量分析
在作者 github 上的一些讨论
- 消融实验更进一步的探索
- Precomputed ControlNet 加快模型推理
- 迁移控制能力到其他 SD1.X 模型上
- 其他

Introduction

提出ControlNet，通过引入该结构微调预训练文生图扩散模型，可以给模型增加空间定位条件.
在Stable Diffusion上使用ControlNet微调，使模型能接受 Canny edges, Hough lines, user scribbles, human key points, segmentation maps, shape normals, depths, cartoon line drawings 图像作为输入条件.
消融实验、定量分析、对比 baseline.

Method

ControlNet

考虑一个预训练好的神经网络 $\mathcal{F}(·;\Theta)$ 表示训练好的神经网络块，它的内部结构可以包括 resnet, conv-bn-relu, muti-head att, transfomer 等. 输入 $x\in\mathbb{R}^{h\times w\times c}$ ，将其转换到 $y$ ，也即

$\large y=\mathcal{F}(x;\Theta)$

使用ControlNet微调神经网络 $\mathcal{F}(·;\Theta)$ ，首先复制 $\mathcal{F}(·;\Theta)$ 的结构和参数，参数命名为 $\Theta_{c}$ ，同时冻结 $\Theta$ . 然后在复制结构的前和后分别引入zero convolution，也即核大小为 $1\times1$ 、初始参数为 $0$ 的卷积层，分别用 $\mathcal{Z}(·;\Theta_{z1})$ 和 $\mathcal{Z}(·;\Theta_{z2})$ 表示. 最后，将 $c$ 作为微调时的条件，将其整合到模型的前向计算中，具体表示为

$\large y_c=\mathcal{F}(x;\Theta)+\mathcal{Z}(\mathcal{F}(x+\mathcal{Z}(c;\Theta_{z1});\Theta_c);\Theta_{z2})$

模型结构如下所示：
暂时无法在飞书文档外展示此内容

在训练的第一步中，zero convolution的参数都为 $0$ ，因此模型输出和未加入ControlNet的输出一样，这样做有助于在刚开始训练时保护微调结构的 backbone，使其免受随机噪声的污染.

ControlNet for Text-to-Image Diffusion

众所周知，Stable Diffusion训练时的网络有这么几个部分构成：

FrozenCLIPEmbedder是一个预训练的 text encoder，将 prompt 嵌入成条件向量，一般情况下参数冻结.
AutoencoderKL是一个预训练的 image encoder，将图像从像素空间转换到隐空间，降低扩散过程中图像向量的尺寸，一般情况下参数冻结.
UNet，主要需要训练的部分，模拟隐空间上图像在数据分布和高斯分布之间转换的过程. 结构上主要包含：
- 若干 encoder 块，主要由 resnet, transformer, avg_pool 组成，用于逐层提取特征.
  - resnet 块融合图像隐向量和扩散时间步的嵌入向量
  - transformer 块融合图像隐向量和 prompt 条件向量
- 一个 middle 块，由 resnet 和 transformer 组成
- 若干 decoder 块，主要由 resnet, transformer, interpolate 组成，用于融合深层特征和浅层特征.

将ControlNet应用于Stable Diffusion做微调，也即应用于其中UNet的 decoder 部分，使这部分网络能进一步融合作为条件的图像。用 $t$ 表示时间步， $c_t$ 表示 prompt 条件， $c_f$ 表示条件图像在隐空间上的表示，修改后的UNet结构为

在这里插入图片描述

Training

用 $z_0$ 表示原始图像的隐向量，经过时间步 $t$ 后加噪的图像表示为 $z_t$ ，应用了ControlNet的UNet表示为 $\epsilon_{\theta}$ ，训练时的损失函数可以表示为

$\large \mathcal{L}=\mathbb{E}_{ z_0,t,c_t,c_f,\epsilon\in\mathcal{N}(0,I)}\left[||\epsilon-\epsilon_{\theta}(z_t,t,c_t,c_f)||_2^2\right]$

在实际训练过程中，作者随机将 $50\%$ 的 prompt 置为空字符串，这种做法能使ControlNet学习到图像条件的语义信息. 由于zero convolution不会引入额外的噪声，因此在训练过程中整个Stable Diffusion模型仍然能生成高质量的图片. 基于这一特性，作者观察到，微调时模型并非逐渐学习到图像条件，而是在训练步数低于 $10\mathrm{K}$ 时的某一步开始突然遵从图像条件. 作者称这其为 “sudden convergence phenomenon”
[图片]

Inference

Stable Diffusion使用CFG控制条件强弱，令 $\epsilon_{uc}$ 表示无 prompt 条件的模型输出， $\epsilon_{c}$ 表示有 prompt 条件的模型输出，超参数 $\beta_{cfg}$ 表示 prompt 条件的强弱，模型最终的输出 $\epsilon_{prd}$ 可以表示为

$\large \epsilon_{\mathrm{prd}}=\epsilon_{\mathrm{uc}}+\beta_{\mathrm{cfg}}(\epsilon_{\mathrm{c}}-\epsilon_{\mathrm{uc}})$

在没有 prompt 条件的极端情况下，如果抽取完深层特征的图像条件同时加到 $\epsilon_{uc}$ 和 $\epsilon_{c}$ 上，这会使CFG完全失去控制条件强弱的作用；如果只加到 $\epsilon_{c}$ 上，又会使控制条件对输出图像的影响过大. 因此，作者提出一种叫做Classifier-free guidance resolution weighting(CFG-RW)的方法. 具体做法，把图像条件加到 $\epsilon_{c}$ 上，在ControlNet每一层输出加回UNet前乘系数 $w_i$ ( $64/h_i$ ， $h_i$ 为第 $i$ 个 decoder 块的尺寸). 下图分别展示了该讨论各种情况下的输出图像：
[图片]

有了上述方法之后，结合不同类别的图像条件，也只需要对应相加即可.

Experiments

消融实验

探索ControlNet其他可能结构

将zero convolution换成随机初始化的卷积层
只使用一个卷积层作为ControlNet

定量分析

作者使用 ADE20K 作为测试集，在 OneFormer 上做语义分割，对比不同方法重构图像和原图像的 IoU .
[图片]

之后，作者评估了不同模型的 FID、CLIP score、CLIP aesthetic score.

[图片]

下图展示了不同模型实际生成的图片
[图片]

在作者 github 上的一些讨论

消融实验更进一步的探索

discussion 链接
将ControlNet简化为ControlNet-lite和ControlNet-mlp两种模型：
[图片]

作者从根据一张房子的图片做了简单地涂鸦风格处理，作为控制条件

在精心设计 prompt 的情况下，发现原版模型和改后的两种模型输出的图像效果都不错.

Professional high-quality wide-angle digital art of a house designed by frank lloyd wright. A delightful winter scene. photorealistic, epic fantasy, dramatic lighting, cinematic, extremely high detail, cinematic lighting, trending on artstation, cgsociety, realistic rendering of Unreal Engine 5, 8k, 4k, HQ, wallpaper

（分别为`ControlNet`、`ControlNet-lite`、`ControlNet-mlp`的输出）

但是当 prompt 为空时，两种改版都很拉胯.

（分别为`ControlNet`、`ControlNet-lite`、`ControlNet-mlp`的输出）

一方面，这样的对比说明更深的 encoder 结构确实拥有更强的识别能力，所以如果你的目标是训练稳健的ControlNet投入到生产环境，这样的识别能力是很重要的. 反之，如果用来做解决特定问题的研究或者训练集足够简单，那可以考虑轻量化的方案.
另一方面，这也解释了ControlNet接受 prompt 条件和时间步输入是重要的，因为这么做可以让使用者仍然能靠 prompt 条件调整模型的输出.

Precomputed ControlNet 加快模型推理

discussion 链接
主要 idea 如下图所示：
[图片]

这样做可以提前计算好ControlNet中每个块的输出，在推理时直接加到原模型的UNet上.
作者观察到这样训练的模型生成的图像更假，并且更不稳健，以失败告终.
评论中有人提到可以尝试使用 NAS (neural architecture search) 探索更好的模型结构，以降低 GPU 消耗.