ResShift: Efficient Diffusion Model for Image Super-resolution by Residual Shifting

ResShift: Efficient Diffusion Model for Image Super-resolution by Residual Shifting (Paper reading)

Zongsheng Yue, S-Lab, Nanyang Technological University, NeurIPS23, Cited:0, Code, Paper

1. 前言

基于扩散的图像超分辨率(SR)方法由于需要数百甚至数千个采样步骤，导致推理速度较低。现有的加速采样技术不可避免地会在一定程度上牺牲性能，导致SR结果过于模糊。为了解决这个问题，我们提出了一种新颖有效的SR扩散模型，该模型显著减少了扩散步骤的数量，从而消除了推理过程中对后加速的需求及其相关的性能下降。我们的方法构建了一个马尔可夫链，通过移动高分辨率图像和低分辨率图像之间的残差来实现高分辨率图像和低分辨率图像之间的转换，大大提高了转换效率。此外，还设计了一套复杂的噪声控制方案，以灵活控制扩散过程中的换档速度和噪声强度。大量的实验表明，即使只有15个采样步骤，所提出的方法在合成和真实数据集上的性能也优于或至少与当前最先进的方法相当。

2. 总体思想

有一个问题，就是本文选择的测试集是自己挑的，而不是常见的benchmarks。我感觉是比不过这些基准数据集，但是扩散模型跟end-to-end的模型比确实没有优势啊。主要想法是将$x_T$换为$y$的某种形式，然后通过扩散模型采样得到$x_0=x$。方法是把这个残差$HR-LR$分配给整个采样阶段，然后构建新的采样过程。

3. 方法

一种常见的方法是将LR图像插入当前扩散模型的输入(例如，DDPM)，并在sr的训练数据上从头开始重新训练模型。另一种流行的方法是使用无条件预训练的扩散模型作为先验，并修改其反向路径以生成预期的HR图像。不幸的是，这两种策略都继承了DDPM底层的马尔可夫链，这在推理中可能是低效的，通常需要数百甚至数千个采样步骤。虽然已经开发了一些加速技术来压缩推理中的采样步骤，但它们不可避免地会导致性能的显著下降，导致结果过于平滑。

在本节中，我们介绍了一个为sr量身定制的扩散模型ResShift。为了便于表示，LR图像和HR图像分别表示为$y_0$和$x_0$。此外，我们假设$y_0$和$x_0$具有相同的空间分辨率，这可以很容易地通过在必要时使用最近邻插值对LR图像$y_0$进行预上采样来实现。

3.1 模型设计

扩散模型的迭代生成范式已被证明在捕获复杂分布方面非常有效，这也激励我们迭代地处理SR问题。我们提出的方法构建了一个马尔可夫链，作为HR和LR图像之间的桥梁。这样，SR任务可以通过对给定任何LR图像的马尔可夫链进行反向采样来完成。接下来，我们将详细介绍专门为SR构建这样一个马尔可夫链的过程。

前向过程。我们将LR和HR图像之间的残差记为$e_0$，即$e_0$ = $y_0-x_0$。我们的核心思想是通过一个长度为$T$的马尔可夫链，逐渐移动它们的残差$e_0$，从$x_0$迁移到$y_0$。首先引入一个位移序列 { η t } t = 1 T \{η_t\}^T_ {t =1} {ηt​}t=1T​，它随时间步长$T$单调增加，满足${\eta }_1\to 0$和${\eta }_t\to 1$。然后根据这一移位顺序推导出过渡分布如下:
$$q(x_t|x_{t-1},y_0)=N(x_t;x_{t-1}+{\alpha }_t{e}_0,k^2{\alpha }_{t}I),t=1,2,...,T\text{\hspace{1em}(1)}$$
其中${\alpha }_t={\eta }_t-{\eta }_{t-1}$，当$t>1$且${\alpha }_1={\eta }_1$时，$k$为控制噪声方差的超参数，$I$为单位矩阵。值得注意的是，我们证明了任意时间步长$t$的边际分布是解析可积的，即
$$q(x_t|x_0,y_0)=N(x_t;x_0+{\eta }_t{e}_0,k^2{\eta }_{t}I),t=1,2,...,T\text{\hspace{1em}(2)}$$
Eq.(1)中给出的过渡分布的设计基于两个主要原则。第一个原理与标准差有关，其目的是促进$x_t$和$x_{t-1}$之间的平滑过渡。这是因为假设图像数据落在[0,1]范围内，则$x_t$与$x_{t-1}$之间的期望距离可以用$\sqrt{{\alpha }_t}$来限定，即:
$$max[(x_0+{\eta }_t{e}_0)-(x_0+{\eta }_{t-1}{e}_0)]=max[{\alpha }_t{e}_0]<{\alpha }_t<\sqrt{{\alpha }_t}\text{\hspace{1em}(3)}$$
其中$max[\cdot ]$表示逐像素最大化操作。引入超参数$k$来增加设计的灵活性。第二个原理是关于均值参数，即$x_0+{\alpha }_t{e}_0$，它可以推导出Eq.(2)中的边际分布。此外，$x_1$和$x_T$的边际分布收敛于$\delta x_0(\cdot )$和$N(\cdot ;y_0，k^{2}I)$，分别作为HR图像和LR图像的两个近似分布。通过以这样一种深思熟虑的方式构造马尔可夫链，可以通过给定LR图像$y_0$的反向采样来处理SR任务。

反向过程。反向过程旨在通过以下公式估计后验分布$p(x_0|y_0)$:
$$p(x_0|y_0)=\int p(x_T|y_0)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t,y_0)d{x}_{1:T}\text{\hspace{1em}(4)}$$
式中$p(x_T|y_0)\approx N(x_T|y_0，k^{2}I)，p_{\theta }(x_{T-1}|x_T,y_0)$为从$x_T$到$x_{T-1}$的具有可学习参数θ的逆转移核。根据大多数关于扩散模型的文献，我们采用$p_{\theta }(x_{t-1}|x_t,y_0)=N(x_{t-1};{\mu }_{\theta }(x_t,y_0,t)，{\Sigma }_{\theta }(x_t,y_0,t))$。θ的优化是通过最小化负证据下界来实现的。由以下推导可以得到采样的后验分布的均值和方差：

考虑到方差参数与$x_t$和$y_0$无关，我们设${\Sigma }_{\theta }(x_t,y_0,t)=k^2\frac{{\eta }_t-1}{\eta t}{\alpha }_{t}I$。对于均值参数${\mu }_{\theta }(x_t,y_0,t)$，重新参数化如下:
$${\mu }_{\theta }(x_t,y_0,t)=\frac{{\eta }_{t-1}}{\eta t}{x}_t+\frac{{\alpha }_t}{\eta t}{f}_{\theta }(x_t,y_0,t)$$
式中，$f_{\theta }$为参数为$\theta$的深度神经网络，目的是预测$x_0$。我们的目标函数简化为$w_t||f_{\theta }(x_t,y_0,t)-x_0|{|}_2^2$。其中$w_t=\frac{{\alpha }_t}{2k^2{\eta }_t{\eta }_{t-1}}$。在实践中，我们通过经验发现，忽略权重$w_t$会导致性能的明显提高。

3.2 噪声策略

该方法采用超参数$k$和移位序列${{\eta }_t}_{T=1}^T$来确定扩散过程中的噪声调度。具体来说，超参数$k$调节了过渡期间的整体噪声强度，其对性能的影响在第4.2节中进行了实证讨论。接下来的论述主要围绕移位序列的构造展开:

由式(2)可知，状态$x_t$下的噪声水平与$\sqrt{{\eta }_t}$成正比，其比例因子为$k$。这一观察结果促使我们专注于设计$\sqrt{{\eta }_t}$而不是${\eta }_t$。Song和Ermon表明，$k\sqrt{{\eta }_1}$应该足够小(例如LDM中的0.04)，以确保$q(x_1|x_0,y_0)\approx q(x_0)$。结合${\eta }_1\to 0$的附加约束，我们将${\eta }_1$设为$(0.04/k{)}^2$和0.001之间的最小值。对于最后一步$T$，我们设置${\eta }_T$为0.999，确保${\eta }_T\to 1$。对于中间时间步，即$t\in [2,t-1]$，我们提出$\sqrt{{\eta }_t}$的非一致几何调度:
$$\begin{gathered} \sqrt{{\eta }_t}=\sqrt{{\eta }_1}\times b_0^{{\beta }_t},t=2,...,T-1 \\ where \\ {\beta }_t=(\frac{t-1}{T-1}{)}^p\times (T-1),b_0=exp\left[\frac{1}{2(T-1)}log\frac{{\eta }_T}{{\eta }_1}\right] \end{gathered}$$
注意，${\beta }_t$和$b_0$的选择是基于${\beta }_1=0，{\beta }_t=T-1，\sqrt{{\eta }_T}=\sqrt{{\eta }_1}\times b_0^{T-1}$的假设。超参数$p$控制着$\sqrt{{\eta }_t}$的增长率，如图2(h)所示。所提出的噪声表在三个关键方面表现出高度的灵活性。首先，对于较小的$k$值，最终状态$x_T$收敛为图2©-(d)所示的LR图像周围的扰动。与高斯噪声下的损坏相比，该设计大大缩短了马尔可夫链的长度，从而提高了推理效率。其次，超参数$p$提供了对换档速度的精确控制，使SR结果的保真度-真实感权衡成为可能，如第4.2节所分析。第三，通过设置$k=40$和$p=0.8$，我们的方法实现了与LDM非常相似的扩散过程。图2(e)-(f)所示扩散过程的视觉结果清楚地证明了这一点，图2(g)所示相对噪声强度的对比也进一步支持了这一点。

4. 实验

本节对ResShift进行了实证分析，并提供了广泛的实验结果，以验证其在一个合成数据集和三个真实数据集上的有效性。我们的研究特别关注更具挑战性的×4 SR任务。

4.1 实验设置

训练细节。我们的训练数据中分辨率为256 × 256的HR图像是根据LDM从ImageNet的训练集中随机裁剪出来的。我们使用RealESRGAN的降解管道合成LR图像。Adam算法使用PyTorch的默认设置和64的mini-batch大小来训练ResShift。在训练期间，我们使用固定的学习率5e-5，并每500K次迭代更新权重参数。在网络架构方面，我们在DDPM中简单地采用UNet结构。为了提高ResShift对任意图像分辨率的鲁棒性，我们用Swin Transformer块替换UNet中的自关注层。

测试数据集。我们基于常用的退化模型合成了一个测试数据集，$即y=(x\ast k)\downarrow +n$，其中$k$为模糊核，$n$为噪声，$y$和$x$分别表示LR图像和HR图像。为了全面评估ResShift的性能，我们考虑了更复杂类型的模糊核、下采样算子和噪声类型。对于HR图像，我们从ImageNet的验证集中随机选择了3000张图像，而不是SR中流行的数据集，例如Set5 ， Set14和Urban100。这是因为这些数据集只包含少量的源图像，无法全面评估各种方法在不同退化类型下的性能。为方便起见，我们将此数据集命名为ImageNet-Test。采用两个真实数据集来评估ResShift的有效性。第一个是RealSR，包含100张由佳能5D3和尼康D810相机拍摄的真实图像。此外，我们收集另一个名为RealSet65的真实世界数据集。它包括最近文献中广泛使用的35张LR图像。剩下的30张图片是我们自己从网上获取的。

比较的方法。我们将ResShift与最近的七种SR方法(ESRGAN、RealSR-JPEG、BSRGAN、RealESRGAN、SwinIR、DASR和LDM)进行了对比，评估了ResShift的有效性。注意，LDM是一种基于扩散的方法，有1000个扩散步骤。为了公平比较，我们将LDM加速到与ResShiftusing DDIM相同的步数，并将其表示为“LDM- a”，其中“a”表示推理步数。DDIM中的超参数$\eta$被设置为1，因为该值产生最真实的恢复图像。使用5个指标评估各种方法的性能，包括PSNR、SSIM、LPIPS、MUSIQ和CLIPIQA。值得注意的是，后两者是非参考指标，专门用于评估图像的真实感。特别是CLIPIQA，它利用了在大规模数据集(即Laion400M)上预训练的CLIP模型，因此表现出很强的泛化能力。在真实数据集上，我们主要依靠CLIPIQA和MUSIQ作为评价指标来比较不同方法的性能。

4.2 模型分析

我们分析了不同设置下ResShift在扩散步数和超参数p和k上的性能。

扩散步骤T和超参数p。在Eq.(1)中提出的过渡分布显著地减少了马尔可夫链中的扩散步骤T。超参数p允许在过渡期间灵活控制剩余移位的速度。表1总结了在不同的T和p配置下ResShift在ImageNet-Test上的性能。我们可以看到，T和p都呈现了一种权衡，在由参考指标(如PSNR、SSIM和LPIPS)测量的保真度与由非参考指标(包括CLIPIQA和MUSIQ)测量的超分辨率结果的真实感之间。以p为例，当p增加时，参考指标会改善，而非参考指标会恶化。此外，从图3的视觉对比可以看出，较大的p值会抑制模型产生更多图像细节的幻觉能力，导致输出模糊。

超参数$k$。由式(2)可知，状态$x_t$时$k$主导噪声强度。我们在表1中报告了$k$对ResShift性能的影响。结合图3的可视化，我们可以发现，过大或过小的$k$值都会使恢复结果平滑，而不管它们的有利指标是PSNR和SSIM。当$k$在[1.0,2.0]范围内时，我们的方法达到了CLIPIQA和MUSIQ所指示的最真实的质量，在实际应用中更理想。因此，我们将$k$设为2.0。

效率的比较。为了提高推理效率，需要限制扩散步数T。然而，这会导致还原的HR图像的真实感降低。为了折衷，可以将超参数$p$设置为较小的值。因此，我们设置T = 15和p = 0.3，并产生我们的模型ResShift。表2总结了在ImageNet-Test上与LDM的效率和性能比较。在LPIPS和CLIPIQA方面，提出的ResShift即使在1000个采样步长下也优于LDM。在这种配置下，ResShift比LDM快40倍，有望解决当前基于扩散的SR方法固有的效率限制。

5. 结论

在这项工作中，我们为SR引入了一种高效的扩散模型ResShift。与现有基于扩散的SR方法需要大量迭代才能获得满意的结果不同，我们提出的方法构建了一个仅15个采样步骤的扩散模型，从而显著提高了推理效率。其核心思想是将HR图像向LR图像腐蚀，而不是高斯白噪声，这样可以有效地切断扩散模型的长度。在合成和真实数据集上的大量实验证明了我们提出的方法的优越性。我们相信，我们的工作将为开发更高效的扩散模型铺平道路，以解决SR问题。