扩散模型逆向方差的选取

$$$$

DDPM: 人为选取方差

首先做一个简要回顾。DDPM 定义前向马尔可夫过程为：

$$\begin{array}{rl}& q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=q({\mathbf{x}}_0)\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})\\ & q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_0,{\beta }_t\mathbf{I})\\ & q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0,(1-{\overline{\alpha }}_t)\mathbf{I})\end{array}$$ 计算可得： $$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_t({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I})$$ 仿照 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$ 的形式，将 $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$​ 定义为： $$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_t({\mathbf{x}}_t,{\mathbf{x}}_{\theta }({\mathbf{x}}_t,t)),{\sigma }_t^2\mathbf{I})$$ 模型通过最大化 ELBO 训练： $$\begin{array}{}\text{(1)}& \text{ELBO}={\mathbb{E}}_{{\mathbf{x}}_1\sim q({\mathbf{x}}_1|{\mathbf{x}}_0)}[\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]-\text{KL}(q({\mathbf{x}}_T|{\mathbf{x}}_0)\Vert p({\mathbf{x}}_T))-\sum _{t=2}^T{\mathbb{E}}_{{\mathbf{x}}_t\sim q({\mathbf{x}}_t|{\mathbf{x}}_0)}[\text{KL}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))]\end{array}$$ 注意 $\text{KL}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))$ 里面的两个分布都是正态分布，前者方差为 ${\tilde{\beta }}_t$，后者方差定义为 ${\sigma }_t^2$；考虑到两个方差相同的正态分布的 KL 散度正比于二者均值的 L2 距离，因此如果直接取 ${\sigma }_t^2={\tilde{\beta }}_t$，那么损失函数就是一个非常简单的 L2 Loss. 但同时，DDPM 也称取 ${\sigma }_t^2={\beta }_t$ 也是可以的，这就是我们之前遗留的问题。

为了理解这个问题，需要注意到 $\text{(1)}$ 式仅仅是针对一个样本 ${\mathbf{x}}_0$ 的推导，实际上真正的优化目标还要套一层对 ${\mathbf{x}}_0$ 的期望： $$\begin{array}{}\text{(2)}& \begin{array}{rl}{\mathbb{E}}_{{\mathbf{x}}_0\sim q({\mathbf{x}}_0)}[\text{ELBO}]& ={\mathbb{E}}_{{\mathbf{x}}_0\sim q({\mathbf{x}}_0)}\left[{\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}]\right]\\ & ={\mathbb{E}}_{{\mathbf{x}}_{0:T}\sim q({\mathbf{x}}_{0:T})}[\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{0:T})}+\log q({\mathbf{x}}_0)]\\ & =-\text{KL}(q({\mathbf{x}}_{0:T})\Vert p_{\theta }({\mathbf{x}}_{0:T}))-H(q({\mathbf{x}}_0))\\ & =-\text{KL}(q({\mathbf{x}}_T)\prod _{t=1}^{T}q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)\Vert p_{\theta }({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))-H(q({\mathbf{x}}_0))\\ & =-\text{KL}(q({\mathbf{x}}_T)\Vert p_{\theta }({\mathbf{x}}_T))-\sum _{t=1}^T\text{KL}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)\Vert p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))-H(q({\mathbf{x}}_0))\end{array}\end{array}$$ 其中 $q({\mathbf{x}}_0)$ 表示数据分布。其中倒数第二步能对 $q$ 展开是因为马尔可夫过程的逆向过程依旧是马尔可夫的： $$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t:T})=\frac{q({\mathbf{x}}_{t-1:T})}{q({\mathbf{x}}_{t:T})}=\frac{q({\mathbf{x}}_{t-1})\prod _{i=t}^{T}q({\mathbf{x}}_i|{\mathbf{x}}_{i-1})}{q({\mathbf{x}}_t)\prod _{i=t+1}^{T}q({\mathbf{x}}_i|{\mathbf{x}}_{i-1})}=\frac{q({\mathbf{x}}_{t-1})q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{q({\mathbf{x}}_t)}=q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$$ 由 $\text{(2)}$ 式可以看见，$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$ 要近似的目标并不是 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$，而是 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$，它们二者之间的关系为： $$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\int q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)q({\mathbf{x}}_0|{\mathbf{x}}_t)\mathrm{d}{\mathbf{x}}_0={\mathbb{E}}_{q({\mathbf{x}}_0|{\mathbf{x}}_t)}[q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)]$$ 因此，仅当 $q({\mathbf{x}}_0|{\mathbf{x}}_t)$ 是 Dirac delta 函数时才有 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$，此时取 ${\sigma }_t^2={\tilde{\beta }}_t$ 才是最优的；否则，直接取 ${\sigma }_t^2={\tilde{\beta }}_t$ 不是最优的。然而，$q({\mathbf{x}}_0|{\mathbf{x}}_t)$ 本身是 intractable 的，所以为了可以计算，DDPM 的作者考虑了两种极端情形：

• 情形一：${\mathbf{x}}_0\sim \mathcal{N}(\mathbf{0},\mathbf{I})$，可知 $q({\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_t;\mathbf{0},\mathbf{I})$，于是 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$，所以应该取 ${\sigma }_t^2={\beta }_t$.
• 情形二：${\mathbf{x}}_0\sim \delta ({\mathbf{x}}_0)$，即数据集只有一个样本 ${\mathbf{x}}_0=\mathbf{0}$，此时 $q({\mathbf{x}}_0|{\mathbf{x}}_t)$ 是 Dirac delta 函数，所以应该取 ${\sigma }_t^2={\tilde{\beta }}_t$.

当然，这些情形的设定本身就不合理——没有人的数据是一堆高斯噪声或者一个样本，只是为了寻找方差可能的选择而做的一些试验性假设罢了。这也为后续的工作埋下了一个伏笔——最优的方差到底应该是什么形式的呢？

iDDPM: 可学习方差

为了解决上述问题，Improved DDPM^[1] 将方差作为可学习参数进行优化而非人为选取。具体而言，考虑到 ${\tilde{\beta }}_t$ 和 ${\beta }_t$ 是两个极端情况，作者把方差参数化为二者在 log domain 的插值： $${\sigma }_{\theta }^2({\mathbf{x}}_t,t)=\exp (v\log {\beta }_t+(1-v)\log {\tilde{\beta }}_t)$$ 损失函数为： $${\mathcal{L}}_{\text{vlb}}=-{\mathbb{E}}_{{\mathbf{x}}_0\sim q({\mathbf{x}}_0)}[\text{ELBO}]=\sum _{t=1}^T{\mathbb{E}}_{{\mathbf{x}}_0\sim q({\mathbf{x}}_0),{\mathbf{x}}_t\sim q({\mathbf{x}}_t|{\mathbf{x}}_0)}[\text{KL}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))]$$ 回忆在 DDPM 中，作者将上式简化为： $${\mathcal{L}}_{\text{simple}}\propto {\mathbb{E}}_{t,{\mathbf{x}}_0,{\mathbf{x}}_t}[\Vert {\mu }_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mu }_t({\mathbf{x}}_t,{\mathbf{x}}_{\theta }({\mathbf{x}}_t,t)){\Vert }^2]$$ 虽然 Improved DDPM 引入可学习方差后，损失函数已经不是 ${\mathcal{L}}_{\text{simple}}$ 了，但它依旧沿用了 ${\mathcal{L}}_{\text{simple}}$ 来训练均值。同时，用 ${\mathcal{L}}_{\text{vlb}}$ 来训练方差，并且 ${\mathcal{L}}_{\text{vlb}}$ 对均值做梯度截断，只用来训练方差。综上，Improved DDPM 采用下面这个混合损失函数： $${\mathcal{L}}_{\text{hybrid}}={\mathcal{L}}_{\text{simple}}+\lambda {\mathcal{L}}_{\text{vlb}}$$ 实验也证明，如果只用 ${\mathcal{L}}_{\text{vlb}}$ 训练，图像生成效果欠佳。这告诉我们理论与实践还是有差距的，最好的设置还是要靠实验探索。

附：两个正态分布的 KL 散度

设 $$ 是两个 $$ 维随机向量，$$，$$，则： $$

Analytic-DPM: 解析最优方差

Improved DDPM 引入可学习方差，缓解了 DDPM 方差设置不合理的问题，但也使得训练更加困难。然而同期工作 Analytic-DPM^[2]却发现，$$ 其实有解析形式的最优解，所以我们压根不需要训练它！

在开始之前，回忆 DDIM 论文将 DDPM 的前向过程扩展为非马尔可夫过程： $$

特别地，取 $$ 就得到了原始 DDPM. 因此下面我们都基于这个更一般的扩散过程叙述。

前文提到，问题实际的优化目标是： $$ 由于： $$ 后三项都与 $$ 无关，因此： $$ 根据 moment matching 理论——在 KL 散度下用高斯分布去拟合任何分布，等价于让二者的均值和方差相等——我们只需要求解 $$​ 的均值和方差即可。

首先计算均值： $$ 其中 $$ 分别是 $$ 中 $$ 和 $$ 前面的系数，即： $$ 注意到 DDPM 中 $$ 正是用于近似 $$​ 的，所以 DDPM 采用均值的确就是最优的均值估计。

接下来计算二阶矩： $$ 所以协方差矩阵为： $$ 考虑到 $$ 的协方差矩阵并不依赖于 $$，所以在上式的基础上对 $$ 取期望： $$ 又考虑到 $$ 的协方差矩阵是一个各向同性对角阵（单位阵的倍数），我们只需要计算上述矩阵的对角线元素的平均，即求迹再除以 $$： $$ 进一步地，可以推导出 $$ 与 score function 有如下关系： $$

代入可得最优方差为： $$ 可以看见，最优方差是在 $$​ 的基础上，增加了一个“修正项”。修正项与 score function 有关，可以代入训练好的模型。不过，由于修正项包含期望，所以需要在采样开始之前做蒙特卡洛估计。

References

1. Nichol, Alexander Quinn, and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171. PMLR, 2021. ↩︎
2. Bao, Fan, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models. In International Conference on Learning Representations. 2021. ↩︎
3. 苏剑林. (Aug. 12, 2022). 《生成扩散模型漫谈（七）：最优扩散方差估计（上） 》[Blog post]. Retrieved from https://kexue.fm/archives/9245 ↩︎