[论文解读] Dual-former: Hybrid Self-attention Transformer for Efficient Image Restoration

Sixiang Chen, Ye Tian+2 more|arXiv (Cornell University)|Oct 3, 2022

Image Enhancement Techniques被引用 11

一句话总结

Dual-former 提出了一种混合自注意力 Transformer 架构，通过在编码器-解码器阶段结合基于卷积的局部特征提取与潜在层中的单一混合 Transformer 块（HTB），高效建模长距离空间与通道依赖关系。该设计在图像去雾、去雨和去雪任务中实现了最先进性能，同时显著降低计算成本——例如，在仅使用其 4.2% GFLOPs（9.28G vs. 216.0G）的情况下，相较于 MAXIM 在室内去雾任务中实现了 1.91 dB 的 PSNR 提升。

ABSTRACT

Recently, image restoration transformers have achieved comparable performance with previous state-of-the-art CNNs. However, how to efficiently leverage such architectures remains an open problem. In this work, we present Dual-former whose critical insight is to combine the powerful global modeling ability of self-attention modules and the local modeling ability of convolutions in an overall architecture. With convolution-based Local Feature Extraction modules equipped in the encoder and the decoder, we only adopt a novel Hybrid Transformer Block in the latent layer to model the long-distance dependence in spatial dimensions and handle the uneven distribution between channels. Such a design eliminates the substantial computational complexity in previous image restoration transformers and achieves superior performance on multiple image restoration tasks. Experiments demonstrate that Dual-former achieves a 1.91dB gain over the state-of-the-art MAXIM method on the Indoor dataset for single image dehazing while consuming only 4.2% GFLOPs as MAXIM. For single image deraining, it exceeds the SOTA method by 0.1dB PSNR on the average results of five datasets with only 21.5% GFLOPs. Dual-former also substantially surpasses the latest desnowing method on various datasets, with fewer parameters.

研究动机与目标

• 解决现有基于视觉 Transformer 的图像恢复模型计算复杂度过高的问题，同时保持或提升性能。
• 在计算效率方面，将卷积的局部建模优势与自注意力的全局建模能力相结合。
• 通过仅将自注意力限制在单一潜在层，减少对所有层中完整注意力机制的依赖。
• 设计一种轻量化但强大的主干网络，在更少参数和 FLOPs 的前提下，优于多种图像恢复基准上的最先进方法。

提出的方法

• 在编码器和解码器中广泛采用基于深度可分离卷积和通道注意力的局部特征提取（LFE）模块，以实现高效的局部表征学习。
• 在潜在层引入混合 Transformer 块（HTB），通过联合通道-空间自注意力机制建模长距离依赖关系，并缓解通道不平衡问题。
• 设计自适应控制模块（ACM），动态融合来自多个注意力分支的特征，提升空间与通道维度之间的特征交互。
• 在 HTB 内部采用多分支前馈网络（MFFN），在保持线性复杂度的同时增强表征能力。
• 采用混合损失函数，结合 PSNR 损失与感知损失，以提升重建保真度与感知质量。
• 仅将自注意力计算限制在单一瓶颈层，与全注意力 Transformer 相比，大幅降低 FLOPs。

实验结果

研究问题

• RQ1 是否可通过仅将自注意力限制在单一潜在层，在图像恢复中保留全局建模能力的同时显著降低计算成本？
• RQ2 将基于卷积的局部特征提取与目标化的混合自注意力块相结合，对性能与效率有何影响？
• RQ3 自适应特征融合机制（ACM）对潜在层中空间与通道注意力之间的交互有何影响？
• RQ4 在 LFE 模块中引入通道注意力在不显著增加 FLOPs 的前提下，对局部表征学习的提升程度如何？
• RQ5 在多种图像恢复任务中，所提出的混合架构在 PSNR、SSIM 和计算效率方面与现有 SOTA 方法相比表现如何？

主要发现

• 在室内去雾数据集上，Dual-former 相较于 SOTA 方法 MAXIM 实现了 1.91 dB 的 PSNR 提升，同时仅消耗其 4.2% 的 GFLOPs（9.28G vs. 216.0G）。
• 在五个去雨数据集上，Dual-former 平均 PSNR 超过 SOTA 方法 0.1 dB，且仅使用 21.5% 的 GFLOPs（9.28G vs. 140.92G）。
• 自适应控制模块（ACM）在特征融合中优于简单拼接与 SK 融合方法，实现 30.64 dB PSNR 与 0.939 SSIM，且参数与 FLOPs 增加有限。
• 在 LFE 模块中引入通道注意力使 PSNR 提升 0.42 dB（从 30.22 提升至 30.64 dB），FLOPs 仅从 9.24G 增至 9.28G，参数从 12.52M 增至 14.23M。
• HTB 中并行的局部特征提取相比非并行设计使 PSNR 提升 0.42 dB，且计算开销极低。
• PSNR 损失与感知损失的组合实现最佳性能（30.64 dB PSNR，0.939 SSIM），优于仅使用 L1 或仅使用 PSNR 的训练方式。