首先回顾一下我们的问题,我们在逆向降噪过程中由于没办法得到 q ( x t − 1 ∣ x t ) q(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}) q(xt−1∣xt),因此我们定义了一个 需要学习的模型模型 p θ ( x t − 1 ∣ x t ) p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) pθ(xt−1∣xt)来对其进行近似,并且在训练阶段我们可以利用后验 q ( x t − 1 ∣ x t , x 0 ) q(\mathbf{x}_{t-1}\vert \mathbf{x}_t,\mathbf{x}_0) q(xt−1∣xt,x0)来对 p θ p_\theta pθ进行优化。
那么现在的问题是我们如何 p θ p_\theta pθ优化得到理想的 μ θ \boldsymbol{\mu}_\theta μθ和 Σ θ \boldsymbol{\Sigma}_\theta Σθ?类似于 VAE ,我们可以最小化在真实数据期望下,模型预测分布的负对数似然,即最小化预测 p d a t a = q ( x 0 ) p_{\mathrm{data}}=q({\mathbf{x}_0}) pdata=q(x0)和 p θ ( x 0 ) p_{\theta}(\mathbf{x}_0) pθ(x0)的交叉熵:
L = E x 0 ∼ q ( x 0 ) [ − log p θ ( x 0 ) ] \begin{equation} \mathcal{L}=\mathbb{E}_{\mathbf{x}_{0} \sim q\left(\mathbf{x}_{0}\right)}\left[-\log p_{\theta}\left(\mathbf{x}_{0}\right)\right] \end{equation} L=Ex0∼q(x0)[−logpθ(x0)]
但是,我们没法得到 p θ ( x 0 ) p_\theta(\mathbf{x}_0) pθ(x0)的表达式,因此公式1的交叉熵是没法计算的。那么可以借助公式Diffusion Model(2):前向扩散过程和逆向降噪过程
2-6 进行一些数学推导。将公式1中的 p θ ( x 0 ) p_\theta(\mathbf{x}_0) pθ(x0)转化为已知的项:
L = − E q ( x 0 ) log p θ ( x 0 ) = − E q ( x 0 ) log ( ∫ p θ ( x 0 : T ) d x 1 : T ) = − E q ( x 0 ) log ( ∫ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) q ( x 1 : T ∣ x 0 ) d x 1 : T ) = − E q ( x 0 ) log ( E q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) q ( x 1 : T ∣ x 0 ) ) ≤ − E q ( x 0 : T ) log p θ ( x 0 : T ) q ( x 1 : T ∣ x 0 ) = E q ( x 0 : T ) [ log q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) ] = L V L B \begin{equation} \begin{aligned} \mathcal{L} &=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log p_{\theta}\left(\mathbf{x}_{0}\right) \\ &=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log \left(\int p_{\theta}\left(\mathbf{x}_{0: T}\right) d \mathbf{x}_{1: T}\right) \\ &=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log \left(\int q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right) \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right)} d \mathbf{x}_{1: T}\right) \\ &=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log \left(\mathbb{E}_{q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right)} \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right)}\right) \\ & \leq-\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)} \log \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right)} \\ &=\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right]=\mathcal{L}_{\mathrm{VLB}} \end{aligned} \end{equation} L=−Eq(x0)logpθ(x0)=−Eq(x0)log(∫pθ(x0:T)dx1:T)=−Eq(x0)log(∫q(x1:T∣x0)q(x1:T∣x0)pθ(x0:T)dx1:T)=−Eq(x0)log(Eq(x1:T∣x0)q(x1:T∣x0)pθ(x0:T))≤−Eq(x0:T)logq(x1:T∣x0)pθ(x0:T)=Eq(x0:T)[logpθ(x0:T)q(x1:T∣x0)]=LVLB
上式中 q ( x 0 ) q(\mathbf{x}_0) q(x0)是真实的数据分布,而 p θ ( x 0 ) p_\theta(\mathbf{x}_0) pθ(x0)是模型,从第四行到第五行使用了Jensen不等式 log E [ f ( x ) ] ≤ E [ log f ( x ) ] \log \mathbb{E}[f(x)] \leq \mathbb{E}[\log f(x)] logE[f(x)]≤E[logf(x)]并结合了对 q ( x 0 ) q(\mathbf{x}_0) q(x0)的期望和对 q ( x 1 : T ∣ x 0 ) q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) q(x1:T∣x0)的期望。
为了最小化这个损失,结合公式2可以将其转化为最小化其上界 L V L B \mathcal{L}_{\mathrm{VLB}} LVLB:
L V L B = E q ( x 0 : T ) [ log q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) ] = E q [ log ∏ t = 1 T q ( x t ∣ x t − 1 ) p θ ( x T ) ∏ t = 1 T p θ ( x t − 1 ∣ x t ) ] = E q [ − log p θ ( x T ) + ∑ t = 1 T log q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) ] = E q [ − log p θ ( x T ) + ∑ t = 2 T log q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) + log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] = E q [ − log p θ ( x T ) + ∑ t = 2 T log ( q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) ⋅ q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) ) + log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] = E q [ − log p θ ( x T ) + ∑ t = 2 T log q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) + ∑ t = 2 T log q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) + log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] = E q [ − log p θ ( x T ) + ∑ t = 2 T log q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) + log q ( x T ∣ x 0 ) q ( x 1 ∣ x 0 ) + log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] = E q [ log q ( x T ∣ x 0 ) p θ ( x T ) + ∑ t = 2 T log q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) − log p θ ( x 0 ∣ x 1 ) ] = E q [ − log p θ ( x 0 ∣ x 1 ) ⏟ L 0 ] + ∑ t = 2 T D K L ( q ( x t − 1 ∣ x t , x 0 ) ∥ p θ ( x t − 1 ∣ x t ) ) ⏟ L t − 1 + D K L ( q ( x T ∣ x 0 ) ∥ p θ ( x T ) ) ⏟ L T \begin{equation} \begin{array}{l} \mathcal{L}_{\mathrm{VLB}}=\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right]\\ =\mathbb{E}_{q}\left[\log \frac{\prod_{t=1}^{T} q\left(\mathbf{x}_{t} \vert \mathbf{x}_{t-1}\right)}{p_{\theta}\left(\mathbf{x}_{T}\right) \prod_{t=1}^{T} p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)}\right]\\ =\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=1}^{T} \log \frac{q\left(\mathbf{x}_{t} \vert \mathbf{x}_{t-1}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)}\right]\\ =\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \frac{\color{blue}q(\mathbf{x}_{t} \vert \mathbf{x}_{t-1})}{p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)}+\log \frac{q\left(\mathbf{x}_{1} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \vert \mathbf{x}_{1}\right)}\right]\\ =\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \left(\frac{\color{blue}q(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_{0})}{p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)} \cdot \frac{\color{blue}q(\mathbf{x}_{t} \vert \mathbf{x}_{0})}{\color{blue}q(\mathbf{x}_{t-1} \vert \mathbf{x}_{0})}\right)+\log \frac{q\left(\mathbf{x}_{1} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \vert \mathbf{x}_{1}\right)}\right]\\ =\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)}+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t} \vert \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{0}\right)}+\log \frac{q\left(\mathbf{x}_{1} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \vert \mathbf{x}_{1}\right)}\right]\\ =\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)}+\log \frac{q\left(\mathbf{x}_{T} \vert \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{1} \vert \mathbf{x}_{0}\right)}+\log \frac{q\left(\mathbf{x}_{1} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \vert \mathbf{x}_{1}\right)}\right]\\ =\mathbb{E}_{q}\left[\log \frac{q\left(\mathbf{x}_{T} \vert \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{T}\right)}+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)}-\log p_{\theta}\left(\mathbf{x}_{0} \vert \mathbf{x}_{1}\right)\right]\\ =\mathbb{E}_{q}[\underbrace{-\log p_{\theta}\left(\mathbf{x}_{0} \vert \mathbf{x}_{1}\right)}_{L_{0}}]+\sum_{t=2}^{T} \underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}\right)\right)}_{L_{t-1}}+\underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{T} \vert \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{T}\right)\right)}_{L_{T}} \end{array} \end{equation} LVLB=Eq(x0:T)[logpθ(x0:T)q(x1:T∣x0)]=Eq[logpθ(xT)∏t=1Tpθ(xt−1∣xt)∏t=1Tq(xt∣xt−1)]=Eq[−logpθ(xT)+∑t=1Tlogpθ(xt−1∣xt)q(xt∣xt−1)]=Eq[−logpθ(xT)+∑t=2Tlogpθ(xt−1∣xt)q(xt∣xt−1)+logpθ(x0∣x1)q(x1∣x0)]=Eq[−logpθ(xT)+∑t=2Tlog(pθ(xt−1∣xt)q(xt−1∣xt,x0)⋅q(xt−1∣x0)q(xt∣x0))+logpθ(x0∣x1)q(x1∣x0)]=Eq[−logpθ(xT)+∑t=2Tlogpθ(xt−1∣xt)q(xt−1∣xt,x0)+∑t=2Tlogq(xt−1∣x0)q(xt∣x0)+logpθ(x0∣x1)q(x1∣x0)]=Eq[−logpθ(xT)+∑t=2Tlogpθ(xt−1∣xt)q(xt−1∣xt,x0)+logq(x1∣x0)q(xT∣x0)+logpθ(x0∣x1)q(x1∣x0)]=Eq[logpθ(xT)q(xT∣x0)+∑t=2Tlogpθ(xt−1∣xt)q(xt−1∣xt,x0)−logpθ(x0∣x1)]=Eq[L0 −logpθ(x0∣x1)]+∑t=2TLt−1 DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))+LT DKL(q(xT∣x0)∥pθ(xT))
上述式子中:
q ( x t ∣ x t − 1 ) = q ( x t ∣ x t − 1 , x 0 ) = q ( x t , x t − 1 ∣ x 0 ) q ( x t − 1 ∣ x 0 ) = q ( x t − 1 ∣ x t , x 0 ) ⋅ q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) \begin{equation} \begin{aligned} q(\mathbf{x}_{t} \vert \mathbf{x}_{t-1}) &= q(\mathbf{x}_{t} \vert \mathbf{x}_{t-1}, \mathbf{x}_0) = \frac{q(\mathbf{x}_t, \mathbf{x}_{t-1} \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} = {\color{red}q(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_0)} \cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}|\mathbf{x}_0)} \end{aligned} \end{equation} q(xt∣xt−1)=q(xt∣xt−1,x0)=q(xt−1∣x0)q(xt,xt−1∣x0)=q(xt−1∣xt,x0)⋅q(xt−1∣x0)q(xt∣x0)
我们可以得到 L V L B \mathcal{L}_{VLB} LVLB实际上是由一个熵( L 0 L_0 L0)以及多个KL散度( ( L t , t ∈ 1 , 2 , 3 , . . . , T ) (L_t,t\in{1,2,3,...,T}) (Lt,t∈1,2,3,...,T))构成。其中最后一项 L T L_T LT中的 x t \mathbf{x}_t xt和 x 0 \mathbf{x}_0 x0分别是数据分布和先验分布,都是固定的,因此它是一个常数,在最小化时可以忽略。转而去研究 L 0 和 L t , t ∈ 1 , 2 , 3 , . . . , T − 1 L_0和L_t, t\in{1,2,3,...,T-1} L0和Lt,t∈1,2,3,...,T−1。
首先来考虑公式2中比较复杂的 L t L_t Lt。
我们的模型 p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}\left(\mathbf{x}_{t-1}; \mu_{\theta}(\mathbf{x}_t, t), \Sigma_\theta(\mathbf{x}_t, t) \right) pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))。
根据公式Diffusion Model(2):前向扩散过程和逆向降噪过程
2-11,我们知道,对于平均值 μ θ ( x t , t ) \mu_{\theta}(\mathbf{x}_t,t) μθ(xt,t)在给定 x t \mathbf{x}_t xt和 t t t的情况下,我们期望它接近于 μ ~ t ( x t , x 0 ) = 1 α t ( x t − β t ( 1 − α ˉ t ) z t ) \tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0)=\frac{1}{{\sqrt{\alpha_t}}}\big(\mathbf{x}_t - \frac{\beta_t}{\sqrt{(1-\bar{\alpha}_{t})}} z_t \big) μ~t(xt,x0)=αt1(xt−(1−αˉt)βtzt)。
因此我们可以通过重参数化,通过学习高斯噪声 z θ ( x t , t ) z_\theta(\mathbf{x}_t,t) zθ(xt,t)来接近于 z t z_t zt:
μ θ ( x t , t ) = 1 α t ( x t − β t 1 − α ˉ t z θ ( x t , t ) ) \begin{equation} \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}\left(\mathbf{x}_{t}, t\right)\right) \end{equation} μθ(xt,t)=αt1(xt−1−αˉtβtzθ(xt,t))
这个过程可以归结为我们的模型实际上学会了估计真正的逆向过程中的噪声。
对于方差 Σ θ ( x t , t ) \Sigma_\theta(\mathbf{x}_t, t) Σθ(xt,t)的处理,DDPM (Ho et al 2020)将其设置成了 σ t 2 I \sigma_t^2\mathbf{I} σt2I,其中 σ t 2 = β t \sigma_t^2=\beta_t σt2=βt或者 σ t 2 = β ~ t = 1 − α ˉ t − 1 1 − α ˉ t β t \sigma_t^2=\tilde{\beta}_t= \frac{1-\bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t σt2=β~t=1−αˉt1−αˉt−1βt。当然此项也可以通过模型学习,此处只是为了简化,并不是唯一的。
现在,结合KL散度的公式Diffusion Model(1):预备知识1-10我们可以写出模型 p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) = 1 α t ( x t − β t 1 − α ˉ t z θ ( x t , t ) ) , Σ θ ( x t , t ) = σ t 2 I ) p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; {\color{green}\mu_{\theta}(\mathbf{x}_t, t)=\frac{1}{\sqrt{\alpha_t}}(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}z_\theta(\mathbf{x}_t,t))}, \Sigma_\theta(\mathbf{x}_t, t) = \sigma^2_t \mathbf{I}) pθ(xt−1∣xt)=N(xt−1;μθ(xt,t)=αt1(xt−1−αˉtβtzθ(xt,t)),Σθ(xt,t)=σt2I)关于 q ( x t − 1 ∣ x t ) = q ( x t − 1 ∣ x t , x 0 ) = N ( x t − 1 ; μ ~ t ( x t , x 0 ) = 1 α t ( x t − β t 1 − α ˉ t z t ) , Σ ~ t = β ˉ t I ) q(\mathbf{x}_{t-1} \vert \mathbf{x}_t) =q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; {\color{blue}\tilde{\mu}_t(\mathbf{x}_t, \mathbf{x}_0)=\frac{1}{\sqrt{\alpha_t}}(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}z_t)}, \tilde{\Sigma}_t =\bar{\beta}_t\mathbf{I}) q(xt−1∣xt)=q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0)=αt1(xt−1−αˉtβtzt),Σ~t=βˉtI)的KL散度:
L t = E x 0 , z [ 1 2 ∥ Σ θ ( x t , t ) ∥ 2 2 ∥ μ ~ t ( x t , x 0 ) − μ θ ( x t , t ) ∥ 2 ] = E x 0 , z [ 1 2 ∥ Σ θ ∥ 2 2 ∥ 1 α t ( x t − β t 1 − α ˉ t z t ) − 1 α t ( x t − β t 1 − α ˉ t z θ ( x t , t ) ) ∥ 2 ] = E x 0 , z [ β t 2 2 α t ( 1 − α ˉ t ) ∥ Σ θ ∥ 2 2 ∥ z t − z θ ( x t , t ) ∥ 2 ] = E x 0 , z [ β t 2 2 α t ( 1 − α ˉ t ) ∥ Σ θ ∥ 2 2 ∥ z t − z θ ( α ˉ t x 0 + 1 − α ˉ t z t , t ) ∥ 2 ] \begin{equation} \begin{array}{l} L_{t}=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}}\left[\frac{1}{2\left\|\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|_{2}^{2}}\left\| \textcolor{blue}{ \tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)}-\textcolor{green}{\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)}\right\|^{2}\right]\\ =\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}}\left[\frac{1}{2\left\|\mathbf{\Sigma}_{\theta}\right\|_{2}^{2}}\left\|\textcolor{blue}{\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{t}\right)}-\textcolor{green}{\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)}\right\|^{2}\right]\\ =\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}}\left[\frac{\beta_{t}^{2}}{2 \alpha_{t}\left(1-\bar{\alpha}_{t}\right)\left\|\mathbf{\Sigma}_{\theta}\right\|_{2}^{2}}\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|^{2}\right]\\ =\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}}\left[\frac{\beta_{t}^{2}}{2 \alpha_{t}\left(1-\bar{\alpha}_{t}\right)\left\|\mathbf{\Sigma}_{\theta}\right\|_{2}^{2}}\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}, t\right)\right\|^{2}\right] \end{array} \end{equation} Lt=Ex0,z[2∥Σθ(xt,t)∥221∥μ~t(xt,x0)−μθ(xt,t)∥2]=Ex0,z[2∥Σθ∥221∥ ∥αt1(xt−1−αˉtβtzt)−αt1(xt−1−αˉtβtzθ(xt,t))∥ ∥2]=Ex0,z[2αt(1−αˉt)∥Σθ∥22βt2∥zt−zθ(xt,t)∥2]=Ex0,z[2αt(1−αˉt)∥Σθ∥22βt2∥ ∥zt−zθ(αˉtx0+1−αˉtzt,t)∥ ∥2]
我们可以发现, L t L_t Lt的训练目标实际上是在使用MSE最小化两个高斯噪声 z t z_t zt和 z θ ( x t , t ) z_\theta(\mathbf{x}_t,t) zθ(xt,t)。
在训练中发现上述带有加权( β t 2 2 α t ( 1 − α ˉ t ) ∥ Σ θ ∥ 2 2 \frac{\beta_{t}^{2}}{2 \alpha_{t}\left(1-\bar{\alpha}_{t}\right)\left\|\mathbf{\Sigma}_{\theta}\right\|_{2}^{2}} 2αt(1−αˉt)∥Σθ∥22βt2)的MSE Loss不太稳定,因此DDPM (Ho et al 2020)使用了不带权重项的简化损失。
L t simple = E x 0 , z t [ ∥ z t − z θ ( α ˉ t x 0 + 1 − α ˉ t z t , t ) ∥ 2 ] L simple = L t simple + C \begin{equation} L_{t}^{\text {simple }}=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_{t}}\left[\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}, t\right)\right\|^{2}\right]\\ L_{\text {simple }}=L_{t}^{\text {simple }} + C \end{equation} Ltsimple =Ex0,zt[∥ ∥zt−zθ(αˉtx0+1−αˉtzt,t)∥ ∥2]Lsimple =Ltsimple +C
其中的 C C C是一项不依赖于 θ \theta θ的常数。
然后来考虑 L 0 L_0 L0的计算。
已知 L 0 = − E x 0 , x 1 log ( p θ ( x 0 ∣ x 1 ) ) L_0=-\mathbb{E}_{\mathbf{x}_0,\mathbf{x}_1}\log(p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)) L0=−Ex0,x1log(pθ(x0∣x1)),而 p θ ( x 0 ∣ x 1 ) = N ( μ θ ( x 1 ) , 1 , σ 1 2 I ) p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)=\mathcal{N}(\mu_\theta(\mathbf{x}_1),1,\sigma_1^2\mathbf{I}) pθ(x0∣x1)=N(μθ(x1),1,σ12I)。因此 L 0 L_0 L0实际上是一个多元高斯分布的负对数似然期望,即其熵。多元高斯分布的熵仅与其协方差有关,即 L 0 L_0 L0仅与 σ 1 2 I \sigma_1^2\mathbf{I} σ12I有关, L 0 L_0 L0是一个常数。
然而,论文DDPM指出,一般而言, x 0 \mathbf{x}_0 x0的分布实际上是离散的,而不是连续的。比如图片数据,像素值取值必须是整数,归一化到 [ − 1 , 1 ] [-1,1] [−1,1]后,依然是离散的点。Diffusion前向过程的第一步实际上是为离散数据添加噪声。那么,逆Diffusion的最后一步,即从 x 1 \mathbf{x}_1 x1到 x 0 \mathbf{x}_0 x0,也不能被简单地看作从 N ( μ θ ( x 1 , 1 ) , σ 1 2 I ) \mathcal{N}(\boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \sigma_1^2\mathbf{I}) N(μθ(x1,1),σ12I)中采样,而是在从 N ( μ θ ( x 1 , 1 ) , σ 1 2 I ) \mathcal{N}(\boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \sigma_1^2\mathbf{I}) N(μθ(x1,1),σ12I)采样的基础上再加上离散化操作。 L 0 L_0 L0也不再是一个常数,而是一个与 μ θ ( x 1 , 1 ) \mu_\theta(\mathbf{x}_1, 1) μθ(x1,1)相关的积分,其具体表达式可以参考DDPM (Ho et al 2020)的Sec3.3。在忽略 σ 1 2 \sigma_1^2 σ12和边缘效应后, L 0 L_0 L0的取值可以被 N ( μ θ ( x 1 , 1 ) , σ 1 2 I ) \mathcal{N}(\boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \sigma_1^2\mathbf{I}) N(μθ(x1,1),σ12I)的密度函数与离散时的分块大小(bin width)相乘所拟合。
另外值得一提的是,逆Diffusion的最后一步,DDPM直接取 μ θ ( x 1 , 1 ) \mu_\theta(\mathbf{x}_1, 1) μθ(x1,1)作为 x 0 \mathbf{x}_0 x0。
训练时,分别从 q ( x 0 ) q(\mathbf{x}_0) q(x0)、 U n i f o r m ( 1 , . . . , T ) Uniform({1,...,T}) Uniform(1,...,T)、 N ( 0 , I ) \mathcal{N}(\mathbf{0},\textbf{I}) N(0,I)中采样得到 x 0 \mathbf{x}_0 x0, t t t和 ϵ \epsilon ϵ(这里的 ϵ \epsilon ϵ就是前面说的 z z z),利用公式Diffusion Model(2):前向扩散过程和逆向降噪过程
2-4计算得到 x t = α ˉ t x 0 + 1 − α ˉ t z \mathbf{x}_t=\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z} xt=αˉtx0+1−αˉtz,将 x t \mathbf{x}_t xt和 t t t送入网络,预测得到一个噪声。结合公式7最小化预测噪声和真实采样的 ϵ \epsilon ϵ之间的距离。重复这一过程直到网络收敛。
采样时,需要从 x T \mathbf{x}_T xT一步一步的变回 x 0 \mathbf{x}_0 x0,其中的每一步都包含三个操作:
在前向扩散过程中,需要确定的超参数有 β t \beta_t βt以及总的步数 T T T。在前向扩散过程中,其值通常会增加。
β t = clip ( 1 − α ˉ t α ˉ t − 1 , 0.999 ) α ˉ t = f ( t ) f ( 0 ) where f ( t ) = cos ( t / T + s 1 + s ⋅ π 2 ) \beta_{t}=\operatorname{clip}\left(1-\frac{\bar{\alpha}_{t}}{\bar{\alpha}_{t-1}}, 0.999\right) \quad \bar{\alpha}_{t}=\frac{f(t)}{f(0)} \quad \text { where } f(t)=\cos \left(\frac{t / T+s}{1+s} \cdot \frac{\pi}{2}\right) βt=clip(1−αˉt−1αˉt,0.999)αˉt=f(0)f(t) where f(t)=cos(1+st/T+s⋅2π)
注意此图的纵坐标是 α ˉ t = 1 − β t \bar{\alpha}_t=1-\beta_t αˉt=1−βt而我们讨论的是 β t \beta_t βt的一个取值。
对于逆向降噪过程中,我们需要选择一个网络来预测一个噪声来从 x t \mathbf{x}_t xt得到 x t − 1 \mathbf{x}_{t-1} xt−1(参照采样的过程)。在这个过程中唯一的要求是我们需要保证输入和输出的维度是一样的。比如输入是 1 × 256 × 256 1\times256\times256 1×256×256那么输出也需要是 1 × 256 × 256 1\times256\times256 1×256×256。因此此处可以选择U-Net。
References: