【文生图系列】基础篇-马尔可夫链
【文生图系列】基础篇-变分推理(数学推导)
扩散模型是扩散概率模型(diffusion probabilistic model)的简称,它是一个参数化的马尔科夫链,使用变分推理进行训练,以在有限的时间后生出与数据匹配的样本。扩散模型包括前向过程(forward process)和反向过程(reverse process),其中前向过程也被成为扩散过程(diffusion process),它逐渐地从方差表(variance schedule) β 1 , ⋯ , β T \beta_{1}, \cdots, \beta_{T} β1,⋯,βT中向数据中添加高斯噪声。反向过程则相反,不断去除噪声,用于生成图片。
如上图所示,前向扩散过程是从 x 0 → x T x_{0} \to x_{T} x0→xT,每一步的条件概率分布定义为 q ( x t ∣ x t − 1 ) q\left( x_{t} | x_{t-1} \right) q(xt∣xt−1),逐步地向人脸图像中添加噪声,直至第 T T T步时,人脸图像彻底为一张噪声图片。反向过程是从 x T → x 0 x_{T} \to x_{0} xT→x0,每一步的条件概率分布定义为 q ( x t − 1 ∣ x t ) q\left(x_{t-1} | x_{t} \right) q(xt−1∣xt),逐步地去除噪声图片中的噪声,以生成人脸图像。但是 q ( x t − 1 ∣ x t ) q\left(x_{t-1} | x_{t} \right) q(xt−1∣xt)却是难以计算的,因此需要学习一个模型 p θ p_{\theta} pθ近似条件概率。
从一个真实数据分布中选择一个样本点 x 0 ∼ q ( x ) x_{0} \sim q\left( x \right) x0∼q(x),每一次向样本中添加少量高斯噪声,执行T次,得到噪音样本序列 x 1 , ⋯ , x T x_{1}, \cdots, x_{T} x1,⋯,xT。前向过程用公式定义如下:
执行步数由方差表 { β t ∈ ( 0 , 1 ) } t = 1 T \{ \beta_{t} \in \left( 0, 1 \right) \}_{t=1}^{T} {βt∈(0,1)}t=1T控制, β t \beta_{t} βt为第 t t t步所采用的方差,它介于 ( 0 , 1 ) \left( 0, 1 \right) (0,1)之间。通常情况下,方差的取值会随着 T T T的增大而越来越大 β 1 < β 2 < ⋯ < β T \beta_{1} < \beta_{2} < \cdots < \beta_{T} β1<β2<⋯<βT。而且随着 T T T的增大,样本数据 x 0 x_{0} x0会逐渐失去它的分布特征,当 T T T趋于无穷时, x T x_{T} xT则等价于一个各向同性高斯分布(isotropic Gaussian distribution)。
已知,扩散过程是一个马尔可夫链,每一步都会生成带噪音的数据 x t x_{t} xt。扩散过程的一个重要特性就是可以基于原始样本数据 x 0 x_{0} x0使用重参数技巧(reparameterization trick)采样任意 t t t步的生成数据 x t x_{t} xt, q ( x t ∣ x 0 ) = N ( x t , α ˉ t x 0 , ( 1 − α ˉ t ) I ) q\left(x_{t} | x_{0} \right)=N\left ( x_{t}, \sqrt{\bar{\alpha }_{t} } x_{0}, \left ( 1 - \bar{\alpha }_{t} \right ) I \right ) q(xt∣x0)=N(xt,αˉtx0,(1−αˉt)I)。
设 α t = 1 − β t , α ˉ t = ∏ s = 1 t = α s \alpha_{t} = 1 - \beta_{t}, \bar{\alpha}_{t}=\prod_{s=1}^{t}=\alpha_{s} αt=1−βt,αˉt=∏s=1t=αs, q ( x t ∣ x 0 ) q\left(x_{t} | x_{0} \right) q(xt∣x0)的推理过程如下所示:
反向过程是前向过程的反向,从一个高斯噪声输入 x T ∼ N ( 0 , 1 ) x_{T} \sim N \left( 0, 1 \right) xT∼N(0,1)中创建真实案例的过程,如果 β t \beta_{t} βt足够小,那么 q ( x t − 1 ∣ x t ) q\left(x_{t-1} | x_{t} \right) q(xt−1∣xt)也符合高斯分布。但是,需要全部的数据计算 q ( x t − 1 ∣ x t ) q\left(x_{t-1} | x_{t} \right) q(xt−1∣xt),这是一件很困难的事情,因此,为了能够运行反向扩撒模型,需要学习一个模型 p θ p_{\theta} pθ近似估计条件概率。
假设 x 1 , ⋯ , x T x_{1}, \cdots, x_{T} x1,⋯,xT为与数据 x 0 ∼ q ( x 0 ) x_{0} \sim q\left( x_{0} \right) x0∼q(x0)有相同维度的隐藏变量,扩散模型就是一个隐藏变量模型,用公式表示为 p θ ( x 0 ) : = ∫ p θ ( x 0 : T ) d x 1 : T p_{\theta}\left( x_{0} \right) := \int p_{\theta} \left( x_{0:T} \right) dx_{1:T} pθ(x0):=∫pθ(x0:T)dx1:T。反向过程则可被定义成这样的一个马尔可夫链,初始状态转移为 p ( x T ) = N ( x T , 0 , I ) p\left( x_{T} \right) = N\left( x_{T}, 0, I \right) p(xT)=N(xT,0,I)。
p θ ( x t − 1 ∣ x t ) p_{\theta}\left( x_{t-1} | x_{t} \right) pθ(xt−1∣xt)为参数化的高斯分布,均值和方差分别是训练的网络 μ θ ( x t , t ) \mu_{\theta}\left( x_{t}, t \right) μθ(xt,t)和 ∑ θ ( x t , t ) \sum_{\theta}\left( x_{t}, t \right) ∑θ(xt,t)。最终的生成模型就是由这些均值和方差网络组成。
虽然条件分布 q ( x t − 1 ∣ x t ) q\left(x_{t-1} | x_{t} \right) q(xt−1∣xt)是不可直接处理的,但是加上条件 x 0 x_{0} x0的后验分布 q ( x t − 1 ∣ x t , x 0 ) q\left(x_{t-1} | x_{t} , x_{0} \right) q(xt−1∣xt,x0)却是可处理的。
推理过程如下图所示。第一步的 C ( c t , c 0 ) C\left( c_{t}, c_{0} \right) C(ct,c0)与 x t − 1 x_{t-1} xt−1无关,可以省略。后验分布 q ( x t − 1 ∣ x t , x 0 ) q\left(x_{t-1} | x_{t} , x_{0} \right) q(xt−1∣xt,x0)的方差是一个定量,均值是一个依赖 x t x_{t} xt的函数。
设 σ 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−αtˉ1−αˉt−1βt,其中 β t = 1 − α t \beta_{t} = 1 - \alpha_{t} βt=1−αt,那么从分布 p θ ( x t − 1 ∣ x t ) p_{\theta}\left( x_{t-1} | x_{t} \right) pθ(xt−1∣xt)中采样 x t − 1 x_{t-1} xt−1,计算公式为 x t − 1 = 1 α t ( x t − β t 1 − α ˉ t ϵ t ) + σ t z x_{t-1} = \frac{1}{\sqrt{\alpha_{t}}} \left ( x_{t} - \frac{\beta_{t}}{\sqrt{1 - \bar{\alpha}_{t} } } \epsilon_{t} \right ) + \sigma_{t}z xt−1=αt1(xt−1−αˉtβtϵt)+σtz,其中 z z z服从正态分布 z ∼ N ( 0 , I ) z \sim N \left( 0, I \right) z∼N(0,I)。采样整体过程伪代码见下图Algorithm 2。
如果将扩散模型中间产生的变量看作隐变量的话,那么扩散模型其实就是包含 T T T个隐变量的隐变量模型,与VAE相似,不同的是扩散模型的隐变量与原始数据同纬度。那么就可以基于变分推断得到ELBO作为最大优化目标。
推导过程如下图所示,最后的优化目标共包含 T + 1 T+1 T+1项。
L T L_{T} LT计算的是噪音分布和先验分布的KL散度,先验 p ( x T ) p\left( x_{T} \right) p(xT)服从正太分布,扩散过程最后得到的随机噪音 q ( x T ∣ x 0 ) q\left( x_{T} | x_{0} \right) q(xT∣x0)也近似正太分布,所以它们两之间的KL散度近似为0。
假设图像数据为的整数,线性缩放到,这样能够保证神经网络反向过程从标准正太分布先验开始,在一致缩放的输入上进行操作。为了能够获得离散对数似然,反向过程中的最后一项设置为一个独立的离散解码器,该解码器源自于高斯分布 N ( x 0 , μ θ ( x 1 , 1 ) , σ 1 2 I ) N \left( x_{0}, \mu_{\theta} \left( x_{1}, 1 \right), \sigma_{1}^{2}I \right) N(x0,μθ(x1,1),σ12I)。
训练扩散模型时,发现忽略权重项,一个简化版的 L t L_{t} Lt会让模型性能更好。简化版的 L t L_{t} Lt如下所示。这就是上述Algorithm 1 训练伪代码第5行公式的由来。
训练部署设置为 T = 1000 T=1000 T=1000,前向过程中方差从 β 1 = 1 0 − 4 \beta_{1} = 10^{-4} β1=10−4线性增长到 β T = 0.02 \beta_{T} = 0.02 βT=0.02。
神经网络backbone是一个U-Net网络。U-Net属于编码-解码架构,编码器分成不同的stages,每个stage都包含下采样模块降低特征的空间大小;解码器将编码器压缩的特征上采样恢复。