M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts

Shiwei Wang; Yongzhen Wang; Bingwen Hu; Liyan Zhang; Xiao-Ping Zhang; Mingqiang Wei

cs.CVMar 16, 2026

M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts

Authors: Shiwei Wang, Yongzhen Wang, Bingwen Hu, Liyan Zhang, Xiao-Ping Zhang, Mingqiang Wei

Organizations: School of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics, Nanjing 210016, China · School of Computer Science and Technology, Anhui University of Technology, Ma’anshan 243032, China · Shenzhen Ubiquitous Data Enabling Key Laboratory, Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen 518055, China · College of Artificial Intelligence, Taiyuan University of Technology, Taiyuan 030024, China

Abstract

While Transformer-based architectures have dominated recent advances in all-in-one image restoration, they remain fundamentally reactive: propagating degradations rather than proactively suppressing them. In the absence of explicit suppression mechanisms, degraded signals interfere with feature learning, compelling the decoder to balance artifact removal and detail preservation, thereby increasing model complexity and limiting adaptability. To address these challenges, we propose M2IR, a novel restoration framework that proactively regulates degradation propagation during the encoding stage and efficiently eliminates residual degradations during decoding. Specifically, the Mamba-Style Transformer (MST) block performs pixel-wise selective state modulation to mitigate degradations while preserving structural integrity. In parallel, the Adaptive Degradation Expert Collaboration (ADEC) module utilizes degradation-specific experts guided by a DA-CLIP-driven router and complemented by a shared expert to eliminate residual degradations through targeted and cooperative restoration. By integrating the MST block and ADEC module, M2IR transitions from passive reaction to active degradation control, effectively harnessing learned representations to achieve superior generalization, enhanced adaptability, and refined recovery of fine-grained details across diverse all-in-one image restoration benchmarks. Our source codes are available at https://github.com/Im34v/M2IR.

Figures & tables

Fig. 1: Performance comparison of recent all-in-one restoration methods under different settings. Here, 3D indicates three degradation types encompassing haze, rain, and noise, while 5D extends 3D by adding blur and low-light. The CDD11 dataset contains four fundamental degradation types, low-light, hazy, rainy, and snowy, as well as seven composite degradation types generated by combining these fundamental types.

Fig. 2: Overview of the proposed M 2 IR. (a) Mamba-Style Transformer (MST): Combining the strengths of Mamba and Transformer, it achieves proactive suppression or regulation of degradation information. (b) Adaptive Degradation Expert Collaboration (ADEC): It utilizes DA-CLIP’s image encoder and text encoder to obtain degradation-aware contextual priors, further deriving routing weights, and then aggregates the results from the selected $K$ most relevant experts and a shared expert to eliminate residual degradations effectively.

Fig. 3: Macro architecture illustrations of (a) Transformer [ 17 ] with MDTA and DGFN, (b) Mamba [ 25 ] with SSM, and (c) the proposed Mamba-Style Transformer (MST) featuring Mamba-Style Attention (MSA).

Fig. 4: The architecture of the Degradation-Aware Contextual Priors (DACP). DACP leverages pre-trained DA-CLIP to extract degradation-aware priors, which guide the generation of more accurate routing weights.

Fig. 5: Architecture of the Degradation-Guided Expert Aggregation (DGEA). For each pixel in the feature map, DGEA selects the $K$ most relevant experts based on routing weights and performs degradation-specific processing in collaboration with the shared expert.

	Method	Venue	Dehazing	Deraining	Denoising on BSD68 [ 29 ]			Average
	Method	Venue	on SOTS [ 30 ]	on Rain100L [ 31 ]	$\sigma=15$	$\sigma=25$	$\sigma=50$	Average
General	Restormer [ 17 ]	CVPR’22	27.78 / 0.958	33.78 / 0.958	33.72 / 0.930	30.67 / 0.865	27.63 / 0.792	30.75 / 0.901
	NAFNet [ 32 ]	ECCV’22	24.11 / 0.928	33.64 / 0.956	33.03 / 0.918	30.47 / 0.865	27.12 / 0.754	29.67 / 0.844
	FSNet [ 33 ]	TPAMI’23	29.14 / 0.968	35.61 / 0.969	33.81 / 0.930	30.84 / 0.872	27.69 / 0.792	31.42 / 0.905
	MambaIR [ 34 ]	ECCV’24	29.57 / 0.970	35.42 / 0.969	33.88 / 0.931	30.95 / 0.874	27.74 / 0.793	31.51 / 0.907
All-in-One	AirNet [ 22 ]	CVPR’22	27.94 / 0.962	34.90 / 0.967	33.92 / 0.933	31.26 / 0.888	28.00 / 0.797	31.20 / 0.910

TABLE I: Quantitative Comparison of recent general and all-in-one methods on three degradations. The best values are highlighted in bold.

Fig. 6: Visual comparisons of M 2 IR with state-of-the-art all-in-one methods on three degradations, with magnified views of red-boxed regions provided for better visualization. The task types from top to bottom are dehazing, deraining, and denoising, respectively.

	Method	Venue	Dehazing	Deraining	Denoising	Deblurring	Low-Light	Average
	Method	Venue	on SOTS [ 30 ]	on Rain100L [ 31 ]	on BSD68 [ 29 ]	on GoPro [ 40 ]	on LOL [ 41 ]	Average
General	Restormer [ 17 ]	CVPR’22	24.09 / 0.927	34.81 / 0.962	31.49 / 0.884	27.22 / 0.829	20.41 / 0.806	27.60 / 0.881
	NAFNet [ 32 ]	ECCV’22	25.23 / 0.939	35.56 / 0.967	31.02 / 0.883	26.53 / 0.808	20.49 / 0.809	27.76 / 0.881
	FSNet [ 33 ]	TPAMI’23	25.53 / 0.943	36.07 / 0.968	31.33 / 0.883	28.32 / 0.869	22.29 / 0.829	28.71 / 0.898
	MambaIR [ 34 ]	ECCV’24	25.81 / 0.944	36.55 / 0.971	31.41 / 0.884	28.61 / 0.875	22.49 / 0.832	28.97 / 0.901
All-in-One	AirNet [ 22 ]	CVPR’22	21.04 / 0.884	32.98 / 0.951	30.91 / 0.882	24.35 / 0.781	18.18 / 0.735	25.49 / 0.846

TABLE II: Quantitative Comparison of recent general and all-in-one methods on five degradations. Denoising results are reported for the noise level

\sigma=25

Method

Single

Double

Triple

Average

L+H

L+R

L+S

H+R

H+S

L+H+R

L+H+S

AirNet [ 22 ]

24.83/0.778

24.21/0.951

26.55/0.891

26.79/0.919

23.23/0.779

22.82/0.710

23.29/0.723

22.21/0.868

23.29/0.901

21.80/0.708

22.24/0.725

23.75/0.814

PromptIR [ 12 ]

26.32/0.805

26.10/0.969

31.56/0.946

31.53/0.960

24.49/0.789

25.05/0.771

24.51/0.761

24.54/0.924

23.70/0.925

23.74/0.752

23.33/0.747

25.90/0.850

WeatherDiff [ 43 ]

23.58/0.763

21.99/0.904

24.85/0.885

24.80/0.888

21.83/0.756

22.69/0.730

22.12/0.707

21.25/0.868

21.99/0.868

21.23/0.716

21.04/0.698

22.49/0.799

WGWS-Net [ 44 ]

24.39/0.774

27.90/0.982

33.15/0.964

34.43/0.973

24.27/0.800

25.06/0.772

24.60/0.765

27.23/0.955

27.65/0.960

23.90/0.772

23.97/0.771

26.96/0.863

OneRestore [ 42 ]

26.48/0.826

32.52/0.990

33.40/0.964

34.31/0.973

25.79/0.822

25.58/0.799

25.19/0.789

29.99/0.957

30.21/0.964

24.78/0.788

24.90/0.791

28.47/0.878

TABLE III: Quantitative Comparison of different methods on CDD11 [ 42 ] Dataset. * denotes the model has been retrained.

Fig. 7: Visual comparisons of M 2 IR with state-of-the-art all-in-one methods on CDD11 [ 42 ] dataset, with magnified views of red-boxed regions provided for better visualization. The composite degradation types from top to bottom are L+H, L+R, L+S, H+R, H+S, L+H+R, and L+H+S, respectively.

Method	Venue	Dehazing				Deraining				Desnowing				Average
Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$
AirNet [ 22 ]	CVPR’22	19.27	0.645	0.3829	134.09	31.56	0.912	0.2236	125.54	20.58	0.737	0.2912	138.57	23.80	0.764	0.2992	132.73
TransWeather [ 47 ]	CVPR’22	18.13	0.621	0.3970	123.21	28.59	0.880	0.2638	149.66	24.06	0.754	0.2250	102.99	23.59	0.752	0.2953	125.29
PromptIR [ 12 ]	NeurIPS’23	19.50	0.658	0.3751	113.55	32.51	0.915	0.1980	111.69	26.35	0.804	0.1951	84.12	26.12	0.792	0.2561	103.12
WGWS-Net [ 44 ]	CVPR’23	11.78	0.532	0.5351	152.76	34.77	0.939	0.1168	60.99	19.39	0.721	0.2481	128.56	21.98	0.731	0.3000	114.10
Histoformer [ 48 ]	ECCV’24	15.82	0.597	0.4371	128.34	28.87	0.876	0.2785	152.42	23.88	0.769	0.2252	105.82	22.86	0.747	0.3136	128.86

TABLE IV: Quantitative Comparison of different methods on WeatherBench [ 46 ] Dataset.

Fig. 8: Visual comparisons of M 2 IR with state-of-the-art all-in-one methods on WeatherBench [ 46 ] dataset. The degradation types, from top to bottom, are haze (rows 1 and 2), rain (rows 3 and 4), and snow (rows 5 and 6).

Method BSD68 [ 29 ] Urban100 [ 50 ] Kodak24 [ 51 ] 15 25 50 15 25 50 15 25 50 DnCNN [ 1 ] 33.90 31.24 27.95 32.98 30.81 27.59 34.60 32.14 28.95 FFDNet [ 2 ] 33.87 31.21 27.96 33.83 31.40 28.05 34.63 32.13 28.98 ADFNet [ 3 ] 34.21 31.60 28.19 34.50 32.13 28.71 34.77 32.22 29.06 MIRNet-v2 [ 52 ] 33.66 30.97 27.66 33.30 30.75 27.22 34.29 31.81 28.55 DGUNet [ 53 ] 33.85 31.10 27.92 33.67 31.27 27.94 34.56 32.10 28.91 Restormer [ 17 ] 34.03 31.49 28.11 33.72 31.26 28.03 34.78 32.37 29.08 NAFNet [ 32 ] 33.67 31.02 27.73 33.14 30.64 27.20 34.27 31.80 28.62 FSNet [ 33 ] 34.09 31.55 28.12 33.88 31.31 28.07 34.75 32.38 29.10 AirNet [ 22 ] 34.14 31.48 28.23 34.40 32.10 28.88 34.81 32.44 29.10 PromptIR* [ 12 ] 34.33 31.70 28.45 34.83 32.59 29.50 35.30 32.86 29.76 Perceive-IR [ 16 ] 34.38 31.74 28.53 34.86 32.55 29.42 34.84 32.50 29.16 AdaIR [ 14 ] 34.36 31.72 28.49 34.96 32.74 29.70 35.39 32.94 29.85 M 2 IR 34.40 31.76 28.54 34.99 32.79 29.78 35.40 32.96 29.89

Method SOTS [ 30 ] DehazeNet [ 4 ] 22.46 / 0.851 AODNet [ 5 ] 20.29 / 0.877 FDGAN [ 6 ] 23.15 / 0.921 DehazeFormer [ 7 ] 31.78 / 0.977 Restormer [ 17 ] 30.87 / 0.969 NAFNet [ 32 ] 30.98 / 0.970 FSNet [ 33 ] 31.11 / 0.971 AirNet [ 22 ] 23.18 / 0.900 PromptIR [ 12 ] 31.31 / 0.973 Perceive-IR [ 16 ] 31.65 / 0.977 AdaIR [ 14 ] 31.80 / 0.981 M 2 IR 32.12 / 0.984 Method Rain100L [ 31 ] UMR [ 8 ] 32.39 / 0.921 MSPFN [ 10 ] 33.50 / 0.948 LPNet [ 9 ] 33.61 / 0.958 DRSformer [ 11 ] 38.14 / 0.983 Restormer [ 17 ] 36.74 / 0.978 NAFNet [ 32 ] 36.63 / 0.977 FSNet [ 33 ] 37.27 / 0.980 AirNet [ 22 ] 34.90 / 0.977 PromptIR [ 12 ] 37.04 / 0.979 Perceive-IR [ 16 ] 38.41 / 0.984 AdaIR [ 14 ] 38.90 / 0.985 M 2 IR 39.33 / 0.986

TABLE V: Quantitative Comparison of recent image restoration methods under the one-by-one settings, with each sub-table dividing methods into task-specific (upper), general-purpose (middle), and all-in-one (lower) categories. * denotes the model has been retrained.

AirNet	PromptIR	AdaIR	M 2 IR
19.32/0.733	20.47/0.764	20.54/0.764	20.83 / 0.770

TABLE VI: Quantitative results on the unseen desnowing task with the CSD [ 54 ] dataset. All models are trained on 3-Degradation settings.

MST	DACP	DGEA		PSNR	SSIM
MST	DACP	special	shared	PSNR	SSIM
✗	✗	✗	✗	38.43	0.9823
✓	✗	✗	✗	38.55	0.9834
✓	✗	✓	✗	38.67	0.9832
✓	✗	✓	✓	38.82	0.9833
✓	✓	✓	✗	38.73	0.9837

TABLE VII: Ablation Studies for each component.

Method	N=2	N=4	N=8	N=16	N=32
PSNR	38.84	38.93	38.81	38.76	38.62
SSIM	0.9834	0.9842	0.9838	0.9836	0.9834

TABLE VIII: Ablation Studies for the number N of specialized experts in each ADEC module, i.e., excluding the shared expert.

Method	K=1	K=2	K=3	K=4
PSNR	38.83	38.93	38.86	38.78
SSIM	0.9840	0.9842	0.9839	0.9837

TABLE IX: Ablation Studies for the number K of the top experts.

$\boldsymbol{\ell_{charb}}$	$\boldsymbol{\ell_{balance}}$	$\boldsymbol{\ell_{FFT}}$	PSNR	SSIM
✗	✗	✗	38.41	0.9822
✓	✗	✗	38.56	0.9834
✓	✗	✓	38.74	0.9836
✓	✓	✗	38.81	0.9839
✓	✓	✓	38.93	0.9842

TABLE X: Ablation Studies for each loss function.

Method	0.001	0.005	0.01	0.05	0.1
PSNR	38.89	38.74	38.93	38.46	38.60
SSIM	0.9841	0.9838	0.9842	0.9823	0.9830

TABLE XI: Ablation Studies for the weight

\lambda_{1}

Method	0.01	0.05	0.1	0.5	1
PSNR	38.63	38.90	38.93	38.71	38.72
SSIM	0.9835	0.9841	0.9842	0.9837	0.9830

TABLE XII: Ablation Studies for the weight

\lambda_{2}

Fig. 9: Gating visualization of degraded input images and their corresponding GT during the encoding stage, illustrating the selective state modulation mechanism in MST. The model is trained on three degradation types.

Panel of Figure 22 — Fig. 10: The t-SNE visualisations of the two encoders under three degeneration types.

Method	Params.	FLOPs	latency	CPU	GPU
IDR	42.3M	1522G	384ms	1387M	6952M
PromptIR	34.1M	745G	282ms	1242M	3127M
GridFormer	34.1M	1941G	746ms	1355M	3830M
Pool-AIO	26.1M	745G	293ms	1246M	3096M
AdaIR	28.8M	778G	333ms	1296M	3110M
MoCE-IR	25.4M	474G	263ms	1312M	1253M

TABLE XIII: Efficiency analysis of model parameters, FLOPs, inference latency, and peak memory usage on both CPU and GPU, evaluated using

720\times 480

inputs on an NVIDIA RTX 4090 GPU.

Degradation-Aware Image Restoration

See all 28 papers

Degradation-Aware Adaptive Context Gating for Unified Image Restoration
May 2, 2026Lei He, Jielei Chu, Fengmao Lv +4All-in-One Image Restoration Image Restoration
the School of Computing and Artificial Intelligence, Southwest Jiaotong University, Chengdu 611756, China · the Engineering Research Center of Sustainable Urban Intelligent Transportation, Ministry of Education, Chengdu, Sichuan 611756, China · School of Computing and Artificial Intelligence, Jiangxi University of Finance and Economics, Nanchang, 330013, China +3
MiM-DiT: MoE in MoE with Diffusion Transformers for All-in-One Image Restoration
Mar 3, 2026Lingshun Kong, Jiawei Zhang, Zhengpeng Duan +6All-in-One Image Restoration Diffusion Models for Image Restoration
Nanjing University of Science and Technology · Nankai University · Harbin Institute of Technology
TaskIR: Task-Driven Image Restoration via Degradation Adaptation and Task Feedback
Sep 25, 2026Yanjie Tu, Qingsen Yan, Axi Niu +4All-in-One Image Restoration Image Restoration
Northwestern Polytechnical University · Shenzhen Research Institute of Northwestern Polytechnical University · Xi’an University of Architecture and Technology

May 2, 2026J/K move · Enter open · S save

Degradation-Aware Adaptive Context Gating for Unified Image Restoration

Lei He, Jielei Chu, Fengmao Lv, Weide Liu+3

the School of Computing and Artificial Intelligence, Southwest Jiaotong University, Chengdu 611756, China · the Engineering Research Center of Sustainable Urban Intelligent Transportation, Ministry of Education, Chengdu, Sichuan 611756, China · School of Computing and Artificial Intelligence, Jiangxi University of Finance and Economics, Nanchang, 330013, China+3

All-in-One Image Restoration Image Restoration Degradation-Aware Image Restoration

Open paper AskPDF arXiv

Unified image restoration using a single model often faces task interference due to diverse degradations. To address this, we propose DACG-IR (Degradation-Aware Adaptive Context Gating), which enables explicit perception of degradation characteristics to dynamically modulate feature representations. Our method constructs degradation-aware contextual representations from the input to modulate attention distribution, frequency-domain features, and feature aggregation. Specifically, a lightweight multi-scale degradation-aware module extracts coarse degradation information and generates layer-wise prompts. These prompts guide attention temperature and output gating in encoder and decoder blocks for adaptive feature extraction. Additionally, a spatial-channel dual-gated adaptive fusion mechanism refines encoder features, suppressing noise propagation from shallow to deep layers. This design effectively suppresses degradation-induced noise while preserving informative structures. Experiments show DACG-IR outperforms state-of-the-art methods in single-task, all-in-one, adverse weather removal, and composite degradation settings. Code: https://github.com/HlHomes/DACG-IR-code

Abstract

Degradation-Aware Image Restoration

Explore similar work

Degradation-Aware Adaptive Context Gating for Unified Image Restoration

MiM-DiT: MoE in MoE with Diffusion Transformers for All-in-One Image Restoration

TaskIR: Task-Driven Image Restoration via Degradation Adaptation and Task Feedback