DDPM 入门

  1. 由于渲染器的限制,多行公式只能分配一个标号,以式 (9.8) 代表式 (9) 第 8 行,以此类推。
  2. 本文的噪声项所使用的符号与原文稍有出入,请读者自行斟酌。

前置

高斯混合模型

一个复杂分布 PθP_{\theta} 可以用 KK 个高斯分布来表示。

Pθ=i=1KP(zi)Pθ(xzi)(1)P_{\theta} = \sum_{i=1}^{K} P (z_i) P_{\theta} (x\vert z_i) \tag{1}

由于 P(z)=1\int P(z) = 1,用一个连续的高斯分布来表示 P(z)N(0,1)P(z) \sim \mathcal{N}(0, 1),因此

Pθ=P(zi)Pθ(xzi)(2)P_{\theta} = \int P (z_i) * P_{\theta} (x\vert z_i) \tag{2}

KL 散度

DKL(pq)=xp(x)logq(x)p(x)dx(m)(3)D_{KL}(p\|q)= -\int\limits_{x} p(x) \log \frac{q(x)}{p(x)} dx \quad (m \to \infty) \tag{3}

重参数

重参数 (Reparameterization) 是处理期望形式目标函数的一种技巧。

Lθ=Ezpθ(z)[f(z)](4)\mathcal{L}_{\theta} = \mathbb{E}_{z \sim p_{\theta}(z)} \left[ f(z) \right] \tag{4}

为使得采样过程能够保留 θ\theta 信息,假设从分布 pθ(z)p_{\theta}(z) 中采样可以分解为两个步骤:

  1. 从无参数分布 q(ϵ)q(\epsilon) 中采样一个 ϵ\epsilon
  2. 通过变换 z=gθ(ϵ)z = g_{\theta}(\epsilon) 生成 zz

那么,式 (4) 就变成了

Lθ=Eεq(ε)[f(gθ(ε))](5)\mathcal{L}_{\theta} = \mathbb{E}_{\varepsilon \sim q(\varepsilon)} \left[ f(g_{\theta}(\varepsilon)) \right] \tag{5}

被采样的分布就没有任何参数了,全部被转移到 ff 内部。

对高斯分布而言,由于

zN(μ,σ2)zμσN(0,1)(6)z \sim \mathcal{N}(\mu, \sigma^2) \Leftrightarrow \frac{z - \mu}{\sigma} \sim \mathcal{N}(0, 1) \tag{6}

则其重参数化

EzN(z;μθ,σθ2)[f(z)]=EεN(ε;0,1)[f(ε×σθ+μθ)](7)\mathbb{E}_{z \sim \mathcal{N}(z; \mu_\theta, \sigma_\theta^2)} \left[ f(z) \right] = \mathbb{E}_{\varepsilon \sim \mathcal{N}(\varepsilon; 0, 1)} \left[ f(\varepsilon \times \sigma_\theta + \mu_\theta) \right] \tag{7}

也就是将「从 N(z;μθ,σθ2)\mathcal{N}(z; \mu_{\theta}, \sigma^2_{\theta}) 中采样一个 zz」转化为:

  1. N(ϵ;0,1)\mathcal{N}(\epsilon; 0, 1) 中采样一个 ε\varepsilon
  2. 计算 z=ε×σθ+μθz = \varepsilon \times \sigma_{\theta} + \mu_{\theta}

优化目标

初衷是希望生成模型所产生的图片所遵从的概率分布尽可能接近真实世界图片的概率分布。使用 KL 散度描述,得到:

argminθKL(PdataPθ)(8)\arg\min_{\theta} \mathrm{KL}(P_{\text{data}} \parallel P_{\theta}) \tag{8}

其中 PθP_{\theta} 是模型所预测的分布,PdataP_{data} 为真实分布。对上式做如下变换

argminθKL(PdataPθ)=argminθxPdata(x)logPθ(x)Pdata(x)dx=argmaxθxPdata(x)logPθ(x)Pdata(x)dx=argmaxθxPdata(x)logPθ(x)dxargmaxθxPdata(x)logPdata(x)dx=argmaxθxPdata(x)logPθ(x)dx=argmaxθExPdata[logPθ(x)]argmaxθi=1mlogPθ(xi)=argmaxθlogi=1mPθ(xi)=argmaxθi=1mPθ(xi)(9)\begin{aligned} \arg\min_{\theta} \mathrm{KL}(P_{\text{data}} \parallel P_{\theta}) &= \arg\min_{\theta} -\int\limits_{x} P_{\text{data}}(x) \log\frac{P_{\theta}(x)}{P_{\text{data}}(x)} dx \\ &= \arg\max_{\theta} \int\limits_{x} P_{\text{data}}(x) \log\frac{P_{\theta}(x)}{P_{\text{data}}(x)} dx \\ &= \arg\max_{\theta} \int\limits_{x} P_{\text{data}}(x) \log P_{\theta}(x) dx - \arg\max_{\theta} \int\limits_{x} P_{\text{data}}(x) \log P_{\text{data}}(x) dx \\ &= \arg\max_{\theta} \int\limits_{x} P_{\text{data}}(x) \log P_{\theta}(x) dx \\ &= \arg\max_{\theta} \mathbb{E}_{x\sim P_{\text{data}}}[\log P_{\theta}(x)] \\ &\approx \arg\max_{\theta} \sum\limits_{i=1}^{m} \log P_{\theta}(x_i) \\ &= \arg\max_{\theta} \log \prod\limits_{i=1}^{m} P_{\theta}(x_i) \\ &= \arg\max_{\theta} \prod\limits_{i=1}^{m} P_{\theta}(x_i) \tag{9} \end{aligned}

式 (9.8) 得到新的优化目标 argmaxθi=1mPθ(xi)\arg\max_{\theta} \prod\limits_{i=1}^{m} P_{\theta}(x_i)

Lower Bound

要使式 (9.8) 成立,即极大似然估计 (Maximum Likelihood Estimate)

MLEθ=arg maxθPθ(x)(10)\text{MLE} \quad \theta^{*} = \argmax_{\theta} P _{\theta}(x) \tag{10}

q(zx)q(z \vert x) 是某个概率分布 (实际上是加噪过程的分布 qϕq_{\phi},通常记作 q(x1..Tx0)q(x_{1..T}|x_0)ϕ\phi 代表加噪过程的参数),则

log(Pθ(x))=log(Pθ(x))q(zx)dz=log(Pθ(x,z)Pθ(zx))q(zx)dz=log(Pθ(x,z)q(zx)q(zx)Pθ(zx))q(zx)dz=log(Pθ(x,z)q(zx))q(zx)dz+KL(q(zx)Pθ(zx))log(Pθ(x,z)q(zx))q(zx)dz=Eq(ZX)[log(Pθ(x,z)q(zx))](11)\begin{aligned} \log(P_\theta(x)) &= \log(P_\theta(x)) \cdot \int q(z|x) \, dz \\ &= \int \log\left(\frac{P_\theta(x,z)}{P_\theta(z|x)}\right) q(z|x) \, dz \\ &= \int \log\left(\frac{P_\theta(x,z)}{q(z|x)} \cdot \frac{q(z|x)}{P_\theta(z|x)}\right) q(z|x) \, dz \\ &= \int \log\left(\frac{P_\theta(x,z)}{q(z|x)}\right) q(z|x) \, dz + \text{KL}(q(z|x) \parallel P_\theta(z|x)) \\ &\geq \int \log\left(\frac{P_\theta(x,z)}{q(z|x)}\right) q(z|x) \, dz \\ &= \mathbb{E}_{q(Z|X)}\left[\log\left(\frac{P_\theta(x,z)}{q(z|x)}\right)\right] \\ \tag{11} \end{aligned}

式 (11.1) → (11.2),由 Pθ(x,z)=Pθ(x)Pθ(zx)P_{\theta}(x, z) = P_{\theta}(x) * P_{\theta}(z | x)

式 (11.4),由 KL(pq)=plogpq0\text{KL}(p\|q) = \int{p\log{\frac{p}{q}}} \ge 0

因此,要优化最大对数似然,就是在优化 Evidence Lower Bound (ELBO). 至此,得到初步的优化目标,下面开始讲解 DDPM.

DDPM

Overview

通俗理解 DDPM。允许多步 (TT) 生成过程,从而分多次估计噪声 ϵ^t\hat{\epsilon}_t 比一步估计图片 xTx^0x_T → \hat{x}_0 容易,或者说逐步学习去噪过程的分布 pθp_{\theta} 比直接学习 PθP_{\theta} 简单。DDPM 的大部分数学框架在更早时候已有之,而该工作的最大贡献其实是用 U-Net 作为噪声估计器调通了这一模型,并得到了不错的效果。

前向扩散过程

给定初始数据分布 x0q(x)x_0 \sim q(x),前向扩散过程向 x0x_0 逐步添加 TT 次高斯噪声,得到一系列带噪声图片 x1,...,xTx_1, ..., x_T。定义 xt1xtx_{t-1} \rightarrow x_t 的加噪过程:

q(xtxt1)=N(xt;1βtxt1,βtI)(12)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}) \tag{12}

q(x1:Tx0)=t=1Tq(xtxt1)(13)q(\mathbf{x}_{1:T} | \mathbf{x}_0) = \prod_{t=1}^T q(\mathbf{x}_t | \mathbf{x}_{t-1}) \tag{13}

从而,加噪过程是一个人为定义的马尔可夫链,满足分布 N(xt;1βtxt1,βtI)\mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t \mathbf{I}),生成 xt1xtx_{t-1} \rightarrow x_t 的过程只由 {βt,xt1}\{\beta_t, x_{t-1}\} 确定。β1,...,βT\beta_1, ..., \beta_T 是超参数,提前确定 β1,...,βT\beta_1, ..., \beta_T 后即可得到加噪过程 x1,...,xTx_1, ..., x_T 任意一步。

关于 qϕq_{\phi} 的设计。为什么选择 1βt\sqrt{1-\beta_t}βt\sqrt{\beta_t} 用于修饰加噪过程的均值和方差,并不知道 DDPM 的前置工作中是否有提及。但在训练过程中,这个设计可以方便地引出 αt\alpha_tαˉt\bar{\alpha}_t 的定义,从而实现对 x0xtx_0 \rightarrow x_t 的一步加噪。

随着加噪过程的进行, βt\beta_t 不断 (线性) 增大,即每一步的噪声比例不断升高。注意,此时我们实际上只定义了加噪过程满足的分布 qϕq_{\phi},尚未定义实际的加噪过程,这部分将在「噪声采样」一节说明。但可以明确的是,当 TT \rightarrow \inftyxTx_T 趋近于一个各向独立的高斯分布,视觉上就是一张接近纯粹高斯噪声的「雪花」图。

噪声采样

上一节提到,正向扩散过程向 x0\mathbf{x}_0 逐步添加 TT 次高斯噪声。这可以写作

xT=x0+η1+η2+...+ηT(14)\mathbf{x}_T = \mathbf{x}_0 + \eta_1 + \eta_2 + ... + \eta_T \tag{14}

反向生成过程则需要通过噪声估计器对噪声进行预测。考虑某一步去噪过程 xt1xt\mathbf{x}_{t-1} \leftarrow \mathbf{x}_t,我们希望用 pθ(xt1xt)p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) 去预测 q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_{t}),由于是期望形式的优化目标,直接采样,梯度信息却无法传递到 θ\theta 上。由式 (7),使用重参数化技巧,每次加噪从 εtN(0,I)\varepsilon_t \sim \mathcal{N}(0, I) 采样,令

xt=1βtxt1+βtεt1(15)\mathbf{x}_t = \sqrt{1 - \beta_t}\mathbf{x}_{t-1} + \sqrt{\beta_t}\varepsilon_{t-1} \tag{15}

此时就实现了将高斯噪声项参数化,得到的 xt\mathbf{x}_t 服从式 (12) 分布。

有同学会问,βt\beta_t 不是超参吗,努力把随机性转移到 εt\varepsilon_t 做什么,为什么不能直接从式 (12) 采样呢。其实加噪过程是可以这样做的,但重参数化同时也面向去噪过程,下文将说明,后验扩散条件概率 q(xt1xt,x0)q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) 是 Markov 的,q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t) 则不是,所以才要用 pθ(xt1xt)p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t) 预测 q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t),这个过程就要用神经网络对 θ\theta 优化,从而需要传递梯度。

实际上,先有一个类似式 (14) (15) 的符合 Intuition 的加噪过程,再将其解释为一个分布,这样更加容易理解。

一步加噪

x0\mathbf{x}_0 和固定序列 {βt(0,1)}t=1T\{\beta_t \in (0, 1)\}_{t=1}^{T} 可以直接得到任意一步 xt\mathbf{x}_t,从而省去前向迭代。这使得训练过程变得容易,同时方便引出逆扩散过程 q(xt1xt,x0)q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) 的定义。

αt=1βt,αt=i=1tαi\alpha_t = 1 - \beta_t,\overline{\alpha}_t = \prod_{i=1}^t \alpha_i,有

xt=1βtxt1+βtεt1=1βt(1βt1xt2+βt1εt2)+βtεt1=αtαt1xt2+αt(1αt1)εt2+1αtεt1=αtαt1xt2+1αtαt1ε...=αtx0+1αtε(16)\begin{aligned} \mathbf{x}_t &= \sqrt{1 - \beta_t}\mathbf{x}_{t-1} + \sqrt{\beta_t}\varepsilon_{t-1} \\ &=\sqrt{1 - \beta_t} (\sqrt{1 - \beta_{t-1}}\mathbf{x}_{t-2} + \sqrt{\beta_{t-1}}\varepsilon_{t-2}) + \sqrt{\beta_t}\varepsilon_{t-1} \\ &= \sqrt{\alpha_t\alpha_{t-1}}\mathbf{x}_{t-2} + \sqrt{\alpha_{t}(1 - \alpha_{t-1})} \varepsilon_{t-2} + \sqrt{1 - \alpha_{t}} \varepsilon_{t-1} \\ &= \sqrt{\alpha_t\alpha_{t-1}}\mathbf{x}_{t-2} + \sqrt{1 - \alpha_t\alpha_{t-1}}\varepsilon \\ &... \\ &= \sqrt{\overline{\alpha}_t}x_0 + \sqrt{1 - \overline{\alpha}_t} \varepsilon \tag{16} \end{aligned}

式 (16.3) → (16.4),由两个独立高斯分布相加仍为高斯分布,且均值为二者均值和,方差为二者方差和。

后文我们会把式 (16.5) 写成 xt=αtx0+1αtεt\mathbf{x}_t = \sqrt{\overline{\alpha}_t}x_0 + \sqrt{1 - \overline{\alpha}_t} \varepsilon_t,注意 εt\varepsilon_t 在这里表示 x0xt\mathbf{x}_0 \rightarrow \mathbf{x}_t 所加噪声。

由式 (6) (16) 可得

q(xtx0)=N(xt;αˉtx0,(1αˉt)I)(17)q\left( \mathbf{x}_{t} \mid \mathbf{x}_{0} \right) = \mathcal{N}\left( \mathbf{x}_{t}; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}, \left(1 - \bar{\alpha}_{t}\right) \mathbf{I} \right) \tag{17}

式 (17) 重参数化就是式 (16.5),现在可以实现一步加噪了。

反向生成过程

如果能运行逆扩散过程,即从 q(xt1xt)q\left( \mathbf{x}_{t-1} \mid \mathbf{x}_{t} \right) 中采样,那么就可以从一个随机的高斯分布 N\mathcal{N} 中重建出一个真实的原始样本。这个过程可以写作一个边缘似然

Pθ(x0)=pθ(x0:T)dx1:T(18)P_\theta(\mathbf{x}_0) = \int p_\theta(\mathbf{x}_{0:T}) \, d\mathbf{x}_{1:T} \tag{18}

上式中 dx1:Td\mathbf{x}_{1:T} 代表对所有 x1..xT\mathbf{x}_1..\mathbf{x}_T 的可能路径积分。其中

pθ(x0:T)=p(xT)pθ(xT1xT)pθ(xT2xT1)pθ(x1x0)(19)p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \cdot p_\theta(\mathbf{x}_{T-1}|\mathbf{x}_T) \cdot p_\theta(\mathbf{x}_{T-2}|\mathbf{x}_{T-1}) \cdots p_\theta(\mathbf{x}_1|\mathbf{x}_0) \tag{19}

定义 pθ(xt1xt)p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) 前,首先考虑 q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t),根据贝叶斯公式,有

q(xt1xt)=q(xtxt1)q(xt1)q(xt)(20)q(\mathbf{x}_{t-1}|\mathbf{x}_t) = \dfrac{q(\mathbf{x}_t|\mathbf{x}_{t-1})q(\mathbf{x}_{t-1})}{q(\mathbf{x}_t)} \tag{20}

研究这个分布有些困难,因为 q(xt),q(xt1)q(\mathbf{x}_{t}), q(\mathbf{x}_{t-1}) 未知。但是 q(xtx0),q(xt1x0)q(\mathbf{x}_t|\mathbf{x}_0), q(\mathbf{x}_{t-1}|\mathbf{x}_0) 已知,可以先由 q(xt1xt,x0)q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) 入手

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)(21)q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) = \dfrac{q(\mathbf{x}_t|\mathbf{x}_{t-1}, \mathbf{x}_0)q(\mathbf{x}_{t-1}|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)} \tag{21}

由式 (17)

q(xt1x0)=N(xt1;αˉt1x0,(1αˉt1)I)(22)q\left( \mathbf{x}_{t-1} \mid \mathbf{x}_{0} \right) = \mathcal{N}\left( \mathbf{x}_{t-1}; \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0}, \left(1 - \bar{\alpha}_{t-1}\right) \mathbf{I} \right) \tag{22}

由式 (12) (17) (21) (22)

q(xtxt1,x0)q(xt1x0)q(xtx0)exp{12[(xtαtxt1)2βt+(xt1αˉt1x0)21αˉt1(xtαˉtx0)21αˉt]}=exp{12[(αtβt+11αˉt1)xt122(αtxtβt+αˉt1x01αˉt1)xt1+C(xt,x0)]}(23)\begin{aligned} \frac{q(\mathbf{x}_t|\mathbf{x}_{t-1}, \mathbf{x}_0)q(\mathbf{x}_{t-1}|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)} &\propto \exp\left\{-\frac{1}{2} \left[ \frac{(\mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0)^2}{1 - \bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}_0)^2}{1 - \bar{\alpha}_t} \right]\right\} \\ &= \exp\left\{-\frac{1}{2} \left[ \left(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}\right)\mathbf{x}_{t-1}^2 - 2\left(\frac{\sqrt{\alpha_t}\mathbf{x}_t}{\beta_t} + \frac{\sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0}{1 - \bar{\alpha}_{t-1}}\right)\mathbf{x}_{t-1} + C(\mathbf{x}_t, \mathbf{x}_0) \right]\right\} \tag{23} \end{aligned}

式 (23.2) 是一个高斯分布。令

{μ~(xt,x0)=αt(1αˉt1)1αˉtxt+αˉt1(1αt)1αˉtx0σt~=βt1αˉt11αˉt(24)\left\{ \begin{aligned} &\tilde{\mu}(\mathbf{x}_t, \mathbf{x}_0) =\frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}} (1-\alpha_t)}{1 - \bar{\alpha}_t} \mathbf{x}_0 \\ &\tilde{\sigma_t} = \beta_t \cdot \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \end{aligned} \right. \tag{24}

因此

q(xtxt1,x0)q(xt1x0)q(xtx0)N(xt1;μ~(xt,x0),σ~tI)(25)\frac{q(\mathbf{x}_t|\mathbf{x}_{t-1}, \mathbf{x}_0)q(\mathbf{x}_{t-1}|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)} \sim \mathcal{N}\left(\mathbf{x}_{t-1}; \tilde{\mu}(\mathbf{x}_t, \mathbf{x}_0), \tilde{\sigma}_t\mathbf{I}\right) \tag{25}

可知,在 x0\mathbf{x}_0 已知的前提下,q(xt1xt,x0)=q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)=q(\mathbf{x}_{t-1}|\mathbf{x}_t),反向生成过程也是一个马尔可夫链,满足分布 N(xt1;μ~(xt,x0),σ~tI)\mathcal{N}\left(\mathbf{x}_{t-1}; \tilde{\mu}(\mathbf{x}_t, \mathbf{x}_0), \tilde{\sigma}_t\mathbf{I}\right).

但实际上,模型在反向过程中并不事先知道 x0\mathbf{x}_0,我们最终希望得到的是 q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t)。考虑 x0=1αˉt(xt1αˉtεt)\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( \mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t} \cdot \varepsilon_t \right),将式 (24) 改写为

{μ~(xt,x0)=1αt(xt1αt1αˉtεt)σt~2=(1αt)1αˉt11αˉt(26)\left\{ \begin{aligned} &\tilde{\mu}(\mathbf{x}_t, \mathbf{x}_0) = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \varepsilon_t \right) \\ &{\tilde{\sigma_t}}^2 = (1 - \alpha_t) \cdot \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \end{aligned} \right. \tag{26}

εt\varepsilon_t 在这里表示 x0xt\mathbf{x}_0 \rightarrow \mathbf{x}_t 所加噪声。这样一来便消去了 x0\mathbf{x}_0,并且实际上已经写出了网络要预测的对象 (即 εt\varepsilon_t),在代码中,我们将用 UNet 作为噪声预测器 εθ\varepsilon_{\theta} 对其分布进行拟合。若写成 εt=εθ(xt,t)\varepsilon_t = \varepsilon_{\theta}(\mathbf{x}_t, t),则生成 xt1xtx_{t-1} \leftarrow x_t 的过程只由 {αt,xt}\{\alpha_t, \mathbf{x}_{t}\} 确定,得到一个近似的遵循高斯分布的马尔可夫链。

根据以上讨论,整理最终表述。前文已经说明,要用 pθp_{\theta} 去拟合真实的 q(xt1xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t),因此,pθ(xt1xt)p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) 定义为高斯分布

pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t))(27)p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}\left(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t), \Sigma_\theta(\mathbf{x}_t, t)\right) \tag{27}

μθ(xt,t)\mu_{\theta}(\mathbf{x}_t, t)Σθ(xt,t)\Sigma_{\theta}(\mathbf{x}_t, t) 在下一节介绍。

整理我们得到的三个分布。前向扩散过程

q(xtxt1)=N(xt;1βtxt1,βtI)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})

x0\mathbf{x}_0 已知的逆扩散过程

q(xt1xt,x0)=N(xt;1αt(xtβt1αˉtεt),βt1αˉt11αˉtI)q(\mathbf{x}_{t-1} | \mathbf{x}_{t}, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \varepsilon_t \right), \beta_t \cdot \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \mathbf{I})

预测的反向生成过程

pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t))p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}\left(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t), \Sigma_\theta(\mathbf{x}_t, t)\right)

再谈优化目标

考虑式 (11),对于 DDPM,现在可将其扩展为链式表达的形式,写成

logPθ(x)Eqϕ(x1:Tx0)[logPθ(x0:T)qϕ(x1:Tx0)](28)\log P_{\theta}(x) \geq \mathbb{E}_{q_{\phi}(x_{1:T}|x_0)} \left[ \log \frac{P_{\theta}(x_{0:T})}{q_{\phi}(x_{1:T}|x_0)} \right] \tag{28}

改写成一个负对数似然,可以使用 Variational Lower Bound 优化,得到的结果是相同的。令

LVLB=Eqϕ(x1:Tx0)[logqϕ(x1:Tx0)pθ(x0:T)](29)L_{\text{VLB}} = \mathbb{E}_{q_{\phi}(\mathbf{x}_{1:T}|\mathbf{x}_0)} \left[ \log \frac{q_{\phi}(\mathbf{x}_{1:T}|\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \right] \tag{29}

可以进一步将式 (29) 展开成几个 KL 散度和熵项的组合 (见附录)

LVLB=Eq[DKL(q(xTx0)p(xT))+t>1DKL(q(xt1xt,x0)pθ(xt1xt))logpθ(x0x1)](30)L_{\text{VLB}} = \mathbb{E}_{q} \left[ D_{\mathrm{KL}}\left(q(\mathbf{x}_{T}|\mathbf{x}_{0}) \parallel p(\mathbf{x}_{T})\right) + \sum_{t>1}D_{\mathrm{KL}}\left(q(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{x}_{0}) \parallel p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})\right) - \log p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1}) \tag{30} \right]

分别标记 LVLBL_{VLB} 中的每个分量,将式 (30) 写成

LVLB=LT+LT1++L0where LT=DKL(q(xTx0)pθ(xT))Lt=DKL(q(xtxt+1,x0)pθ(xtxt+1)),1tT1L0=logpθ(x0x1)(31)\begin{aligned} L_{\text{VLB}} &= L_T + L_{T-1} + \cdots + L_0 \\ \text{where } L_T &= D_{\text{KL}}\left(q(\mathbf{x}_T|\mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T)\right) \\ L_t &= D_{\text{KL}}\left(q(\mathbf{x}_t|\mathbf{x}_{t+1},\mathbf{x}_0) \parallel p_\theta(\mathbf{x}_t|\mathbf{x}_{t+1})\right), \quad 1 \leq t \leq T - 1 \\ L_0 &= -\log p_\theta(\mathbf{x}_0|\mathbf{x}_1) \tag{31} \end{aligned}

式 (31) 形式的 LVLBL_{VLB} 中,LTL_T 为常数项,L0L_0 为重建项,训练过程中通常可忽略。最终优化目标由 LtL_t 决定,观察 LtL_t,由 KL 散度含义知,是要让 qqpθp_{\theta} 尽量接近。这等价于 μ~(xt,x0)μθ(xt,t),σ~tΣθ(xt,t)\tilde{\mu}(\mathbf{x}_t, \mathbf{x}_0) \to \mu_{\theta}(\mathbf{x}_t, t), \tilde{\sigma}_t \to \Sigma_{\theta}(\mathbf{x}_t, t)。注意到 σ~t\tilde{\sigma}_t 只由 α\alpha 决定,在 DDPM 中,为了简化过程并保证训练稳定性,直接将 Σθ\Sigma_{\theta} 固定为 σ~\tilde{\sigma}。此时,只需 μ~(xt,x0)μθ(xt,t)\tilde{\mu}(\mathbf{x}_t, \mathbf{x}_0) \to \mu_{\theta}(\mathbf{x}_t, t),令

μθ(xt,t)=1αt(xt1αt1αˉtεθ(xt,t))(32)\mu_{\theta}(\mathbf{x}_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \varepsilon_{\theta}(\mathbf{x}_t, t)\right) \tag{32}

对比式 (26.1) (32),现在优化目标转化为 εθ(xt,t)εt\varepsilon_{\theta}(\mathbf{x}_t, t) \to \varepsilon_t,使用 MSE 描述

Lsimple(θ):=Et,x0,ε[εtεθ(xt,t)2](33)L_{\text{simple}}(\theta) := \mathbb{E}_{t,\mathbf{x}_0,\varepsilon} \left[ \| \varepsilon_t - \varepsilon_\theta(\mathbf{x}_t, t) \|^2 \right] \tag{33}

训练和推理

训练过程的重参数化采样由式 (16.5) 给出。实际训练过程中,为每个训练样本 x0q(x0)\mathbf{x}_0 \sim q(\mathbf{x}_0) 选择一个随机时间步 tt,一步加噪后使用式 (33) 计算 Loss 并梯度下降。

推理过程。由式 (26) (27) (32),对 pθ(xt1xt)p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t) 进行重参数化采样

xt1=1αt(xt1αt1αˉtεθ(xt,t))+σtε(t>1)(34)\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \varepsilon_{\theta}(\mathbf{x}_t, t)\right) + \sigma_t\varepsilon \quad (t>1)\tag{34}

注意 σtε\sigma_t\varepsilon 项,这一项使用 εN(0,I)\varepsilon \sim \mathcal{N}(0, I) 对方差进行重参数化,与 εθ\varepsilon_{\theta} 含义不同。在 x1x0\mathbf{x}_1 \rightarrow \mathbf{x}_0 则直接输出 μθ(x1,1)\mu_{\theta}(\mathbf{x}_1, 1),不添加这一噪声项。

σtε\sigma_t\varepsilon 稍微起到了正则项的作用,可以从结果上近似认为是一种风格损失 Lstyle\mathcal{L}_{style},也可以认为它补强了高频。去噪过程中不加这一项,则重建出的 x^0\hat{\mathbf{x}}_0 是过于平滑、缺乏纹理的色块。

伪代码

Summary

断断续续数日,终于是整理完了 DDPM 公式,下一步辅以代码强化理解。虽如此,还是有一些基础知识,如流匹配等需要补充。

References

DDPM

  1. 【较真系列】讲人话-Diffusion Model 全解(原理+代码+公式)_哔哩哔哩_bilibili(讲得很不错,完整看下来基本明白 DDPM 前后向过程都在干什么)
  2. VictorYuki 视频专辑-VictorYuki 视频合集-哔哩哔哩视频(手写过程,可作为参考)
  3. 【大白话01】一文理清 Diffusion Model 扩散模型 | 原理图解+公式推导_哔哩哔哩_bilibili(入门 DDPM 看的第一个精讲,还可以)
  4. 54、Probabilistic Diffusion Model 概率扩散模型理论与完整 PyTorch 代码详细解读_哔哩哔哩_bilibili | 64、扩散模型加速采样算法 DDIM 论文精讲与 PyTorch 源码逐行解读_哔哩哔哩_bilibili(主讲代码)
  5. 深入浅出扩散模型 (Diffusion Model) 系列:基石 DDPM(人人都能看懂的数学原理篇) - 知乎
  6. DDPM 解读(一)| 数学基础,扩散与逆扩散过程和训练推理方法 - 知乎
  7. 漫谈重参数:从正态分布到 Gumbel Softmax - 科学空间 | Scientific Spaces
  8. Denoising Diffusion Probabilistic Models(DDPM)保姆级解析——附代码实现-CSDN博客
  9. 此外还看了李宏毅的视频,和之前看老师其他视频的感受一样,需要有些基础再来看,作为巩固增强。否则很多推导都是一语带过,理解起来会有问题。

DDIM

  1. 【串讲系列】讲人话-Stable Diffusion 全解(原理+代码+公式)之 DDIM + SDXL Turbo_哔哩哔哩_bilibili

DDNM

  1. denoising_diffusion_nullspace_model_ddnm__method_explained

附录

LVLBL_{VLB} 展开