基于 Mamba-Transformer 聚合的多维视觉提示增强图像修复
计算机视觉与模式识别
2024-12-23 v1
摘要
近期的图像修复研究聚焦于开发能够在单一模型中处理不同退化类型和程度的“全能”模型。然而,大多数基于Transformer的模型在模型能力与计算负担之间面临困境,因为自注意力机制随图像尺寸呈二次增长,且不足以捕捉长程依赖。大多数基于Mamba的方法仅在空间维度上扫描特征图进行全局建模,未能充分利用通道维度的信息。为此,本文提出一种充分利用Mamba和Transformer互补优势且不牺牲计算效率的方法。具体而言,采用Mamba的选择性扫描机制专注于空间建模,实现线性复杂度下的长程空间依赖捕获;采用Transformer的自注意力机制专注于通道建模,避免随图像空间维度呈二次增长的高计算负担。此外,为丰富有效的图像修复提示,本文提出多维提示学习模块,从多尺度编码器/解码器层学习提示流,促进从空间和通道视角揭示各种退化特征,从而增强“全能”模型解决各种修复任务的能力。大量实验结果表明,我们的方法在多个图像修复基准任务(包括图像去噪、去雾和去雨)上实现了新的最先进性能,优于众多主流方法。相关源代码和预训练参数将公开于 github https://github.com/12138-chr/MTAIR。
引用
@article{arxiv.2412.15845,
title = {Multi-dimensional Visual Prompt Enhanced Image Restoration via Mamba-Transformer Aggregation},
author = {Aiwen Jiang and Hourong Chen and Zhiwen Chen and Jihua Ye and Mingwen Wang},
journal= {arXiv preprint arXiv:2412.15845},
year = {2024}
}