VmambaIR: Visual State Space Model for Image Restoration (2403.11423v1)

Abstract: 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.

VmambaIR: A Visual State Space Model for Image Restoration

The paper "VmambaIR: Visual State Space Model for Image Restoration" introduces an innovative approach to image restoration leveraging state space models (SSMs). The authors present a detailed exposition of the limitations of current deep learning methodologies in image restoration, such as convolutional neural networks (CNNs), generative adversarial networks (GANs), vision transformers, and diffusion models. They highlight that these models face distinct challenges like difficulty in capturing long-range dependencies, high computational burdens, and complexities related to the input sequence sizes.

Method and Contributions

The core proposition of the paper is the VmambaIR model which utilizes the Mamba block—a novel variant of SSM that is adept at capturing high-frequency information in data sequences. VmambaIR stands out due to its incorporation of the Omni Selective Scan (OSS) mechanism designed to overcome the unidirectional data flow limitation inherent in traditional SSMs. VmambaIR boasts a linear complexity feature which increases its efficiency compared to transformers with quadratic input sequence complexity.

The architecture of VmambaIR is built upon a Unet-like structure, utilizing OSS blocks which include an OSS module and an Efficient Feed-Forward Network (EFFN). Key innovations presented include:

1. Omni Selective Scan (OSS): This mechanism allows for the comprehensive modeling of information in multidimensional spaces, expanding directional scanning beyond the conventional approaches seen in earlier SSM methodologies. It employs bidirectional scanning across the spatial and channel dimensions of image data, effectively capturing complex image patterns with minimal computational overhead.
2. Efficient Feed-Forward Network (EFFN): The EFFN aids in effectively modulating the information flow, ensuring maintained high performance while leveraging linear structures for reduced computational cost.

Experimental Evaluation

VmambaIR’s effectiveness was validated across various image restoration tasks such as single image super-resolution, real-world image super-resolution, and image deraining. The experimental results indicate competitive performance, with VmambaIR achieving state-of-the-art accuracy while requiring notably fewer computational resources and parameters. For instance, in real-world image super-resolution tasks, the method demonstrated superior performance with merely 26% of the computational expense incurred by existing leading methods.

Quantitative comparisons using metrics like LPIPS, PSNR, and SSIM across multiple datasets underscore VmambaIR's capabilities in maintaining high image fidelity and perceptual quality. These experiments manifest the strong potential of SSMs in practical image restoration scenarios, correcting issues traditional models fail at overcoming efficiently.

Implications and Future Directions

The research underscores the viability of state space models as alternatives to CNNs and transformers, particularly in image restoration applications where efficiency and high fidelity are paramount. With the prevalence and criticality of image restoration in various applications like photography, remote sensing, and medical imaging, models such as VmambaIR present significant practical implications for reducing resource consumption while retaining or improving output quality.

From a theoretical perspective, successful integration of SSMs into visual tasks expands the perceptions of their applicability beyond time-series data—where they traditionally thrive—into domains necessitating spatial reasoning and multidimensional data processing. Future research can build upon these findings, exploring further enhancements in omni-domain scanning techniques and their generalizations to more complex networks and larger datasets.

The introduction of VmambaIR marks a pivotal stride towards unlocking the potential of state space models in domains previously dominated by more conventional deep learning architectures. The contributions encapsulated in this paper illustrate substantial advancements towards more efficient, robust, and versatile solutions in the field of image restoration.