Source-linked AI summary
VmambaIR: Visual State Space Model for Image Restoration
Yuan Shi, Bin Xia, Xiaoyu Jin, Xing Wang, Tianyu Zhao, Xin Xia, Xuefeng Xiao, Wenming Yang
TL;DR
Image restoration requires models that handle long-range image dependencies without the computational burdens of diffusion models or quadratic-complexity transformers. VmambaIR addresses this with a Unet-based state space architecture using OSS blocks and six-direction image scanning, achieving strong results across multiple restoration tasks with reduced resource requirements.
Problem
CNNs have limited long-range modeling, diffusion models require costly priors and denoising, and transformers face quadratic complexity as image size grows.
Method
VmambaIR combines a Unet architecture with OSS blocks containing Omni Selective Scan and EFFN modules to model multi-scale image information efficiently.
Results
Across image deraining and super-resolution tasks, VmambaIR surpasses baseline accuracy with fewer computational resources; real-world super-resolution uses 26% of the computational cost while achieving higher reconstruction accuracy.
Takeaways & Limitations
The results highlight state space models as potential alternatives to transformer and CNN architectures for image restoration.
Abstract
from arXiv · showhide
Image restoration is a critical task in low-level computer vision, aiming to restore high-quality images from degraded inputs. Various models, such as convolutional neural networks (CNNs), generative adversarial networks (GANs), transformers, and diffusion models (DMs), have been employed to address this problem with significant impact. However, CNNs have limitations in capturing long-range dependencies. DMs require large prior models and computationally intensive denoising steps. Transformers have powerful modeling capabilities but face challenges due to quadratic complexity with input image size. To address these challenges, we propose VmambaIR, which introduces State Space Models (SSMs) with linear complexity into comprehensive image restoration tasks. We utilize a Unet architecture to stack our proposed Omni Selective Scan (OSS) blocks, consisting of an OSS module and an Efficient Feed-Forward Network (EFFN). Our proposed omni selective scan mechanism overcomes the unidirectional modeling limitation of SSMs by efficiently modeling image information flows in all six directions. Furthermore, we conducted a comprehensive evaluation of our VmambaIR across multiple image restoration tasks, including image deraining, single image super-resolution, and real-world image super-resolution. Extensive experimental results demonstrate that our proposed VmambaIR achieves state-of-the-art (SOTA) performance with much fewer computational resources and parameters. Our research highlights the potential of state space models as promising alternatives to the transformer and CNN architectures in serving as foundational frameworks for next-generation low-level visual tasks.
1 Introduction
VmambaIR applies state space modeling to comprehensive image restoration, using an Unet with OSS blocks to address visual-data modeling challenges. Across deraining and super-resolution tasks, it reports higher accuracy with lower computational requirements.
- Image restoration reconstructs high-quality images from low-quality inputs across tasks including deblurring, super-resolution, and deraining.
- Existing approaches face distinct limitations: CNNs struggle with long-range dependencies, diffusion models require costly priors and denoising, and transformers scale quadratically with image size.
- VmambaIR uses a Unet architecture with OSS blocks containing an OSS module and an EFFN to model multi-scale image information efficiently.
- Omni Selective Scan models image information flow in six directions, addressing the unidirectional modeling limitation of Mamba-based state space models.
- Across image deraining, single-image super-resolution, and real-world super-resolution, VmambaIR surpasses baseline accuracy while using fewer computational resources.
- 26% computational cost accompanies higher reconstruction accuracy for VmambaIR in real-world super-resolution compared with the existing SOTA method.
2 Related work
Related work progresses from CNN-based restoration through transformers and diffusion models to state space models. Transformers capture global dependencies but incur quadratic complexity, while SSMs and Mamba offer linear scalability for long sequences.
- CNN-based methods became prominent in image restoration after early compact networks demonstrated strong performance, especially for denoising.
- Transformers capture global dependencies and complex relationships, but self-attention complexity grows quadratically with input image size.
- Diffusion models use powerful data fitting and pretrained prior knowledge to generate visually appealing restored images, but resource demands and fidelity limitations hinder adoption.
- Structured state space models were developed to address long-range dependencies while reducing the computational complexity associated with earlier LSSL models.
- Mamba introduces a data-dependent SSM layer with linear scalability in sequence length and computational efficiency relevant to large-scale image processing.
3 Preliminaries: State Space Models
State space models map input signals to output responses through continuous linear systems and can be discretized for machine learning. Their parameters represent hidden-state dynamics, input/output mappings, and a skip connection.
- SSMs are linear time-invariant systems that map an input stimulation x(t) ∈ R^L to an output response y(t) ∈ R^L.
- The continuous SSM formulation uses linear ordinary differential equations to describe the system dynamics and output response.
- The hidden state h(t) has size N, while A, B, and C parameterize the state dynamics and input/output mappings; D represents a skip connection.
- Discretization converts the continuous-time ODE into a discrete function aligned with the sample rate of the input signal.
- The zeroth-order hold rule uses a timescale parameter Δ to convert continuous parameters A and B into discrete parameters.
- Omni Mamba adds Mamba’s selection mechanism to capture long-sequence characteristics with computational efficiency and accuracy.
4 Method
VmambaIR uses a multi-scale Unet built from OSS blocks to restore images while modeling information flow across image dimensions. Its Omni Selective Scan addresses Mamba’s unidirectional limitation by scanning image features from six directions with linear complexity.
- Model Architecture: VmambaIR uses a multi-scale Unet architecture with OSS blocks to extract, reconstruct, and refine image features at different scales.Encoder and decoder features are connected through skip connections before iterative OSS refinement.
- Model Architecture: The tail block uses convolution and pixel shuffle for super-resolution, while same-resolution restoration tasks use vanilla convolution to produce the residual output.This task-dependent design changes the image size only when upsampling is required.
- OSS Block: Each OSS block combines an Omni Selective Scan module with an Efficient Feed-Forward Network for feature-flow modeling and regulation.The OSS module processes multiple information streams, while EFFN uses convolution, depth-wise convolution, and gating after normalization.
- Omni Selective Scan Mechanism: OSS scans feature information bidirectionally across spatial dimensions and incorporates channel-dimension scanning to address Mamba’s unidirectional and incomplete multidimensional modeling.The mechanism applies forward and backward scans on the image plane and scans feature channels from front to back.
- Omni Selective Scan Mechanism: Omni Selective Scan models image features from six directions while maintaining linear computational complexity compared with self-attention.The six-direction design extends Mamba’s long-range modeling to three-dimensional image features without the quadratic complexity of self-attention.
5 Experiments and Analysis
VmambaIR is evaluated on single-image super-resolution, real-world super-resolution, and deraining, achieving strong restoration quality with reduced computational cost and parameter count. Ablations show that omni-directional scanning and EFFN contribute to accuracy and efficiency.
- Experimental Setup: VmambaIR is evaluated on single-image super-resolution, real-world image super-resolution, and image deraining.The experiments use task-specific encoder-decoder configurations and training objectives.
- Single Image Super-Resolution: 0.48 dB PSNR improvement over BebyGAN is achieved on Urban100 while maintaining a lower LPIPS score.VmambaIR also outperforms existing methods on PSNR and LPIPS across all test datasets.
- Real-World Image Super-Resolution: VmambaIR outperforms existing methods across all reported real-world super-resolution quality metrics, achieving lower LPIPS and higher PSNR and SSIM.Evaluation uses benchmark datasets and LPIPS, SSIM, and PSNR.
- Real-World Image Super-Resolution: 26% of the computational resources and the lowest parameter count are reported for VmambaIR compared with existing state-of-the-art methods in real-world super-resolution.The comparison does not rely on distillation or pruning.
- Image Deraining: Over 0.1 dB PSNR improvement is achieved on Rain100H, Rain100L, and Test1200 while using fewer parameters and lower computational complexity.VmambaIR also reports higher SSIM scores across all deraining datasets.
- Ablation Analysis: Removing bidirectional channel scanning decreases accuracy by approximately 0.14 dB without materially changing computational complexity.Replacing omni selective scan with one-way scanning reduces complexity by approximately 7% but decreases accuracy by approximately 0.43 dB.
- Ablation Analysis: EFFN is used to improve computational accuracy and efficiency because selective scan operations can be slower than vanilla convolution at comparable complexity.The ablation discussion attributes this issue to data-type and dimension conversions in selective scan operations.
- Ablation Analysis: Linear layers and reshape operations increase computational speed by approximately 8.6% while maintaining comparable complexity, parameter count, and accuracy.This comparison concerns detailed design choices within the OSS module.
6 Conclusion
The paper presents VmambaIR as a simple state-space-model image restoration network combining UNet multi-scale processing with OSS and EFFN modules. Across deraining and super-resolution tasks, it reports state-of-the-art performance without elaborate auxiliary techniques.
- Conclusion: VmambaIR leverages mamba’s linear complexity and high-frequency modeling within a UNet architecture for image restoration.The network is designed to model images at different scales.
- Conclusion: The OSS block combines an OSS module for comprehensive image-feature modeling with EFFN for mapping and modulating information flow.The paper describes these components as improving network accuracy and efficiency.
- Conclusion: Extensive experiments report state-of-the-art performance on image deraining, image super-resolution, and real-world image super-resolution.The model is intentionally presented without distillation, teacher networks, or hybrid network structures.
7 Appendix
The appendix describes training configurations and additional qualitative comparisons for the three restoration tasks. These details cover degradation synthesis, optimization, losses, datasets, and visual comparisons with prior methods.
- Training Configuration: Real-world super-resolution data uses a dual degradation process with resizing, Gaussian noise, gray noise, and blur, following Real-ESRGAN settings.The configuration is intended to maintain comparable results.
- Training Configuration: Super-resolution training uses GAN, perceptual, and L1 losses with equal weights and a U-Net discriminator.The same optimizer settings are applied to the generator and discriminator in the described configuration.
- Training Configuration: Single-image super-resolution training uses bicubic downsampling and the 3450-image DF2K dataset.The low-quality image size is 64 × 64 for 4× super-resolution.
- Training Configuration: Deraining uses blocks across network layers, two refinement blocks, L1 loss, and AdamW optimization.Progressive training uses input sizes from 128 to 384 and eight V100 GPUs for approximately six days.
- Qualitative Results: Additional visual comparisons include VmambaIR against prior state-of-the-art methods for real-world and single-image super-resolution.The cited comparisons use Figures 8 and 9 and emphasize fidelity and detail.
- Qualitative Results: Deraining comparisons report near-perfect restoration across diverse images relative to Restormer, RESCAN, and DerainNet.The appendix presents these results as visual advantages over previous methods.