论文链接：High-Resolution Image Synthesis with Latent Diffusion Models
官方实现：CompVis/latent-diffusion、CompVis/stable-diffusion
这一篇文章的内容是 Latent Diffusion Models（LDM），也就是大名鼎鼎的 Stable Diffusion。先前的扩散模型一直面临的比较大的问题是采样空间太大，学习的噪声维度和图像的维度是相同的。当进行高分辨率图像生成时，需要的计算资源会急剧增加，虽然 DDIM 等工作已经对此有所改善，但效果依然有限。Stable Diffusion 的方法非常巧妙，其把扩散过程转换到了低维度的隐空间中，解决了这个问题。
方法介绍
本方法的整体结构如下图所示，主要分为三部分：最左侧的红框对应于感知图像压缩，中间的绿框对应 Latent Diffusion Models，右侧的白框表示生成条件，下面将分别介绍这三个部分。
感知图像压缩
LDM 把图像生成过程从原始的图像像素空间转换到了一个隐空间，具体来说，对于一个维度为 x ∈ R H × W × 3 \mathbf{x}\in\mathbb{R}^{H\times W\times 3} x∈RH×W×3 的 RGB 图像，可以使用一个 encoder E \mathcal{E} E 将其转换为隐变量 z = E ( x ) \mathbf{z}=\mathcal{E}(\mathbf{x}) z=E(x)，也可以用一个 decoder D \mathcal{D} D 将其从隐变量转换回像素空间 x ~ = D ( E ( x ) ) \tilde{\mathbf{x}}=\mathcal{D}(\mathcal{E}(\mathbf{x})) x~=D(E(x))。在转换时会将图像下采样，作者测试了一系列下采样倍数 f ∈ { 1 , 2 , 4 , 8 , 16 , 32 } f\in\{1, 2, 4, 8, 16, 32\} f∈{1,2,4,8,16,32}，发现下采样 4-16 倍的时候可以比较好地权衡效率和质量。
在进行图像压缩时，为了防止压缩后的空间是某个高方差的空间，需要进行正则化。作者使用了两种正则化，第一种是 KL-正则化，也就是将隐变量和标准高斯分布使用一个 KL 惩罚项进行正则化；第二种是 VQ-正则化，也就是使用一个 vector quantization 层进行正则化。
Latent Diffusion Models
实际上 latent diffusion models 和普通的扩散模型没有太大区别，只是因为从像素空间变到了隐空间，所以维度降低了。训练的优化目标也没有太大变化，普通的扩散模型优化目标为：
L
DM
=
E
x
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∣
∣
ϵ
−
ϵ
θ
(
x
t
,
t
)
∣
∣
2
2
]
L_\textrm{DM}=\mathbb{E}_{\mathbf{x},\epsilon\sim\mathcal{N}(0,1),t}\left[||\epsilon-\epsilon_\theta(\mathbf{x}_t,t)||_2^2\right]
LDM=Ex,ϵ∼N(0,1),t[∣∣ϵ−ϵθ(xt,t)∣∣22]
而 Latent Diffusion Models 的优化目标只是套了一层 autoencoder：
L
LDM
=
E
E
(
x
)
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∣
∣
ϵ
−
ϵ
θ
(
x
t
,
t
)
∣
∣
2
2
]
L_\textrm{LDM}=\mathbb{E}_{\textcolor{red}{\mathcal{E}(\mathbf{x})},\epsilon\sim\mathcal{N}(0,1),t}\left[||\epsilon-\epsilon_\theta(\mathbf{x}_t,t)||_2^2\right]
LLDM=EE(x),ϵ∼N(0,1),t[∣∣ϵ−ϵθ(xt,t)∣∣22]
在采样时，首先从隐空间随机采样噪声，在去噪后再用 decoder 转换到像素空间即可。
条件生成
为了进行条件生成，需要学习 ϵ θ ( x t , t , y ) \epsilon_\theta(\mathbf{x}_t,t,y) ϵθ(xt,t,y)，这里使用的方法是在去噪网络中加入 cross attention 层，条件通过交叉注意力注入。在计算注意力时， z \mathbf{z} z 为 Query、 y y y 为 Key 和 Value，具体的内容已经在 Classifier-Free Guidance 的文章中介绍过了，对具体细节感兴趣的读者可以去看一下。
代码解读
Stable Diffusion 有两套主流的代码实现，第一种是 CompVis 的官方实现，第二种是 huggingface 的实现。这里的代码解读都以文生图任务为例。
CompVis 的实现
这个实现的代码比较分散，层次结构不太好梳理，不过可以照着配置文件看各部分都在哪里。这个配置文件有点类似 openmmlab 的那套框架的写法，例如文生图的配置文件 models/ldm/text2img256/config.yaml：
model:
target: ldm.models.diffusion.ddpm.LatentDiffusion
params:
unet_config:
target: ldm.modules.diffusionmodules.openaimodel.UNetModel
first_stage_config:
target: ldm.models.autoencoder.VQModelInterface
cond_stage_config:
target: ldm.modules.encoders.modules.BERTEmbedder
无关的内容都略去，可以看到顶层的模块是 LatentDiffusion，去噪网络是 UNetModel、encoder 是 VQModelInterface、文本编码器是 BERTEmbedder。
这里主要还是关注 LatentDiffusion 的采样过程。具体的采样代码位于 LatentDiffusion.sample：
@torch.no_grad()
def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
verbose=True, timesteps=None, quantize_denoised=False,
mask=None, x0=None, shape=None, **kwargs):
# 一些数据的封装以及格式转换等等
if shape is None:
shape = (batch_size, self.channels, self.image_size, self.image_size)
if cond is not None:
if isinstance(cond, dict):
cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
else:
cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
# 实际的采样过程
return self.p_sample_loop(cond,
shape,
return_intermediates=return_intermediates, x_T=x_T,
verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
mask=mask, x0=x0)
可以看到实际的采样过程并不是在这一层进行，这一层只进行了一些封装，例如采样的大小以及条件的数据格式等等，具体的采样则是在 p_sample_loop 中进行的：
@torch.no_grad()
def p_sample_loop(self, cond, shape, timesteps=None):
iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(range(0, timesteps))
for i in iterator:
ts = torch.full((b,), i, device=device, dtype=torch.long)
img = self.p_sample(img, cond, ts,
clip_denoised=self.clip_denoised,
quantize_denoised=quantize_denoised)
return img
去掉一堆杂七杂八的代码之后可以发现在 p_sample_loop 中是一个循环，也就对应于一步步进行降噪的过程，具体的降噪在 p_sample 中实现：
@torch.no_grad()
def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
return_codebook_ids=False, quantize_denoised=False, return_x0=False,
temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
b, *_, device = *x.shape, x.device
outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
return_codebook_ids=return_codebook_ids,
quantize_denoised=quantize_denoised,
return_x0=return_x0,
score_corrector=score_corrector,
corrector_kwargs=corrector_kwargs)
model_mean, _, model_log_variance = outputs
noise = noise_like(x.shape, device, repeat_noise) * temperature
if noise_dropout > 0.:
noise = torch.nn.functional.dropout(noise, p=noise_dropout)
# no noise when t == 0
nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
在 p_sample 中，首先用模型预测出了均值和方差（也就是 p_mean_variance，这里就不展开讲了），然后进行了去噪。
综合上述分析来看，如果看原始代码，可能会觉得非常混乱，但是其实去掉不重要的内容之后，核心的代码并不算非常多。这里没有展开具体的 p_mean_variance 内部的内容，在 CompVis 的框架中，定义了很多 diffusion 中常用的常量（例如 alphas_cumprod、sqrt_recipm1_alphas_cumprod 等）和方法（例如 q_mean_variance、p_mean_variance 等），后续我应该还会写一篇文章专门介绍这些内容，这里暂时略过，只需要知道最顶层的 p_mean_variance 是预测了均值和方差即可。
huggingface 的实现
相比于 CompVis 的实现，huggingface 的实现更加工程化一点，相关的在 diffusers 库中。这个库主要包括三大类元素：models（各种神经网络的实现，unet、vae 等）、schedulers（diffusion 相关的操作，加噪去噪等）、pipelines（high level 封装，相当于 models+schedulers，这个应该是方便用户直接用的）。
这里直接看 diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py 的采样过程，定义在 __call__ 函数中：
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
timesteps: List[int] = None,
sigmas: List[float] = None,
guidance_scale: float = 7.5,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
**kwargs,
):
可以看到参数实在是非常的多，我们在这里不关注工程的部分，只关注核心的逻辑。这里的第一个需要关注的点是对生成条件进行编码：
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
prompt,
device,
num_images_per_prompt,
self.do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
)
# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
这里实际上还有 LoRA 和 IP-Adaptor 相关的处理，暂时省略。可以看到这里对生成的 prompt 进行了编码，并且不仅有正常的 prompt，还有 negative 的 prompt，这是为了做 classifier-free guidance。并且由于两个 prompt 需要分别推理，这里还将其在 batch 维度拼接，来进行并行化。随后获取 timesteps：
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, device, timesteps, sigmas
)
然后初始化噪声，这个就相当于 x T \mathbf{x}_T xT：
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
上边准备了 x \mathbf{x} x、timestep 以及 condition，现在就可以正式进行生成了：
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs=self.cross_attention_kwargs,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
可以看到有一些为了 classifier-free guidance 进行的处理，其他的都是正常 diffusion 的操作。最后将隐变量解码回像素空间得到生成结果：
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0]
总结
最近看了这么多文章，感觉比较成功的 researcher 的工作都是连贯的。就像宋飏研究 sliced score matching，然后紧随其后做出了 score-based generative model；又如 OpenAI 训出 CLIP 然后基于 CLIP 做了一系列文生图的工作。今天这篇文章看起来也是 CompVis 把 VQGAN 迁移到 diffusion models 上的成果，感觉对平时做研究的启发还是很大的，我个人一直以来研究方向都比较摇摆不定，也应该反思学习一下。
参考资料：
diffusion model(五)：LDM: 在隐空间用diffusion model合成高质量图片 扩散模型（六）| Stable Diffusion本文原文以 CC BY-NC-SA 4.0 许可协议发布于 笔记｜扩散模型（七）：Latent Diffusion Models（Stable Diffusion）理论与实现，转载请注明出处。
总结