Masked Image Training for Generalizable Deep Image Denoising 论文阅读笔记

• CVPR2023 港科大（广州）发的一篇denoising的论文，作者里面有上海AILab的董超老师（看introduction的时候看到有一段很像董超老师 Networks are slaching off 的论文的思想，说网络overfitting的时候学习了训练集的噪声模式而非图片内容，翻回去作者看，果然有董超老师），文章提出了一种提高现有denoising方法泛化性的方法，动机就是迫使网络学习图像内容而不要学习噪声模式，主要通过两种mask的方法，一个是input mask，一个是attention mask。通过提出的方法，网络用高斯噪声训练，可以泛化到各种各样的噪声类型，如speckle noise、poisson noise等等。

方法介绍

• 文章做了一个很有意思的动机实验：在非自然图像上训练高斯去噪，然后放到自然图像加高斯噪声上去测试，如果去噪成功，说明忘了学习的是训练集上的高斯噪声的模式，从而可以泛化到具有相同噪声却有不同图像内容的自然图像上。从下图的denoised by swinir可以看出来，确实去噪成功，也证实了网络在这次实验中只是在学习高斯噪声的模式，从而无法泛化到其它的噪声上。如果希望网络训练出来之后能够泛化到各种各样的噪声上，应该让网络学习自然图像是怎么样的而非学习噪声是怎么样的。
  
• 从而，文章提出，对输入的图像进行mask操作，使得网络在学习去噪的同时学习重建，这样就能迫使网络学习图像内容而非噪声模式，这样训练的泛化结果（以mask的图片作为输入）即上图的右下角，可以看出网络确实是企图去恢复一张细胞图像，从而确认网络学习到了图像内容。
• 文章使用了SwinIR这个backbone，但是做了一点小修改（正文没有提，github提了issue回答了我，baseline就是swinIR去掉global residual），input mask被加在最初的特征图上，也就是第一层1x1的卷积之后，把随机数量的特征更换成mask token，这一步有点像drop-out操作，只是不置0而是用一个mask token。mask token被初始化为0 vector并且是可学习的参数。(文章提到，其实不学习，直接用0 vector也ok)
• 然而，直接mask input的训练方式会使得测试的时候图片亮度被提高。这是由于train和test的做法的不同，test的时候是没有mask的，其实相当于，如果train的时候用了drop out，test的时候不用drop out，卷积核一次见到的非零值就会变多，最终的结果就会变得比train的时候大。针对这个问题有一个ECCV2018 inpainting的工作提出了mask的卷积（“Image Inpainting for Irregular Holes Using Partial Convolutions”），在卷积的输出乘以一个该卷积核当前区域被mask掉的像素数量成比例的因子，从而解决了这个问题，但这里是transformer，所以文章提出对特征进行mask。其实就是把网络中间的特征，也用随机的概率mask掉一些位置，换成另一个mask token。
  

其它

• 文章很多细节没有披露，好在github的issue有进行详细的解答，比如我一开始比较疑惑为什么attention mask可以解决训练和测试阶段不统一导致的亮度问题，看了issue才明白，原来他们训练和测试都会使用attention mask，而input mask仅在训练使用。至于测试使用attention mask导致信息丢失问题，作者也大方承认了，并说这是feature work的内容。
  
• 当被问及为什么不在其它网络结构上进行实验是，作者说只是没做，不代表不可以，也说这是feature work的内容，我觉得回答得很诚恳！
  

实验

• 在高斯噪声上训练，在各种其它噪声上测试。

• 实验发现使用随机的ratio比使用固定的ratio好
  

• 另一个有意思的点是，随着mask ratio的增加，高斯噪声的去噪效果只有极缓慢的下降而其它噪声的去噪效果却可以增加，直到75%，所以选75%左右是最佳选择，存在一个泛化和拟合的trade off

• 与SOTA方法的比较在下面：
  

• 这里的baseline model和dropout到底是什么，文章并没有明确指示。我猜测baseline是文章正文中描述的transformer网络结构去掉两个mask机制的模型，也就是说文章用的其实并不是swinir，而是一点修改，从结果也看出baseline比swinir效果好一点，而baseline加上两个mask机制效果又比baseline效果好。github的回答证实了我这一点猜想：
  

• 文章做了很多泛化实验，这些实验有很好的借鉴性。比如，文章可视化了训练过程中在不同噪声上的psnr和ssim曲线，可以看出，随着训练进行，baseline对其它噪声的psnr下降了，而mask training的模型对其它噪声的psnr提高了，由此看出baseline在过拟合而mask training在有效拟合。

• 

• 再比如下面这个图：
  

个人观点

• 我觉得实验结果的展示上可以借鉴，当图像的亮度有变化却不明显时，可以借鉴直方图来展示使得读者能更直观地感受亮度的变化。

• dropout模型根据github的回复，是放在整个网络输出层的前面，虽说是引用了另一篇文章的做法，但我觉得奇怪，从本文相同的目的来说，为什么不是放在网络的中间以及图片的输入处（不过这样好像就和本文提出的方法一模一样了）