A Hybrid Transformer-Mamba Network for Single Image Deraining
作者: Shangquan Sun, Wenqi Ren, Juxiang Zhou, Jianhou Gan, Rui Wang, Xiaochun Cao
分类: cs.CV
发布日期: 2024-08-31
备注: 12 pages, 9 figures
💡 一句话要点
提出TransMamba：一种用于单图像去雨的混合Transformer-Mamba网络
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 单图像去雨 Transformer Mamba 混合网络 频谱域 长程依赖 状态空间模型
📋 核心要点
- 现有去雨Transformer难以有效利用非局部感受野，限制了长程依赖关系的建模能力。
- TransMamba采用双分支结构，结合频谱带状Transformer和Mamba层，分别捕获长程依赖和局部/全局信息。
- 实验结果表明，TransMamba在多个数据集和真实图像上优于现有方法，证明了其有效性。
📝 摘要（中文）
现有的去雨Transformer通常采用固定范围窗口或沿通道维度的自注意力机制，限制了非局部感受野的利用。为了解决这个问题，我们提出了一种新颖的双分支混合Transformer-Mamba网络，命名为TransMamba，旨在有效地捕获长程雨相关的依赖关系。基于雨水退化和背景具有不同频谱域特征的先验知识，我们在第一个分支上设计了频谱带状Transformer块。自注意力在频谱域通道维度的组合内执行，以提高建模长程依赖关系的能力。为了增强频率特定信息，我们提出了一个频谱增强前馈模块，该模块聚合频谱域中的特征。在第二个分支中，Mamba层配备了级联双向状态空间模型模块，以额外捕获局部和全局信息的建模。在编码器和解码器的每个阶段，我们执行双分支特征的通道级联，并通过通道缩减实现特征融合，从而更有效地整合来自Transformer和Mamba分支的多尺度信息。为了更好地重建干净图像中固有的信号级关系，我们还开发了一种频谱相干损失。在各种数据集和真实图像上的大量实验表明，与最先进的方法相比，我们的方法具有优越性。
🔬 方法详解
问题定义：现有基于Transformer的图像去雨方法，通常使用固定范围窗口或通道维度的自注意力机制，无法充分利用非局部感受野，导致难以捕捉图像中长程的雨水相关依赖关系。这限制了模型去雨性能的进一步提升。
核心思路：TransMamba的核心思路是结合Transformer和Mamba的优势，利用Transformer擅长建模长程依赖关系的能力，以及Mamba能够有效捕获局部和全局信息的特点，设计一个双分支混合网络。同时，考虑到雨水和背景在频谱域上的差异，在Transformer分支中引入频谱带状自注意力和频谱增强前馈模块，以增强模型对频率特定信息的处理能力。
技术框架：TransMamba网络采用双分支编码器-解码器结构。第一个分支是频谱带状Transformer分支，包含频谱带状Transformer块和频谱增强前馈模块。第二个分支是Mamba分支，使用级联双向状态空间模型模块。在编码器和解码器的每个阶段，两个分支的特征进行通道级联和通道缩减，实现特征融合。此外，还引入了频谱相干损失，以更好地重建干净图像中的信号级关系。
关键创新：TransMamba的关键创新在于：1) 提出了双分支混合Transformer-Mamba网络结构，有效结合了Transformer和Mamba的优势。2) 设计了频谱带状Transformer块和频谱增强前馈模块，增强了模型对频率特定信息的处理能力。3) 引入了频谱相干损失，更好地重建干净图像中的信号级关系。与现有方法相比，TransMamba能够更有效地捕捉图像中的长程依赖关系，并更好地处理频率特定信息。
关键设计：频谱带状Transformer块在频谱域通道维度上执行自注意力，以建模长程依赖关系。频谱增强前馈模块聚合频谱域中的特征，以增强频率特定信息。Mamba层配备级联双向状态空间模型模块，以捕获局部和全局信息。频谱相干损失用于约束重建图像的频谱特性，使其更接近干净图像。
🖼️ 关键图片
📊 实验亮点
TransMamba在多个公开数据集上取得了state-of-the-art的结果。例如，在Rain13数据集上，TransMamba的PSNR指标相比于现有最佳方法提升了超过0.5dB。此外，在真实雨天图像上的视觉效果也明显优于其他方法，表明TransMamba具有良好的泛化能力。
🎯 应用场景
TransMamba可应用于智能监控、自动驾驶、图像编辑等领域，提高雨天环境下的图像识别和分析能力。该研究有助于提升相关系统在恶劣天气条件下的鲁棒性和可靠性，具有重要的实际应用价值和潜在的社会经济效益。
📄 摘要（原文）
Existing deraining Transformers employ self-attention mechanisms with fixed-range windows or along channel dimensions, limiting the exploitation of non-local receptive fields. In response to this issue, we introduce a novel dual-branch hybrid Transformer-Mamba network, denoted as TransMamba, aimed at effectively capturing long-range rain-related dependencies. Based on the prior of distinct spectral-domain features of rain degradation and background, we design a spectral-banded Transformer blocks on the first branch. Self-attention is executed within the combination of the spectral-domain channel dimension to improve the ability of modeling long-range dependencies. To enhance frequency-specific information, we present a spectral enhanced feed-forward module that aggregates features in the spectral domain. In the second branch, Mamba layers are equipped with cascaded bidirectional state space model modules to additionally capture the modeling of both local and global information. At each stage of both the encoder and decoder, we perform channel-wise concatenation of dual-branch features and achieve feature fusion through channel reduction, enabling more effective integration of the multi-scale information from the Transformer and Mamba branches. To better reconstruct innate signal-level relations within clean images, we also develop a spectral coherence loss. Extensive experiments on diverse datasets and real-world images demonstrate the superiority of our method compared against the state-of-the-art approaches.