原文：The Annotated Diffusion Model (huggingface.co)
作者: Niels Rogge, Kashif Rasul
译者: 蓝方_Blusque
注:
本文主要用于个人学习记录用, 在翻译的准确性与对扩散模型理解的准确性上都有很大的欠缺和偏颇, 还请不吝赐教.
另: 本文中的 "正态分布" 一词因为一次笔误写成了 "正太分布", 但感觉还蛮可爱的索性不改了. ^_^
结论
先说结论, 扩散模型与 VAE 有些相似, 只不过扩散模型的编码过程与解码过程 (在扩散模型中被称为加噪过程和采样过程) 都是多次从正太分布中采样的过程. 扩散模型中应用神经网络学习的是在采样过程中每一步的正太分布中的均值 , 并通过重参数化转化为学习每一步的加噪
. 对于每一步的方差, 在最原始的论文中采用了定值, 而在此后的改进Improved diffusion models中, 采用了可学习的方差.
开篇
在这篇文章中，我们将会深入了解 Denoising Diffusion Probabilistic Models (DDPM, 就是diffusion model, 亦或是 autoencoders), 在这个模型上，研究人员已经在条件（非条件）图像 / 音频 / 视频生成任务中取得了较好的结果。比较广为人知的成果包括 OpenAI 的 GLIDE 以及 DALL-E 2, 海德堡大学的 Latent Diffusion 以及 Google Brain 的 ImageGen.
我们将会回顾DDPM的原论文 (Ho et al., 2020), 并一步一步在 PyTorch 中基于 Phil Wang 的代码重新实现这个模型 -- 他的代码是基于原始的TensorFlow 实现完成的. 将扩散思想用于生成模型最早在2015年就被提出了(Sohl-Dickstein et al., 2015). 然而，直到2019年 (Song et al., 2019) 以及此后的 (Ho et al., 2020) 分别改进了这个方法，才得以得到广泛应用.
值得注意的是，对于扩散模型存在很多种分类 (原文为 "perspectives" ). 在这里我们采用的是离散时间版本 (discrete-time, 也称 latent variable model), 也请注意查看其他的版本.
好的, 让我们继续吧!
from IPython.display import Image
Image(filename='assets/78_annotated-diffusion/ddpm_paper.png')
配置环境的部分就略过了, 毕竟不是PyTorch安装教程, 其他的库可以参考原文安装. 需要的库如下所示:
import math
from inspect import isfunction
from functools import partial
%matplotlib inline
import matplotlib.pyplot as plt
from tqdm.auto import tqdm
from einops import rearrange, reduce
from einops.layers.torch import Rearrange
import torch
from torch import nn, einsum
import torch.nn.functional as F
什么是扩散模型?
扩散模型相比于其他的生成模型, 如标准化流 (Normalizing Flows), 生成式对抗网络 (GANs) 或 变分自编码器 (VAEs), 并不复杂: 这些模型都是将从简单分布中采样的噪声转化为数据样本. 扩散模型也一样, 是一个从纯噪声中学习如何去噪的神经网络.
关于更具体的细节, 整个过程可以分为两个部分:
- 一个固定 (或预定义) 的前向扩散过程
, 这个过程逐步向原图像加噪, 直到得到一个单纯的噪声 (关于可以得到纯噪声的证明过程见原论文).
- 一个可学习的反向去噪过程
, 其中一个神经网络被训练用于从纯噪声中逐步去噪, 直到得到一个 "真实" 的图像.
前向过程和后向过程都在有限的时间步 T 内发生, 并且其中的每一步都有标记 t (DDPM 的作者采用的 T = 1000). 假设我们从第 0 步, 即从数据集中抽取一张图像 , 开始 (这里我们用 ImageNet 数据集中一张猫的照片来举例), 之后在前向过程中的每一个时间步向前一个时间步的图像中加入从正太分布中采样的噪声. 在给定一个足够大的 T 以及一个良好的参数分布序列 (schedule) 来对每一个时间步的图像加噪后, 你就可以通过这个分步过程得到一个各向同性高斯分布 (isotropic Gaussion distribution).
更加数学的表述
让我们从更加正式的角度描述这个问题, 因为最终我们需要一个易于处理的损失函数以供我们的神经网络更新之用.
使 作为真实的数据分布, 即 "真实图像". 我们可以从这个分布中采样从而得到一个图像,
. 我们定义了前向扩散过程
来在每一步 t 向前一步的图像加入高斯噪声, 这个噪声分布的方差
服从
. 即,
记得标准正太分布 (也叫高斯分布) 是通过两个参数来定义的: 一个均值 以及一个方差
. 基本上来说, 每一步 t 的新图像都是从一个
并且
的条件高斯分布中采样得到的, 这样的一张图像我们可以通过采样一个噪声
并代入
得到.
注意这里的 并不是每一步 t 都相同的常量 (根据下标变化) --- 实际上是通过定义一个 "variance schedule" 来对
进行采样的. 这个过程 (或分布) 可以是线性的, 平方的, 余弦的或者其他的, 有点类似于学习率的变化.
之后, 我们可以从真实图像 开始, 随着
的加噪, 如果我们的参数设置正确的话, 最终得到的
为纯高斯噪声.
现在, 如果过我们可以直到解码时的条件概率分布 , 那么我们就可以跑通这个反向过程: 通过采样随机高斯噪声
, 并逐步对它进行 "去噪", 我们最终会得到原始的真是图像
.
然而可惜的是, 我们现在还不知道 , 实际上这个分布是不可处理的, 因为它要求我们知道所有可能会出现的图像的分布才能求得这个条件概率. 因此, 我们需要一个神经网络来估计 (学习) 这个条件概率分布. 让我们把这个需要学习的分布叫做
, 其中
是神经网络的参数, 可以通过梯度下降更新.
现在, 我们需要一个神经网络来表示后向过程的 (条件) 概率分布. 如果我们假定这个反向过程也是高斯分布, 那么由于任何一个高斯分布都定义自以下两个变量:
- 一个均值参数
- 一个方差参数
因此我们可以将这个分布参数化为
其中均值和方差都是根据前一步 t 作为条件确定的.
自此, 我们的神经网络需要学习的就是分布的均值和方差. 然而实际上, DDPM 的作者决定将方差固定, 使得整个神经网络仅用来学习 (表示) 这个条件概率分布的均值 . 在文中写道:
"First, we set
to untrained time dependent constants. Experimentally, both
and
(about the sign
please refer to the paper) had similar results."
在之后的文章 Improved diffusion models中, 这一点 (固定 ) 得到了改进, 其中的神经网络在均值之外, 也会学习反向过程的方差.
在这里, 我们仅仅按照 DDPM 的描述, 假定我们的模型只学习条件概率分布的均值.
