Liu Lei, Ma Luan, Wang Shuai, Wang Jun, Melo Silas N
School of Computer Science and Technology, Huaibei Normal University, Huaibei, 235000, China.
Huaibei Key Laboratory of Digital Multimedia Intelligent Information Processing, Huaibei Normal University, Huaibei, 235000, China.
Sci Rep. 2025 Jul 1;15(1):22265. doi: 10.1038/s41598-025-07032-3.
Image restoration is fundamental in computer vision to restore high-quality images from degraded ones. Recently, models such as the transformer and diffusion have shown notable success in addressing this challenge. However, transformer-based methods face high computational costs due to quadratic complexity, while diffusion-based methods often struggle with suboptimal results due to inaccurate noise estimation. This study proposes Diff-Mamba, a two-stage adaptive Mamba-based diffusion model for image restoration. Diff-Mamba integrates the linear complexity state space model (SSM, also known as Mamba) into image restoration, expanding its applicability to visual data generation. Diff-Mamba mainly consists of two parts: the diffusion state space model (DSSM) and the diffusion feedforward neural network (DFNN). DSSM combines Mamba's high efficiency with the representative power of diffusion models, enhancing both inference and training. DFNN regulates the information flow, enabling each depthwise convolutional layer to focus on the details of image, thus learning more effective local structures for image restoration. The study's findings, verified through extensive experiments, indicate that Diff-Mamba outperforms both diffusion-based and transformer-based methods in image deraining, denoising, and deblurring, demonstrating competitive restoration performance with various commonly used datasets. Code is available at https://github.com/maluan-ml/Diff-Mamba.
图像恢复是计算机视觉中的基础任务，旨在从退化图像中恢复高质量图像。最近，诸如Transformer和扩散模型等在应对这一挑战方面取得了显著成功。然而，基于Transformer的方法由于二次复杂度而面临高计算成本，而基于扩散的方法则常常因噪声估计不准确而难以获得最优结果。本研究提出了Diff-Mamba，一种基于Mamba的两阶段自适应扩散模型用于图像恢复。Diff-Mamba将线性复杂度状态空间模型（SSM，也称为Mamba）集成到图像恢复中，扩展了其在视觉数据生成中的适用性。Diff-Mamba主要由两部分组成：扩散状态空间模型（DSSM）和扩散前馈神经网络（DFNN）。DSSM将Mamba的高效性与扩散模型的代表性能力相结合，增强了推理和训练效果。DFNN调节信息流，使每个深度卷积层能够专注于图像细节，从而学习更有效的局部结构用于图像恢复。通过大量实验验证的研究结果表明，Diff-Mamba在图像去雨、去噪和去模糊方面优于基于扩散和基于Transformer的方法，在各种常用数据集上展现出具有竞争力的恢复性能。代码可在https://github.com/maluan-ml/Diff-Mamba获取。
