BLURRING DIFFUSION MODELS (Paper reading)

BLURRING DIFFUSION MODELS

Emiel Hoogeboom, Google Research, Brain Team, ICLR2023, Cited:11, Code: 无, Paper.

1. 前言

最近，提出了一种基于散热或模糊的生成建模的新型扩散过程，作为各向同性高斯扩散的替代方案。在这里，我们表明模糊可以等价地通过具有非各向同性噪声的高斯扩散过程来定义。在建立这种联系的过程中，我们弥合了逆散热和去噪扩散之间的差距，并阐明了这种建模选择导致的归纳偏差。最后，我们提出了一类广义扩散模型，它提供了标准高斯去噪扩散和逆散热的最佳效果，我们称之为模糊扩散模型。

2. 背景

2.1 扩散模型

扩散模型的扩散过程可以定义为：$q(z_t|x)=N(z_t|{\alpha }_{t}x,{\sigma }_t^2\mathbf{I})$，这里$x$表示的是数据data，$z_t$是噪声隐变量（中间变量）。其中${\alpha }_t$单调递减，${\sigma }_t$单调递增，数据的信息会被逐渐的被噪声摧毁。假设上述定义是马尔可夫的，那么转移分布可以写为$q(z_t|z_s)=N(z_t|{\alpha }_{t|s}{z}_s,{\sigma }_{t|s}^{2}I)$，其中${\alpha }_{t|s}={\alpha }_t/{\alpha }_s$，${\sigma }_{t|s}^2={\sigma }_t^2-{\alpha }_{t|s}^2{\sigma }_s^2$。

扩散模型的逆过程（去噪过程）在给定数据$x$的情况下可以写为去噪分布
$$\begin{array}{c}q(z_s|z_t,x)=N(z_s|{\mathbf{\mu }}_{t\to s},{\sigma }_{t\to s}^2\mathbf{I})\\ \\ {\sigma }_{t\to s}^2={\left(\frac{1}{{\sigma }_s^2}+\frac{{\alpha }_{t\mid s}^2}{{\sigma }_{t\mid s}^2}\right)}^{-1}\text{ and }{\mathbf{\mu }}_{t\to s}={\sigma }_{t\to s}^2\left(\frac{{\alpha }_{t\mid s}}{{\sigma }_{t\mid s}^2}{\mathbf{z}}_t+\frac{{\alpha }_s}{{\sigma }_s^2}\mathbf{x}\right)\end{array}$$
显然数据$x$未知，通常利用一个神经网络来估计$\hat{x}(z_t)\approx x$。此外，像DDPM中直接预测噪声${\widehat{ϵ}}_t=f_{\theta }(z_t,t)$来估计$x$，或直接预测均值($x_{t-1}$)都是可以的。

最优化：模型对数似然的连续时间变分下界由以下对平方重构误差的期望给出，经验可得权重项为1效果不错。:
$$\mathcal{L}={\mathbb{E}}_{t\sim \mathcal{U}(0,1)}{\mathbb{E}}_{{\mathbf{ϵ}}_t\sim \mathcal{N}(0,\mathbf{I})}\left[w(t){\Vert f_{\theta }\left({\mathbf{z}}_t,t\right)-{\mathbf{ϵ}}_t\Vert }^2\right],$$

2.2 Inverse Heat Dissipation（反向散热）

通俗的讲，散热算法，如上图，可以理解为沸水的冷却，分子剧烈运动趋于平淡，或是人慢慢冷静的过程。上图可以看到的是图片中的高频成分都被平滑了，也就是模糊的过程。

逆散热模型（IHDM）使用散热来破坏信息，而不是增加越来越多的高斯噪声。他们观察到用于散热的拉普拉斯偏微分方程
$$\frac{\partial }{\partial t}z(i,j,t)=△z(i,j,t)$$
可以用余弦变换频域的对角矩阵求解。让$z_t$表示在时间步$t$的Laplace方程的解：
$$z_t={\mathbf{A}}_t{z}_0=\mathbf{V}{\mathbf{D}}_t{\mathbf{V}}^T{z}_0$$
其中${\mathbf{V}}^T$表示离散余弦变换（DCT），$\mathbf{V}$表示逆DCT。$z_0,z_t$应该考虑在空间维度上向量化，以允许矩阵乘法。对角线矩阵${\mathbf{D}}_t$是频率$\mathbf{\Lambda }$和时间$t$的加权矩阵的指数，因此${\mathbf{D}}_t=exp(-\mathbf{\Lambda }t)$。这里是个结论，直接用就可以。因此可以定义扩散过程的边缘分布：
$$q(z_t|x)=N(z_t|{\mathbf{A}}_{t}x,{\sigma }^2\mathbf{I})$$
中间扩散状态$z_t$是通过对data逐步添加模糊和固定的噪声而构建的。对于生成过程，用一个可学习的生成模型近似地反转的散热过程：
$$p(z_{t-1}|z_t)=N(z_{t-1}|f_{\theta }(z_t),{\sigma }^2\mathbf{I})$$
既$z_{t-1}$的均值是通过神经网络$f_{\theta }$直接学习的，并且具有固定的标量方差。类似于DDPMs，IHDM模型是通过采样随机时间步长$t$的前向过程$z_t\sim q(z_t|x)$，然后最小化模型之间的平方重建误差$f_{\theta }(zt)$和真实目标$\mathbb{E}(z_{t-1}|x)={\mathbf{A}}_{t-1}x$，产生以下训练损失：
$$\mathcal{L}={\mathbb{E}}_{t\sim \mu (1,...,T)}{\mathbb{E}}_{z_t\sim q(z_t|x)}[||{\mathbf{A}}_{t-1}x-f_{\theta }(z_t,t)|{|}^2]$$
有些问题仍然存在： (1)散热过程是否是马尔可夫的，如果是，什么是$q(z_t|z_s)$?(2)真正的逆加热过程是否也是各向同性的，如生成过程？(3)最后，除了预测中前一个时间步长的均值之外，是否还有其他的方法呢？在下面的章节中，我们会发现： (1)是的，这个过程可以是马尔可夫的。(2)不，生成过程不是各向同性的，虽然它在频域上是对角性的。(3)是的，像散热这样的过程可以被参数化，类似于标准扩散模型中的环境参数化。

2.3 散热形式为高斯扩散

我们将散热过程重新解释为一种高斯扩散的形式，类似于DDPM所使用的形式；Score-based Models等。本文将两个向量之间的乘法和除法定义为元素化的。我们从边际分布定义开始：
$$q(z_t|x)=N(z_t|{\mathbf{A}}_{t}x,{\sigma }^2\mathbf{I})$$
在本节中，我们让${\mathbf{V}}^T$表示正交的DCT，这是DCT的一个特定的归一化设置。把变量$u={\mathbf{V}}^T{z}_t$写成这种形式，我们可以写出频率空间的扩散过程：
$$q({\mathbf{u}}_t|{\mathbf{u}}_x)=N({\mathbf{u}}_t|{\mathbf{d}}_t\cdot {\mathbf{u}}_x,{\sigma }^2\mathbf{I})$$
其中$u_x={\mathbf{V}}^{T}x$是$x$的频率相应，$d_t=exp(-\mathbf{\lambda }{\tau }_t)$是对角线。这里${\tau }_t$表示关于时间步的噪声策略：线性或余弦等。上式表明，在每个维度$i$上，频率$u_t$的边缘分布在其标量元素$u_t^{(i)}$上被完全分解。同样，逆向散热模型$p_{\theta }({\mathbf{u}}_s|{\mathbf{u}}_t)$也被完全分解。我们可以等价地用每个维$i$的标量形式来描述散热过程（及其逆过程）：
$$q\left(u_t^{(i)}\mid u_0^{(i)}\right)=\mathcal{N}\left(u_t^{(i)}\mid d_t^{(i)}{u}_0^{(i)},{\sigma }^2\right)\iff u_t^{(i)}=d_t^{(i)}{u}_0^{(i)}+\sigma {ϵ}_t\text{, with }{ϵ}_t\sim \mathcal{N}(0,1)$$
这个方程可以看作是2.1节中介绍的标准高斯扩散过程的一个特殊情况。让$s_t$表示为频域空间中的标准高斯扩散过程，那么$s_t={\alpha }_t{u}_0^{(i)}+{\sigma }_t{ϵ}_t$。从概率论的角度来看，这里只有${\alpha }_t/{\sigma }_t$的比率很重要，而不是个体${\alpha }_t$，${\sigma }_t$的特殊选择。这是正确的，因为所有的值都可以简单地重新缩放，而无需以一种有意义的方式改变分布。这意味着，与其进行模糊处理和添加固定噪声，散热过程可以等效地定义为一个相对标准的高斯扩散过程，尽管是在频率空间。

马尔可夫转移分布：上述分析说明了与高斯扩散的等下性，因此散热模型同样也存在马尔可夫过程：
$$q({\mathbf{u}}_t|{\mathbf{u}}_s)=N({\mathbf{u}}_t|{\mathbf{\alpha }}_{t|s}{\mathbf{u}}_s,{\mathbf{\sigma }}_{t|s}^2\mathbf{I})$$
这与2.1中定义的是一样的，只是其中${\mathbf{\alpha }}_{t|s}={\mathbf{d}}_t/{\mathbf{d}}_s$，${\mathbf{\sigma }}_{t|s}^2=(1-({\mathbf{d}}_t/{\mathbf{d}}_s{)}^2){\sigma }^2$。

去噪过程（逆过程）：
$$q({\mathbf{u}}_s|{\mathbf{u}}_t,x)=N({\mathbf{u}}_s|{\mathbf{\mu }}_{t\to s},{\mathbf{\sigma }}_{t\to s}^2)$$
其中：
$${\mathbf{\sigma }}_{t\to s}^2={\left(\frac{1}{{\mathbf{\sigma }}_s^2}+\frac{{\mathbf{\alpha }}_{t\mid s}^2}{{\mathbf{\sigma }}_{t\mid s}^2}\right)}^{-1}\text{ and }{\mathbf{\mu }}_{t\to s}={\mathbf{\sigma }}_{t\to s}^2\left(\frac{{\mathbf{\alpha }}_{t\mid s}}{{\mathbf{\sigma }}_{t\mid s}^2}{\mathbf{u}}_t+\frac{{\mathbf{\alpha }}_s}{{\mathbf{\sigma }}_s^2}{\mathbf{u}}_x\right)\text{\hspace{1em}(19)}$$

3. Blurring Diffusion Models （模糊扩散模型）

利用第2.3节的分析，我们可以将该模型在频率空间中定义为一个高斯扩散模型，具有不同的维度schedules。如何参数化模型以及参数化${\alpha }_t$和${\sigma }_t$的具体schedules是很重要的。与传统模型不同，扩散过程是在频率空间中定义的:
$$q({\mathbf{u}}_t|{\mathbf{u}}_x)=N({\mathbf{u}}_t|{\mathbf{\alpha }}_t\mathbf{u},{\mathbf{\sigma }}_t^2\mathbf{I})$$
不同的频率可能以不同的速率扩散，这是由向量${\mathbf{\alpha }}_t$和${\mathbf{\sigma }}_t$中的值控制的。直接估计${\mathbf{\mu }}_{t\to s}$对于神经网络来说是困难的。类似与DDPM中，学习$x$的近似并带入去噪分布，通过参数化间接进行：
$${\mathbf{u}}_t={\mathbf{\alpha }}_t{\mathbf{u}}_x+{\mathbf{\sigma }}_t{\mathbf{u}}_{ϵ,t}\text{where}{\mathbf{u}}_x={\mathbf{V}}^{T}xand{\mathbf{u}}_{ϵ,t}={\mathbf{V}}^T{\mathbf{ϵ}}_t\text{\hspace{1em}(c)}$$
估计的${\hat{\mathbf{u}}}_x=\frac{1}{{\mathbf{\alpha }}_t}({\mathbf{u}}_t-{\mathbf{\sigma }}_t{\hat{\mathbf{u}}}_{ϵ,t})$。尽管在频率空间中表达扩散和去噪过程很方便，但神经网络往往在标准像素空间上运行良好。正是由于这个原因，神经网络$f_{\theta }$将$z_t=\mathbf{V}{u}_t$作为输入，并预测${\widehat{ϵ}}_t$。通过DCT矩阵就可以获得${\hat{u}}_{ϵ,t}={\mathbf{V}}^T{\widehat{ϵ}}_t$。将估计的${\hat{u}}_x$带入到去噪分布的均值中可以得到：
$${\hat{\mathbf{\mu }}}_{t\to s}={\mathbf{\sigma }}_{t\to s}^2\left(\frac{{\mathbf{\alpha }}_{t\mid s}}{{\mathbf{\sigma }}_{t\mid s}^2}{\mathbf{u}}_t+\frac{1}{{\mathbf{\alpha }}_{t|s}{\mathbf{\sigma }}_s^2}({\mathbf{u}}_t-{\mathbf{\sigma }}_t{\hat{\mathbf{u}}}_{ϵ,t})\right)\text{\hspace{1em}(24)}$$
在像素空间无权重的目标函数可写为：
$$\mathcal{L}={\mathbb{E}}_{t\sim u(0,1)}{\mathbb{E}}_{{ϵ}_t\sim N(0|I)}[||{ϵ}_t-f_{\theta }(z_t,t)|{|}^2],\text{where}{z}_t=\mathbf{V}({\mathbf{\alpha }}_t{\mathbf{V}}^T{\mathbf{x}}_t+{\mathbf{\sigma }}_t{\mathbf{V}}^T{\mathbf{ϵ}}_t)\text{\hspace{1em}(d)}$$

3.1 噪声和模糊Schedules

为了精确地指定模糊过程，需要为$t\in [0,1]$定义${\mathbf{\alpha }}_t$和${\mathbf{\sigma }}_t$的schedules。对于${\mathbf{\sigma }}_t$，我们为所有频率选择相同的值，因此给标量值${\mathbf{\sigma }}_t$一个schedule就足够了。这些shcedules是通过结合一个典型的高斯噪声扩散schedule（由标量在${\alpha }_t$和${\sigma }_t$处指定）和一个模糊schedule（由向量${\mathbf{d}}_t$指定）来构造的。

噪声schedule:选择Variance Preserving Cosine Schedule，既${\sigma }_t^2=1-{\alpha }_t^2,\text{where}{\alpha }_t=cos(t\pi /2)\text{for}t\in [0,1]$。为了避免在$t\to 0$和$t\to 1$时的不稳定，对数信噪比$(log{\alpha }_t^2/{\sigma }_t^2)$在$t=0$时的最大值为$+10$，在$t=1$时至少为$-10$。

模糊schedule: 没有理由使模型的概念时间步长$t$与耗散时间完美匹配。因此，${\mathbf{D}}_t=exp(-\Lambda {\tau }_t)$被重新定义，其中${\tau }_t$随$t$单调增加。变量${\tau }_t$在噪声扩散方面与${\alpha }_t$和${\sigma }_t$具有非常相似的函数：它允许相对于模型的概念时间步长$t$的任意耗散调度。

尺度为${\sigma }_B$的高斯模糊对应于随时间$\tau ={\sigma }_B^2/2$的耗散。根据经验，下面的策略work well：
$${\sigma }_{B,t}={\sigma }_{B,max}sin(t\pi /2{)}^2$$
其中，${\sigma }_{B,max}$是可调的超参数，对应于图像最大模糊的程度。因此这个模糊schedule可以定义为${\tau }_t={\sigma }_{B,t}^2/2$。如果仅仅使用$d_t=exp(-\lambda {\tau }_t)$用于${\alpha }_t$，并同样用于步骤$s$，那么$d_t/d_s$可能包含非常小的高频值。因此，一个不希望出现的副作用是，小的误差可能会被去噪过程中的许多步骤放大。因此，我们稍微修改了这个过程，让：
$${\mathbf{d}}_t=(1-d_{min})\cdot exp(-\mathbf{\lambda }{\tau }_t)+d_{min}\text{\hspace{1em}(b)}$$
这里令$d_{min}=0.001$。这种模糊变换将频率抑制到一个小的值dmin，同时去噪过程放大高频的侵略性较低。至此，结合高斯噪声schedule和模糊schedule，可以得到：
$${\mathbf{\alpha }}_t={\alpha }_t\cdot {\mathbf{d}}_t\text{and}{\mathbf{\sigma }}_t=1{\sigma }_t\text{\hspace{1em}(a)}$$
其中$1$是值都为1的向量。

3.4 伪代码

# 计算alpha_t和sigma_t，见公式(a)
def get_alpha_sigma(t):
	dt = get_dt(t)
	a, sigma_t= get_noise_scaling_cosine(t)
	alpha_t = a * dt.
return alpha_t, sigma_t

# 计算d_t，见公式(b)
def get_dt(t, min_d=0.001):
	sigma_blur = sigma_blur_max * sin(t * pi / 2)^2
	dissipation_t = sigma_blur^2 / 2
	freq = pi * linsapce(0, img_dim-1, img_dim) / img_dim
	labda = freqs[None, :, None, None]^2 + freqs[None, None, :, None]^2
	dt = (1-min_d) * exp(-labda * dissipation_t) + min_d
	return dt

# t~[0,1], 见3.1噪声Schedule部分
def get_noise_schaling_cosine(t, logsnr_min=-10, logsnr_max=10):
	limit_max = arctan(exp(-0.5 * logsnr_max))
	limit_min = arctan(exp(-0.5 * logsnr_min)) - limit_max
	logsnr = -2 * log(tan(limit_min * t + limit_max))
	# Transform logsnr to a, sigma .
	return sqrt(sigmoid(logsnr)), sqrt(sigmoid(-logsnr))

# 前向过程, 见公式(c)
def diffuse(x, t):
	x_freq=DCT(x)
	alpha, sigma = get_alpha_sigma(t)
	eps = random_normal_like(x)
	z_t = IDCT(alpha * x_freq) + sigma * eps
return z_t , eps

# 损失函数，见公式(d)
def loss(x):
	t = random_uniform(0, 1)
	z_t, eps = diffuse(x, t)
	error = (eps - neural_net(z_t, t))^2
	return mean(error)

# 采样过程，t=T,T-1,...,1/T, 见公式(24)(19)
def denoise(z_t, t, delta =1e-8):
	alpha_s, sigma_s = get_alpha_sigma(t-1 / T)
	alpha_t, sigma_t = get_alpha_sigma (t)
	# Compute helpful coefficients .
	alpha_ts = alpha_t / alpha_s
	alpha_st = 1 / alpha_ts
	sigma2_ts = ( sigma ^2 - alpha_ts ^2 * sigma_s ^2)
	# Denoising variance .
	sigma2_denoise = 1 / clip (
	1 / clip ( sigma_s ^2 , min= delta ) +
	1 / clip ( sigma_t ^2 / alpha_ts ^2 - sigma_s ^2 , min= delta ) ,
	min = delta )
	# The coefficients for u_t and u_eps .
	coeff_term1 = alpha_ts * sigma2_denoise / ( sigma2_ts + delta ) # eq.24
	coeff_term2 = alpha_st * sigma2_denoise / clip(sigma_s ^2, min= delta)
	# Get neural net prediction .
	hat_eps = neural_net(z_t, t)
	# Compute terms .
	u_t = DCT(z_t)
	term1 = IDCT(coeff_term1 * u_t)
	term2 = IDCT(coeff_term2 * (u_t - sigma_t * DCT(hat_eps)))
	mu_denoise = term1 + term2
	# Sample from the denoising distribution .
	eps = random_normal_like(mu_denoise)
return mu_denoise + IDCT (sqrt(sigma2_denoise) * eps)

4. 实验

CIFAR10 and LSUN Churches: 为了测量生成样本的视觉质量，在经过200万步训练后，从模型中采样的50000个样本计算FID评分。从这些分数中可以看出（表1），模糊的扩散模型能够生成比ihdm的质量更高的图像，以及文献中其他类似的方法。我们的模糊扩散模型也优于标准的去噪扩散模型，尽管在这种情况下，性能上的差异不那么明显。

4.1 在不同的噪音水平和Schedules之间的比较

在本节中，我们分析了在最大模糊${\sigma }_{B,max}$和两种不同的噪声schedule（$si{n}^2$和$sin$）方面有不同的设置。其中${\sigma }_{B,max}=0$等价于一个标准的去噪扩散模型。对于CIFAR10，性能最好的模型使用的模糊值为${\sigma }_{B,max}=20$，其FID为3.17，优于没有应用模糊时的FID3.60，如表3所示。对于LSUN，表现最好的模型使用了更少的模糊${\sigma }_{B,max}=10$，性能再次相对接近于${\sigma }_{B,max}=20$的模型。当将$si{n}^2$和$sin$ Schedule进行比较时，对于更高的最大模糊，$si{n}^2$表现得更好。我们的假设是，$sin$ 模糊过于激进，而$si{n}^2$的图在$t=0$附近的扩散过程开始时逐渐增加模糊。

有趣的是，具有较高最大模糊的模型收敛较慢，但当训练时间足够长时，比具有较少模糊的模型表现更好。比较最大模糊1和20的设定时，只有当20的设定训练超过一定步数才会比1的设定要好。似乎更高的模糊需要更多的时间来训练，但可以学会更好地匹配数据。

结论

模糊扩散模型的一个局限性是，模糊的使用具有正则化效果：当使用模糊时，将生成模型训练到收敛需要更长的时间。例如，正则化效应通常是有益的，并可以导致样本质量的提高。