到了第十二天,今天学习的内容是扩散模型,记录一下:
学习内容
扩散模型是目前最流行的文生图模型,这篇学习笔记就不写代码了,只记录一下实现原理。
什么是Diffusion Model?
如果将Diffusion与其他生成模型(如Normalizing Flows、GAN或VAE)进行比较,它并没有那么复杂,它们都将噪声从一些简单分布转换为数据样本,Diffusion也是从纯噪声开始通过一个神经网络学习逐步去噪,最终得到一个实际图像。
Diffusion对于图像的处理包括以下两个过程:
-
我们选择的固定(或预定义)正向扩散过程 q q q :它逐渐将高斯噪声添加到图像中,直到最终得到纯噪声
-
一个学习的反向去噪的扩散过程 p θ p_\theta pθ :通过训练神经网络从纯噪声开始逐渐对图像去噪,直到最终得到一个实际的图像
由 t t t 索引的正向和反向过程都发生在某些有限时间步长 T T T(DDPM作者使用 T = 1000 T=1000 T=1000)内。从 t = 0 t=0 t=0开始,在数据分布中采样真实图像 x 0 \mathbf{x}_0 x0(本文使用一张来自ImageNet的猫图像形象的展示了diffusion正向添加噪声的过程),正向过程在每个时间步长 t t t 都从高斯分布中采样一些噪声,再添加到上一个时刻的图像中。假定给定一个足够大的 T T T 和一个在每个时间步长添加噪声的良好时间表,您最终会在 t = T t=T t=T 通过渐进的过程得到所谓的各向同性的高斯分布。
扩散模型实现原理
Diffusion 前向过程
所谓前向过程,即向图片上加噪声的过程。虽然这个步骤无法做到图片生成,但这是理解diffusion model以及构建训练样本至关重要的一步。
首先我们需要一个可控的损失函数,并运用神经网络对其进行优化。
设 q ( x 0 ) q(x_0) q(x0) 是真实数据分布,由于 x 0 ∼ q ( x 0 ) x_0 \sim q(x_0) x0∼q(x0) ,所以我们可以从这个分布中采样以获得图像 x 0 x_0 x0 。接下来我们定义前向扩散过程 q ( x t ∣ x t − 1 ) q(x_t | x_{t-1}) q(xt∣xt−1) ,在前向过程中我们会根据已知的方差 0 < β 1 < β 2 < . . . < β T < 1 {0}<\beta_{1}<\beta_{2}< ... <\beta_{T}<{1} 0<β1<β2<...<βT<1 在每个时间步长 t 添加高斯噪声,由于前向过程的每个时刻 t 只与时刻 t-1 有关,所以也可以看做马尔科夫过程:
q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) 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}) q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)
回想一下,正态分布(也称为高斯分布)由两个参数定义:平均值 μ \mu μ 和方差 σ 2 ≥ 0 \sigma^2 \geq 0 σ2≥0 。基本上,在每个时间步长 t t t 处的产生的每个新的(轻微噪声)图像都是从条件高斯分布中绘制的,其中
q ( μ t ) = 1 − β t x t − 1 q(\mathbf{\mu}_t) = \sqrt{1 - \beta_t} \mathbf{x}_{t-1} q(μt)=1−βtxt−1
我们可以通过采样 ϵ ∼ N ( 0 , I ) \mathbf{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) ϵ∼N(0,I) 然后设置
q ( x t ) = 1 − β t x t − 1 + β t ϵ q(\mathbf{x}_t) = \sqrt{1 - \beta_t} \mathbf{x}_{t-1} + \sqrt{\beta_t} \mathbf{\epsilon} q(xt)=1−βtxt−1+βtϵ
请注意, β t \beta_t βt 在每个时间步长 t t t (因此是下标)不是恒定的:事实上,我们定义了一个所谓的“动态方差”的方法,使得每个时间步长的 β t \beta_t βt 可以是线性的、二次的、余弦的等(有点像动态学习率方法)。
因此,如果我们适当设置时间表,从 x 0 \mathbf{x}_0 x0 开始,我们最终得到 x 1 , . . . , x t , . . . , x T \mathbf{x}_1, ..., \mathbf{x}_t, ..., \mathbf{x}_T x1,...,xt,...,xT,即随着 t t t 的增大 x t \mathbf{x}_t xt 会越来越接近纯噪声,而 x T \mathbf{x}_T xT 就是纯高斯噪声。
那么,如果我们知道条件概率分布 p ( x t − 1 ∣ x t ) p(\mathbf{x}_{t-1} | \mathbf{x}_t) p(xt−1∣xt) ,我们就可以反向运行这个过程:通过采样一些随机高斯噪声 x T \mathbf{x}_T xT,然后逐渐去噪它,最终得到真实分布 x 0 \mathbf{x}_0 x0 中的样本。但是,我们不知道条件概率分布 p ( x t − 1 ∣ x t ) p(\mathbf{x}_{t-1} | \mathbf{x}_t) p(xt−1∣xt) 。这很棘手,因为需要知道所有可能图像的分布,才能计算这个条件概率。
Diffusion 逆向过程
为了解决上述问题,我们将利用神经网络来近似(学习)这个条件概率分布 p θ ( x t − 1 ∣ x t ) p_\theta (\mathbf{x}_{t-1} | \mathbf{x}_t) pθ(xt−1∣xt) , 其中 θ \theta θ 是神经网络的参数。如果说前向过程(forward)是加噪的过程,那么逆向过程(reverse)就是diffusion的去噪推断过程,而通过神经网络学习并表示 p θ ( x t − 1 ∣ x t ) p_\theta (\mathbf{x}_{t-1} | \mathbf{x}_t) pθ(xt−1∣xt) 的过程就是Diffusion 逆向去噪的核心。
现在,我们知道了需要一个神经网络来学习逆向过程的(条件)概率分布。我们假设这个反向过程也是高斯的,任何高斯分布都由2个参数定义:
-
由 μ θ \mu_\theta μθ 参数化的平均值
-
由 μ θ \mu_\theta μθ 参数化的方差
综上,我们可以将逆向过程公式化为
p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_\theta (\mathbf{x}_{t-1} | \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1};\mu_\theta(\mathbf{x}_{t},t), \Sigma_\theta (\mathbf{x}_{t},t)) pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))
其中平均值和方差也取决于噪声水平 t t t ,神经网络需要通过学习来表示这些均值和方差。
-
注意,DDPM的作者决定保持方差固定,让神经网络只学习(表示)这个条件概率分布的平均值 μ θ \mu_\theta μθ 。
-
本文我们同样假设神经网络只需要学习(表示)这个条件概率分布的平均值 μ θ \mu_\theta μθ 。
-
在这里, x 0 \mathbf{x}_0 x0 是初始(真实,未损坏)图像, ϵ \mathbf{\epsilon} ϵ 是在时间步长 t t t 采样的纯噪声, ϵ θ ( x t , t ) \mathbf{\epsilon}_\theta (\mathbf{x}_t, t) ϵθ(xt,t)是我们的神经网络。神经网络是基于真实噪声和预测高斯噪声之间的简单均方误差(MSE)进行优化的。
训练算法现在如下所示:
换句话说:
-
我们从真实未知和可能复杂的数据分布中随机抽取一个样本 q ( x 0 ) q(\mathbf{x}_0) q(x0)
-
我们均匀地采样 1 1 1和 T T T之间的噪声水平 t t t(即,随机时间步长)
-
我们从高斯分布中采样一些噪声,并使用上面定义的属性在 t t t 时间步上破坏输入
-
神经网络被训练以基于损坏的图像 x t \mathbf{x}_t xt 来预测这种噪声,即基于已知的时间表 x t \mathbf{x}_t xt 上施加的噪声
实际上,所有这些都是在批数据上使用随机梯度下降来优化神经网络完成的。
U-Net神经网络预测噪声
神经网络需要在特定时间步长接收带噪声的图像,并返回预测的噪声。请注意,预测噪声是与输入图像具有相同大小/分辨率的张量。因此,从技术上讲,网络接受并输出相同形状的张量。那么我们可以用什么类型的神经网络来实现呢?
这里通常使用的是非常相似的自动编码器,您可能还记得典型的"深度学习入门"教程。自动编码器在编码器和解码器之间有一个所谓的"bottleneck"层。编码器首先将图像编码为一个称为"bottleneck"的较小的隐藏表示,然后解码器将该隐藏表示解码回实际图像。这迫使网络只保留bottleneck层中最重要的信息。
在模型结构方面,DDPM的作者选择了U-Net,出自(Ronneberger et al.,2015)(当时,它在医学图像分割方面取得了最先进的结果)。这个网络就像任何自动编码器一样,在中间由一个bottleneck组成,确保网络只学习最重要的信息。重要的是,它在编码器和解码器之间引入了残差连接,极大地改善了梯度流(灵感来自于(He et al., 2015))。
可以看出,U-Net模型首先对输入进行下采样(即,在空间分辨率方面使输入更小),之后执行上采样。