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

TL;DR

CU-Mamba is a novel image restoration model that integrates spatial and channel state space models into a U-Net architecture, achieving superior performance with efficient global context encoding.

Contribution

It introduces a dual State Space Model framework within U-Net, combining spatial and channel SSMs for improved global context modeling in image restoration.

Findings

Outperforms existing state-of-the-art methods in image restoration tasks.

Efficiently encodes global context with linear complexity.

Effectively preserves channel correlation features.

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.