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
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
on SOTS [ 30 ]
on Rain100L [ 31 ]
σ=15
σ=25
σ=50
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
on SOTS [ 30 ]
on Rain100L [ 31 ]
on BSD68 [ 29 ]
on GoPro [ 40 ]
on LOL [ 41 ]
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 σ=25 .
Method
Single
Double
Triple
Average
L
H
R
S
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
PSNR ↑
SSIM ↑
LPIPS ↓
FID ↓
PSNR ↑
SSIM ↑
LPIPS ↓
FID ↓
PSNR ↑
SSIM ↑
LPIPS ↓
FID ↓
PSNR ↑
SSIM ↑
LPIPS ↓
FID ↓
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).
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
special
shared
✗
✗
✗
✗
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.
ℓcharb
ℓbalance
ℓ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 λ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 λ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.
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×480 inputs on an NVIDIA RTX 4090 GPU.
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
Lei He, Jielei Chu, Fengmao Lv +4
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 is challenging because different degradation types, such as haze, blur, noise, and low-light, impose diverse requirements on restoration strategies, making it difficult for a single model to handle them effectively. In this paper, we propose a unified image restoration framework that integrates a dual-level Mixture-of-Experts (MoE) architecture with a pretrained diffusion model. The framework operates at two levels: the Inter-MoE layer adaptively combines expert groups to handle major degradation types, while the Intra-MoE layer further selects specialized sub-experts to address fine-grained variations within each type. This design enables the model to achieve coarse-grained adaptation across diverse degradation categories while performing fine-grained modulation for specific intra-class variations, ensuring both high specialization in handling complex, real-world corruptions. Extensive experiments demonstrate that the proposed method performs favorably against the state-of-the-art approaches on multiple image restoration task.
Lingshun Kong, Jiawei Zhang, Zhengpeng Duan +6
Nanjing University of Science and Technology · Nankai University · Harbin Institute of Technology
Task-driven image restoration aims to improve both image quality and downstream task performance. However, existing methods predominantly focus on single degradation type and struggle to handle the diverse degradations encountered in real-world scenarios. Different degradations impose distinct restoration demands, and insufficient restoration may leave residual degradations and artifacts that impair object boundaries and semantic cues, thereby compromising downstream task performance. To address these challenges, we propose TaskIR, a two-stage task-driven unified image restoration framework that integrates degradation-adaptive restoration with task feedback refinement. In Stage I, a Degradation Representation Module (DRM) extracts degradation representations, enabling a Degradation-Guided Transformer Block (DGTB) to dynamically modulate feature transformations for adaptive restoration. In Stage II, a Task-to-Restoration Feedback Generation module (TRFG) transforms heterogeneous task features into restoration feedback by modeling task-representation discrepancies associated with the current restoration. Subsequently, a Selective Task Feedback Refinement module (STFR) assesses feedback relevance and selectively refines intermediate restoration features to mitigate interference with well-restored content. Extensive experiments demonstrate that TaskIR achieves competitive restoration quality and downstream task performance across diverse degradations and tasks.
Yanjie Tu, Qingsen Yan, Axi Niu +4
Northwestern Polytechnical University · Shenzhen Research Institute of Northwestern Polytechnical University · Xi’an University of Architecture and Technology