MiM-DiT: MoE in MoE with Diffusion Transformers for All-in-One Image Restoration
Organizations: Nanjing University of Science and Technology · Nankai University · Harbin Institute of Technology
Abstract
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.
Figures & tables
| Degradation | Metrics | AirNet [ 15 ] | DGUNet [ 29 ] | TransWeather [ 37 ] | PromptIR [ 34 ] | DiffIR [ 40 ] | DiffUIR [ 50 ] | DA-CLIP [ 26 ] | AutoDIR* [ 12 ] | FoundIR [ 17 ] | DiT4SR [ 8 ] | Ours |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Blur | LPIPS | 0.4340 | 0.3855 | 0.3911 | 0.3900 | 0.3692 | 0.2043 | 0.3067 | 0.3291 | 0.1722 | 0.3019 | 0.2372 |
| FID | 42.4248 | 35.3296 | 34.3489 | 36.2908 | 33.8067 | 9.0676 | 16.9121 | 17.8850 | 10.5639 | 29.0668 | 16.9659 | |
| NIQE | 7.6429 | 7.7016 | 7.2645 | 7.8885 | 7.2844 | 5.5757 | 6.4046 | 6.6792 | 5.5645 | 2.9719 | 2.9546 | |
| LIQE | 1.0102 | 1.0089 | 1.0042 | 1.0086 | 1.0127 | 1.8270 | 1.0274 | 1.0140 | 2.0964 | 3.0184 | 3.9440 | |
| MUSIQ | 22.9102 | 22.8828 | 21.1029 | 23.1877 | 23.5724 | 50.8270 | 31.3108 | 30.8041 | 53.4637 | 66.8670 | 68.5028 | |
| CLIP-IQA | 0.3677 | 0.3199 | 0.3121 | 0.2929 | 0.2981 | 0.4806 | 0.4772 | 0.3789 | 0.5083 | 0.6242 | 0.6405 |
| Dataset | Metrics | TransWeather [ 37 ] | AirNet [ 15 ] | PromptIR [ 34 ] | PromptIR* | DiffIR [ 40 ] | DiffUIR [ 50 ] | DiffUIR* | AutoDIR* [ 12 ] | FoundIR [ 17 ] | DiT4SR [ 8 ] | Ours |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 4KRD [ 6 ] | LPIPS | 0.4468 | 0.4837 | 0.4421 | 0.3130 | 0.2373 | 0.2669 | 0.2421 | 0.3328 | 0.2347 | 0.3283 | 0.2896 |
| FID | 27.5878 | 32.6022 | 25.6506 | 34.3458 | 14.6920 | 21.1369 | 15.5659 | 15.2305 | 16.6398 | 23.4762 | 18.1328 | |
| NIQE | 7.9220 | 8.1894 | 8.8772 | 6.1589 | 6.4251 | 6.0133 | 5.7992 | 7.3755 | 6.0495 | 4.5789 | 3.6536 | |
| LIQE | 1.0119 | 1.0115 | 1.0192 | 1.0531 | 1.1162 | 1.0711 | 1.0795 | 1.0413 | 1.1014 | 2.2047 | 2.5894 | |
| MUSIQ | 17.6972 | 18.7134 | 18.9286 | 23.9315 | 27.6394 | 26.1672 | 28.0036 | 22.1480 | 27.3308 | 40.8834 | 42.0540 | |
| CLIP-IQA | 0.3305 | 0.4515 | 0.4172 | 0.3531 | 0.3895 | 0.3849 | 0.3842 | 0.3516 | 0.4010 | 0.5136 | 0.5587 |
| Methods | LPIPS | FID | NIQE | LIQE | MUSIQ | CLIPIQA |
|---|---|---|---|---|---|---|
| w/o Intra-MoE | 0.1873 | 34.5619 | 4.7352 | 2.3143 | 50.5890 | 0.5511 |
| w/ Intra-MoE | 0.1636 | 26.8815 | 4.2502 | 2.9133 | 53.4975 | 0.5763 |
| Methods | LPIPS | FID | NIQE | LIQE | MUSIQ | CLIPIQA |
|---|---|---|---|---|---|---|
| (a) Spatial-only | 0.1924 | 35.8375 | 4.3574 | 2.6470 | 52.4738 | 0.5673 |
| (b) Channel-only | 0.1801 | 31.3849 | 4.2863 | 2.5394 | 53.4217 | 0.5526 |
| (c) Swin-only | 0.2199 | 38.1274 | 4.9472 | 2.5539 | 51.6261 | 0.5395 |
| (d) SE-only | 0.2577 | 40.3904 | 5.1836 | 2.1044 | 49.1557 | 0.4972 |
| (e) Heterogeneity | 0.1636 | 26.8815 | 4.2502 | 2.9133 | 53.4975 | 0.5763 |
| Methods | LPIPS | FID | NIQE | LIQE | MUSIQ | CLIPIQA |
|---|---|---|---|---|---|---|
| Sparse Inter + Sparse Intra | 0.2457 | 43.5841 | 5.4490 | 2.0107 | 48.8349 | 0.5183 |
| Sparse Inter + Dense Intra | 0.2348 | 44.8341 | 5.1905 | 2.3049 | 46.4920 | 0.5532 |
| Dense Inter + Sparse Intra | 0.1636 | 26.8815 | 4.2502 | 2.9133 | 53.4975 | 0.5763 |
| Attention Type | Blur | Noise | Haze | Rain | Low-light |
|---|---|---|---|---|---|
| Spatial | 0.3622 | 0.3953 | 0.1630 | 0.3321 | 0.2081 |
| Channel | 0.1399 | 0.2784 | 0.4008 | 0.1406 | 0.3716 |
| Swin | 0.3592 | 0.1514 | 0.2006 | 0.3673 | 0.1423 |
| SE | 0.1387 | 0.1749 | 0.2356 | 0.1600 | 0.2780 |
All-in-One Image Restoration
All-in-one image restoration trains unified models to remove diverse and co-occurring degradations while preserving scene content.
See all 87 papers