离散变量扩散模型

笔者最近在研究图上的扩散模型，阅读了离散变量扩散模型的经典文献 Discrete Denoising Diffusion Probabilistic Model (D3P , NeurIPS, 2021)。D3P 是连续变量上的去噪扩散模型 Denoising Diffusion Probabilistic Model (DDPM ，NeurIPS, 2020) 在离散变量上的扩展，遵循 DDPM 的总体思路，但技术细节上有所不同。作者在原论文中对部分推导过程未做展开，网络上对离散变量扩散模型的介绍较少，本博文旨在填补这部分空白。

去噪扩散模型回顾

连续变量扩散模型有较多优秀的博客介绍，在此简单重述。如上图所示，DDPM 包含一个加噪过程和一个去噪过程。一个复杂分布通过加噪过程逐渐变成一个无意义的噪声，然后通过去噪过程将该噪声恢复为原始数据分布。去噪过程将用于最后的生成。DDPM 的损失函数是原始数据负对数似然的变分上界，即

\[\mathbb{E}\left[-\log p_\theta(\mathbf{x}_0)\right] \leq \mathbb{E}_q \left[\text{KL}\left[q(\mathbf{x}_T|\mathbf{x}_0)||p(\mathbf{x}_T)\right] + \sum_{t > 1} \text{KL}\left[q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)||p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)\right] -\log p_\theta(\mathbf{x}_0|\mathbf{x}_1) \right] \]

模型优化的关键在于损失函数有封闭表达式。记 \(\mathbf{x}_0\) 为原始样本，\(\mathbf{x}_t\) 为加噪过程第 \(t\) 步得到的样本，\(\mathbf{x}_s\) 为第 \(t\) 步的上一步的样本。则损失函数具有封闭表达式需要满足以下五个条件：

1. 单步转移分布 \(q(\mathbf{x}_t|\mathbf{x}_s)\) 有封闭表达式；
2. 多步分布 \(q(\mathbf{x}_t|\mathbf{x}_0)\) 有闭式解；
3. 稳态分布 \(q(\mathbf{x}_\infty|\mathbf{x}_0)\) 已知，表达式与 \(\mathbf{x}_0\) 无关；
4. 辅助分布 \(q(\mathbf{x}_s|\mathbf{x}_t,\mathbf{x}_0)\) 有闭式解；
5. \(q(\mathbf{x}_s|\mathbf{x}_t,\mathbf{x}_0)\) 和 \(p_\theta(\mathbf{x}_s|\mathbf{x}_t)\) 的 KL 散度存在闭式解。

根据 DDPM 原文的推导，以上条件发挥了以下作用：

• 条件 1 和条件 2 是条件 4 的前置条件；
• 根据条件 3，损失函数第一项可视为常数；
• 条件 2，条件 4 和条件 5 确保损失函数第二项有闭式解；

记一个离散随机变量（独热向量）为 \(\mathbf{x}\in\{0,1\}^K\)，其服从类别分布，概率密度函数为 \(\text{Cat}(\mathbf{x};\mathbf{p})\)，其中 \(\mathbf{p}\) 描述了随机变量不同取值的概率。如何定义离散变量的加噪过程和去噪过程，使其满足以上 5 个条件？

离散变量扩散模型

满足条件 1 的构造是容易想象的，只需定义离散状态之间的概率转移矩阵即可。条件 3 也是容易想象的，可以假定稳态分布为均匀分布，虽然实际上可以有其他选择。条件 5 是容易满足的，因为类别分布的 KL 散度可以直接按定义计算。相比之下，条件 2 是不容易想象的，条件 4 看起来是困难的，下面将分点讲解。

条件 1

由于 \(\mathbf{x}\) 是离散变量，容易通过状态转移矩阵 \(\mathbf{Q}\) 定义转移分布 \(q(\mathbf{x}_t|\mathbf{x}_s)\)，满足 \(\mathbf{Q}_{ij}=q(\mathbf{x}_{t,j}=1|\mathbf{x}_{s,i}=1)\)。写成矩阵形式得到以下表达式

\[\begin{equation} q(\mathbf{x}_t|\mathbf{x}_s)=\text{Cat}\left(\mathbf{x}_t|\mathbf{p}=\mathbf{Q}_t\mathbf{x}_s\right). \end{equation} \]

此处要特别注意到，对于取定的 \(\mathbf{x}_t,\mathbf{x}_s\)，

\[\begin{equation} q(\mathbf{x}_t|\mathbf{x}_s)=\mathbf{x}_t^T\mathbf{Q}_t\mathbf{x}_s. \end{equation} \]

条件 2

由本博客公式 2 及加噪过程的马尔可夫性，

\[\begin{equation} q(\mathbf{x}_t|\mathbf{x}_{s-1})=\mathbf{x}_t^T\mathbf{Q}_t\mathbf{Q}_s\mathbf{x}_{s-1}. \end{equation} \]

以此类推得到多步转移分布

\[\begin{equation} \begin{aligned} q(\mathbf{x}_t|\mathbf{x}_0)&=\mathbf{x}_t^T\prod^t_{\tau=1}\mathbf{Q}_\tau\mathbf{x}_{0}\\ &=\mathbf{x}_t^T\bar{\mathbf{Q}}_t\mathbf{x}_{0}\\ &= \text{Cat}\left(\mathbf{x}_t|\mathbf{p}=\bar{\mathbf{Q}}_t\mathbf{x}_{0}\right). \end{aligned} \end{equation} \]

问题在于如何恰当选取单步转移分布，使得多步转移分布具有闭式解。

下面对 D3P 附录 A 2.2.1 和 A 2.2.2 的设计进行解读。两个设计分别导致稳态分布为均匀分布和单点分布，可以有统一的描述形式。

定义转移矩阵为

\[\begin{equation} \mathbf{Q}_t = (1-\beta_t)\mathbf{I}+\beta_t\mathbf{Q}_\infty, \end{equation} \]

其中 \(\mathbf{Q}_\infty\) 为稳态分布矩阵，满足幂等性质 \(\mathbf{Q}_\infty^2=\mathbf{Q}_\infty\)。

只需计算 \(\mathbf{Q}_t\mathbf{Q}_{s}\)。

\[\begin{equation} \begin{aligned} \mathbf{Q}_t\mathbf{Q}_{s} &= \left((1-\beta_t)\mathbf{I}+\beta_t\mathbf{Q}_\infty\right)\left((1-\beta_s)\mathbf{I}+\beta_s\mathbf{Q}_\infty\right)\\ &= (1-\beta_t)(1-\beta_s)\mathbf{I}+\left((1-\beta_t)\beta_s+(1-\beta_t)\beta_s\right)\mathbf{Q}_\infty+\beta_t\beta_s\mathbf{Q}_\infty^2\\ &= (1-\beta_t)(1-\beta_s)\mathbf{I}+\left((1-\beta_t)\beta_s+(1-\beta_t)\beta_s\right)\mathbf{Q}_\infty+\beta_t\beta_s\mathbf{Q}_\infty\\ &= (1-\beta_t)(1-\beta_s)\mathbf{I}+\left(1-(1-\beta_t)(1-\beta_s)\right)\mathbf{Q}_\infty\\ \end{aligned} \end{equation} \]

记 \(\alpha=1-\beta,\bar{\alpha}=\alpha_t\alpha_s\)，则有

\[\mathbf{Q}_t\mathbf{Q}_{s} = (1-\bar{\alpha})\mathbf{I}+\bar{\alpha}\mathbf{Q}_\infty \]

即两步的转移矩阵与单步的转移矩阵具有相同形式。因此，多步转移矩阵为

\[\bar{\mathbf{Q}}_t = (1-\bar{\alpha}_t)\mathbf{I}+\bar{\alpha}_t\mathbf{Q}_\infty, \]

其中

\[\bar{\alpha}_t=\prod^t_{\tau=1}\alpha_\tau. \]

可以看出，当 \(t\to\infty\)，\(\bar{\alpha}_t\to 0\)，此时 \(\bar{\mathbf{Q}}_t\to\mathbf{Q}_\infty\)。

附录 A 2.2.1，A 2.2.2，A 2.2.6 提供的稳态分布分别为：

• 均匀分布：\(\mathbf{Q}_\infty=\frac{1}{K}\mathbf{1}\mathbf{1}^T\)；
• 单点分布：\(\mathbf{Q}_\infty=\mathbf{1}\mathbf{e}^T_m\)；
• 均匀分布与单点分布的混合：\(\mathbf{Q}_\infty=\frac{\alpha}{K}\mathbf{1}\mathbf{1}^T+\beta\mathbf{1}\mathbf{e}^T_m,\alpha,\beta \geq 0, \alpha+\beta\leq 1\)。

条件 3

通过适当选取 \(\mathbf{Q}_t\)，可以使得稳态分布

\[\begin{equation} \underset{t\to\infty}{\text{lim}}\bar{\mathbf{Q}}_t\mathbf{x_0}=\boldsymbol{\pi}, \end{equation} \]

存在且与初值 \(\mathbf{x}_0\) 无关。根据 D3P 在其论文附录 A.1 的推导，如果转移矩阵 \(\mathbf{Q}_t\) 是双随机矩阵（行和列和均为 1），满足不可约性和非周期性，则稳态分布为均匀分布。

条件 4

辅助分布

\[\begin{equation} \begin{aligned} q(\mathbf{x}_s|\mathbf{x}_t,\mathbf{x}_0) &=\frac{q(\mathbf{x}_t|\mathbf{x}_s,\mathbf{x}_0)q(\mathbf{x}_s|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)}\\ &= \frac{q(\mathbf{x}_t|\mathbf{x}_s)q(\mathbf{x}_s|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)}\\ &= \frac{\left(\mathbf{x}_t^T\mathbf{Q}_t\mathbf{x}_s\right)\cdot\left(\mathbf{x}_s^T\bar{\mathbf{Q}}_s\mathbf{x}_0\right)}{\mathbf{x}_t^T\bar{\mathbf{Q}}_t\mathbf{x}_0}\\ &= \frac{\left(\mathbf{x}_s^T\mathbf{Q}^T_t\mathbf{x}_t\right)\cdot\left(\mathbf{x}_s^T\bar{\mathbf{Q}}_s\mathbf{x}_0\right)}{\mathbf{x}_t^T\bar{\mathbf{Q}}_t\mathbf{x}_0}\\ &= \text{Cat}\left(\mathbf{x}_s;\mathbf{p}=\frac{\left(\mathbf{Q}^T_t\mathbf{x}_t\right)\circ\left(\bar{\mathbf{Q}}_s\mathbf{x}_0\right)}{\mathbf{x}_t^T\bar{\mathbf{Q}}_t\mathbf{x}_0}\right). \end{aligned} \end{equation} \]

其中：

• 第一个等号是由贝叶斯公式得到，第二个等号是由加噪过程的马尔可夫性得到，与原文推导一致；
• 第三个等号来自本博客的公式 2，即直接计算类别分布概率密度函数的值；
• 第四个等号对分子的第一项转置得到；
• 第五个等号根据定义，考虑 \(\mathbf{x}_s\) 的全部可能取值得到其类别分布的参数。

可以检验分母确实为为归一化常数，可以证明对分子求和得到分母，因为

\[\begin{aligned} \mathbf{1}^T\left(\mathbf{Q}_t\mathbf{x}_t \circ \bar{\mathbf{Q}}_{t-1}^T\mathbf{x}_0\right) &= \left(\bar{\mathbf{Q}}_{t-1}^T\mathbf{x}_0\right)^T\left(\mathbf{Q}_t\mathbf{x}_t\right)\\ &= \mathbf{x}_0^T\bar{\mathbf{Q}}_{t-1}\mathbf{Q}_t\mathbf{x}_t\\ &= \mathbf{x}_0^T\bar{\mathbf{Q}}_t\mathbf{x}_t. \end{aligned} \]

因此，在计算概率分布时，只需计算分子，对分子求和即可得到分母，这也是大部分实现所采用的方法。

条件 5

记 \(q(\mathbf{x}_s|\mathbf{x}_t,\mathbf{x}_0)=\text{Cat}(\mathbf{x}_s|\mathbf{p}_1)\) 和 \(p_\theta(\mathbf{x}_s|\mathbf{x}_t)=\text{Cat}(\mathbf{x}_s|\mathbf{p}_2)\)，根据 KL 散度的定义，

\[\begin{equation} \begin{aligned} \text{KL}\left[\text{Cat}(\mathbf{x}_s|\mathbf{p}_1)||\text{Cat}(\mathbf{x}_s|\mathbf{p}_2)\right] &=\mathbf{p}_1^T\ln\frac{\mathbf{p}_1}{\mathbf{p}_2}. \end{aligned} \end{equation} \]

总结

下面列表格对比说明连续变量和离散变量扩散模型的关键设计：

项目	连续变量	离散变量
变量范围	\(\mathbf{x}\in\mathbb{R}^d\)	\(\mathbf{x}\in\{0,1\}^K\)
概率密度函数	\(\mathcal{N}(\mathbf{x};\boldsymbol{\mu},\sigma^2\mathbf{I})\)	\(\text{Cat}(\mathbf{x};\mathbf{p})\)
单步转移分布 \(q(\mathbf{x}_t|\mathbf{x}_s)\)	\(\mathcal{N}(\mathbf{x}_t;\sqrt{1-\beta_t}\mathbf{x}_s,\beta_t\mathbf{I})\)	\(\text{Cat}(\mathbf{x}_t;\mathbf{p}=\mathbf{Q}_t\mathbf{x}_s)\)
多步转移分布 \(q(\mathbf{x}_t|\mathbf{x}_0)\)	\(\mathcal{N}(\mathbf{x}_t;\sqrt{\bar{\alpha}_t}\mathbf{x}_s,(1-\bar{\alpha}_t)\mathbf{I})\)	\(\text{Cat}(\mathbf{x}_t;\mathbf{p}=\bar{\mathbf{Q}}_t\mathbf{x}_0)\)
稳态分布 \(q(\mathbf{x}_\infty|\mathbf{x}_0)\)	\(\mathcal{N}(\mathbf{x}_\infty;\mathbf{0},\mathbf{I})\)	\(\text{Cat}(\mathbf{x}_\infty;\mathbf{p}=K^{-1}\mathbf{1})\)
辅助分布 \(q(\mathbf{x}_s|\mathbf{x}_t,\mathbf{x}_0)\)	\(\mathcal{N}(\mathbf{x}_s;\widetilde{\boldsymbol{\mu}}(\mathbf{x}_t,\mathbf{x}_0),\widetilde{\beta}_t\mathbf{I})\)	\(\text{Cat}\left(\mathbf{x}_s;\mathbf{p}=\frac{\left(\mathbf{Q}^T_t\mathbf{x}_t\right)\circ\left(\bar{\mathbf{Q}}_s\mathbf{x}_0\right)}{\mathbf{x}_t^T\bar{\mathbf{Q}}_t\mathbf{x}_0}\right)\)
KL 散度	多元高斯分布间的 KL 散度	类别分布间的 KL 散度

连续变量扩散模型的具体公式可以参考 DDPM 原文。

思考

D3P 附录提供了多种转移矩阵的构造，其中附录 A 2.2.3-A 2.2.5 构造的双随机矩阵结构很精细，那么转移矩阵的设计空间是什么，不同的设计有什么适用场景，不同的设计存在本质不同吗？

参考文献：

1. Denoising diffusion probabilistic models. In NeurIPS, 2020.

2. Structured Denoising Diffusion Models in Discrete State-Spaces. In NeurIPS, 2021.