arXiv:2412.15845cs.CV2024-12被引 6

用Mamba+Transformer融合提升图像修复效果,兼顾速度与精度。

Multi-dimensional Visual Prompt Enhanced Image Restoration via Mamba-Transformer Aggregation

  • Mamba处理空间信息,线性复杂度捕捉长距离依赖。
  • Transformer专注通道建模,避免计算量二次增长。
  • 多维提示学习增强对各类退化的感知能力,适合多任务修复。

近期图像修复研究致力于开发可处理多种退化类型和程度的“全功能”模型。然而,主流基于Transformer的方法在模型能力与计算开销间面临困境:自注意力机制随图像尺寸呈平方级增长,且难以捕捉长程依赖。多数Mamba相关方法仅沿空间维度扫描特征图进行全局建模,未能充分利用通道维度信息。为此,本文提出充分融合Mamba与Transformer的优势,在不牺牲计算效率的前提下实现互补。具体而言,Mamba的选通扫描机制用于空间建模,以线性复杂度捕获长程空间依赖;Transformer的自注意力机制则用于通道建模,避免随图像空间维度平方增长的高计算负担。此外,为丰富有效提示,提出多维提示学习模块,从多尺度编码器/解码器层学习提示流,从空间与通道双重视角揭示各类退化的潜在特征,从而增强“全功能”模型解决多样化修复任务的能力。在图像去噪、去雾、去雨等多个基准任务上的大量实验表明,该方法显著优于众多主流方法,达到新SOTA。相关代码与预训练参数将开源于GitHub:https://github.com/12138-chr/MTAIR。

原文摘要 · Abstract (English)

Recent efforts on image restoration have focused on developing "all-in-one" models that can handle different degradation types and levels within single model. However, most of mainstream Transformer-based ones confronted with dilemma between model capabilities and computation burdens, since self-attention mechanism quadratically increase in computational complexity with respect to image size, and has inadequacies in capturing long-range dependencies. Most of Mamba-related ones solely scanned feature map in spatial dimension for global modeling, failing to fully utilize information in channel dimension. To address aforementioned problems, this paper has proposed to fully utilize complementary advantages from Mamba and Transformer without sacrificing computation efficiency. Specifically, the selective scanning mechanism of Mamba is employed to focus on spatial modeling, enabling capture long-range spatial dependencies under linear complexity. The self-attention mechanism of Transformer is applied to focus on channel modeling, avoiding high computation burdens that are in quadratic growth with image's spatial dimensions. Moreover, to enrich informative prompts for effective image restoration, multi-dimensional prompt learning modules are proposed to learn prompt-flows from multi-scale encoder/decoder layers, benefiting for revealing underlying characteristic of various degradations from both spatial and channel perspectives, therefore, enhancing the capabilities of "all-in-one" model to solve various restoration tasks. Extensive experiment results on several image restoration benchmark tasks such as image denoising, dehazing, and deraining, have demonstrated that the proposed method can achieve new state-of-the-art performance, compared with many popular mainstream methods. Related source codes and pre-trained parameters will be public on github https://github.com/12138-chr/MTAIR.

图像修复MambaTransformer提示学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。