基于对抗性混频的鲁棒图像去噪

论文链接：https://openaccess.thecvf.com/CVPR2024/Ryou_Robust_Image_Denoising_through_Adversarial_Frequency_Mixup

项目链接：https://github.com/dhryougit/AFM

Abstract

基于深度神经网络的图像去噪方法经常与训练数据中存在的特定噪声分布的过拟合作对抗。这一挑战存在于现有的现实世界去噪网络中，这些网络使用有限的真实噪声分布谱进行训练，因此对分布外的真实噪声类型表现出较差的鲁棒性。为了缓解这个问题，我们开发了一种新的训练框架，称为对抗性混频(AFM)。AFM利用频域的混合产生具有独特和具有挑战性的噪声特征的噪声图像，同时保留真实世界噪声的属性。随后，将这些噪声图像纳入训练pipeline中，增强去噪网络对噪声分布变化的鲁棒性。在广泛的真实噪声基准上进行的大量实验和分析表明，使用我们提出的框架训练的去噪网络在对看不见的噪声分布的鲁棒性方面表现出显着的改进。

1. Introduction

基于深度神经网络的图像去噪[5,6,22,26 - 28,30,31]由于问题表述的简单性和新数据集的构建，取得了前所未有的成功。传统上，去噪网络使用合成噪声模型进行训练，例如高斯或泊松噪声，这些模型被人为地添加到图像中以用于训练和评估目的。然而，现实世界的噪声受图像信号处理(ISP)pipeline中的各种因素(如去马赛克和伽马校正)的影响，表现出明显的信号依赖性，并且通常遵循与合成对应物不同的分布。当将去噪模型应用于真实的噪声图像时，合成噪声和真实噪声分布之间的差异提出了重要的泛化问题。

虽然人们一直致力于在现实世界中创建具有干净和噪声图像对的数据集，但收集这些数据集是一项重大挑战。因此，各种自监督方法[10,12,18]已经成为图像去噪的有前途的解决方案，旨在减少对成对的无噪声图像数据集的依赖。然而，与监督学习方法相比，这些方法在性能上往往不足。因此，在实际场景中，对真实噪声数据集的监督仍然是有利的。

然而，在真实噪声数据集上进行监督训练的图像去噪器有其自身的缺陷。最值得注意的是，我们观察到，即使是最先进的去噪器也难以泛化到噪声分布的变化，这些变化是由不同的相机传感器类型、拍摄环境和ISP过程等因素引起的。我们在图1中说明了这种现象，并在图2中提供了频率分析。虽然缺乏真实的噪声泛化对去噪模型的广泛使用产生了不利影响，但在图像去噪研究领域，这是一个相对未被探索的问题。

在这项工作中，我们提出了对抗混频(AFM)，这是一种模型不可知的训练框架，可以提高去噪网络对实际噪声分布变化的泛化能力。为此，AFM构建了具有独特且难以去噪的噪声的图像，同时保持了真实世界噪声的特性。然后，通过将这些噪声图像整合到训练pipeline中，去噪网络对实际噪声分布的变化具有鲁棒性。更具体地说，AFM的工作原理是根据频域的混合掩码将有噪声和去噪的图像混合在一起。这个混合掩码是由一个单独的轻量级神经网络生成的，它产生了我们方法的两个变体：AFM-E和AFM-B。AFM-E以元素方式生成混频掩码，而AFM-B通过为每个频带分配一个混频值来生成掩码。此外，AFM使用对抗性损失进行训练，并可以生成对抗性混合掩码。

总体而言，我们的贡献总结如下：

• 我们提出了AFM，这是一种模型不可知的训练框架，通过在频域中的对抗性混合生成新的噪声图像，提高了去噪网络对现实世界分布外噪声图像的泛化和鲁棒性。
• 我们设计了两个对抗性掩码生成网络，AFM-E和AFM-B。AFM-E以元素方式生成混合掩码，而AFM-B为单个波段生成混合掩码。
• 我们在多个真实世界的分布外图像去噪基准上证明了AFM的有效性。我们的方法在各种去噪架构甚至最先进的去噪模型上始终优于普通训练。

2. Related Works

本节回顾了现有的基于有监督和无监督的深度神经网络的图像去噪模型，并讨论了最近在鲁棒图像去噪方面的努力。

2.1. 有监督图像去噪

近年来，有监督图像去噪领域取得了重大进展，其中有噪声和干净的图像配对可用于训练。最初的突破主要是由基于CNN的模型推动的，其中去噪卷积神经网络(Denoising Convolutional Neural Network, DnCNN)[30]为该领域的进一步创新奠定了基础。在此基础上，基于U-Net的架构[4,5,26,27]已经成为突出的模型，利用跳过连接来有效地结合局部和全局上下文信息。此外，基于Transformer的模型的引入[22,28]标志着去噪策略的范式转变。配备了注意力机制，基于Transformer的模型擅长识别复杂的依赖关系，有效地减少了过程中的噪声失真。这些模型在降低噪声方面显示出卓越的技能，大大减轻了它们使用成对数据集训练的特定噪声分布的影响。然而，他们在有效管理与他们所训练的噪声分布不同的噪声分布方面遇到了挑战，这对他们在现实场景中的应用构成了障碍。

2.2. 自监督图像去噪

从现实世界中收集干净和有噪声的对数据是成本很高的。为了缓解这些问题，人们积极研究探索无监督和自监督学习技术，如N2N[12]、N2S[2]、N2V[10]和R2R[18]。这些方法证明了使用仅包含噪声对或单个噪声图像的数据集训练去噪网络的可行性，而不需要干净的噪声对数据集。此外，DIP[21]、Self2Self[19]和Neighbor2Neighbor[7]引入了新的方法来训练去噪网络，即使在没有训练数据的情况下，也可以使用单个带噪声的图像来产生干净的对应图像。然而，由于网络需要对每个不同的图像进行重新训练，因此这些方法的效率低下且耗费大量时间。重要的是，与监督方法相比，它们的性能往往不足。

2.3. 去噪泛化

现有的去噪模型存在泛化问题，对训练噪声分布存在过拟合问题。只有少数作品解决了这一挑战[3,15]。Mohan等[15]观察到噪声过拟合是由偏置项引起的，并从网络中去除所有的偏置。然而，这种方法的目的是增强在不同程度的噪声中的鲁棒性，而不是噪声类型。Chen等人[3]在训练过程中引入输入像素和注意特征掩码，专注于重建干净的图像本身，而不是去噪，但其在真实噪声上的表现仍然不足。

3. Methods

我们的目标是训练一个去噪网络，在各种看不见的真实噪声的例子中表现出强大的泛化性能。为了实现这一目标，我们提出了对抗混频(AFM)，这是一种基于具有真实噪声分布的频域增强图像的新型训练框架。

3.1. 背景

去噪网络的目的是生成干净的图像，而不考虑输入图像上的噪声模式。这种步骤表述为：
$${\mathcal{D}}_{\theta }(\mathbf{x}+\mathbf{n})=\mathbf{x},\text{\hspace{1em}(1)}$$
式中，${\mathcal{D}}_{\theta }$表示用$\theta$参数化的去噪网络，$(x,n)$表示干净图像及其噪声的一对。

对于监督学习，一种流行的方法是使用真实的噪声数据集(如SIDD[1])训练去噪模型，该数据集由干净图像$(x)$和噪声图像$(x+n)$对组成。问题的目标是通过最小化以下损失来优化模型参数$\theta$：
$$\mathcal{L}=\Vert {\mathcal{D}}_{\theta }(\mathbf{x}+\mathbf{n})-\mathbf{x}\Vert .\text{\hspace{1em}(2)}$$
然而，传统的方法容易学习干净图像和噪声图像之间的固定映射，这往往导致较差的泛化。在实际场景中，由于相机传感器和图像信号处理(ISP)pipeline的不同特性，存在很大范围的噪声变化。因此，真实的噪声去噪模型必须能够推广到这种噪声分布的变化。

3.2. 对抗性混频

图3(a)说明了所提出的AFM方法的概述，该方法旨在提高对看不见的真实噪声图像的泛化。给定一个带噪声的输入图像$I\in {\mathbb{R}}^{C\times H\times W}$，我们首先使用去噪网络${\mathcal{D}}_{\theta }$预测去噪图像$I$，该网络由公式（3）给出：
$$\hat{I}={\mathcal{D}}_{\theta }(I).\text{\hspace{1em}(3)}$$

然后，我们将原始输入$I$与频域预测$\hat{I}$混合如下：
$$I_{\mathrm{hard}}={\mathcal{F}}^{-1}\left(\mathcal{F}(I)\odot \mathbf{m}+\mathcal{F}(\hat{I})\odot (1-\mathbf{m})\right),\text{\hspace{1em}(4)}$$
其中$\mathcal{F}$是快速傅里叶变换(FFT)，$\odot$是元素乘法算子，$m\in [0,1{]}^{1\times H\times W}$是频域中的任意掩码，$I_{hard}$是得到的混合图像。由于掩码$m$以[0,1]为界，公式（4）的右侧表示频域中$I$和$\hat{I}$之间的元素插值。给定$I$和$\hat{I}$，这里的目标是生成带有一些噪声的新图像，其特征是与原始噪声图像的分布不同，同时也类似于现实场景中遇到的噪声。我们在公式(4)中的设计选择—特别是频率混频的使用—直接受到这个目标的启发。

$I$和$\hat{I}$上的傅里叶变换将每个图像映射到频域，与空间域相比，在频域操作图像的噪声分布相对容易。图2显示了支持这一主张的直观可视化，其中噪声和干净信号的区域在频域中可以清晰区分。此外，频率分量的变化对图像起全局作用，从而最大限度地降低了对底层图像内容进行不必要操作的风险。从本质上讲，结果$I_{hard}$停留在真实噪声图像的流形中，因为它是$I$和$\hat{I}$的插值。因此，混频可以防止构造带有任意和人为噪声的合成图像。

通过推导出公式(4)的简化表达式，我们对混合操作有了更直观的理解：
$$\begin{array}{rl}{I}_{\mathrm{hard}}& ={\mathcal{F}}^{-1}\left(\mathcal{F}(\hat{I})+\mathbf{m}\odot (\mathcal{F}(I)-\mathcal{F}(\hat{I}))\right)\\ & =\hat{I}+{\mathcal{F}}^{-1}\left(\mathbf{m}\odot \mathcal{F}(I-\hat{I})\right)\\ & =\hat{I}+{\mathcal{F}}^{-1}(\mathbf{m})\ast {\mathcal{F}}^{-1}\left(\mathcal{F}(\hat{\mathbf{n}})\right)\\ & =\hat{I}+{\mathcal{F}}^{-1}(\mathbf{m})\ast \hat{\mathbf{n}}\\ & =\hat{I}+h\ast \hat{\mathbf{n}},\end{array}\text{\hspace{1em}(5)}$$
其中$h$表示滤波器，$\ast$表示卷积算子，$\hat{n}=I-\hat{I}$表示预测噪声。本质上，混频并不是随机改变图像，而是根据$h$变换噪声分量$\hat{n}$， $h$由掩码$m$​的傅里叶反变换定义。

剩下的是选择$m$，或者更具体地说，如何为混频创建适当的掩码。正如将在3.3节中更详细地讨论的那样，我们选择$m$的对抗性设计，这导致去噪网络认为是hard(因此称为$I_{hard}$)的真实噪声图像。然后，我们可以将$I_{hard}$纳入到训练pipeline中，以提高去噪网络在现实世界中的泛化程度。
$$\begin{array}{rl}\tilde{I}& ={\mathcal{F}}^{-1}\left(\mathcal{F}(I_1)\odot \mathbf{m}+\mathcal{F}(I_2)\odot (1-\mathbf{m})\right),\\ & ={\mathcal{F}}^{-1}\left(\mathcal{F}(I_2)+\mathbf{m}\odot (\mathcal{F}(I_1)-\mathcal{F}(I_2))\right)\\ & =I_2+{\mathcal{F}}^{-1}\left(\mathbf{m}\odot \mathcal{F}(I_1-I_2)\right)\\ & =I_2+{\mathcal{F}}^{-1}(\mathbf{m})\ast (I_1-I_2)\\ & =I_2+\mathbf{h}\ast (I_1-I_2)\\ & =I_2+\mathbf{h}\ast \Delta I_{1,2},\end{array}\text{\hspace{1em}(S.1)}$$

3.3. 掩码设计

掩码$m$是对抗性地生成的，使得$I_{hard}$的去噪输出${\mathcal{D}}_{\theta }(I_{hard})$与ground truth $I_{gt}$之间的差值最大化。如图3(b)所示，这是通过掩码生成网络$G_{\varphi }$完成的，该网络可以由AFM-E或AFM-B实现。

无论实现选择如何，m表示为：
$$\mathbf{m}=G_{\varphi }\left(\left[I,\hat{I},|\mathcal{F}(I)|,|\mathcal{F}(\hat{I})|\right]\right),\text{\hspace{1em}(6)}$$
式中$\left|\mathcal{F}(I)\right|$和$\left|\mathcal{F}(\hat{I})\right|$分别是$I$和$\hat{I}$的频率幅度，$[\cdot ,\cdot ]$是沿通道维度的级联运算符。给定掩码$m$，我们根据公式(4)构造$I_{\mathrm{hard}}$。此外，为了稳定$G_{\varphi }$的训练，我们还使用相反的掩码$1-m$构造$I_{\mathrm{easy}}$：
$$I_{\text{easy}}={\mathcal{F}}^{-1}\left(\begin{array}{c}\mathcal{F}(I)\odot (1-\mathbf{m})+\mathcal{F}(\hat{I})\odot \mathbf{m}\end{array}\right).\text{\hspace{1em}(7)}$$
我们制定了以下损失函数，使$I_{\mathrm{hard}}$对抗并训练掩码生成网络$G_{\varphi }$：
$${\mathcal{L}}_{\mathrm{AFM}}=\mathrm{PSNR}({\hat{I}}_{\mathrm{hard}},I_{\mathrm{gt}})-\gamma \mathrm{PSNR}({\hat{I}}_{\mathrm{easy}},I_{\mathrm{gt}}),\text{\hspace{1em}(8)}$$
其中，${\hat{I}}_{\mathrm{hard}}={\mathcal{D}}_{\theta }(I_{\mathrm{hard}}),{\hat{I}}_{\mathrm{easy}}={\mathcal{D}}_{\theta }(I_{\mathrm{easy}})$， $\gamma$是平衡这两项的超参数。第一项满足我们的主要目的：通过最小化${\hat{I}}_{\mathrm{hard}}$和$I_{\mathrm{gt}}$之间的峰值信噪比(PSNR)来生成对抗情况$I_{\mathrm{hard}}$。第二项的目的是最大化${\tilde{\hat{I}}}_{\mathrm{easy}}$和$I_{\mathrm{gt}}$之间的PSNR，因为$I_{\mathrm{easy}}$是$I_{\mathrm{hard}}$的补项，应该作为去噪网络相对容易的输入来生成。虽然在公式(8)中，训练一个对抗性的$G_{\varphi }$并不是严格必要的，但我们从经验上发现，它有助于防止生成不重要的的掩码。

3.3.1 AFM-E：逐元素掩码生成

AFM-E通过像素分割来构造与输入图像大小相同的掩码。因此，每个元素在$\mathcal{F}(I)$和$\mathcal{F}(\hat{I})$中，可以用不同的混合比例混合。如图4(a)所示，合并后的输入通过基于UNet的具有跳跃连接的编码器-解码器架构[20]。关于架构的更多细节可以在附录中找到。

3.3.2 AFM-B：逐频带掩码生成

虽然AFM-E是生成混合掩码的有效方法，但我们观察到AFM-E生成的掩码自然表现出一定程度的旋转不稳定性，即对应于相同频率的掩码元素具有相似的混合值。因此，我们引入了一种简化但功能强大的替代方案：逐频带掩码生成。如图4(b)所示，该网络使用几个卷积和全连接层构建，这产生了一组N个混合值。所得到的掩码由沿极轴均匀分布的N个圆形带组成。一个频带区域内的所有元素对应于网络产生的混频值。

3.4. 训练过程

在算法1中详细介绍的训练过程中，我们同时训练去噪网络和AFM网络。

在每次迭代中，我们采用两步方法。在第一步中，我们保持掩码生成网络不变，更新去噪网络，损失如下：
$$\begin{array}{rl}{\mathcal{L}}_{\mathrm{D}}& =-\mathrm{PSNR}(\hat{I},I_{\mathrm{gt}})-\lambda \mathrm{PSNR}({\hat{I}}_{\mathrm{hard}},I_{\mathrm{gt}})\\ & ={\mathcal{L}}_{\mathrm{rec}}+\lambda {\mathcal{L}}_{\mathrm{hard}},\end{array}\text{\hspace{1em}(9)}$$
其中$\lambda$是一个超参数。在第二步中，我们保持去噪网络固定，并使用公式(8)中定义的对抗损失函数更新AFM模块。

值得注意的是，我们的AFM框架只影响训练过程。在推理中，我们简单地省略AFM模块，并使用去噪网络进行预测，而不需要任何额外的内存或计算成本。

4. Experiments

我们将我们的框架AFM应用于各种图像去噪架构，并在多个分布外(OOD)基准上进行评估。我们的结果表明，无论结构如何，AFM都显著提高了图像去噪网络的OOD泛化(表1)，甚至在OOD泛化方面优于最先进的去噪网络(表2)。最后，我们给出了一些定性结果。

4.1. 实验设置

数据集 在所有实验中，我们只在Smartphone Image Denoising Dataset (SIDD)(SIDD) Medium[1]数据集上训练去噪网络。我们在SIDD验证集上测量了分布(ID)性能，并在五个真实噪声基准上评估了分布外(OOD)性能，以确保在各种数据域上进行鲁棒的评估。五个OOD基准包括Poly[23]、CC[16]、HighISO[24]、iPhone[9]和Huawei[9]。Poly、CC和HighISO的图片大小为512×512, iPhone和Huawei的图片大小为1024×1024。

训练细节 我们训练所有模型进行200K次迭代，批大小为32，训练patch大小为256×256。对于去噪网络和掩码生成网络，我们使用AdamW[13]优化器，初始学习率为10⁻³，在余弦退火计划后降低到10⁻⁶。最后，我们设$\lambda =0.8$， $\gamma =0.3$。

4.2. 结果

在表1中，我们比较了使用和不使用所提出框架的现有去噪架构的ID和OOD性能。我们在四种不同的架构上进行了测试，包括DnCNN[30]、CBDNet[6]、NAFNet[5]和MPRNet[27]。结果表明，结合AFM可以在几乎所有OOD基准和架构中持续提高性能，同时保持ID性能的竞争力。例如，集成AFM后，在DnCNN、CBDNet、MPRNet和NAFNet上，OOD的平均PSNR分别提高了+0.75 dB、+0.33 dB、+0.54 dB和+1.00 dB。这表明AFM可以应用于各种去噪网络，以提高网络的泛化性和对未知噪声分布的鲁棒性。

我们还与最先进的(SOTA)监督去噪网络(如MIRNetv2[26]、Uformer[22]和Restormer[28])进行了比较，并将结果显示在表2中。我们使用官方发布的权重评估所有SOTA网络，并将这些网络与使用AFMB框架训练的DnCNN架构进行比较。值得注意的是，尽管在推理时需要更少的内存和计算，但具有AFM-B的DnCNN在PSNR方面至少优于所有SOTA模型+0.39 dB (Restormer)，在SSIM方面优于所有SOTA模型+0.009 (MIRNet-v2)。

我们还比较了其他基于SOTA数据集的自监督方法的去噪性能。如表2所示，使用SIDD数据集的基于数据集的自监督方法训练的模型在分布内和分布外的情况下都表现不佳，而我们的方法在这两种情况下都表现出色。

定性结果 图5显示了用AFM-B训练的DnCNN在两个OOD基准样本上的去噪输出：CC和HighISO。为了比较，我们还可视化了其他去噪网络的输出，这些网络没有经过AFM的训练。这些定性比较清楚地表明，与未经AFM训练的模型相比，使用AFM训练的模型产生更清晰的输出。

5. Analysis and Discussions

分析了AFM与其他泛化方法的有效性，并进行了讨论。

5.1. 与泛化技术的比较

我们将AFM框架与其他一些旨在提高泛化的技术进行了比较。我们简要总结每种技术：

• Dropout[8]在网络的最后一个卷积层之前沿着通道维度随机丢弃特征。
• 输入掩码(Input mask)[3]在空间维度上随机掩码第一个卷积层之后的特征。
• CutMix[25]是一种常见的数据增强方案，通过在训练数据集中剪切和粘贴来自不同图像的补丁来创建新的数据样本。
• 对抗性训练(Adversarial training) 首先使用PGD攻击在输入图像中产生对抗性扰动[14]，然后训练网络去噪对抗性扰动。
• 随机混频(Random frequency mixup) 与AFM-B相同，除了每个频带的混合值是从均匀分布中采样。
• ASM-E(Adversarial Spatial Mixup) 与AFM-E相同，除了混是在空间域而不是频域进行的。

我们用上面提到的每一种技术训练一个DnCNN网络，并在表3中给出了ID和平均OOD结果。我们首先观察到Dropout和Input mask并不能提高泛化性能，这可以通过平均OOD性能的下降来证明。此外，我们注意到CutMix表现出较小的增益，并且优于AFM-E和AFM-B，它们在OOD性能方面表现出显着改善。

在Adversarial training的结果中有一个有趣的观察。尽管与我们的AFM框架有一些相似之处，但我们发现使用PGD攻击的朴素对抗性训练未能提高对OOD图像的鲁棒性(-0.63 dB)。我们假设，差异在于噪声是如何产生的；虽然PGD攻击会产生合成的对抗噪声，但我们的AFM旨在产生逼真的对抗噪声。这意味着保持真实属性的噪声增强在提高泛化性能方面起着至关重要的作用。此外，Random frequency mixup导致OOD性能的适度改善(+0.36 dB)，但不如AFM-E (+0.70 dB)或AFM-B (+0.75 dB)，这突出了AFM模块中对抗性掩码生成网络的重要性。最后，ASM-E不能提高泛化性能，这加强了在频域使用混合的理由。

5.2. 讨论

频率分析 图6可视化了对抗掩码$\mathbf{m}$的示例，以及来自训练数据集的相应图像$I_{\mathrm{hard}}$，${\hat{I}}_{\mathrm{hard}}$， 和 $\hat{I}$, 的频率值。由于输入是分布内样本，去噪后的图像$\hat{I}$显示干净的频率。然而，去噪网络$I_{\mathrm{hard}}$很难去噪；事实上，图6中的$|\mathcal{F}({\hat{I}}_{\mathrm{hard}})|$与图2中的输出频率非常相似，这是从一个OOD数据集中采样的。这意味着我们的AFM模块构建了与OOD图像共享许多特征的训练图像。

图像去噪的泛化 表2中的结果揭示了即使在性能最好的去噪网络（如Restormer）中，也缺乏对看不见的真实噪声分布的泛化。这引起了一个重大关切；尽管泛化是广泛使用图像去噪网络的必要条件，但对不同实噪声分布的泛化常常被忽视。因此，我们希望我们的工作能够启发未来的图像去噪研究，以探索这个方向，而不是仅仅专注于提高分布性能。

6. Conclusion

我们提出了对抗性混频(AFM)，这是一种新的图像去噪网络训练框架，有助于更好的泛化和鲁棒性，以适应各种现实世界的噪声分布。通过利用频域中的对抗性混合，AFM模块生成新的噪声图像，这些图像保留了现实世界中遇到的噪声特性。然后，将这些图像整合到训练数据集中，学习具有分布外鲁棒性的改进图像去噪网络。值得注意的是，与使用传统监督训练方法训练的图像去噪器相比，使用AFM框架训练的图像去噪器表现出明显优越的泛化能力。最后，我们发现图像去噪网络缺乏鲁棒性是一个重要但经常被忽视的研究领域，并希望未来的工作将积极地朝着这个方向发展。

Appendix

https://openaccess.thecvf.com/content/CVPR2024/supplemental/Ryou_Robust_Image_Denoising_CVPR_2024_supplemental