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

Lingshun Kong; Jiawei Zhang; Zhengpeng Duan; Xiaohe Wu; Yueqi Yang; Xiaotao Wang; Dongqing Zou; Lei Lei

cs.CVMar 3, 2026

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

Authors: Lingshun Kong, Jiawei Zhang, Zhengpeng Duan, Xiaohe Wu, Yueqi Yang, Xiaotao Wang, Dongqing Zou, Lei Lei,

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

Figure 1: Homogeneous MoE vs. the proposed MiM. The symbols in the figure are consistent with those defined in our method section and share the same meanings. Our hierarchical MiM adopts a dual-level routing mechanism, which ensures dynamic and specialized processing by routing inputs to appropriate architectural priors and fine-grained specialists. Compared with homogeneous MoE, our method enables better restoration for blur, haze, and low-light conditions through adaptive structural selection.

Figure 2: Overview of the proposed Hierarchical MoE in MoE (MiM) integrated into the DiT backbone. The framework processes low-quality (LQ) images through a series of MiM-DiT blocks. Given the LQ input, the MiM module extracts degradation-specific features and processes them through a hierarchical MoE architecture composed of two levels: Inter-MoE and Intra-MoE. At the Inter-MoE level, four expert groups based on distinct attention mechanisms—spatial self-attention channel, self-attention, Swin attention, and SE attention—are combined via a dense router that computes adaptive weights over all groups. This dense fusion enables the model to leverage complementary inductive biases. Within each expert group, Intra-MoE captures fine-grained variations within each degradation category via sparse routing. These processed features are then injected as conditional input into the DiT backbone through a Zero-Linear pathway, dynamically guiding the diffusion process to generate restoration results.

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 $\downarrow$	0.4340	0.3855	0.3911	0.3900	0.3692	0.2043	0.3067	0.3291	0.1722	0.3019	0.2372
	FID $\downarrow$	42.4248	35.3296	34.3489	36.2908	33.8067	9.0676	16.9121	17.8850	10.5639	29.0668	16.9659
	NIQE $\downarrow$	7.6429	7.7016	7.2645	7.8885	7.2844	5.5757	6.4046	6.6792	5.5645	2.9719	2.9546
	LIQE $\uparrow$	1.0102	1.0089	1.0042	1.0086	1.0127	1.8270	1.0274	1.0140	2.0964	3.0184	3.9440
	MUSIQ $\uparrow$	22.9102	22.8828	21.1029	23.1877	23.5724	50.8270	31.3108	30.8041	53.4637	66.8670	68.5028
	CLIP-IQA $\uparrow$	0.3677	0.3199	0.3121	0.2929	0.2981	0.4806	0.4772	0.3789	0.5083	0.6242	0.6405

Table 1: Quantitative evaluations of the proposed method against state-of-the-art ones on FoundIR [ 17 ] benchmarks. Models marked with an asterisk ( ∗ ) are evaluated using officially released pre-trained weights, while those without are retrained following the FoundIR [ 17 ] on the same dataset for fair comparison. The best and second performances are marked in red and blue , respectively.

Dataset	Metrics	TransWeather [ 37 ]	AirNet [ 15 ]	PromptIR [ 34 ]	PromptIR*	DiffIR [ 40 ]	DiffUIR [ 50 ]	DiffUIR*	AutoDIR* [ 12 ]	FoundIR [ 17 ]	DiT4SR [ 8 ]	Ours
4KRD [ 6 ]	LPIPS $\downarrow$	0.4468	0.4837	0.4421	0.3130	0.2373	0.2669	0.2421	0.3328	0.2347	0.3283	0.2896
	FID $\downarrow$	27.5878	32.6022	25.6506	34.3458	14.6920	21.1369	15.5659	15.2305	16.6398	23.4762	18.1328
	NIQE $\downarrow$	7.9220	8.1894	8.8772	6.1589	6.4251	6.0133	5.7992	7.3755	6.0495	4.5789	3.6536
	LIQE $\uparrow$	1.0119	1.0115	1.0192	1.0531	1.1162	1.0711	1.0795	1.0413	1.1014	2.2047	2.5894
	MUSIQ $\uparrow$	17.6972	18.7134	18.9286	23.9315	27.6394	26.1672	28.0036	22.1480	27.3308	40.8834	42.0540
	CLIP-IQA $\uparrow$	0.3305	0.4515	0.4172	0.3531	0.3895	0.3849	0.3842	0.3516	0.4010	0.5136	0.5587

Table 2: Quantitative evaluations of the proposed method against state-of-the-art ones on additional public benchmarks, including 4KRD [ 6 ] , RealRain-1K [ 19 ] , HazeRD [ 49 ] , and UHD-LL [ 16 ] . Models marked with an asterisk ( ∗ ) are evaluated using officially released pre-trained weights, while those without are retrained following the FoundIR [ 17 ] on the same dataset for fair comparison. The best and second performances are marked in red and blue , respectively.

Panel of Figure 3 — Figure 3: Deblurred results on the FoundIR dataset [ 17 ] . The deblurred results in (c)-(g) still contain significant blur effects. In contrast, our method generates clear results.

Panel of Figure 4 — Figure 4: Dehazed results on the FoundIR dataset [ 17 ] . The results in (c) to (g) fail to fully restore the original scene content while removing haze. In contrast, our method generates clear and faithful reconstructions.

Panel of Figure 5 — Figure 5: Low-light enhanced results on the FoundIR dataset [ 17 ] . Results from (c) to (g) suffer from color casts and detail smearing. In contrast, our method recovers accurate colors and fine structures.

Methods	LPIPS $\downarrow$	FID $\downarrow$	NIQE $\downarrow$	LIQE $\uparrow$	MUSIQ $\uparrow$	CLIPIQA $\uparrow$
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

Table 3: Ablation study on the effectiveness of the Intra-MoE module, comparing performance with and without Intra-MoE on the FoundIR dataset [ 17 ] . The variant with Intra-MoE achieves consistently better results, validating the importance of Intra-MoE.

Panel of Figure 6 — Figure 6: Effect of the Intra-MoE. The model with Intra-MoE produces a sharper and clearer result.

Methods	LPIPS $\downarrow$	FID $\downarrow$	NIQE $\downarrow$	LIQE $\uparrow$	MUSIQ $\uparrow$	CLIPIQA $\uparrow$
(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

Table 4: Ablation study on structural heterogeneity in Inter-MoE. Heterogeneous Inter-MoE (fusing all four attention types) outperforms single-structure variants across all metrics on FoundIR dataset [ 17 ] , proving that combining multiple attention mechanisms is essential for all-in-one restoration.

Methods	LPIPS $\downarrow$	FID $\downarrow$	NIQE $\downarrow$	LIQE $\uparrow$	MUSIQ $\uparrow$	CLIPIQA $\uparrow$
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

Table 5: Ablation studies of the Router Design, Dense Inter-MoE + Sparse Intra-MoE outperforms alternative configurations on FoundIR dataset [ 17 ] .

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

Table 6: Average adaptive weights across degradation types. Values represent normalized mean routing weights.

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

M2IR: Proactive All-in-One Image Restoration via Mamba-style Modulation and Mixture-of-Experts
Mar 16, 2026Shiwei Wang, Yongzhen Wang, Bingwen Hu +3Degradation-Aware Image Restoration All-in-One Image Restoration
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 +1
Divide and Restore: A Modular Task-Decoupled Framework for Universal Image Restoration
Mar 30, 2026Joanna Wiekiera, Martyna ZurImage Restoration All-in-One Image Restoration
Faculty of Automatic Control, Electronics and Computer Science, Silesian University of Technology, Gliwice, Poland
Leveraging Multimodal Large Language Models for All-in-One Image Restoration via a Mixture of Frequency Experts
May 12, 2026Eunho Lee, Rei Kawakami, Youngbae HwangMixture of Experts All-in-One Image Restoration
Chungbuk National University, Cheongju 28644, Republic of Korea · Institute of Science Tokyo, Tokyo 152-8550, Japan

Mar 16, 2026J/K move · Enter open · S save

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

Shiwei Wang, Yongzhen Wang, Bingwen Hu, Liyan Zhang+2

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+1

Degradation-Aware Image Restoration All-in-One Image Restoration Image Restoration Mixture of Experts

Open paper AskPDF arXiv

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.

Abstract

All-in-One Image Restoration

Explore similar work

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

Divide and Restore: A Modular Task-Decoupled Framework for Universal Image Restoration

Leveraging Multimodal Large Language Models for All-in-One Image Restoration via a Mixture of Frequency Experts