基于Mamba-Transformer聚合的多维视觉提示增强图像恢复
Multi-dimensional Visual Prompt Enhanced Image Restoration via Mamba-Transformer Aggregation
AI总结:
本文提出一种基于Mamba-Transformer聚合的多维视觉提示图像恢复方法,结合Mamba空间建模与Transformer通道建模,并引入多维提示学习模块,在多个基准任务上实现了最先进性能。
摘要:
近期图像恢复的研究致力于开发“多合一”模型,以在单一模型内处理不同的退化类型和程度。然而,大多数主流基于Transformer的模型面临着模型能力与计算负担之间的困境,因为自注意力机制的计算复杂度随图像大小呈二次增长,且在捕捉长程依赖方面存在不足。大多数与Mamba相关的模型仅在空间维度扫描特征图进行全局建模,未能充分利用通道维度的信息。为解决上述问题,本文提出充分利用Mamba和Transformer的互补优势,同时不牺牲计算效率。具体而言,采用Mamba的选择性扫描机制专注于空间建模,在线性复杂度下捕捉长程空间依赖。Transformer的自注意力机制被应用于通道建模,避免了随图像空间维度呈二次增长的高计算负担。此外,为丰富有效图像恢复的信息提示,提出了多维提示学习模块,从多尺度编码器/解码器层学习提示流,有助于从空间和通道视角揭示各种退化的潜在特征,从而增强“多合一”模型解决各种恢复任务的能力。在图像去噪、去雾和去雨等多个图像恢复基准任务上的大量实验结果表明,与许多流行的主流方法相比,所提出的方法能够取得新的最先进性能。相关源代码和预训练参数将在github公开。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。