论文题目:Masked Image Training for Generalizable Deep Image Denoising(面向可泛化深度图像去噪的遮蔽图像训练)
会议:CVPR 2023
摘要:在图像采集与存储过程中,设备不可避免地会引入噪声,去除这些噪声是一项重要的图像去噪任务。深度学习已成为图像去噪的主流方法,特别是随着基于 Transformer 的模型出现,它们在各种图像任务上取得了显著的先进结果。然而,基于深度学习的方法往往缺乏泛化能力。例如,在高斯噪声上训练的深度模型,面对其他噪声分布时可能表现不佳。为解决这一问题,本文提出一种称为遮蔽训练(Masked Training)的新方法,以增强去噪网络的泛化性能。训练时,该方法随机遮蔽输入图像中的像素,并重建缺失的信息;同时在自注意力层中遮蔽特征,以避免训练与测试不一致带来的影响。相比其他深度学习模型,所提方法表现出更好的泛化能力,并能够直接应用于真实场景。此外,可解释性分析也展示了该方法的优势。
源码:https://github.com/haoyuc/MaskedDenoising
一、研究背景与核心问题
深度图像去噪的常见做法是:向干净图像加入预先设定的噪声,用“含噪图像—干净图像”成对数据训练网络。只要测试噪声与训练噪声一致,DnCNN、SwinIR、Restormer 等模型通常能够恢复较好的视觉结果。真正困难的是分布外噪声(Out-of-Distribution Noise):模型只见过标准差 σ = 15 的高斯噪声,测试时却遇到泊松噪声、散斑噪声、空间相关噪声或几种噪声的混合。
论文 Figure 1:相同 SwinIR 在训练噪声与未知混合噪声上的去噪差异
Figure 1 给出了论文的出发点:在训练分布内,普通 SwinIR 去噪表现良好;换成混合噪声,性能明显下降,而采用 Masked Training 的模型仍能清除相当一部分噪声。这里要区分两个目标:拟合已知噪声与恢复不同噪声下的图像内容,两者并不等价。
作者认为,传统监督训练虽然以干净图像为目标,却允许网络通过熟悉训练噪声的统计模式来降低重建误差;这种“识别特定噪声”的捷径可能比真正学习纹理、边缘和结构更加容易。为验证这一判断,作者用与自然图像差异很大的免疫组织化学图像训练 SwinIR,再在相同高斯噪声条件下测试自然图像。
论文 Figure 2:跨图像内容的诊断实验及遮蔽重建结果
普通 SwinIR 即使没有见过自然图像,也能对自然图像去噪;遮蔽训练模型则倾向于按照训练集中熟悉的纹理补全陌生图像。这个实验并非严格证明网络“只学习噪声”,但为作者的机制假设提供了支持:传统方法更依赖噪声模式,遮蔽训练更强调内容先验。论文要解决的核心问题不是拟合更多种噪声,而是怎样改变网络的学习倾向,让恢复图像内容成为必经之路。
二、方法整体框架
作者借鉴 Masked Image Modeling 的思想,将原本直接从含噪图像恢复干净图像的任务,改造成一个更难的“遮蔽—补全—去噪”问题。其直觉是:当大量位置的输入信息被拿走,网络就无法只根据局部噪声模式做简单映射,而必须利用剩余像素之间的结构关系进行重建。
论文 Figure 3:不同遮蔽比例下的图像补全示意
Figure 3 展示了遮掉 70% 或 90% 的像素信息后,模型仍可在一定程度上重建图像。这不是说极端遮蔽没有损失,而是说明图像存在可利用的空间冗余与结构规律。
论文 Figure 4:基于 SwinIR 的 Masked Training 整体架构
整体流程可以概括为:
含噪图像 → 1 × 1 卷积特征映射 → Input Mask → Masked Swin Transformer Blocks → 图像重建 → L1 监督
方法没有重新设计骨干网络,而是在 SwinIR 式 Transformer 上增加两种遮蔽:Input Mask在输入端制造信息缺失,改变网络要学的内容;Attention Mask在自注意力内部再次遮蔽特征,缓解训练与测试不一致。两种 Mask 承担的任务不同。
三、Input Mask:让网络必须学习内容重建
Input Mask 的操作非常直接。先用 1 × 1 卷积将 RGB 输入映射为 H × W × C 的特征 Token,再以概率 p_IM 随机将部分 Token 替换为 Mask Token。论文指出,Mask Token 可以设为可学习向量,但实验发现直接使用零向量已经合适。之所以强调 1 × 1 卷积,是因为它在特征嵌入阶段不会把相邻像素信息混合到一起,后续遮蔽才能真正移除对应位置的信息。
训练仍然使用成对的干净图像,并最小化网络输出与 Ground Truth 之间的 L1 像素误差,即 L = ‖ŷ − y‖₁;其中 ŷ 为重建结果、y 为干净目标图像。这不是不需要干净标签的自监督去噪,也不是只在被遮蔽位置计算误差的 MAE 式预训练。变化主要发生在输入构造上:网络被迫根据有限观测推断缺失区域,同时输出完整的去噪图像。
为什么这样可能更泛化?传统网络能从带噪像素中找到与训练分布相关的局部线索;遮蔽多数 Token 后,这些线索不再充分,纹理连续性、边缘走向及图像结构的重要性随之提高。不过,遮蔽越多不一定越好:可利用的真实细节也会减少,必须在泛化与信息保留之间权衡。
四、Attention Mask:解决训练与测试的分布差异
如果只有 Input Mask,训练时模型始终面对大量缺失信息;测试时为了保留细节,却要送入未做输入遮蔽的完整含噪图像。两种输入状态不一致,作者观察到模型可能出现整体亮度偏移,导致量化结果和视觉结果都变差。
论文 Figure 5:Attention Mask 对亮度分布与重建结果的影响
Figure 5 的直方图和花朵示例说明,去掉 Attention Mask 后,输出亮度分布偏向高亮区域;加入后,亮度分布更接近参考图像。换句话说,Input Mask 负责制造有价值的学习任务,但也产生新的输入分布问题,Attention Mask 则用来修正这个副作用。
具体做法是在 Swin Transformer 的自注意力计算中,以独立比例 p_AM 将部分特征 Token 替换为另一种 Mask Token,让注意力从其余位置提取信息。标准注意力仍遵循 A = softmax(QKᵀ / √D + B),其中 Q、K 是查询与键,D 是投影维度,B 是相对位置偏置;本文修改的是参与计算的特征,而不是把注意力矩阵中的权重手工置零。两种遮蔽共同作用,使模型在输入完整的测试情形下也能稳定工作。
五、实验结果与消融分析
5.1 实验设置与比较范围
为检验真正的跨噪声泛化,作者将各模型统一训练在 σ = 15 的高斯噪声上,干净训练图像来自 DIV2K、Flickr2K、BSD500 和 WED;训练裁块为 64 × 64,批大小 64,总迭代 200K。测试覆盖 CBSD68、Kodak24、McMaster、Urban100 等图像数据集上的多种未见噪声,包括散斑、泊松、空间相关高斯、椒盐、混合噪声,以及 Monte Carlo 渲染噪声与合成 ISP 噪声。指标采用 PSNR、SSIM(越高越好)和 LPIPS(越低越好)。
关键在于:训练阶段没有混入各种测试噪声。因此实验考察的是只学过一种高斯噪声之后,能否处理新噪声,而非多噪声联合训练的覆盖范围。
5.2 两种 Mask 是否都必要
论文 Table 1:Input Mask 与 Attention Mask 的组合消融
Table 1 在 CBSD68 混合噪声上分别考察两种机制。只使用 Attention Mask 时,PSNR 为 29.17 dB;只使用 Input Mask 时降至 26.96 dB;两者同时使用则达到 29.74 dB,SSIM 为 0.8672。数值说明,单独依赖 Input Mask 会受到训练与测试差异的干扰,而加入 Attention Mask 能恢复稳定性。
论文 Figure 6:去掉不同 Mask 后的混合噪声可视化结果
Figure 6 区分了两种模块的职责:缺少 Input Mask,模型难以清理陌生噪声;缺少 Attention Mask,虽然仍能消噪,却出现亮度问题。前者是泛化能力的主要来源,后者使训练方式可用,并非两种重复的正则化操作。
论文 Table 2:不同 Attention Mask 比例的消融
在 Attention Mask 比例为 65%、75%、85% 时,PSNR 分别为 29.57、29.76、28.84 dB,最优点约在 75%。比例过大也会破坏有效上下文,说明 Attention Mask 同样不是越强越好。
论文 Figure 7:Input Mask 比例与已知及未知噪声性能的权衡
Figure 7 的关键现象是:提高 Input Mask 比例,在一定范围内会改善散斑与混合噪声表现,但已知高斯噪声的性能最终会下降;超过约 85% 后,信息损失更加明显。实验支持的结论是存在一个有利于未知噪声、又不过度丢失细节的遮蔽区间,而非无限增大遮蔽率。
5.3 未见噪声上的主实验与应用场景
论文 Figure 8:空间相关噪声、散斑噪声和椒盐噪声的视觉对比
Figure 8 中,普通模型在陌生噪声下留下明显残留或伪影;Masked Training 对长颈鹿纹理、建筑线条和椒盐噪点的处理更稳定。这表明收益不只体现在 PSNR 上,但局部纹理仍可能被平滑。
论文 Figure 11:Kodak24 上四类未知噪声随强度变化的 PSNR 与 SSIM 曲线
Figure 11 比较四类噪声的多个强度。Masked Training 的曲线整体更高,随噪声增强下降也较慢,改进并不局限于某个固定等级。但这些仍是论文设定的测试分布,不能推出任意未知噪声下都最优。
论文 Table 3:Input Mask 比例及 Monte Carlo 渲染与 ISP 噪声的定量结果
Table 3a 进一步比较固定遮蔽率和区间随机遮蔽率。在 CBSD68 混合噪声上,随机取 75%~85% 达到 PSNR 30.04 dB、SSIM 0.8756,高于表中几组固定比例;这也是论文最终偏向随机比例的直接证据。Table 3b 则显示,面对 128 spp 的 Monte Carlo 渲染噪声,本文方法取得 30.62 dB,普通 Baseline 为 29.68 dB;在更难的 64 spp 下,两者分别为 28.25 dB 和 25.91 dB,提升达到 2.34 dB。
论文 Figure 9:Monte Carlo 渲染噪声的细节恢复对比
Figure 9 使用渲染图像中的局部金色装饰细节进行对比,部分传统模型仍有强烈噪点,本文方法的结构更完整。这里的 Monte Carlo 噪声来自渲染采样过程,与训练阶段的高斯噪声并不一致,因此是颇有说服力的迁移测试。
论文 Figure 10:ISP 噪声场景中的文字与纹理恢复
Table 3c 对应按 ISP 成像过程合成的噪声:本文方法与 DnCNN 的 PSNR同为 29.44 dB,并非这一项上独占第一;但其 SSIM 为 0.7920、LPIPS 为 0.2368,优于 DnCNN 的 0.7857、0.3083。结合 Figure 10 可以看出,本文方法在文字纹理和视觉质量上更有优势。要注意,这里报告的是合成 ISP 噪声,不宜直接表述为已经在所有真实相机噪声上验证成功。
5.4 泛化提升是否有机制解释
论文 Figure 13:不同训练迭代下的跨噪声测试曲线
Figure 13 展示出普通训练模型在未知噪声上的表现可能随训练继续而变差,呈现过拟合趋势;遮蔽训练模型则逐步提高并维持较好性能。这从训练动态上支持了“改变学习偏好”而不仅是“多训一会儿”的解释。
论文 Figure 12:普通训练与遮蔽训练网络的 CKA 特征相似性
作者用 CKA(Centered Kernel Alignment)比较不同层和不同噪声条件下的内部表征。Figure 12a 显示两种训练策略的深层表示明显不同;Figure 12b 显示普通模型在高斯与泊松噪声下的深层表示相似度较低,而遮蔽训练模型维持了更高的一致性。这不能直接证明网络“理解图像语义”,但与它更少依赖特定噪声统计的假设相吻合。
论文 Figure 14:使用 SRGA 指标比较不同噪声下的泛化能力
SRGA 衡量网络内部表示跨退化条件的一致程度,数值越低表示泛化越好。Figure 14 中,遮蔽训练在散斑、泊松、椒盐及空间相关噪声上普遍低于普通训练,给出了与 PSNR、SSIM 不同维度的证据。
论文 Figure 15:不同噪声输入下的深层特征分布
Figure 15 进一步可视化了这种差异:普通训练模型的深层特征分布随噪声类型变化更明显,而 Masked Training 的多种噪声特征分布更加集中。结合 CKA、SRGA 和训练曲线,作者建立了一条相对完整的机制证据链:强迫内容补全 → 减少对训练噪声的依赖 → 跨噪声表征更一致 → 未见噪声下的恢复表现更稳。
六、总结与思考
这篇论文最值得记住的,不是“给输入随机打洞”这个操作本身,而是它对去噪训练目标的重新审视:监督损失相同,并不意味着网络会学到我们想让它学到的东西。当直接映射允许模型依赖固定噪声模式时,制造大量信息缺失可以迫使它更多借助纹理和结构;但这一改变又引入训练与测试分布差异,因此需要 Attention Mask 做配套修正。
实验结果支持该设计在多类未见噪声、Monte Carlo 渲染与合成 ISP 噪声上的有效性,也通过消融、训练曲线和表征分析解释了各模块的分工。边界同样明确:遮蔽必然丢失部分输入信息,作者在结论中承认仍需探索如何更好地保留细节。从工程角度看,这也意味着实际使用时应结合目标噪声强度与细节保真要求选择遮蔽比例,不能把跨噪声泛化的增益理解为所有条件下无代价的提升。
如果用一句话概括本文:与其让去噪网络记住一种噪声长什么样,不如在训练时迫使它学会图像原本应该长什么样。