[ICCV 2025] MiM-DiT: Mastering All-in-One Image Restoration with Hierarchical MoE and Diffusion Transformers

Executive Summary

TL;DR: MiM-DiT is a unified image restoration framework that solves the "one-size-fits-all" limitation of current models. By nesting a sparse Mixture-of-Experts (MoE) within a dense heterogeneous MoE and plugging it into a Diffusion Transformer (DiT), the model dynamically adjusts its architecture to fit the specific degradation (blur, haze, noise) of the input.

Positioning: This work moves beyond simple prompt-based restoration or uniform diffusion processes. It represents a shift toward Architectural Adaptation, where the model doesn't just change its "features" but changes its "wiring" based on the task at hand.

The "Conflict of Priors" Problem

Traditional restoration models face a fundamental trade-off. Handling motion blur requires sharp local filters, while low-light enhancement requires global illumination modeling. When a single network is trained on all tasks, it often settles for a "mean" solution, resulting in images that are either physically inaccurate or perceptually "mushy" (over-smoothed).

The authors identify two failures in current SOTA:

1. Homogeneous Experts: Most MoE models use the same block structure for all experts, ignoring that different tasks need different types of attention.
2. Uniform Diffusion: Standard diffusion models apply the same denoising logic to every degradation, often losing fine textures in the process.

Methodology: The Hierarchical "MoE in MoE" (MiM)

The core innovation is the MiM block, which replaces standard Feed-Forward Networks or Attention layers with a two-tiered hierarchy.

1. Inter-MoE (Structural Diversity)

Instead of identical experts, the Inter-MoE level uses four distinct attention mechanisms:

• Spatial Self-Attention: For long-range geometric dependencies.
• Channel Self-Attention: For inter-channel relationship recalibration.
• Swin Attention: For efficient local-window processing.
• SE (Squeeze-and-Excitation): For global illumination and transmission modeling (critical for dehazing).

A dense router fuses these outputs, allowing the model to "blend" these different logical priors.

2. Intra-MoE (Fine-grained Specialization)

Inside each of the four groups above, an Intra-MoE layer uses sparse Top-k routing. This activates only the sub-experts best suited for the specific intensity of degradation (e.g., distinguishing between a light drizzle and a heavy rainstorm).

Integration with DiT

The MiM features are injected into a pretrained Stable Diffusion 3.5 (DiT) backbone. To prevent the new modules from "shocking" the pretrained weights during early training, the authors use a Zero-Linear pathway. This ensures the model starts by relying on its strong generative priors and gradually learns to use the restoration guidance.

Experimental Results & SOTA Comparison

MiM-DiT was tested against 10+ SOTA methods (including AirNet, PromptIR, and DiffUIR) across five major restoration tasks.

• Quantitative Success: The model achieved the best LPIPS (fidelity) and MUSIQ/CLIPIQA (perceptual quality) scores on nearly all benchmarks.
• Qualitative Superiority: While competing diffusion models like AutoDIR or DiT4SR sometimes introduce artifacts or color distortions, MiM-DiT maintains structural integrity while fully recovering high-frequency textures.

In the deblurring task, MiM-DiT effectively removes motion blur that other models leave behind.

Ablation Insight: Why Heterogeneity Matters

The authors proved that using four different types of attention (Heterogeneous) significantly outperformed using four experts of the same type (Homogeneous). For example, the "SE-only" variant failed at deblurring, while the "Spatial-only" variant struggled with low-light enhancement.

Critical Analysis & Conclusion

Takeaways

MiM-DiT proves that heterogeneity is the next frontier for MoE. Simply adding more parameters isn't enough; we need to add different kinds of processing logic. By acting as a "surgical" conditioner for the Diffusion Transformer, MiM-DiT preserves the "hallucination" power of diffusion while grounding it in physical reality.

Limitations

• Computational Weight: While sparse routing helps, a hierarchical MoE embedded in a 3D-attention DiT (like SD3.5) remains computationally expensive for real-time mobile applications.
• Training Data: The performance relies heavily on the FoundIR dataset, which might not cover every niche real-world degradation.

Future Work: We expect to see this MiM architecture distilled into smaller, "student" DiT models or extended into video restoration where temporal consistency provides an even greater challenge for MoE routers.