ControlMambaIR: Conditional Controls with State-Space Model for Image Restoration

TL;DR

ControlMambaIR introduces a new image restoration approach combining the Mamba network with diffusion models, achieving superior perceptual quality across tasks like deraining, deblurring, and denoising, with robust performance on multiple datasets.

Contribution

The paper presents a novel integration of Mamba architecture with diffusion models for enhanced conditional control in image restoration tasks.

Findings

Outperforms existing methods in perceptual metrics like LPIPS and FID.

Maintains comparable performance in PSNR and SSIM metrics.

Mamba architecture outperforms CNN and Attention-based controls in diffusion models.

Abstract

This paper proposes ControlMambaIR, a novel image restoration method designed to address perceptual challenges in image deraining, deblurring, and denoising tasks. By integrating the Mamba network architecture with the diffusion model, the condition network achieves refined conditional control, thereby enhancing the control and optimization of the image generation process. To evaluate the robustness and generalization capability of our method across various image degradation conditions, extensive experiments were conducted on several benchmark datasets, including Rain100H, Rain100L, GoPro, and SSID. The results demonstrate that our proposed approach consistently surpasses existing methods in perceptual quality metrics, such as LPIPS and FID, while maintaining comparable performance in image distortion metrics, including PSNR and SSIM, highlighting its effectiveness and adaptability. Notably, ablation experiments reveal that directly noise prediction in the diffusion process achieves better performance, effectively balancing noise suppression and detail preservation. Furthermore, the findings indicate that the Mamba architecture is particularly well-suited as a conditional control network for diffusion models, outperforming both CNN- and Attention-based approaches in this context. Overall, these results highlight the flexibility and effectiveness of ControlMambaIR in addressing a range of image restoration perceptual challenges.