Refusion:Enabling Large-Size Realistic Image Restoration withLatent-Space Diffusion Models-CVPRW2023-洪萨配资

Paper：https://arxiv.org/pdf/2304.08291
Code：https://github.com/Algolzw/image-restoration-sde
Baseline：IR-SDE

文章目录

前言

3. Preliminaries: Mean-Reverting SDE

4. Improving the Diffusion Model

4.1. Latent Diffusion under U-Net

4.2. Modified NAFBlocks for Noise Prediction

4.3. Improved Training Strategies

5. Experiments

5.1. Tasks and Datasets

5.2. Implementation Details

5.3. Experimental Results

前言

问题：IR-SDE的一个缺点是它在测试时计算量很大，因为它需要对整个图像进行多步去噪以恢复最终输出。这对于现实世界的应用程序来说可能会有问题，特别是对于高分辨率图像。

动机：本文的目的是改进扩散模型，以提高其在处理各种现实世界图像恢复任务中的有效性。IR-SDE由于其简单和灵活的适应不同的问题，成为Refusion的基础。

贡献：通过探索不同的噪声网络架构，我们表明使用非线性无激活网络(NAFNet)[7]可以在噪声/分数预测方面取得良好的性能，同时具有更高的计算效率。

此外，我们还说明了不同噪声水平、去噪步骤、训练图像大小和优化器/调度程序选择的有效性。

为了进一步处理大图像，我们提出了一种基于U-Net的潜在扩散策略。这使我们能够在压缩和低分辨率的潜在空间中执行图像恢复，从而加快了训练和推理。

3. Preliminaries: Mean-Reverting SDE

我们的方法利用扩散模型进行逼真的图像恢复。具体来说，我们使用IR-SDE[41]作为基本扩散框架，它可以自然地将高质量图像转换为降级的低质量图像，而不管退化有多复杂(即使对于现实世界的退化，参见图2)。IR-SDE的前向过程定义为：

其中，和分别是表征平均回归速度和随机波动率的时变正参数。如果设(1)中的SDE系数满足，则边际分布可根据[41]计算。

其中。注意，随着的增加，平均值和方差分别收敛于和。因此，初始状态被迭代地转换为带有附加噪声的，其中噪声水平固定为。

IR-SDE前向过程(1)是SDE的前向时间过程，其逆时表示为[57]：

请注意，在训练期间，我们可以访问HQ图像，这意味着我们可以使用(2)来计算真实值分数函数：

重新参数化技巧现在允许我们根据：

对进行采样，其中是标准高斯噪声。然后我们可以把(4)写成。通常训练CNN网络来估计噪声，然后在测试时模拟反向SDE将低质量图像转换为高质量版本，类似于其他基于扩散的模型。

4. Improving the Diffusion Model

4.1. Latent Diffusion under U-Net

在具有高分辨率图像的任务上迭代地运行扩散模型(即使只有几个去噪步骤)是出了名的耗时。特别是对于HR去雾，所有图像都以6000 × 4000 × 3像素捕获，这远远超出了现有任何扩散模型的输入尺寸。

为了处理大的输入尺寸，我们建议在低分辨率潜在空间中执行恢复，通过结合预训练的U-Net网络。所提出的基于U-Net的潜在扩散模型的总体架构如图3所示。编码器将LQ图像压缩为潜在表示，通过IR-SDE后向处理将其转换为HQ潜在表示。由此，解码器然后重建HQ图像。与使用VAE-GAN作为压缩模型的潜在扩散[52]相比，一个重要的区别是，所提出的U-Net通过跳跃连接保持从编码器到解码器的多尺度细节流。这样可以更好地捕获输入图像的信息，并为解码器提供额外的细节，以重建更准确的HQ图像。

在训练U-Net模型时，我们需要确保压缩的潜在表示具有判别性，并且包含主要的退化信息。U-Net解码器还必须能够从转换后的LQ潜在表示中重建HQ图像。因此，我们采用了一种潜在替换训练策略，如图4所示。

每个LQ图像首先由U-Net编码和解码，并应用重构L1损失。然后，U-Net也被训练来重建相应的HQ图像，方法是用HQ图像的潜在表示替换LQ的潜在表示，并再次运行解码器。重要的是，我们提出的训练策略不涉及任何对抗性优化。因此，模型训练比潜在扩散[52-LDM]更稳定。

4.2. Modified NAFBlocks for Noise Prediction

一种常用的噪声/分数预测架构是U-Net[53]，带有残差块[25]和注意机制，如通道注意和自注意[28,57]。最近提出的DiT[49]利用基于 Transformer 的结构，在低分辨率潜在空间中模拟扩散，在FID方面，它在类条件ImageNet 512 × 512生成方面设置了新的最先进的技术。但是，即使在潜在扩散框架[52]下，纯 Transformer 架构仍然会产生比传统图像恢复应用更大的计算成本。

为了解决上述模型效率问题，我们探索了一种新的噪声预测架构。具体来说，我们的噪声网络是基于稍微修改的非线性激活自由块(NAFBlocks)[7]。无非线性激活意味着我们用“SimpleGate”取代所有非线性激活函数，这是一种元素操作，将特征通道分成两个部分，然后将它们相乘以产生输出。如图2(b)所示，我们为注意层和前馈层添加了一个额外的多层感知器，将时间嵌入处理为通道尺度和移位参数γ和β。

为了适应不同的任务，我们还使用特定于任务的架构对网络进行了轻微的修改，例如散景效果变换中的镜头信息和立体图像超分辨率中的双输入。基于U-Net和NAFNet的扩散学习曲线如图6所示。注意，修改后的NAFNet在阴影去除任务上明显优于U-Net骨干网。

4.3. Improved Training Strategies

（1）Noise levels

当涉及到扩散模型的性能时，噪声水平(即第3节中的平稳方差 λ)可以发挥重要作用[41,46]。在图像恢复中，我们经常直接从LQ图像中恢复HQ图像，而不是从纯噪声中恢复HQ图像，这意味着终端状态的标准高斯函数是不必要的。如图5a所示，我们比较了四种不同的噪声水平σ ={10,30,50,70}对阴影去除任务的影响。训练曲线表明，σ = 50或σ = 70比σ = 10时更稳定。

（2）Denoising steps

一些研究提出使用长步预训练权值，但使用更少的步数生成图像[5,40,56]，这确实提高了样本效率，但也以降低图像质量为代价。在图像恢复中，我们必须为所有任务从头开始重新训练扩散模型。由于IR-SDE具有稳定的鲁棒的学习过程，我们考虑在保持性能的同时，在训练中直接调整去噪步骤。图5b对比了100步和1000步去噪时IR-SDE的训练曲线。我们发现使用更少的去噪步骤可以获得相当的，有时甚至更好的恢复性能。

（3）Training patch sizes

一种常见的做法是使用大 patch 进行训练可以提高图像恢复性能[36,37]。但是现有的研究都没有讨论 patch 大小对训练扩散模型的影响。这里我们比较了patch大小为128 × 128和256 × 256的训练IR-SDE，如图5c所示。可以观察到，使用大 patch 的训练效果要好得多，这与其他基于CNN/Transformer的图像恢复方法一致。

（4）Optimizer/scheduler

具有适当学习率调度器的优秀优化器对性能也很重要。例如，在ImageNet分类任务[26]上，简单地添加一个余弦衰减调度器可以将ResNet-50的准确率提高0.5%。为了找出哪个优化器更适合扩散模型，我们提供了三个比较：

1)Adam +多步衰减

2)AdamW[39] +余弦衰减

3)Lion[9] +余弦衰减。

图5d的结果显示，在多步学习率衰减的情况下，AdamW和Lion的性能都略好于Adam优化器。

5. Experiments

5.1. Tasks and Datasets

（1）Image Shadow Removal

按照NTIRE 2023 Shadow Removal Challenge[61,62]中的说明，使用1000对阴影和无阴影图像进行训练，使用100对阴影图像进行验证。

（2）Stereo Image Super-Resolution

为了训练和评估我们的模型，我们使用了NTIRE 2023立体图像超分辨率挑战赛[64]提供的数据集，该数据集由来自Flickr1024[66]数据集的800张训练立体图像和112张验证立体图像组成。所有的低分辨率图像都是通过双三次下采样生成的。

（3）Bokeh Effect Transformation

我们考虑在NTIRE 2023 Bokeh Effect Transformation挑战[15]中提出的新数据集，其中使用10000对具有不同镜头信息的合成图像进行训练，并使用500对具有源和目标镜头信息的图像进行验证。

（4）HR Non-Homogeneous Dehazing

在这里，我们使用了在NTIRE 2023 HR非均匀除雾竞赛[3]中提出的新数据集。该数据集内容丰富，收集方式类似于NH-HAZE[1,2]，但所有图像像素均为6000×4000×3。此外，它只包含40个模糊/无模糊图像对用于训练，10个图像用于验证和测试。

5.2. Implementation Details

对于所有的实验，我们使用与NAFNet相同的设置。批量大小设置为8，训练 patch 为256 × 256像素。我们使用Lion优化器[9]，β1 = 0.9， β2 = 0.99。初始学习率设置为3 × 10−5，并通过余弦调度器衰减为1e-7。对于所有任务，噪声级别固定为50，扩散去噪步骤数设置为100。我们还增加了随机水平翻转和90度旋转的训练数据。所有模型都使用PyTorch[48]实现，并在单个A100 GPU上训练约3天。

对于阴影去除和立体超分辨率，我们使用正常扩散策略，并将训练迭代设置为50万次。另一方面，对于HR去雾和散景效果转换，我们采用了潜在扩散策略。具体来说，我们首先在每个数据集上训练U-Net进行30万次迭代，然后基于U-Net训练Refusion模型进行40万次迭代。在这里，所有的输入 patch 被裁剪为1024 × 1024，同时使用UNet将它们压缩为128 × 128像素。