Source-linked AI summary
MambaIR: A Simple Baseline for Image Restoration with State-Space Model
Hang Guo, Jinmin Li, Tao Dai, Zhihao Ouyang, Xudong Ren, Shu-Tao Xia
TL;DR
Image restoration backbones face a trade-off between global receptive fields and efficient computation, while standard Mamba suffers local pixel forgetting and channel redundancy in low-level vision. MambaIR adapts Mamba with local enhancement and channel attention, and is reported to outperform other strong baselines across multiple restoration tasks while combining global receptive fields with linear complexity.
Problem
Image restoration backbones face a trade-off between global receptive fields and efficient computation, while standard Mamba suffers local pixel forgetting and channel redundancy in low-level vision.
Method
MambaIR adapts Mamba with stacked Residual State Space Blocks that use local convolution, channel attention, and learnable residual scaling.
Results
MambaIR outperforms other strong baselines across multiple image restoration tasks while providing a global effective receptive field with linear computational complexity.
Takeaways & Limitations
MambaIR provides a simple state-space alternative for image restoration backbones that combines long-range modeling with restoration-specific local and channel mechanisms.
Abstract
from arXiv · showhide
Recent years have seen significant advancements in image restoration, largely attributed to the development of modern deep neural networks, such as CNNs and Transformers. However, existing restoration backbones often face the dilemma between global receptive fields and efficient computation, hindering their application in practice. Recently, the Selective Structured State Space Model, especially the improved version Mamba, has shown great potential for long-range dependency modeling with linear complexity, which offers a way to resolve the above dilemma. However, the standard Mamba still faces certain challenges in low-level vision such as local pixel forgetting and channel redundancy. In this work, we introduce a simple but effective baseline, named MambaIR, which introduces both local enhancement and channel attention to improve the vanilla Mamba. In this way, our MambaIR takes advantage of the local pixel similarity and reduces the channel redundancy. Extensive experiments demonstrate the superiority of our method, for example, MambaIR outperforms SwinIR by up to 0.45dB on image SR, using similar computational cost but with a global receptive field. Code is available at \url{https://github.com/csguoh/MambaIR}.
1 Introduction
Image restoration backbones face a trade-off between global receptive fields and efficient computation. MambaIR adapts Mamba with local enhancement and channel attention, achieving global receptive fields with linear complexity for restoration.
- Motivation: Large receptive fields help restoration networks use wider contextual information and capture higher-level image structures.The passage connects larger receptive fields with reconstruction and structure preservation, and reports that Transformer-based methods outperform CNN-based methods experimentally.
- Motivation: Current CNN backbones offer efficient convolution but limited receptive fields, whereas standard Transformers provide global receptive fields at quadratic computational complexity.This creates an inherent choice dilemma between global context and efficient computation.
- Motivation: Mamba offers long-range dependency modeling with efficient parallel processing, suggesting a route toward balancing global receptive fields and computational efficiency.Its recursive state-space formulation models very long-range dependencies, while parallel scan processes tokens efficiently on GPUs.
- Challenges: Standard Mamba can cause local pixel forgetting on flattened images and channel redundancy because it must memorize long sequence dependencies with many hidden states.These issues make vanilla Mamba poorly suited to low-level image restoration without task-specific adaptation.
- Method: MambaIR uses shallow feature extraction, stacked Residual State Space Blocks, and high-quality reconstruction to adapt Mamba for image restoration.RSSBs combine local convolution, channel attention, and a learnable residual scale to address restoration-specific challenges.
- Contribution: MambaIR is presented as a simple state-space restoration backbone that combines a global effective receptive field with linear computational complexity.The introduction describes it as an alternative to CNN- and Transformer-based restoration methods and reports strong results across multiple tasks.
2 Related Work
Image restoration research progressed from CNN-based methods toward Transformers and state-space models. These model families differ in how they represent dependencies, computational scaling, and application scope across vision tasks.
- CNNs and Transformers: Early image restoration methods used CNNs for super-resolution, denoising, and JPEG compression artifact reduction, with residual and dense connections improving representation ability.CNN-based restoration methods nevertheless face challenges in effectively modeling global dependencies.
- CNNs and Transformers: Transformers became promising for image restoration after demonstrating effectiveness across several broader machine-learning tasks.The supplied related-work passage introduces Transformer-based restoration as a response to CNN limitations in global dependency modeling.
- State Space Models: State Space Models originated in control theory and entered deep learning as backbones for state-space transformation and long-range dependency modeling.Their linearly scaling sequence processing attracted interest as an alternative to attention-based architectures.
- State Space Models: The SSM literature progressed from S4 and S5 through H3 and gated state-space layers to Mamba, which uses selective mechanisms and efficient hardware design.The passage also notes applications of Mamba and related state-space models in classification, video understanding, biomedical segmentation, and other vision tasks.
3 Methodology
MambaIR adapts state-space modeling to image restoration by combining long-range dependency modeling with local enhancement and channel attention. Its architecture uses stacked residual state-space blocks, a vision state-space module, and 2D selective scanning.
- State-Space Preliminaries: Mamba uses structured state-space equations and parallel scanning to model ultra-long dependencies efficiently.Its input-dependent parameters enable dynamic feature representation, while parallel scan processes tokens in parallel.
- Overall Architecture: MambaIR processes images through shallow feature extraction, stacked residual state-space groups, and high-quality reconstruction.A 3 × 3 convolution produces shallow features, while group outputs are refined and combined with the shallow features before reconstruction.
- Residual State-Space Block: The Residual State-Space Block applies LayerNorm and VSSM for spatial long-term dependencies, with a learnable scale controlling the skip connection.RSSB is introduced as a restoration-specific block rather than a direct replacement of Transformer attention with SSM.
- Residual State-Space Block: Local convolution after VSSM preserves neighborhood information that flattening can separate in the 1D sequence.The four-direction unfolding strategy provides only four nearest neighbors, so spatially close pixels may become distant in the flattened sequence.
- Residual State-Space Block: Channel attention selects critical channels after SSM processing to reduce redundancy caused by the large hidden-state capacity.The authors motivate channel attention by observing that many channels are inactive without it.
- Vision State-Space Module: VSSM uses two parallel expanded-channel branches, combines them with a Hadamard product, and projects the result back to the input channel count.One branch includes depth-wise convolution, SiLU, 2D-SSM, and LayerNorm; the other uses linear projection and SiLU.
4 Experiences
Across super-resolution and denoising experiments, MambaIR generally achieves strong restoration quality while retaining efficient scaling and a global receptive field. Ablations attribute these results to restoration-specific local interaction, channel attention, and multi-directional scanning.
- Ablation Study: MambaIR improves Mamba by combining local convolution and channel attention within the Residual State Space Block.The ablation identifies local convolution as improving local interaction and channel attention as part of the effective RSSB design.
- Ablation Study: Four-direction scanning improves image processing compared with one- and two-direction alternatives.The scan-mode ablation evaluates one-, two-, and four-direction unfolding strategies for flattened image features.
- Classic Image Super-Resolution: 0.41dB on Manga109 at ×2 is MambaIR’s reported advantage over SwinIR in classic image super-resolution.MambaIR achieves the best performance on almost all five benchmark datasets across scale factors and reconstructs sharp edges and natural textures qualitatively.
- Model Complexity Comparison: Linear complexity with input resolution makes MambaIR more efficient than the full-attention baseline while retaining a global receptive field.The comparison scales input resolution from 48 × 48 to 84 × 84 and positions MambaIR’s scaling alongside efficient attention methods such as SwinIR.
- Lightweight Image Super-Resolution: 0.34dB PSNR is MambaIR-light’s reported improvement over SwinIR-light on Manga109 at ×4 with similar parameters and MACs.The lightweight model is presented as evidence of scalability and efficiency.
- Image Denoising: 0.48dB over SwinIR at σ=50 on Urban100 is reported for Gaussian color image denoising.MambaIR achieves the best performance on most denoising datasets, while real-image denoising is comparable to Restormer and exceeds Uformer by 0.12dB PSNR on SIDD.
5 Conclusion
The paper presents MambaIR as a simple state-space baseline for image restoration that addresses the trade-off between efficient computation and a global effective receptive field. It uses local enhancement and channel attention to mitigate neighborhood pixel forgetting and channel redundancy, with experiments across multiple restoration tasks supporting its effectiveness.
- Conclusion: MambaIR combines local enhancement with channel attention to address neighborhood pixel forgetting and channel redundancy in Mamba-based image restoration.The method adapts Mamba for restoration while targeting these two stated limitations.