MambaIRv2: Attentive State Space Restoration (2411.15269v2)

Abstract: The Mamba-based image restoration backbones have recently demonstrated significant potential in balancing global reception and computational efficiency. However, the inherent causal modeling limitation of Mamba, where each token depends solely on its predecessors in the scanned sequence, restricts the full utilization of pixels across the image and thus presents new challenges in image restoration. In this work, we propose MambaIRv2, which equips Mamba with the non-causal modeling ability similar to ViTs to reach the attentive state space restoration model. Specifically, the proposed attentive state-space equation allows to attend beyond the scanned sequence and facilitate image unfolding with just one single scan. Moreover, we further introduce a semantic-guided neighboring mechanism to encourage interaction between distant but similar pixels. Extensive experiments show our MambaIRv2 outperforms SRFormer by even 0.35dB PSNR for lightweight SR even with 9.3\% less parameters and suppresses HAT on classic SR by up to 0.29dB. Code is available at https://github.com/csguoh/MambaIR.

• The paper shows that integrating non-causal modeling via an Attentive State-space Equation significantly enhances global image restoration performance.
• It employs Semantic Guided Neighboring to reorganize spatial neighborhoods into semantically unified 1D sequences for effective long-range interactions.
• Prompt-based learning transforms MambaIRv2 into a globally aware architecture, achieving a 0.35 dB PSNR boost on Urban100 with fewer parameters.

Overview of MambaIRv2: Attentive State Space Restoration

The paper under discussion introduces MambaIRv2, a state-of-the-art approach in the domain of image restoration that addresses the limitations inherent in Mamba-based architectures. Mamba, a selective state-space model, has shown potential in various image restoration tasks, yet it is restricted by causal modeling that hinders global image restoration capabilities. This research posits that integrating non-causal modeling capabilities into Mamba can effectively enhance its performance.

Key Innovations

1. Attentive State-space Equation (ASE): Central to the innovation is the re-engineering of the state-space equation to enable non-causal modeling similar to Vision Transformers (ViTs). The proposed ASE modifies the matrix traditionally analogous to an attention mechanism's query component, allowing it to query and incorporate information from the complete image rather than a single scanned sequence.
2. Semantic Guided Neighboring (SGN): This mechanism manages the potential degradation of long-range dependencies by restructuring conventional spatial neighborhoods into semantically unified 1D sequences. By repositioning semantically similar pixels to be proximate in the 1D space, SGN enables effective interactions even over long distances in the original spatial configuration.
3. Prompt-based Learning: ASE utilizes a pool of prompts, effectively combining additional input learned representations into the state-space model, akin to a query function in attention mechanisms. This prompt integration effectively transforms Mamba's awareness from a linear, sequence-based model to a globally aware nonlinear model, mitigating the constraints of multiple scanning sequences.

Methodology

The methodology intricately combines these components to form an overarching framework geared at maximizing both computational efficiency and restoration fidelity. MambaIRv2 employs a combination of Attentive State-space Modules (ASSMs) and window-based multi-head self-attention (MHSA) layers to craft a network that affords both local and global context learning. Through an iterative learning process that trains on extensive datasets, MambaIRv2 is optimized for diverse image restoration tasks beyond just super-resolution—specifically, JPEG artifact reduction and image denoising are also addressed effectively.

Experimental Evaluation

The results from the experimental evaluation are impressive and speak to the potency of MambaIRv2 across several standard datasets, including Set5, Set14, BSDS100, Urban100, and Manga109. For instance, the model demonstrates an increase of 0.35dB in PSNR on the Urban100 dataset against the previous best method, despite using fewer parameters, underscoring its efficiency. The weak performance in past architectures arose from excessive parameter demands due to multiple sequential scans. MambaIRv2 negates this issue through its single-scan approach aided by attentive modeling and prompt learning.

Implications and Future Directions

MambaIRv2's outcomes pose substantial implications for both practical and theoretical strides in image restoration. Practically, the ability to process images with fewer computational resources while maintaining high fidelity holds promise in real-world applications such as medical imaging and satellite image processing. Theoretically, bridging state-space models with attention mechanisms marks an important area for exploration, especially in reducing the overhead traditionally associated with models like ViTs without compromising performance.

The future landscape could see further exploration into multi-modal restorations beyond single-image tasks, broadening the applied scope of the attentive state-space framework. Moreover, exploring interactions with other architectures and further reducing computational costs could yield net gains across the spectrum of neural-based frameworks in this space.

Conclusion

MambaIRv2 emerges as a substantial iteration over previous work, capitalizing on the integration of non-causal modeling within a state-space framework. Its architecture embodies a fine balance between leveraging the inherent advantages of attention mechanisms while maintaining the computational agility of Mamba. As a research contribution, it lays the groundwork for subsequent advances in efficient and effective image restoration techniques.