Masked Training:用遮蔽训练提升图像去噪的跨噪声泛化)
论文题目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 5Attention 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 1Input Mask 与 Attention Mask 的组合消融Table 1 在 CBSD68 混合噪声上分别考察两种机制。只使用 Attention Mask 时PSNR 为 29.17 dB只使用 Input Mask 时降至 26.96 dB两者同时使用则达到 29.74 dBSSIM 为 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 7Input Mask 比例与已知及未知噪声性能的权衡Figure 7 的关键现象是提高 Input Mask 比例在一定范围内会改善散斑与混合噪声表现但已知高斯噪声的性能最终会下降超过约 85% 后信息损失更加明显。实验支持的结论是存在一个有利于未知噪声、又不过度丢失细节的遮蔽区间而非无限增大遮蔽率。5.3 未见噪声上的主实验与应用场景论文 Figure 8空间相关噪声、散斑噪声和椒盐噪声的视觉对比Figure 8 中普通模型在陌生噪声下留下明显残留或伪影Masked Training 对长颈鹿纹理、建筑线条和椒盐噪点的处理更稳定。这表明收益不只体现在 PSNR 上但局部纹理仍可能被平滑。论文 Figure 11Kodak24 上四类未知噪声随强度变化的 PSNR 与 SSIM 曲线Figure 11 比较四类噪声的多个强度。Masked Training 的曲线整体更高随噪声增强下降也较慢改进并不局限于某个固定等级。但这些仍是论文设定的测试分布不能推出任意未知噪声下都最优。论文 Table 3Input 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 9Monte Carlo 渲染噪声的细节恢复对比Figure 9 使用渲染图像中的局部金色装饰细节进行对比部分传统模型仍有强烈噪点本文方法的结构更完整。这里的 Monte Carlo 噪声来自渲染采样过程与训练阶段的高斯噪声并不一致因此是颇有说服力的迁移测试。论文 Figure 10ISP 噪声场景中的文字与纹理恢复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 特征相似性作者用 CKACentered 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 噪声上的有效性也通过消融、训练曲线和表征分析解释了各模块的分工。边界同样明确遮蔽必然丢失部分输入信息作者在结论中承认仍需探索如何更好地保留细节。从工程角度看这也意味着实际使用时应结合目标噪声强度与细节保真要求选择遮蔽比例不能把跨噪声泛化的增益理解为所有条件下无代价的提升。如果用一句话概括本文与其让去噪网络记住一种噪声长什么样不如在训练时迫使它学会图像原本应该长什么样。