CU-Mamba: Selective State Space Models with Channel Learning for Image Restoration

Rui Deng,Tianpei Gu

2024-04-18

Abstract:Reconstructing degraded images is a critical task in image processing. Although CNN and Transformer-based models are prevalent in this field, they exhibit inherent limitations, such as inadequate long-range dependency modeling and high computational costs. To overcome these issues, we introduce the Channel-Aware U-Shaped Mamba (CU-Mamba) model, which incorporates a dual State Space Model (SSM) framework into the U-Net architecture. CU-Mamba employs a Spatial SSM module for global context encoding and a Channel SSM component to preserve channel correlation features, both in linear computational complexity relative to the feature map size. Extensive experimental results validate CU-Mamba's superiority over existing state-of-the-art methods, underscoring the importance of integrating both spatial and channel contexts in image restoration.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address issues in the field of image restoration, particularly the limitations of existing methods such as Convolutional Neural Networks (CNNs) and Transformer-based models in handling long-range dependencies and computational costs. Specifically, the paper proposes a new model named Channel-Aware U-Shaped Mamba (CU-Mamba), which combines a bidirectional Selective State Space Model (SSM) to achieve global context encoding and channel feature preservation with linear computational complexity.

The main contributions of the paper include:
1. Proposing the CU-Mamba model, which introduces spatial SSM modules and channel SSM modules into the traditional U-Net architecture to enhance global context and channel features in image restoration tasks.
2. Validating the effectiveness of the spatial and channel SSM modules through detailed ablation experiments.
3. Experimental results show that the CU-Mamba model achieves superior performance on various image restoration datasets, surpassing existing state-of-the-art methods while maintaining lower computational costs.

Overall, the CU-Mamba model addresses the limitations of existing CNN and Transformer models in image restoration tasks, particularly in capturing long-range dependencies and handling channel information.