Random Noise Injection on Feature Statistics for Generalizable Deep Image Denoising

ICLR2025—Random Is All You Need: 在特征统计上随机注入噪声，用于泛化深度图像去噪

论文链接（Paper under double-blind review）：https://openreview.net/forum?id=z8PcUSKXXN

ABSTRACT

广义深度图像去噪的最新进展促进了鲁棒噪声处理模型的发展。目前最先进的Masked Training（MT）构建了一个蒙面SwinIR模型，该模型仅对高斯噪声（σ=15）进行训练，但可以在各种噪声类型（即斑点噪声，泊松噪声）中实现令人满意的去噪性能。然而，该方法虽然侧重于内容重建，但往往会产生过度平滑的图像，并且在蒙版比例优化方面存在挑战，使其与其他方法的集成变得复杂。作为回应，本文介绍了RNINet，一种基于流线型编码器-解码器框架的新型架构，以提高效率和整体性能。最初，我们在单个噪声类型上训练纯RNINet（仅简单的编码器-解码器），观察特征统计量（如均值和方差）在响应不同噪声条件时的移位。利用这些见解，我们结合了一个噪声注入块，将随机噪声注入到我们框架内的特征统计中，显著提高了对看不见的噪声类型的泛化。我们的框架不仅简化了在MT中发现的架构复杂性，而且提供了卓越的性能。综合实验评估表明，我们的方法在各种看不见的噪声条件下，在去噪效果和计算效率（更低的MACs和GPUs内存使用）方面优于机器翻译，实现了高达10倍的推理速度，并强调了它的大规模部署能力。

1 INTRODUCTION

图像去噪是低水平图像处理的一个重要研究领域，旨在从噪声中恢复干净的图像。深度学习的快速发展激发了许多研究，提出了专门的图像去噪网络。这些网络通常在预定义的噪声分布上进行训练，在去噪方面表现出显著的性能。然而，它们对其他噪声类型的泛化是有限的，限制了它们在实际场景中的应用，因为实际场景中的噪声分布通常与训练阶段的噪声分布不同。

在目前对图像去噪任务的研究趋势中，大多数现有的作品(即SwinIR （Liang et al ., 2021), Restormer (Zamir et al ., 2022)）在被高斯噪声破坏的图像上训练和评估模型，这限制了它们对特定噪声分布的性能。为了解决这一限制，一些方法（Zhang et al .，2017）对特定噪声类型假设未知的噪声水平，而其他方法(Brooks et al .，2019b；Wei et al ., 2020)试图通过合成或收集更接近目标噪声的训练数据或直接对目标噪声进行无监督训练来提高现实场景中的性能(Chen et al ., 2018；Yuan et al, 2018)。尽管做出了这些努力，但Chen等人（Chen et al .，2023）最近的工作认为，这些方法都没有从本质上提高去噪网络的泛化性能，而且当噪声分布不匹配时，它们仍然很困难（Abdelhamed et al .，2018b）。为此，他们提出了掩模训练，并构建了一个掩模SwinIR模型，该模型学习图像纹理和结构的重建，而不是过度拟合到特定的噪声类型，他们的模型是在高斯噪声σ = 15上训练的，但可以很好地推广到其他不同的看不见的噪声类型。然而，我们注意到，尽管增强了性能，但他们的模型也引入了不必要的副作用，这些副作用倾向于过度平滑图像内容，导致高频细节的丢失和PSNR的下降（参见图1）。深度去噪的泛化挑战仍然是广泛应用的重大障碍。

在本文中，我们提出了RNINet，这是一种基于流线型编码器-解码器框架的新架构，可提高广义深度图像去噪的效率和整体性能。最初，我们在单个噪声类型上训练一个纯RNINet（仅简单的编码器-解码器），并观察到特征统计量（如均值和方差）随着不同噪声条件的变化而变化（参见图2）。最近的一些研究(Liu et al ., 2021；2023;Chen et al ., 2023)进行了基于特征统计分布的泛化分析实验，但都没有直接对学习到的特征统计进行操作。利用这些见解，我们在RNINet中加入噪声注入块，在特征统计上注入随机噪声，从而创建影响模型学习的噪声特征。而特征统计可以包含特定于领域的信息(Huang & Belongie, 2017；Li等人，2021)，这种噪声注入操作允许带噪特征统计量表示潜在的未见噪声域，显著增强模型的泛化能力。本工作的主要贡献总结如下：

• 我们提出了RNINet，这是一种利用流线型编码器-解码器框架来提高效率和改善可泛化深度图像去噪性能的新架构。这种方法简化了现有一般化去噪方法中常见的架构复杂性，促进了在实际部署环境中的更广泛应用。

• 我们引入了一个噪声注入块，将随机噪声注入到特征统计中，针对潜在的看不见的噪声域。这一发展显著提高了泛化能力，将我们的方法与现有的主要关注泛化分析而没有直接操作干预的研究区分开来。

• 综合实验表明，RNINet在各种看不见的噪声条件下超越了最先进的MT方法的性能，提供了卓越的去噪效果和计算效率（更低的mac和GPU内存使用率），实现了高达10倍的推理速度，并强调了其大规模部署的能力。

2 RELATED WORK

2.1 图像去噪

图像去噪技术主要分为两类：传统的基于模型的方法和数据驱动的深度学习方法。传统方法（Buades等人，2005；Dabov等人，2007；Elad&Aharon，2006；Gu等人，2014；Mairal等人，2009）通常依赖于对图像先验进行建模，从受噪声影响的图像中恢复内容，在各种噪声类型中表现出一定程度的灵活性和泛化能力（Abdelhamed等人，2018b）。然而，这些方法往往难以重建精细的图像细节和实现高PSNR。相比之下，数据驱动的深度学习模型取得了显著的去噪性能。CNN模型(Zhang et al ., 2017；2022;Lefkimmiatis, 2018;2017;Mao et al ., 2016；Divakar & Venkatesh Babu, 2017；Jia et al ., 2019；Zhang等人，2018)曾经是去噪模型的主流，比传统方法提供了实质性的性能改进。Vision Transformer（Dosovitskiy等人，2020）的出现标志着重大的范式转变，它将像素视为token并利用自我关注来解析token交互。基于Vision Transformer的变体(Zamir et al, 2022；Liang等，2021；赵等，2023；张等，2023a；Wang et al ., 2022；Chen等，2021；2022b)已经在很大程度上取代了CNN模型作为主流解决方案，因为它们捕获全球依赖关系的能力增强了。尽管有这些进步，一个普遍的问题是在噪声模式上的模型训练与测试过程中遇到的相同，其中主要的性能指标成为网络对训练噪声的过拟合能力。

2.2 泛化问题

在图像去噪等低层次视觉任务中，训练降噪与测试降噪之间存在差异往往会导致泛化困境。传统上，模型是在高斯噪声上训练的，这种做法与现实世界中遇到的主要非高斯噪声不一致，导致性能下降。为了解决这个问题，解决方案分为两种主要方法：一种是在训练期间更紧密地模拟现实世界的噪音(Brooks等人，2019b；Wei et al ., 2020；Chen et al ., 2018；郭等人，2019；Plotz & Roth, 2017；Krull等人，2019；Abdelhamed等人，2018b)，而另一个则开发了“盲”去噪模型，假设噪声水平未知或在大量噪声类型上进行训练(Krull等人，2019；Yue等人，2019；张等，2023b；2017)。Chen等人最近的工作（Chen等人，2023）指出，这些努力没有充分研究泛化的不足；这些方法仍然不能推广到训练数据集中没有表示的噪声类型。一些研究(Liu et al ., 2021；2023)试图分析超分辨率模型泛化能力差的原因，并发现传统的训练方法倾向于使模型过拟合特定的退化类型以获得更高的PSNR。在这些见解的基础上，Chen等人（Chen等人，2023）引入了掩膜训练并构建了一个masked SwinIR (Liang等人，2021)，旨在专注于内容重建，而不是过度拟合特定的噪声类型。虽然他们的方法取得了值得称赞的结果，但它也引入了一些不必要的副作用：图像内容倾向于过度平滑，导致高频细节的丢失和PSNR性能的下降。最近，Cheng等人（Cheng et al, 2024）通过结合预训练的CLIP模型提高了泛化性能。然而，由于CLIP模型已经暴露于包含不同噪声类型的数十亿图像的先验信息中，因此他们的方法被分类到另一个独立的轨道中，并且通常被排除在与仅使用高斯噪声训练的基准方法的直接比较之外，以确保公平性。

2.3 神经网络的特征统计

特征统计（即均值和方差）通常用于各种神经网络的分析，研究(Huang & Belongie, 2017；Li et al ., 2021)，表明它们可以捕获特定领域的信息特征（例如，颜色、纹理和对比度）。在涉及非分布数据的场景中，由于域特征不同，特征统计往往与训练域的统计不一致(Wang et al ., 2019；Gao et al ., 2021)，而一些归一化方法(Ioffe & Szegedy, 2015；Li et al ., 2022)也可以通过操纵特征统计来提高模型性能。近年来，在低层次视觉领域的研究也在探索基于特征统计的模型泛化能力。例如（Liu et al ., 2023）提出了一种基于特征统计的度量来评估超分辨率模型的泛化能力，（Liu et al ., 2021）引入了基于特征统计的降维聚类后的CHI评价分数。随后，MT （Chen et al ., 2023）利用（Liu et al ., 2021）的度量来验证其可泛化去噪模型的泛化性能。然而，这些研究通常将特征统计作为从特征中获得的确定性值，并仅依靠对这些值的统计分析来验证其方法的泛化能力和有效性。相比之下，我们的方法提供了一个新的视角，即注入随机噪声张量来改变特征统计，从而增强我们去噪模型的泛化能力。

3 METHOD

3.1 整体

RNINet的总体结构如图3所示。在一般推理流形中，给定噪声输入图像$\mathbf{I}\in {\mathbb{R}}^{H\times W\times 3}$，我们提出的RNINet首先通过卷积操作提取低级特征${\mathbf{F}}_0\in {\mathbb{R}}^{H\times W\times C}$，然后是ReLU激活函数，其中$H\times W$表示空间分辨率，$C$表示通道数。随后，这些特征嵌入${\mathbf{F}}_0$通过四层分层编码器-解码器结构进行处理，转化为深度特征${\mathbf{F}}_{\mathbf{d}}\in {\mathbb{R}}^{H\times W\times C}$。每个编码器-解码器级别包含多个基本块，每个基本块由卷积层、批处理规范化层和ReLU激活层组成。编码器逐步降低空间分辨率，同时增强信道容量，最终达到低分辨率潜在表示${\hat{\mathbf{F}}}_{\mathbf{n}}$。为了方便编码过程，在编码器的每两个级别上都有策略地集成了噪声注入块。这些块通过注入随机噪声张量获得带噪特征统计量。解码器的目标是逐步重建来自${\hat{\mathbf{F}}}_{\mathbf{n}}$的高分辨率干净输出。特征内的下采样和上采样分别使用卷积和转置卷积执行。将改进后的深度特征${\mathbf{F}}_{\mathbf{d}}$经过最后的卷积层处理，得到残差图像$\mathbf{R}\in {\mathbb{R}}^{H\times W\times 3}$，将残差图像与退化后的输入相加，得到恢复后的图像：$\overline{\mathbf{I}}=\mathbf{I}+\mathbf{R}$。接下来，我们详细介绍了包括基本和噪声注入块的特定模块。

3.2 Basic Block

为了降低过拟合的风险，从而提高我们去噪模型的泛化能力，现有研究表明(Liu et al ., 2023；Chen et al ., 2023；Liu等人，2021)，我们在RNINet框架中为基本块选择了一个简单的结构。如图3所示，基本块由三层组成，在网络的编码器和解码器阶段都充当特征提取器。给定输入特征$\mathbf{F}\in {\mathbb{R}}^{B\times H\times W\times C}$，基本块促进的变换过程定义为：
$${\mathbf{F}}_{\mathbf{e}}=\text{ReLU(BN(Conv(}\mathbf{F}\text{)))}\text{\hspace{1em}(1)}$$
式中Conv为卷积运算，BN为批归一化，ReLU为整流线性单元激活函数。这种流线型结构确保了高效和有效的特征提取，降低了复杂性，防止了过拟合问题。

3.3 噪声注入块

噪声注入块旨在生成具有改变统计属性的噪声特征，增强模型对未知噪声类型的泛化。给定特征${\mathbf{F}}_{\mathbf{e}}$，该块首先通过卷积进行下采样（不包括最后的块），然后对特征进行批处理归一化和ReLU激活，得到${\mathbf{F}}_{\mathbf{e}}\in {\mathbb{R}}^{B\times H^s\times W^s\times C^s}$，其中$H^s$，$W^s$，$C^s$