Lazy loaded image
Logs for Foundations of Diffusion Models
00 min
Feb 9, 2024
Mar 19, 2024
type
status
date
slug
summary
tags
category
icon
password
The content of this blog page only has Simplified Chinese version.
想写关于 Diffusion 的博客很久了,但苦于没有什么很好的教程遂搁置至今。最近适逢过年回家,在科研的夹缝中终于觅得宝贵时间,遂参考如下 surveys 及 blog page [1-5] 而得此文,希望能写完。
 

Sect 0. 前置的前置知识

在深入 Diffusion 之前,我们有必要对 Generative Models 有一个了解。Diffusion Models 的纯技术文章有非常多,但它们几乎都集中在模型本身,[2] 首次将 VAE 与 Diffusion 的底层逻辑统一,这是很难得可贵的。本部分逻辑将根据 [2] 进行撰写。
 
给定若干观测样本 , 生成式模型(Generative Model)试图对样本背后的数据分布 进行建模。但凡学到,我们就可以使用这个模型生成更多的样本。
 
生成模型种类非常多,但在简单回顾之前,有一个数学概念名词需要回顾。

似然 (likelihood): 这是一个非常文邹邹的词,但实际上可以把它理解成“可能性”。如果你有一个带参概率模型 ,模型参数用 表示。当 固定,模型退化成 ,亦即描述样本 概率分布(Probability of );当 固定,模型退化成 ,亦即描述参数 似然分布(likelihood of
 
一个例子:
有一个硬币,它有 的概率会正面向上,有 的概率反面向上。 是存在的,但是你不知道它是多少。为了获得 的值,你做了一个实验:将硬币抛10次,得到了一个正反序列:
无论 的值是多少,已知这个序列的概率值为
在这个例子里面, 概率模型的观测值 已经被固定,故以上推导的即为似然函数。
 
在实际运用中,经常使用 log-likelihood 的形式表示某数据的似然函数。原因见此

Sect 1. 前置知识:VAE 和 ELBO

VAE 模型的基本结构此处不讲细节。感兴趣的读者可以参考这篇 post 以及下面这张图:
VAE 通过刻画数据 与低维隐变量 的关系对数据分布进行建模。在模型中我们很容易能观测 以及它们的联合概率分布 。我们的目标是极大化数据分布 似然函数让我们能够更好的还原源数据 。为了符号描述简单起见,以下推导省略参数部分将数据分布视作似然函数
当存在联合概率分布 ,我们有以下两种方式还原出单变量分布:
1.直接对隐变量积分:
2.使用概率论的链式法则:
但我们既不可能在方程 1 中对无穷无尽的隐变量 积分,也不可能在还不知道 的条件概率 情况下代入方程 2 求解(注意 是 Learnable 的,我们同样也不知道是什么)。故我们需要新的数学工具进行刻画,亦即标题中的证据下界 (Evidence Lower Bound, ELBO)。我们这里直接给出对应使用形式:
推导过程如下:由方程 2 开始:
从这个推导中,我们可以清楚地观察到,借助一个 Learnable 的条件似然函数 ,证据等于 ELBO 加上近似的条件似然函数 和真实分布 之间的 KL 散度。由以上推导,我们可以发现证据下界充当了很好的优化目标,因为它保证了似然函数的下界。并且,轻松使用 Learnable 的条件似然函数 与联合分布函数 ,我们即可获得 ELBO。
 
在继续以下内容之前,我们引入三个名词:
后验概率(知果求因):某事已经发生,求相关关联事件的概率
先验概率(由历史求因):某事未发生,由过往的经验猜测该事的发生概率
似然估计(由因求果):先定下来原因,根据原因来估计结果的概率分布。
 
因此我们可以说,条件似然函数 似然估计。

notion image
 
进一步地,从 ELBO 出发,再引入一个 Learnable 的后验概率函数 ,我们可以得到:
由该式推导结果,我们可以发现似然估计函数 可以被视作一个编码器!它将输入 转换为低维隐变量 上的一个分布。同时,我们还发现这个形式引入的后验概率函数 可以视作将给定的低维隐变量 转换为原始数据 ,亦即一个解码器。KL 散度形式鼓励编码器实际学习一个分布,而不是塌缩成一个狄拉克 函数。
 
为了方便起见,VAE 的编码器 通常选择来建模一个具有对角协方差的多元高斯分布,而先验(低维隐变量 分布)通常选为标准多元高斯分布。
然后,ELBO 的 KL 散度项可以解析地计算出来,而重构项可以使用蒙特卡洛估计来近似。于是目标函数可以重新写为:
其中 是从编码器 里面的采样个数。然后,网络通过重参数化(Reparameterization)实现梯度连续,此处不展开,详情请参阅此链接

在 VAE 之后,拥有层次结构的层级 VAE(HVAE, Hierarchical Variational Autoencoder)的实现往 Diffusion 的推导更进了一步:
notion image
 
由图可知,这种 VAE 变体实质上把生成过程变成了一个马尔可夫链;也就是说,层次结构中的每一个转换都是马尔可夫的,其中解码每一个低维潜变量 仅依赖于前一个低维潜变量 ,亦即把 T 个 VAE 进行层叠。于是乎对应的先验/后验分布形式变成了:
又于是,ELBO 变成了:
至此,前置知识的部分就结束了。我们将在后面的部分进一步对上面的式子进行解析。
 

Sect 2. Variational Diffusion Probabilistic Model

Sect 2.1 三个限制条件

扩散模型(Diffusion Model, DM)最简单的理解方式就是将其看作一个具有三个关键限制的 HVAE:
  • 隐变量 维度与输入数据 维度完全相等。于是在接下来的篇幅中,我们将直接以 表示数据以及隐变量。
  • 每个时间步的隐变量编码器 结构不是学习得到的;它预定义为线性高斯模型。换句话说,它是 以 前一时间步 的输出 为中心的高斯分布。
  • 潜在编码器的高斯参数随时间变化,在最终时间步 时,潜在变量的分布是一个标准高斯分布
notion image
 
然后我们再一步步针对刚刚得到的前置条件一步步进行解析。
由限制条件 1,HVAE 的后验
可以被重写为:
由限制条件 2,我们知道在编码器中每个潜在变量的分布是一个以其前一个层次潜在变量为中心的高斯分布。与 HVAE 不同的是,每个时间步骤 的编码器结构不是学习得到的;它被固定为一个线性高斯模型,其中均值和标准差可以事先作为超参数设置,或作为参数学习。我们使用均值 和方差 来参数化高斯编码器。这样的设置方式只是最为广泛的一种,也允许使用其他高斯参数化形式。甚至 亦可作为 Learnable Parameter 以提供灵活性。请注意, 是可以随时间步骤 变化的。结合以上分析我们有每个时间步骤的编码器描述如下:
由限制条件 3,我们知道使得最终潜在变量的分布 一个标准高斯分布。于是,结合限制条件 1,2 和 HVAE 的先验/后验分布
就变成了:
总的来说,这一系列假设描述的是随时间对图像输入进行稳定的噪声化过程;我们通过逐渐添加高斯噪声来逐步破坏图像,直到最终它完全变成纯粹的高斯噪声。请注意,编码器部分 已经不依赖参数 学习由于这一部分已经是纯粹的线性高斯过程了。所以在 Diffusion Model 里面我们感兴趣的是如何去噪
 

Sect 2.2 Diffusion ELBO 以及其含义

有了编码器/解码器的表达形式,我们自然对如何优化感兴趣,在这里我们还是利用 ELBO对优化过程进行刻画。根据前述 ELBO 的推导过程,我们立刻能得到如下过程:
 
推导出的ELBO(证据下界)可以通过其各个组成部分来解释:
1. 可被视为一个重构项,预测给定第一步潜变量 时,原始数据样本 的对数概率。这一项也出现在标准 VAE 中,并且可以以类似的方式进行训练。
2. 是一个先验匹配项;当最终潜在分布 与高斯先验 匹配时,该项被最小化。这一项不需要优化,因为它没有可训练的参数;此外,由于我们假设了足够大的 ,使得最终分布是高斯的,于是这个项实际上无限趋近于零。
3. 是一个一致性约束项;它试图使 处的分布从前向和后向过程中保持一致。也就是说,对于每一个中间时间步骤,从更嘈杂的图像进行的去噪步骤应该与从更清晰的图像进行的加噪步骤匹配;这是通过 KL 散度实现的。当我们训练 以匹配高斯分布 时,该项被最小化。
这些组成部分共同定义了模型的训练目标,通过最小化这些项来优化模型参数,以此来学习数据的生成过程。

在这种推导下计算出来的 ELBO 的所有项都以期望的形式表示,因此可以使用蒙特卡洛估计来近似。但这种近似的稳定性无法得到保证:在该式中,一致性约束项利用了每个时间步骤上的两个随机变量 计算此项对应的期望,进一步地加剧了结果的随机不稳定性。通过这种方式的计算在 个时间步骤上求和无疑会加剧结果的随机性。故我们需要推导出另一种形式的 ELBO,以尽力较小随机性带来的高方差影响。
 

Sect 2.3 更稳定的 ELBO

上节说到,由于在一致性约束项中采用两个随机变量估计方差可能造成结果不稳定,本节会尝试解决该问题。
观察到编码器可以再额外添加原始输入 作为约束条件,同时再利用链式贝叶斯法则可得:
于是乎我们再次重新尝试推导 ELBO:
通过链式贝叶斯公式,我们导出了仅含一个随机变量的一致性约束项。在这种推导下,ELBO的各个组成部分可以这样理解:
 
1. 重构项与之前一致。该项可以使用蒙特卡洛估计来近似和优化。
2. 先验匹配项依旧与之前一致,表示最终噪声化输入的分布与标准高斯先验的接近程度。我们再次重申它没有可训练的参数,并且在我们的假设下也等于零
3. 在 去噪匹配项上。Learnable 的去噪步骤 近似了真实的(ground-truth)去噪过渡步骤 可以作为监督信号的原因是它定义了如何在最终、完全去噪的图像 的限制下去噪噪声图像 。因此,通过衡量可学习项分布与监督信号分布的 KL 散度即可达成优化目标。同时减小了推导过程中包含的随机变量数。
 
注:到现在我们能看到 Diffusion 的整个前向过程没有要学习的参数,而 VAE 的加噪过程是需要训练的,因为它希望通过训练将不同类的图片特征在 latent space 里区别开,也就是说 VAE 压缩出来的噪声图片实际上也是包含信息的,因此它需要进行图片 -> 对应噪声 -> 对应图片这样一对一的转换才能保证正确性。而 Diffusion 认为压缩出来的噪声图片完全不包含任何信息,它学习的是噪声 -> 图片的每一步转换过程。
 
 

Sect 2.4 Diffusion ELBO 与前后向推导

在给出了合适的 ELBO 之后,我们终于可以开始推导 ELBO 的详细表达式了。Sect 2.3 给出的 ELBO 的最终式子是:
首先结合限制条件 2 和贝叶斯定理的结果,我们有:
所以一致性约束项我们可以能够推导出来,而后向过程 可以通过网络生成。剩下需要推导的项即是 以及 ,由于前向过程是线性高斯过程,故我们中均可推导给出。

假设我们的加噪过程已经进行了 步,现在要进行 的加噪。在线性高斯过程(限制 2)下根据重参数化技巧,有:
其中 。写成通式是:
不妨假定 相互独立,根据正态分布的可加性,上式可以写成:
其中 。我们尝试减少变量数量,以进一步化简。故我们尝试把这两个参数归一化,亦即令 。这样能行吗?既然我们希望加噪结束最终得到的图片 ,那么归一化显然不会影响我们达到这个目标,故我们尝试归一化。此时有:
于是我们令 ,从而 ,于是有(此处将 记为 ,考虑到 的相互独立性):
至此,我们成功推导出了一个时间步骤的前向过程。进一步的,通式有:
于是前向过程中的 , 我们均可导出。

至此 ELBO 中所有项我们都能导出了,我们现在只剩最后一个问题:在 ELBO 的一致性约束项 中,我们不知道 的分布究竟为何,需要推导出分布,我们才能使用 KL 散度进一步优化,故这一小节就用来分析。
由贝叶斯过程,
我们有:
 
所以我们证明了这个分布仍然是正态分布。

现在我们剩下最后一项 的去噪建模项 。我们期望让模型通过神经网络的参数 去学习每一步去噪过程的期望和方差:
为什么要纠结分布是正态分布?因为如果是正态分布,我们就可以用下面的公式计算两个正态分布的 KL 散度:
把之前的所有推导带入进 KL 散度计算,我们有:
于是乎我们发现算了这么一堆,整个 diffusion 的关键是让去噪过程与加噪高斯过程中的均值尽可能靠近!
 

Sect 2.5 均值的三种学法

现在考虑怎么学习这个 ,当然我们可以直接硬学(方法 1),但是如果能做一点分解然后学其中的一个部分或许会更简单。

我们立刻能知道
在后向过程中,我们学习的 的函数,所以上面的 部分可以直接拿下来用,需要通过神经网络 拟合的只有 那部分,亦即:
于是优化可以这么写:
亦即我们学的最终还是如何拟合原图 (方法 2)!

那么以上是式子还能不能够被进一步化简?答案是可以的,因为我们有 之间的关系: ,于是乎我们又有:
那么我们又有了一种新的预测方法:预测噪声(方法 3)!也就是建模
然后优化目标变成了:
从数学本质上来看,预测噪声和预测原始图像应该是一样的,但是实际操作中发现预测噪声比预测原始图像表现要更好。
 

Sect 3. Guidance

在把生成模型的原理讲清楚之后(我希望是讲清楚了),我们开始探讨如何控制生成模型。条件控制生成的方式分两种:事后修改(Classifier-Guidance)和事前训练(Classifier-Free)。
Classifier-Guidance:
对于大多数人来说,一个 SOTA 级别的扩散模型训练成本太大了,而分类器(Classifier)的训练还能接受,所以直接复用别人训练好的无条件扩散模型,用一个分类器来调整生成过程以实现控制生成。
Classifier-Free:
往扩散模型的训练过程中就加入条件信号,达到更好的生成效果。
以上推导请参考文首给出的参考 blog page[4]。

结语

这整篇东西… 可以说是给 diffusion 系列仅仅开了个头,要写的东西实在太多,况且已经有人写了一整个系列,只能说在这个方面还是比较任重道远… Maybe 这个东西可以作为我未来某个大长篇系列的一部分吧。
 
Fin.
 
Prev Post
博三某夜之静夜思
Next Post
Logs for Casual Deep Learning

Comments
Loading...