离散状态、离散时间扩散

前置

• index
• 离散时间马尔可夫链

1. 问题设定

本文讨论有限离散状态、离散时间的扩散模型，即 D3PM（Discrete Denoising Diffusion Probabilistic Model）。

1.1 可测空间

状态空间取有限集并配备幂集 $\sigma$-代数：

$$\mathcal{X}:=\{1,\dots ,K\},(\mathcal{X},\mathcal{P}(\mathcal{X})).$$

给 $\mathcal{X}$ 离散拓扑后，Borel $\sigma$-代数 $\mathcal{B}(\mathcal{X})$ 与 $\mathcal{P}(\mathcal{X})$ 重合，故 $(\mathcal{X},\mathcal{P}(\mathcal{X}))$ 是标准 Borel 空间。每个子集都可测，每个 $\mathcal{X}\to \mathbb{R}$ 的函数都可测；$\mathcal{X}$ 上的概率测度因此与概率向量一一对应。

时间集取

$$\mathbb{T}:=\{0,1,\dots ,T\}.$$

一条轨迹是 $\mathbb{T}\to \mathcal{X}$ 的映射，路径空间为有限笛卡尔积

$$\Omega :={\mathcal{X}}^{\mathbb{T}}=\{(x_0,\dots ,x_T):x_t\in \mathcal{X}\}.$$

由 积空间sigma代数，$\Omega$ 上使全部坐标投影可测的乘积 $\sigma$-代数为 $\mathcal{P}(\mathcal{X}{)}^{\otimes \mathbb{T}}$。由于 $\mathbb{T}$ 与 $\mathcal{X}$ 都有限，

$$\mathcal{F}:=\mathcal{P}(\mathcal{X}{)}^{\otimes \mathbb{T}}=\mathcal{P}(\Omega ).$$

证明。 任一 $\sigma$-代数都含于幂集，故
$\mathcal{P}(\mathcal{X}{)}^{\otimes \mathbb{T}}\subseteq \mathcal{P}(\Omega )$。
反方向只需证明每个单点集属于乘积 $\sigma$-代数。记
${\pi }_t:\Omega \to \mathcal{X}$ 为第 $t$ 个坐标投影。对任意轨迹
$x=(x_t{)}_{t\in \mathbb{T}}$，

$$\{x\}=\bigcap _{t\in \mathbb{T}}{\pi }_t^{-1}(\{x_t\}).$$

每个 $\{x_t\}$ 属于 $\mathcal{P}(\mathcal{X})$，故对应柱集属于乘积
$\sigma$-代数。$\mathbb{T}$ 有限，有限交仍可测。又因 $\Omega$ 有限，
任意 $A\subseteq \Omega$ 都是有限个单点集之并，所以
$\mathcal{P}(\Omega )\subseteq \mathcal{P}(\mathcal{X}{)}^{\otimes \mathbb{T}}$。$■$

坐标投影

$$X_t:\Omega \to \mathcal{X},X_t(x_0,\dots ,x_T):=x_t$$

都是可测映射，从而 $X=(X_t{)}_{t\in \mathbb{T}}$ 是该可测空间上的随机过程。

1.2 概率空间

前向过程与生成模型共用可测空间 $(\Omega ,\mathcal{F})$，但使用不同的概率测度：

$$(\Omega ,\mathcal{F},Q),(\Omega ,\mathcal{F},P_{\varphi }).$$

$\Omega$ 有限，二者由单点质量 $Q(x_{0:T})$、$P_{\varphi }(x_{0:T})$ 完全确定。数据分布 $P_{\mathrm{data}}$ 与参考分布 $P_T$ 则是 $(\mathcal{X},\mathcal{P}(\mathcal{X}))$ 上的概率测度。

边缘分布是坐标映射的像测度。例如

$$(Q\circ X_t^{-1})(\{j\})=Q(X_t=j).$$

要求

$$Q\circ X_0^{-1}=P_{\mathrm{data}},Q\circ X_T^{-1}\approx P_T.$$

概率核在有限状态空间上可由矩阵表示。对每个 $t=1,\dots ,T$，一步前向核为

$$q_t:\mathcal{X}\times \mathcal{P}(\mathcal{X})\to [0,1],q_t(i,\{j\}):=Q(X_t=j\mid X_{t-1}=i).$$

于是两条概率路径为

$$X_0\stackrel{Q}{\to }{X}_1\stackrel{Q}{\to }\cdots \stackrel{Q}{\to }{X}_T,X_T\stackrel{P_{\varphi }}{\to }{X}_{T-1}\stackrel{P_{\varphi }}{\to }\cdots \stackrel{P_{\varphi }}{\to }{X}_0.$$

前向过程逐渐丢失 $X_0$ 的信息，逆向过程尝试逐步恢复数据。

2. 前向扩散过程

从本节起额外假设前向过程是一阶非齐次马尔可夫链：对 $t=1,\dots ,T$，

$$Q(X_t=j\mid X_0,\dots ,X_{t-1})=Q(X_t=j\mid X_{t-1})$$

在条件概率有定义处成立。

2.1 一步转移

用行随机矩阵记转移核：

$${\mathbf{Q}}_t\in [0,1{]}^{K\times K},({\mathbf{Q}}_t{)}_{ij}:=q_t(i,\{j\}),\sum _{j=1}^K({\mathbf{Q}}_t{)}_{ij}=1.$$

若 $Q(X_{t-1}=i)=0$，对应行的条件概率版本可任取为一个概率向量。本文约定起点对应行、终点对应列；若 ${\mathbf{e}}_i$ 是第 $i$ 个标准基列向量，则从状态 $i$ 出发的一步分布是行向量 ${\mathbf{e}}_i^{\top }{\mathbf{Q}}_t$。

链式法则与马尔可夫性给出

$$\begin{array}{rl}Q(x_{0:T})& =Q(x_0)\prod _{t=1}^{T}Q(x_t\mid x_{0:t-1})\\ & =P_{\mathrm{data}}(x_0)\prod _{t=1}^{T}Q(x_t\mid x_{t-1})\\ & =P_{\mathrm{data}}(x_0)\prod _{t=1}^T({\mathbf{Q}}_t{)}_{x_{t-1},x_t},\end{array}$$

故

$$Q(x_{1:T}\mid x_0)=\prod _{t=1}^T({\mathbf{Q}}_t{)}_{x_{t-1},x_t}.$$

2.2 任意时刻的边缘分布

定义累积转移矩阵

$${\overline{\mathbf{Q}}}_t:={\mathbf{Q}}_1{\mathbf{Q}}_2\cdots {\mathbf{Q}}_t,{\overline{\mathbf{Q}}}_0:=\mathbf{I}.$$

由 Chapman—Kolmogorov 方程，

$$Q(X_t=j\mid X_0=i)=({\overline{\mathbf{Q}}}_t{)}_{ij},$$

即

$$X_t\mid (X_0=i)\sim \mathrm{Cat}\left({\mathbf{e}}_i^{\top }{\overline{\mathbf{Q}}}_t\right).$$

训练时可以从 $X_0$ 直接采样任意 $X_t$，无需模拟前 $t-1$ 步。

3. 噪声矩阵

3.1 随机替换噪声

给定类别分布

$$\mathbf{\pi }\in {\Delta }^{K-1},{\Delta }^{K-1}:=\left\{\mathbf{\pi }\in {\mathbb{R}}^K:{\pi }_j\ge 0,\sum _{j=1}^K{\pi }_j=1\right\}.$$

令

$${\mathbf{Q}}_t=(1-{\beta }_t)\mathbf{I}+{\beta }_{t}1{\mathbf{\pi }}^{\top },0\le {\beta }_t\le 1.$$

每一步以概率 $1-{\beta }_t$ 保留原类别，以概率 ${\beta }_t$ 从 $\mathbf{\pi }$ 重新采样。记

$${\alpha }_t:=1-{\beta }_t,{\overline{\alpha }}_t:=\prod _{s=1}^t{\alpha }_s,$$

则

$${\overline{\mathbf{Q}}}_t={\overline{\alpha }}_t\mathbf{I}+(1-{\overline{\alpha }}_t)1{\mathbf{\pi }}^{\top },$$

以及

$$Q(X_t=j\mid X_0=i)={\overline{\alpha }}_t{1}_{\{i=j\}}+(1-{\overline{\alpha }}_t){\pi }_j.$$

当 ${\overline{\alpha }}_T\approx 0$ 时，$X_T$ 几乎不依赖 $X_0$，可取
$P_T=\mathrm{Cat}(\mathbf{\pi })$。均匀噪声是 ${\pi }_j=1/K$ 的特例：

$${\mathbf{Q}}_t=(1-{\beta }_t)\mathbf{I}+\frac{{\beta }_t}{K}11^{\top }.$$

3.2 吸收态或掩码噪声

令 $m=[\mathrm{MASK}]$ 为吸收态。若前 $K-1$ 个状态是普通类别，最后一个状态是 $m$，则

$${\mathbf{Q}}_t=\left[\begin{array}{cc}(1-{\beta }_t){\mathbf{I}}_{K-1}& {\beta }_{t}1\\ 0^{\top }& 1\end{array}\right].$$

普通类别以概率 ${\beta }_t$ 变成 $m$，而 $m$ 不会再离开。此时

$$Q(X_t=m\mid X_0\ne m)=1-{\overline{\alpha }}_t.$$

4. 精确后验

训练逆向模型所需的关键量是

$$Q(X_{t-1}\mid X_t,X_0).$$

固定 $X_0=k$、$X_t=j$。对任意中间状态 $i$，由贝叶斯公式和马尔可夫性，

$$\begin{array}{rl}Q(X_{t-1}=i\mid X_t=j,X_0=k)& =\frac{Q(X_t=j\mid X_{t-1}=i)Q(X_{t-1}=i\mid X_0=k)}{Q(X_t=j\mid X_0=k)}\\ & =\frac{({\mathbf{Q}}_t{)}_{ij}({\overline{\mathbf{Q}}}_{t-1}{)}_{ki}}{({\overline{\mathbf{Q}}}_t{)}_{kj}}.\end{array}$$

因此后验概率向量为

$${\mathbf{r}}_{t-1}=\frac{({\mathbf{Q}}_t{)}_{:,j}\odot ({\overline{\mathbf{Q}}}_{t-1}{)}_{k,:}^{\top }}{({\overline{\mathbf{Q}}}_t{)}_{kj}},$$

其中 $\odot$ 是逐元素乘法。分母就是分子元素之和；实现时通常先算未归一化分子再归一化。该公式只在 $({\overline{\mathbf{Q}}}_t{)}_{kj}>0$ 的支持集上使用。

5. 逆向模型

生成模型定义为

$$P_{\varphi }(x_{0:T})=P_T(x_T)\prod _{t=1}^T{P}_{\varphi }(x_{t-1}\mid x_t,t).$$

最直接的方法是让网络输出 $K$ 个 logits，直接参数化
$P_{\varphi }(X_{t-1}\mid X_t,t)$。D3PM 中也常让网络预测干净状态：

$${\tilde{P}}_{\varphi }(X_0=k\mid X_t=j,t).$$

再用已知前向核构造未归一化权重

$$\begin{array}{rl}{w}_i& :=\sum _{k=1}^{K}Q(X_{t-1}=i,X_t=j\mid X_0=k){\tilde{P}}_{\varphi }(X_0=k\mid X_t=j,t)\\ & =({\mathbf{Q}}_t{)}_{ij}\sum _{k=1}^K({\overline{\mathbf{Q}}}_{t-1}{)}_{ki}{\tilde{P}}_{\varphi }(X_0=k\mid X_t=j,t),\end{array}$$

并在候选前态上归一化：

$$P_{\varphi }(X_{t-1}=i\mid X_t=j,t)=\frac{w_i}{{\sum }_{r=1}^K{w}_r}.$$

这相当于离散模型中的“预测 $x_0$”参数化：网络负责识别干净数据，已知前向过程负责计算一步去噪分布。若 ${\tilde{P}}_{\varphi }$ 把全部质量放在真实 $X_0=k$ 上，上式就还原为第 4 节的精确后验。

6. 训练目标

6.1 负 ELBO

对数据样本 $x_0$，

$$-\log P_{\varphi }(x_0)\le {\mathcal{L}}_{\mathrm{VLB}}(x_0):={\mathbb{E}}_{Q(x_{1:T}\mid x_0)}\left[\log \frac{Q(x_{1:T}\mid x_0)}{P_{\varphi }(x_{0:T})}\right].$$

它分解为

$${\mathcal{L}}_{\mathrm{VLB}}={\mathcal{L}}_T+\sum _{t=2}^T{\mathcal{L}}_{t-1}+{\mathcal{L}}_0,$$

其中

$${\mathcal{L}}_T=D_{\mathrm{KL}}\left(Q(X_T\mid x_0)\Vert P_T(X_T)\right),$$ $${\mathcal{L}}_{t-1}={\mathbb{E}}_{Q(X_t\mid x_0)}\left[D_{\mathrm{KL}}\left(Q(X_{t-1}\mid X_t,x_0)\Vert P_{\varphi }(X_{t-1}\mid X_t,t)\right)\right],$$

以及

$${\mathcal{L}}_0={\mathbb{E}}_{Q(X_1\mid x_0)}\left[-\log P_{\varphi }(x_0\mid X_1,1)\right].$$

最后再对 $x_0\sim P_{\mathrm{data}}$ 取期望。由于状态空间有限，这些 KL 散度可由分类概率精确计算，不需要连续变量模型中的重参数化技巧。

6.2 辅助交叉熵

采用 $X_0$ 参数化时，还可加入

$${\mathcal{L}}_{\mathrm{CE}}={\mathbb{E}}_{\begin{array}{c}{x}_0\sim P_{\mathrm{data}},\\ t\sim \mathrm{Unif}\{1,\dots ,T\},\\ X_t\sim Q(\cdot \mid x_0)\end{array}}\left[-\log {\tilde{P}}_{\varphi }(x_0\mid X_t,t)\right].$$

实践中常使用

$$\mathcal{L}={\mathcal{L}}_{\mathrm{VLB}}+\lambda {\mathcal{L}}_{\mathrm{CE}}.$$

$\lambda$ 控制辅助目标权重。交叉熵通常更易优化，但只有完整 VLB 直接对应似然界。

7. 训练与采样流程

一次训练迭代：

1. 采样 $x_0\sim P_{\mathrm{data}}$；
2. 采样 $t\in \{1,\dots ,T\}$；
3. 通过 ${\overline{\mathbf{Q}}}_t$ 直接采样
   $x_t\sim Q(x_t\mid x_0)$；
4. 网络根据 $(x_t,t)$ 输出逆向分布或
   ${\tilde{P}}_{\varphi }(x_0\mid x_t,t)$；
5. 用闭式后验计算对应 KL 项，并按时间采样分布做重要性加权。

生成流程：

1. 从参考分布采样 $x_T\sim P_T$；
2. 对 $t=T,T-1,\dots ,1$，采样
   $x_{t-1}\sim P_{\varphi }(x_{t-1}\mid x_t,t)$；
3. 返回 $x_0$。

与 DDPM 一样，标准采样需要 $T$ 次串行网络调用，噪声步数会直接影响生成延迟。

8. 从单个类别推广到序列

文本、离散图像等数据通常是

$${\mathbf{X}}_t=(X_t^{(1)},\dots ,X_t^{(D)})\in {\mathcal{X}}^D.$$

状态可测空间换成乘积空间
$({\mathcal{X}}^D,\mathcal{P}(\mathcal{X}{)}^{\otimes D})$。有限 $\mathcal{X}$ 下，它仍等于 $({\mathcal{X}}^D,\mathcal{P}({\mathcal{X}}^D))$。路径空间相应为

$${\Omega }_D:=({\mathcal{X}}^D{)}^{\mathbb{T}},{\mathcal{F}}_D:=\mathcal{P}({\Omega }_D).$$

最常见的前向过程对各位置独立加噪：

$$Q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=\prod _{d=1}^D({\mathbf{Q}}_t{)}_{x_{t-1}^{(d)},x_t^{(d)}},$$

从而

$$Q({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\prod _{d=1}^D({\overline{\mathbf{Q}}}_t{)}_{x_0^{(d)},x_t^{(d)}}.$$

逆向分布通常也按位置分解，

$$P_{\varphi }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,t)=\prod _{d=1}^D{P}_{\varphi }(x_{t-1}^{(d)}\mid {\mathbf{x}}_t,t),$$

但每个因子都可依赖完整噪声序列 ${\mathbf{x}}_t$。“条件独立”不等于“模型忽略位置间依赖”：Transformer 等去噪网络仍会利用全局上下文预测各位置。数据依赖主要由逆向模型学习，前向核保持简单、可采样且后验可计算。

9. 常见误区

• 离散指状态而不是概率。 更准确的说法是“离散状态空间上的扩散”，概率值仍在 $[0,1]$ 中连续变化。
• $P_T$ 不会自动等于终端边缘。 日程必须让 $Q(X_T\mid X_0)$ 充分接近参考分布，否则 ${\mathcal{L}}_T$ 会变大。
• 矩阵乘法顺序取决于约定。 本文使用行随机矩阵，因此 ${\overline{\mathbf{Q}}}_t={\mathbf{Q}}_1\cdots {\mathbf{Q}}_t$。
• 多步掩码扩散不等于 BERT。 特定的一步 D3PM 可还原 BERT 的去噪目标；生成式 D3PM 还需要终端先验和完整逆向链。
• 类别不一定适合均匀替换。 若类别有序、稀疏或具有图结构，结构化转移通常比任意类别间跳转更合理。

参考资料

1. Austin et al., Structured Denoising Diffusion Models in Discrete State-Spaces, 2021.
2. Hoogeboom et al., Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions, 2021.
3. Sohl-Dickstein et al., Deep Unsupervised Learning using Nonequilibrium Thermodynamics, 2015.