《Masked Image Training for Generalizable Deep Image Denoising》——CVPR23论文阅读笔记

Project page: https://github.com/haoyuc/MaskedDenoising

前提：在捕获和存储图像时，设备不可避免地会引入噪声。减少这种噪声是一项关键任务，称为图像去噪。深度学习已经成为图像去噪的事实方法，尤其是随着基于Transformer的模型的出现，这些模型在各种图像任务上都取得了显著的最新成果。

核心问题：基于深度学习的方法去噪缺乏泛化能力。如何提高深度学习去噪泛化能力，使适应更广泛的场景。

方法：提出一种新的方法来提高去噪网络的泛化性能，称为掩码训练。其包括在训练期间掩蔽输入图像的随机像素并重建丢失的信息，屏蔽了自我注意层中的特征，以避免训练-测试不一致性的影响。

该图说明了去噪网络的泛化问题。在σ = 15的高斯噪声上训练SwinIR模型。在相同噪声条件下测试时，SwinIR表现出出色的性能。然而，当应用于分布外噪声时，例如，各种噪音的混合。SwinIR遭受了巨大的性能下降。

在这种场景下，更具体地说，大多数现有的去噪工作在被高斯噪声破坏的图像上训练和评估模型，将它们的性能限制在单个噪声分布上。当这些模型被应用于去除从其他分布中提取的噪声时，它们的性能急剧下降。

近年来，针对深度模型的泛化问题，一些方法假设特定噪声类型的噪声水平是未知的，而另一些方法试图通过合成或收集更接近目标噪声的训练数据或直接对目标噪声执行无监督训练来改善真实世界场景中的性能。然而，这些方法都没有实质性地提高去噪网络的泛化性能，并且当噪声分布不匹配时，它们仍然很难实现。深度去噪的泛化问题仍然对这些方法的广泛应用提出了挑战。

模型在噪声上的表现与训练期间观察到的不同。作者认为深度去噪的泛化问题是由于训练噪声的过拟合。现有的训练策略直接优化去噪图像和地面真实值之间的相似性。这背后的意图是，网络应该学会正确地重建自然图像的纹理和语义。然而，经常被忽视的是，网络也可以简单地通过过拟合噪声模式来减少损失，这比学习图像内容更容易。这是泛化问题的核心。作者的研究试图通过指导深度去噪网络学习图像内容重建而不是过度拟合训练噪声来提高深度去噪网络的泛化性能。从最近的掩蔽建模方法中汲取灵感，采用掩蔽训练策略来显式学习图像内容重建的表示，而不是训练噪声。利用图像处理变换器的特性，引入了两种掩蔽机制：输入掩码和注意力掩码。在训练过程中，输入掩码随机删除输入图像像素，网络重建删除的像素。注意力掩码在Transformer的每个自注意力层中实现，使其能够动态地学习掩码特征的完成，并减轻掩码学习中训练和测试之间的分布偏移。

动机：

目标是创建能够推广到训练集中未遇到的噪声的去噪模型。当在大量图像上训练深度网络时，期望网络能够学习从受噪声污染的测试用例中识别自然图像的丰富语义。作者认为去噪模型的泛化能力差是由训练方法造成的，这导致模型专注于过拟合训练噪声，而不是学习图像重建。

使用与自然图像差异很大的图像训练了SwinIR去噪网络，使用高斯噪声合成训练数据对，然后评估模型在具有高斯噪声的自然图像上的性能。根据假设，如果模型从训练集中学习图像语义的内容和重建，它在自然图像上的表现应该不会很好，因为它没有接触过任何图像。如果模型只是过拟合噪声，即使图像不同，模型也可以去除噪声，因为模型主要依赖于检测噪声进行去噪。结果如图所示。很明显，该方法试图从测试自然图像的训练集重建免疫组织化学图像纹理，而不是依赖于噪声过拟合去噪。这证明了这种想法在提高泛化性能方面的潜力。通过在自然图像上训练该方法，它将专注于重建自然图像的内容，这与将深度学习用于低级视觉任务的核心概念相一致。

正如所观察到的，在免疫组织化学图像上训练的SwinIR仍然可以去噪和再现自然图像。这支持了作者关于泛化能力的猜想，表明大多数现有方法通过过拟合训练噪声来进行去噪。因此，当噪声偏离训练条件时，这些模型的去噪性能显著下降。我们的目标是让模型学习图像纹理和结构的重建，而不是只关注噪声。

在本文中提出了一种新的去噪网络的掩蔽训练策略。在训练过程中，屏蔽掉一部分输入像素，然后训练深度网络来完成它们，如上图所示。强调重建自然的图像纹理和图像中观察到的边缘，而不是过拟合噪声。

模型：

首先进行噪声输入，使用1*1卷积层作为特征嵌入模块，1*1卷积层可确保像素在特征嵌入过程中不会相互影响，这有利于后续的掩蔽操作。特征提取进行掩蔽后送入Mask STLB中(4个Mask STLB,每个Mask STLB包含6个Masked STL和1个卷积层），最后一个Mask STLB后也有一个卷积层。

imput mask：随机屏蔽第一个卷积层嵌入的特征令牌，并鼓励网络在训练期间完成屏蔽信息。输入掩码显式地构造了一个非常具有挑战性的修复问题，即使高达90%的像素信息被破坏，网络仍然可以在一定程度上重建目标图像。

给定特征token张量，以概率pIM将token随机替换为[mask token]∈ RC，其中pIM称为输入掩码比。该网络是在重建图像的L1范数和地面实况的监督下训练的。[mask token]可以是可学习的，并可以用0向量初始化。但实际上发现0向量本身已经是一个合适的选择。输入掩码的存在迫使网络学习从非常有限的信息中识别和重建图像的内容。

attention mask：

不能仅仅依靠输入掩码操作来构建可用的图像处理网络。因为在测试过程中将输入未损坏的图像以保留足够的信息。此时，由于训练和测试之间的不一致性，网络会倾向于增加输出图像的亮度。例如图中的示例。由于Transformer使用自注意操作来处理空间信息，可以通过在自注意过程中执行相同的掩码操作来缩小训练和测试之间的差距。具体的掩码操作类似于输入掩码，但使用不同的注意掩码比率pAM和[mask token]。当自我注意中的某些标记被屏蔽时，注意操作将适应这些标记的信息不再可靠的事实。自我注意力将集中在每一层中未屏蔽的标记上，并完成屏蔽信息。这种操作很难在卷积网络上实现。上图显示了注意力遮罩的效果。可以看出，注意力掩模成功地使掩蔽的训练网络在未掩蔽的输入图像上工作。

实验：

消融实验

表显示了不同输入掩码比的影响，从均匀分布中测试固定比率和随机比率。从实验来看，固定的比率对于训练来说不如从一个范围内随机选择的稳定，并且性能也更差。在75%至85%之间的随机抽样率下，可以实现最佳的定量性能。这是去噪泛化能力和图像细节保留之间的权衡。

如图所示，较小的比率不足以让网络学习图像的分布，因为保留了更多的噪声模式。较大的比率提高了模型的泛化能力，因为模型更关注重建。但与此同时，一些图像细节可能会丢失。对于注意掩模比率在表中显示了效果。最佳比例在75%左右。

结果：

不同噪声类型和声级的测试曲线：

局限性：

总之，掩码训练方法提供了一种很有前途的方法来提高基于深度学习的图像去噪模型的泛化性能。局限性在于掩码操作不可避免地丢失信息。如何保留更多的细节需要在今后的工作中探索。