【ICCV2023】Masked Image Training for Generalizable Deep Image Denoising

论文摘要：

🧐

本文提出了一种提高去噪网络泛化性能的新方法，称为掩蔽训练。该方法涉及在训练过程中对输入图像的随机像素进行掩蔽，并在训练过程中重建缺失的信息。同时，还本文还在自注意力层中掩蔽特征，以避免训练和测试不一致性的影响。

🔥

本文方法具有比其他深度学习模型更好的泛化能力，并且可以直接应用于实际场景。

动机：

目前深度去噪网络在泛化能力方面存在一些问题。例如，当训练集中的噪声类型与测试集中的噪声类型不一致时，模型的性能会受到很大的影响，表现出较差的泛化能力。目前，大多数方法都是在高斯噪声上训练和测试的，当遇到训练过程中未见过的噪声时，这些方法的效果会急剧下降。

本文进行了一个小实验，在高斯噪声(σ=15)上训练了一个SwinIR模型。如图a所示，在同样的高斯噪声(σ=15)下进行测试，SwinIR展现出了出色的性能。然而，当应用于分布不同的噪声（各种噪声的混合）时，如图b所示，SwinIR的性能显著下降。尽管如此，本文方法在高斯噪声上训练的模型仍然能够保持较好的去噪效果。

作者认为，出现这种问题的原因是目前的模型倾向于过度拟合训练噪声，通过简单地过拟合噪声模式来降低损失，而这比学习图像内容更加容易，从而导致模型没有学习图像自身的内在结构。

本文进行了实验，以验证其猜想的正确性。作者在仅带有高斯噪声的免疫组织化学图像上训练了一个SwinIR模型，并使用该模型对自然图像进行去噪，结果仍然表现出很好的效果。这个观察结果说明了：大多数现有方法主要通过过拟合训练噪声来进行去噪，而不是真正理解图像内容。

本文所提出的方法强调在自然图像上重建训练集中观察到的自然图像纹理和语义边缘，而不是依靠噪声过拟合来进行去噪。这种区别强调了我们的方法与以前方法之间的根本区别。如图所示，在免疫组织化学图像训练的本文模型上，对自然图像进行去噪的效果较差。

本文要解决的核心问题是：如何提高深度学习去噪模型的泛化能力，使其不仅能去除训练使用的噪声类型，还能够很好地处理其他未知的噪声，适应更广泛的场景。

本文方法：

为了解决上述问题，我们需要通过改进训练策略，让模型学习重构图像内容。作者提出了遮挡训练的方法，该方法可以使模型学习理解和重构图像的内容，而不仅仅依赖于噪声特征，从而获得更好的泛化能力。

如图所示，展示了本文所提出的 mask-and-complete 训练策略，每组图片左侧是 mask 后的效果，右侧是重建后的效果。可以看到，即使有大量的像素被遮罩，模型仍然可以在一定程度上重建输入。

整体网络架构基于了Swin Transformer，作者提出了一种 masked training 的方法来提高模型的泛化能力。

在本文方法中，使用一个核大小为1的卷积层作为特征嵌入模块，将三通道像素值投影到C维的特征标记中。卷积层确保了像素在特征嵌入过程中不会相互影响，便于后续的掩码操作。特征标记的形状为，其中、和是高度、宽度和特征尺寸。

移动窗口机制首先通过将输入划分为不重叠的本地窗口，其中是窗口的总数，将每个帧的特征映射重塑为特征。在同一窗口内的特征标记上计算自注意力。因此，标记涉及到每个标准的自注意力操作，我们产生了局部窗口特征。

在每个self-attention层中，查询、键和值计算为，其中是权重矩阵，是投影向量的维数。然后，我们使用查询生成注意力图，其中是可学习的相对位置编码。然后，这个注意力图用于中的向量的加权和。

本文的掩码训练主要涉及两个方面：输入掩码（Input Mask）和注意力掩码（Attention Mask），尽管它们都涉及掩码操作，但它们的目的是不相同的，下面将对这两个部分做详细介绍。

输入掩码 Input Mask

在特征提取之后, 会对输入图像进行随机大比例遮盖（input mask），比如遮盖75%~85%的像素。这将构造一个非常具有挑战性的图像修复问，迫使网络学习重构被遮盖的内容，增强对图像本身分布的建模能力，从而对未见过的噪声类型也有好的去噪效果。

给定特征标记张量，我们以概率随机将标记替换为一个 [mask token] ，其中被称为输入掩码比例。网络在重建图像的L1范数和真实图像的监督下进行训练。[mask token] 可以是可学习的，并且可以用一个 0 向量初始化。

注意力掩码 Attention Mask

在self-attention层也进行类似的随机遮盖。这是为了减轻训练和测试的不一致性。由于 input mask 只在训练使用，测试时的输入是完整的图像。因此，我们可以通过在自注意力过程中执行相同的掩码操作来缩小训练和测试之间的差距。

由于Transformer使用自注意力操作处理空间信息，具体的掩码操作类似于输入掩码，但使用了不同的注意力掩码比例和 [mask token]。当自注意力中的某些标记被屏蔽时，注意力操作将调整到这些标记的信息不再可靠的情况。自注意力将专注于每层中的未掩码标记，并补全掩码信息。

实验结果：

消融实验：

如果没有输入掩码，模型将失去其泛化能力，并且不能有效地去除训练集之外的噪声。在没有注意力掩码的情况下，由于训练-测试不一致，量化性能显著下降，输出图像会出现错误的亮度。此外，即使没有注意力屏蔽，模型的泛化能力也没有明显影响，大部分噪声仍然被有效去除。

遮盖比例参数：

图 7 展示了在不同的输入掩码比例时，模型在训练集噪声（高斯 15）和训练集外的噪声（Speckle, Mixture noise）上的性能。可以看到，在掩码比率为75%之前，训练噪声的性能下降并不明显。而在掩码比例为 50%左右时，就已经体现出较高的性能。而在训练集之外的噪声上的性能提升大于在训练集上的性能损失。

较小的比率不足以使网络学习到图像的分布，因为更多的噪声模式被保留下来。较大的比率提高了模型的泛化能力，因为模型更加关注重构。但与此同时，一些图像细节可能会丢失。

除噪效果展示：

Figure 8 展示了不同去噪算法在各种未出现在训练中的噪声下的视觉效果。测试噪声类型包括Spatially correlated Gaussian、Speckle noise和 Salt-and-pepper noise，都不属于训练使用的高斯噪声。对比算法包括DnCNN、RIDNet、RNAN、SwinIR、Restormer等当前主流算法。

实验结果显示，这些对比算法完全失效，无法有效去除测试噪声，图像效果很差。而本文提出的遮挡训练方法在所有的测试噪声下都获得了很好的去噪视觉效果。这直观地反映了作者方法相比其他算法在泛化性上的明显提高。尤其是在其他方法完全失败的情况下，遮挡训练仍能有效去噪，突出了方法的优势。这验证了遮挡训练可以减少对训练噪声的依赖，提高模型对复杂未知噪声的适应性。