Structured Denoising Diffusion Models in Discrete State-Spaces

目录

• 概
• 符号说明
• Motivation
• 基于转移概率矩阵的 D3PM
• 转移概率矩阵的设计
  • Uniform diffusion
  • Diffusion with an absorbing state
• 代码

Austin J., Johnson D. D., Ho J., Tarlow D. and van den Berg R. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems (NIPS), 2021.

概

DPM 在离散空间上的探讨.

符号说明

• forward process: $q(\bm{x}_{1:T}|\bm{x}_0) = \prod_{t=1}^T q(\bm{x}_t|\bm{x}_{t-1}), \bm{x}_0 \sim q(\bm{x}_0)$;
• reverse process: $p_{\theta}(\bm{x}_{0:T}) = p(\bm{x}_T) \prod_{t=1}^T p_{\theta} (\bm{x}_{t-1}|\bm{x}_t)$.
• 这里向量 $\bm{x}$ 都表示行向量.

Motivation

• DDPM 已经在图像生成领域取得了巨大成功, 它分为上述的前向和后向过程, 并通过如下损失进行优化:

  $$\tag{1} \begin{array}{l} L_{\text{vb}} = \mathbb{E}_{q(\bm{x}_0)} \Bigg [ \underbrace{D_{KL}[q(\bm{x}_T|\bm{x}_0 \| p(\bm{x}_T))]}_{L_T} + \sum_{t=2}^T \underbrace{\mathbb{E}_{q(\bm{x}_t|\bm{x}_0)} [D_{KL} [q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0)\| p_{\theta}(\bm{x}_{t-1}|\bm{x}_t)]]}_{L_{t-1}} \\ \quad \quad \quad \quad \quad \quad \underbrace{-\mathbb{E}_{q(\bm{x}_1|\bm{x}_0)}[\log p_{\theta}(\bm{x}_0| \bm{x}_1)]}. \Bigg] \end{array}$$

• 之前的扩散模型主要研究连续空间的情形 (故而通常采用比较好采样的高斯分布), 那么如何把 DPM 推广到离散空间中呢.

• 首先, 合适的前向扩散过程需要满足:

  1. $q(\bm{x}_t|\bm{x}_0)$ 对任意的时间戳 $t$ 的采样都是容易的;
  2. $q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0)$ 最好是有容易求解的显示表达, 这样有利于求解 $L_{t-1}$.

基于转移概率矩阵的 D3PM

• 假设离散空间中有 $K$ 个元素, 记为 $1,2, \ldots, K$, 则对于任意的时间戳 $t$ 的状态 $x_t$ 都可以取遍这些元素. 我们知道, DPM 通常假设前向以及反向过程都是满足马氏性的, 而马氏性通常可以用转移概率矩阵 $Q$ 来描述.

• 令 $Q_t \in \mathbb{R}^{K \times K}$, $[Q_t]_{ij} = q(x_t = j | x_{t-1} = i)$, 则

  $$q(x_t = j|x_{t-1} = i) = [\underbrace{0, 0, \cdots, 0}_{i-1}, 1, 0, \cdots, 0] Q_t [\underbrace{0, 0, \cdots, 0}_{j-1}, 1, 0, \cdots, 0]^T.$$

• 不妨假设 $\bm{x}_t \in \mathbb{R}^{1 \times K}$ 为一 one-hot 向量 (注意和之前的定义区分), 即 $[\underbrace{0, 0, \cdots, 0}_{j-1}, 1, 0, \cdots, 0]$, 则,

  $$q(\bm{x}_t| \bm{x}_{t-1}) = \bm{x}_{t-1} Q_t \bm{x}_t^T \sim \text{Cat}(\bm{x}_t; \bm{p} = \bm{x}_{t-1}Q_t).$$

• 进而注意到

  $$\begin{array}{ll} q(\bm{x}_t|\bm{x}_0) &= \sum_{\bm{x}_{1:t-1}} \prod_{k=1}^{t} q(\bm{x}_k|\bm{x}_{k=1}) \\ &= \sum_{\bm{x}_{1:t-1}} \prod_{k=1}^{t} \bm{x}_{k-1}Q_k \bm{x}_{k}^T \\ &= \sum_{\bm{x}_{1:t-1}} \bm{x}_0 Q_1 \bm{x}_1^T \cdots \bm{x}_{k-1}Q_k \bm{x}_{k}^T \cdots \bm{x}_{t-1} Q_t \bm{x}_t^T\\ &= \bm{x}_0 Q_1 (\sum_{\bm{x}_{1}} \bm{x}_1^T \bm{x}_1) \cdots (\sum_{\bm{x}_{k-1}} \bm{x}_{k-1}^T \bm{x}_{k-1}) Q_k (\sum_{\bm{x}_k}\bm{x}_{k}^T \bm{x}_k) \cdots (\sum_{\bm{x}_{t-1}}\bm{x}_{t-1}^T \bm{x}_{t-1}) Q_t \bm{x}_t^T\\ &= \bm{x}_0 Q_1 I Q_2 \cdots I Q_k I \cdots I Q_t \bm{x}_t^T\\ &= \bm{x}_0 \underbrace{Q_1 Q_2 \cdots Q_k \cdots Q_t}_{\bar{Q}_t} \bm{x}_t^T\\ &= \bm{x}_0 \bar{Q}_t \bm{x}_t^T \sim \text{Cat}(\bm{x}_t; \bm{p} = \bm{x}_{0}\bar{Q}_t). \end{array}$$

• 以及

  $$\begin{array}{ll} q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) &= \frac{q(\bm{x}_t|\bm{x}_{t-1}, \bm{x}_0) q(\bm{x}_{t-1}|\bm{x}_0)}{q(\bm{x}_t|\bm{x}_0)} \\ &= \frac{q(\bm{x}_t|\bm{x}_{t-1}) q(\bm{x}_{t-1}|\bm{x}_0)}{q(\bm{x}_t|\bm{x}_0)} \\ &= \frac{\bm{x}_{t-1}Q_t \bm{x}_t^T \cdot \bm{x}_0 \bar{Q}_{t-1} \bm{x}_{t-1}^T}{\bm{x}_0 \bar{Q}_t \bm{x}_{t}^T} \\ &= \frac{(\bm{x}_{t}Q_t^T \bm{x}_{t-1}^T) \cdot (\bm{x}_0 \bar{Q}_{t-1} \bm{x}_{t-1}^T)}{\bm{x}_0 \bar{Q}_t \bm{x}_{t}^T} \\ &= \frac{(\bm{x}_{t}Q_t^T) \odot (\bm{x}_0 \bar{Q}_{t-1}) (\bm{x}_{t-1}^T)}{\bm{x}_0 \bar{Q}_t \bm{x}_{t}^T} \: \leftarrow (?) \\ &\sim \text{Cat}(\bm{x}_{t-1}; \bm{p} = \frac{(\bm{x}_{t}Q_t^T) \odot (\bm{x}_0 \bar{Q}_{t-1})}{\bm{x}_0 \bar{Q}_t \bm{x}_{t}^T} ). \end{array}$$

  需要注意的是, (?) 处成立完全是因为 $\bm{x}_t$ 是 one-hot 的向量, 否则一般情况下是不成立的.

• 现在, 我们以及搭建好了基于转移概率矩阵的一个整体框架, 下面我们就需要讨论如何设计 $Q_t$ 使得 $\bar{Q}_t$ 是易求得. 此外, $Q_t$ 的行和应当为 $1$, 且 $\bar{Q}_t, t \rightarrow +\infty$ 应当是收敛的.

转移概率矩阵的设计

• 当 $Q$ 是 double stochastic 的时候, 即 行和, 列和 均为1, 此时倘若 $Q_t = Q$, 则该分布必收敛到均匀分布:
  $$[\bm{\pi}Q]_j = \sum_{i=1}^K \frac{1}{K} [Q]_{ij} = \frac{1}{K} = \pi_j.$$

接下来, 我们具体介绍几种设计.

Uniform diffusion

• 构造:

  $$[Q_t]_{ij} = \left \{ \begin{array}{ll} 1 - \frac{K - 1}{K} \beta_t & \text{ if } j = i, \\ \frac{1}{K} \beta_t & \text{ if } j \not = i. \end{array} \right .$$

• 用人话说就是, 知道了上一个状态 $x_{t-1}$, $x_t$ 仍为 $x_{t-1}$ 的概率为 $1 - \frac{K-1}{K} \beta_t$, 否则它等概率地成为其它状态.

• 矩阵表示为

  $$Q_t = (1 - \beta_t) I + \beta_t \bm{1}^T \bm{1} / K.$$

• 容易发现:

  $$\bar{Q}_t = Q_{t}Q_{t-1} \cdots Q_1 = \prod_{k=1}^t (1 - \beta_k) I + (1 - \prod_{k=1}^t (1 - \beta_k)) \bm{1}^T \bm{1} / K.$$

Diffusion with an absorbing state

• 构造

  $$[Q_t]_{ij} = \left \{ \begin{array}{ll} 1 & \text{ if } j=i=m \\ 1 - \beta_t & \text{ if } j = i \not= m, \\ \beta_t & \text{ if } j = m, i \not= m. \end{array} \right .$$

  这里 $m$ 是一个特殊的吸收态, 比如语言模型里的 [mask].

• 用人话说就是,

  1. 如果 $x_{t-1}$ 已经是 $m$ 了, 那么 $x_t$ 必为 $m$;
  2. 如果 $x_{t-1}$ 不为 $m$, 则 $x_t$ 以 $\beta_t$ 的概率成为 $m$, 否则仍为 $x_{t-1}$.

• 矩阵表示为

  $$Q_t = (1 - \beta_t) I + \beta_t \bm{1}^T \bm{e}_m,$$

  这里 $\bm{e}_m$ 表示第 $m$ 个元素为 1, 其余均为 0 的行向量.

• 显然, 这不是一个 double stochastic 的转移矩阵.

• 容易发现:

  $$\bar{Q}_t = Q_{t}Q_{t-1} \cdots Q_1 = \prod_{k=1}^t (1 - \beta_k) I + (1 - \prod_{k=1}^t (1 - \beta_k)) \bm{1}^T \bm{e}_m.$$

注: 作者还介绍了如何把高斯离散化, 以及利用 embedding 的相似度来刻画转移矩阵的方法, 这里不作介绍了.

注: 作者对于 noise schedule 的一些见解非常有趣, 但是说实话, 我只能看懂一半, 写在这里就不懂装懂了.

代码

[official]