[论文解读] U2-Former: A Nested U-shaped Transformer for Image Restoration

Haobo Ji, Xin Feng+3 more|arXiv (Cornell University)|Dec 4, 2021

Image Enhancement Techniques被引用 14

一句话总结

U2-Former 提出了一种深度嵌套的 U 形 Transformer 架构用于图像恢复，通过多尺度特征交互和特征过滤机制提升计算效率。通过整合多视角对比学习，该方法能够解耦噪声与干净图像，实现在去反射、去雨和去雾任务上的最先进性能。

ABSTRACT

While Transformer has achieved remarkable performance in various high-level vision tasks, it is still challenging to exploit the full potential of Transformer in image restoration. The crux lies in the limited depth of applying Transformer in the typical encoder-decoder framework for image restoration, resulting from heavy self-attention computation load and inefficient communications across different depth (scales) of layers. In this paper, we present a deep and effective Transformer-based network for image restoration, termed as U2-Former, which is able to employ Transformer as the core operation to perform image restoration in a deep encoding and decoding space. Specifically, it leverages the nested U-shaped structure to facilitate the interactions across different layers with different scales of feature maps. Furthermore, we optimize the computational efficiency for the basic Transformer block by introducing a feature-filtering mechanism to compress the token representation. Apart from the typical supervision ways for image restoration, our U2-Former also performs contrastive learning in multiple aspects to further decouple the noise component from the background image. Extensive experiments on various image restoration tasks, including reflection removal, rain streak removal and dehazing respectively, demonstrate the effectiveness of the proposed U2-Former.

研究动机与目标

• 解决由于计算成本高和跨尺度通信能力差，导致 Transformer 在图像恢复中应用深度有限的问题。
• 通过增强层间尺度交互，克服标准编码器-解码器 Transformer 在图像恢复中效率低下的问题，实现更深层次的特征学习。
• 通过在块级、内容级和噪声级表示之间引入多视角对比学习，提升噪声与背景的分离能力。
• 通过特征过滤机制优化自注意力模块的计算效率，压缩低质量 token 以减少计算负载。
• 在不依赖 CNN 先验的前提下，实现在去反射、去雨和去雾等多种图像恢复任务中的优越性能。

提出的方法

• 采用嵌套 U 形架构，外层为编码器-解码器框架，内层为 U 形 Transformer 块，以实现深度多尺度特征交互。
• 引入特征过滤机制，在自注意力计算前剪枝低质量 token，降低计算负载，同时保留关键特征。
• 实施多视角对比学习，配置三种正样本对：(1) 同一恢复图像中的图像块，(2) 恢复图像与真实干净图像，(3) 恢复图像与随机真实图像，以学习噪声不变特征。
• 结合真实图像的监督信号与对比学习，联合优化图像恢复与特征解耦。
• 设计模型在编码和解码路径中深度堆叠 Transformer 块，实现跨尺度的全局上下文建模。
• 在所有层中利用自注意力机制捕捉长距离依赖关系，提升恢复图像的空间一致性。

实验结果

研究问题

• RQ1 具有嵌套 U 形结构的深度 Transformer 架构是否能有效改善图像恢复中跨尺度的特征通信？
• RQ2 Transformer 块中的特征过滤机制是否能显著降低计算成本，同时保持性能？
• RQ3 多视角对比学习是否能有效解耦图像恢复中噪声成分与干净图像内容？
• RQ4 在多种图像恢复任务中，所提出的 U2-Former 在 PSNR 和 SSIM 指标上与最先进方法相比表现如何？
• RQ5 与标准 U-Net 或 Uformer 风格架构相比，嵌套 U 形设计在性能提升方面有多大贡献？

主要发现

• 在真实世界去反射基准上，U2-Former 在 Real20 数据集上达到 23.67 dB PSNR 和 0.835 SSIM，优于 Uformer（22.02 dB PSNR）及其他最先进方法。
• 在 Rain100L 数据集上，U2-Former 达到 39.31 dB PSNR，较 Uformer（37.68 dB）提升 1.63 dB，且可视化结果中细节恢复更优。
• 在 Rain100H 数据集上，U2-Former 达到 30.87 dB PSNR，较 Uformer 提升 1.50 dB，展现出在强降雨条件下的优异性能。
• 在图像去雾任务中，U2-Former 在室内数据集上达到 36.42 dB PSNR 和 0.988 SSIM，在室外数据集上达到 31.10 dB PSNR 和 0.976 SSIM，分别优于 Uformer 4.51 dB 和 4.58 dB。
• 视觉结果表明，U2-Former 有效缓解了色彩失真和过度平滑问题，相比 DCP、DA 等基线方法更好地保留了细节。
• 消融实验确认，嵌套 U 形结构与多视角对比学习均对性能提升有显著贡献，特征过滤机制在不降低性能的前提下提升了推理效率。