diffusion model(十二)： StableCascade技术小结

	info
paper	https://arxiv.org/abs/2306.00637
github	https://github.com/Stability-AI/StableCascade/tree/master
个人blog位置	http://myhz0606.com/article/stablecascade
http://stability.AI blog	https://stability.ai/news/introducing-stable-cascade

前置知识：

latent diffusion model（stable diffusion）知乎 个人博客 paper

classifier-free guided: 知乎 个人博客 paper

背景

Diffusion model(DM)的提出将图像生成任务推到了前所未有的高度。随着应用人数的增多，人们对文生图模型有了更高的期望：生成质量更高、生成分辨率更大、推理速度更快、生成的图片与文本的alignment更强等。

更大的分辨率意味着更大的算力需求。现在的常用做法是先用DM先生成一个分辨率相对较低的图片，再结合图片超分的方法使其更加清晰。目前所用的超分算法一般采用GAN的架构，就效果而言，超分生成的图片细节不如DM直接在高分辨率生成的图片效果。

提高DM的生成分辨率，意味着更多的训练成本。虽然Latent diffusion model(LDM)[1]使得diffusion过程从像素空间转化为隐空间，大大减少训练和推理成本。但LDM方法在不明显降低图片生成质量的条件下，只能支持到4-16x的压缩比率，训练成本依旧很大。如stable diffusion(SD) v1.4(512分辨率，隐空间下采样factor为8) 的训练用了150000小时的A100 PCIe 40GB[

2]，SD v2.1(768分辨率，隐空间下采样factor为8)的训练用了200000小时的A100 PCIe 40GB[3]。

能不能在采用更大采样率（意味着更低的训练和推理成本）的同时不降低生成图片质量呢？ 这就是**Wuerstchen （StableCascade）**想要解决的问题。

方法大意

区别于二阶段(two-stage) SD，StableCascade是一个三阶段（three-stage）的生成方法。如下所示。可以看到StableCascade他有两个不同压缩比的text-condition LDM。下文将详细介绍StableCascade的推理过程和训练过程是如何利用这两个不同压缩比的text-condition LDM，以及为何这样的架构能够加速推理的同时保证图片质量。

	SD	StableCascade
stageA	VAGAN	VAGAN
stageB	text-condition LDM(压缩比8x)	text-condition LDM（低压缩比4x) Unet架构
stageC	无	text-condition LDM （高压缩比42x）。非Unet架构，堆叠了16个ConvNext block, 内部没有downsample

推理过程

推理过程的pipeline如下图所示，从stageC开始，stageB依赖stageC的输出，stageA依赖stageB的输出。为了方便表述，不妨记stageC的text-condition LDM为$f_{{\theta }_c}^C$,stageB的text-condition LDM为$f_{{\theta }_b}^B$,stageA的VAGAN的decode为$\mathcal{D}$

STEP1: 通过stageC预测Semantic Compressor的输出$C_{sc}$。stageC的推理步长${\tau }_C=60$

$$C_{sc}^t=f_{{\theta }_c}^C({ϵ}_t^c,C_{\mathrm{text}},t)\text{\hspace{1em}(1)}$$

式中$C_{sc}^t$为时间步为$t$时的输出；$f_{{\theta }_c}^C$ 为stageC的text-condition LDM ；${ϵ}_t^c$ 为时间步为$t$时预测的噪声, ${ϵ}_0^c\sim \mathcal{N}(0,\mathbf{I})$；$C_{\mathrm{text}}$为prompt的text embedding。

Semantic Compressor可以理解为是一个特征提取器，如resnet、vit…。其作用是给定一张图片，提取其特征，维度变化为：${\mathbb{R}}^{3\times 768\times 768}\stackrel{\mathrm{Semantic}\mathrm{Compressor}+1x1conv}{⟶}{\mathbb{R}}^{16\times 24\times 24}$ 。论文中，作者用EfficientV2(S)作为Semantic Compressor。

STEP2: 通过StageB来预测图片的隐表征$z$。stageB的推理步长${\tau }_B=12$

$$z_t=f_{{\theta }_b}^B({ϵ}_t^b,C_{\mathrm{text}},C_{sc},t)$$

式中$z_t$为时间步为$t$时预测的隐表征；$f_{{\theta }_b}^B$ 为stageB的text-condition LDM ；${ϵ}_t^b$ 为时间步为$t$时预测的噪声,其中 ${ϵ}_0^b\sim \mathcal{N}(0,\mathbf{I})$；$C_{sc}$为stageC的输出，为了增强模型的robust，会对$C_{sc}$添加一些noise使其non-perfect；$C_{\mathrm{text}}$为prompt的text embedding （与stageC中的是一致的）。

这里可能会有一个疑问：$\ast \ast C_{sc}^t$如何融入？**在StableCascade源码实现和Wuerstchen paper不同。paper中是将$C_{sc}$ reshape后和$C_{\mathrm{text}}$组合一起作为condition融入。而StableCascade源码的实现是先对$C_{sc}$进行transform，随后resize为模型内部特征的shape进行相加。源码位置：https://github.com/Stability-AI/StableCascade/blob/master/modules/stage_b.py#L227C22-L227C35

STEP3: 通过StageA的decode $\mathcal{D}$来将隐空间变换到像素空间

$$x=\mathcal{D}(z)$$

式中$x$为生成的图片；$\mathcal{D}$为VAGAN的decode；$z$为stageB的输出。

就此StableCascade的推理过程结束。回顾一下：**Wuerstchen （StableCascade）**想要解决的问题是希望用一个更大采样率（意味着更低的训练和推理成本）的同时不降低生成图片质量呢。 从上面对推理pipeline描述中我们发现StableCascade相比LDM还多了一个stage，为什么还能降低推理成本提升生成质量呢？
个人理解：StableCascade仅在高压缩率的StageC（42x压缩率）采用较大的推理步长(${\tau }_C=60$)，而在低压缩率的stageB用了较低的推理步长${\tau }_B=12$，从而实现加速。之所以能这样做，是由于stageC给stageB所提供的图片特征先验$C_{sc}$能让stageB仅用较少步长就能实现较好的效果。

训练过程

StableCascade需要训练3个模型

1. StageA的VAGAN 。由一个encoder $\mathcal{E}$和一个decoder $\mathcal{D}$ 构成。参数量18M
2. stageB的text-condition LDM$f_{{\theta }_b}^B$,参数量1B，架构为Unet
3. stageC的text-condition LDM$f_{{\theta }_c}^C$，参数量为1B；架构：堆叠了16个ConvNext block,没有进行下采样（压缩比已经很大了，作者发现再下采样会恶化生成质量）。

前面提到的Semantic Compressor采用的是开源预训练的EfficientV2,只额外训练一个1x1卷积调整维度。

训练数据格式为图文对数据，来源于improved-aesthetic LAION-5B。数据格式为: $\{({\mathrm{img}}^{(i)},{\mathrm{text}}^{(i)})|i=1,2,\cdots N\}$
stageA VQGAN的训练目标可参考[4]

StageB和stageC的训练目标和LDM [1]中的一致。需要注意的是stageB的预测目标是VQGAN encoder$\mathcal{E}$的输出$z=\mathcal{E}(x)$，StageC的预测目标是Semantic Compressor+1x1 conv的feature map $C_{sc}$。

结果

下图展示了StableCascade 在不同batch下和SD2.1和SDXL的推理速度对比结果。

下图展示了不同架构生成质量对比结果。

生成图片展示

小结

StableCascade的核心在于新增了一个图片隐特征的先验$C_{sc}$，从而带来推理和训练的加速。在训练阶段这个先验来源于pre-training model。在推理阶段这个先验来源于stageC对这个先验信息的估计。

reference

[1] High-Resolution Image Synthesis with Latent Diffusion Models [BLOG](http://myhz0606.com/article/ldm)

[2] https://huggingface.co/CompVis/stable-diffusion-v-1-4-original

[3] https://huggingface.co/stabilityai/stable-diffusion-2-1

[4] (VQGAN) Taming transformers for high-resolution image synthesis