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

This paper proposes MiM-DiT, a unified image restoration framework that integrates a dual-level Mixture-of-Experts architecture with pretrained diffusion transformers to effectively handle diverse and fine-grained degradation types through adaptive coarse-grained and fine-grained expert selection.

The Gist

Imagine you have a photo that is ruined. Maybe it's blurry because you moved the camera, maybe it's foggy because of a storm, maybe it's too dark, or maybe it's covered in static noise.

For a long time, fixing these photos was like having a Swiss Army Knife with only one blade. You could try to use that one blade to cut, screw, or scrape, but it never did any of them perfectly. If you tried to fix a blurry photo with a tool designed for dark photos, the result looked weird and smudged.

The paper you shared introduces a new tool called MiM-DiT. Think of it not as a single knife, but as a high-tech, magical repair shop that knows exactly which tool to grab for every specific problem.

Here is how it works, broken down into simple concepts:

1. The Problem: One Size Does Not Fit All

The authors realized that fixing a "blurry" photo is totally different from fixing a "foggy" one.

• Blur needs you to sharpen edges and fix motion.
• Fog needs you to brighten the image and remove a white haze.
• Darkness needs you to boost the light without making it look fake.

Old AI models tried to be "generalists," doing everything with the same brain. This often led to results that looked "over-smoothed" (like a plastic doll) or lost important details.

2. The Solution: A "Shop Within a Shop" (MoE in MoE)

The core idea of this paper is MiM-DiT. Let's break that name down using a Restaurant Analogy:

Imagine a massive, world-class restaurant (the Diffusion Transformer). This restaurant is famous for making delicious, high-quality food (generating perfect images). But the customers (the damaged photos) have very specific, weird orders.

To handle this, the restaurant builds a two-level expert system:

Level 1: The "Inter-MoE" (The Head Chefs)

This is the Main Kitchen. Instead of having one chef, they have four different Head Chefs, each with a totally different style of cooking:

• Chef A (Spatial): Great at fixing shapes and lines (good for blur).
• Chef B (Channel): Great at fixing colors and lighting (good for fog or darkness).
• Chef C (Swin): A mix of local and global views (good for complex scenes).
• Chef D (SE): Specialized in global illumination (good for low-light).

When a damaged photo arrives, a Smart Manager (The Router) looks at the damage and decides: "This photo is blurry, so let's ask Chef A and Chef C to help." It doesn't just pick one; it mixes their advice together perfectly. This is the Inter-MoE (Mixture of Experts).

Level 2: The "Intra-MoE" (The Specialized Sous-Chefs)

But wait! Not all blurs are the same. Some are "fast motion blur" (like a car speeding by), and some are "slow motion blur" (like a shaky hand).

Inside Chef A's kitchen, there isn't just one person. There are three Sous-Chefs who all know how to fix blur, but each specializes in a different type of blur.

• Sous-Chef 1: Handles fast motion.
• Sous-Chef 2: Handles slow motion.
• Sous-Chef 3: Handles focus issues.

The Smart Manager looks at the specific photo again and says, "This is fast motion blur, so let's activate Sous-Chef 1." This is the Intra-MoE (Mixture of Experts inside the experts).

3. The Magic Ingredient: The "Pre-trained Brain"

Why is this so good? Because the restaurant (the Diffusion Transformer) is already a Master Chef trained on millions of perfect photos. It already knows what a "perfect tree" or a "perfect face" looks like.

The MiM system doesn't try to teach the restaurant how to cook from scratch. Instead, it just guides the Master Chef.

• It says: "Hey, this photo is foggy. Use your knowledge of trees, but apply the 'Fog Removal' technique from Chef B."
• It says: "This photo is dark. Use your knowledge of faces, but apply the 'Light Boost' technique from Chef D."

Because the Master Chef already knows what a perfect image looks like, the final result is incredibly realistic, with sharp textures and natural colors, avoiding that "plastic" look of older AI.

4. Why This Matters

• Flexibility: It handles any mix of problems (blur + fog + noise) at the same time.
• Efficiency: It doesn't waste energy. It only "wakes up" the specific chefs needed for the job, rather than making the whole kitchen work on every single photo.
• Quality: By combining the "Master Chef's" knowledge with "Specialized Chefs," the results are sharper and more detailed than anything else currently available.

The Bottom Line

MiM-DiT is like a super-intelligent repair team that doesn't just use a hammer for everything. Instead, it has a team of specialists who can instantly recognize if your photo needs a "sharpening tool," a "lighting tool," or a "color tool," and they work together seamlessly to restore your photo to its original, beautiful state.

The paper proves this works better than all previous methods on tests involving blur, fog, rain, darkness, and noise, making it a huge step forward for fixing real-world photos.

Technical Summary

1. Problem Statement

All-in-One Image Restoration aims to reconstruct high-quality images from inputs degraded by various types of corruptions (e.g., haze, blur, noise, low-light, rain) using a single unified model.

• The Core Challenge: Different degradation types impose conflicting structural requirements. For instance, removing haze requires modeling global illumination and transmission, while removing motion blur requires reconstructing sharp local edges. A single homogeneous network often struggles to satisfy these divergent needs simultaneously, leading to compromised performance (e.g., over-smoothing or structural distortion).
• Limitations of Existing Methods:
  • Task-Specific Models: Require separate training for each degradation type, lacking generalization.
  • Unified Deterministic Models: Often rely on pixel-level regression losses, resulting in over-smoothed outputs and a lack of texture generation.
  • Standard Diffusion Models: While powerful, they typically apply a uniform processing strategy to all degradations, ignoring the distinct inductive biases required for specific corruption types.
  • Standard Mixture-of-Experts (MoE): Existing MoE approaches often use homogeneous expert designs (same architecture for all experts), failing to leverage the distinct structural priors needed for different degradation categories.

2. Methodology: MiM-DiT

The authors propose MiM-DiT, a framework that integrates a hierarchical Mixture-of-Experts (MoE) architecture into a pretrained Diffusion Transformer (DiT) backbone (specifically based on Stable Diffusion 3.5).

A. Hierarchical MoE-in-MoE (MiM) Architecture

The core innovation is a dual-level routing mechanism designed to handle both coarse-grained structural differences and fine-grained variations within degradation types.

1. Inter-MoE Level (Coarse-Grained Adaptation):

   • Goal: Select the appropriate architectural prior for the specific degradation type.
   • Structure: Four distinct expert groups are constructed, each based on a different attention mechanism:
     • Spatial Self-Attention: For long-range spatial dependencies.
     • Channel Self-Attention: For channel-wise feature recalibration.
     • Swin Attention: For balancing local and global context efficiently.
     • SE (Squeeze-and-Excitation) Attention: For global illumination modeling (crucial for haze/low-light).
   • Routing: A Dense Router adaptively fuses the outputs of all four groups. Unlike sparse selection, this dense fusion allows the model to simultaneously leverage complementary inductive biases (e.g., combining spatial and channel features) based on the input's degradation characteristics.

2. Intra-MoE Level (Fine-Grained Adaptation):

   • Goal: Handle variations within a specific degradation category (e.g., light vs. heavy haze, short vs. long motion blur).
   • Structure: Inside each of the four expert groups defined above, multiple sub-experts are deployed. These sub-experts share the same attention mechanism as their parent group but have independently trained parameters.
   • Routing: A Sparse Router (Top-k gating) dynamically selects and activates the most relevant sub-experts for the specific input, ensuring efficiency and specialization.

B. Integration with Diffusion Transformer (DiT)

• Backbone: The method utilizes a pretrained DiT (from SD3.5) which employs a rectified flow-matching formulation. This provides strong generative priors and the ability to model diverse image distributions.
• Conditional Generation: The MiM module acts as a conditioning mechanism.
  1. The Low-Quality (LQ) image is processed through the hierarchical MiM to extract degradation-specific features.
  2. These features are passed through a Zero-Linear layer (initialized with zero weights) to project them into the latent space.
  3. The resulting features are concatenated with text embeddings and noise tokens before entering the DiT attention layers.
• Training Stability: The Zero-Linear initialization ensures that the conditioning signal has minimal influence at the start of training, preventing instability, while allowing the model to gradually learn to leverage these signals for high-fidelity restoration.

3. Key Contributions

1. Unified Framework: Proposes the first all-in-one restoration framework that successfully integrates the MoE paradigm with a pretrained Diffusion Transformer, preserving generative capacity while enabling task-specific specialization.
2. Hierarchical MoE-in-MoE Design: Introduces a novel dual-level routing mechanism:
   • Inter-MoE: Routes between heterogeneous architectural priors (Spatial, Channel, Swin, SE) via dense fusion.
   • Intra-MoE: Routes to specialized sub-experts within each architecture via sparse selection.
   • This design explicitly decouples structural adaptation from intra-class specialization.
3. State-of-the-Art Performance: Demonstrates superior performance across multiple benchmarks, outperforming existing universal restoration methods and specialized diffusion models.

4. Experimental Results

The method was evaluated on the FoundIR dataset and four real-world benchmarks (4KRD, RealRain-1K, HazeRD, UHD-LL).

• Quantitative Performance:

  • Achieved State-of-the-Art (SOTA) results on most metrics, particularly in no-reference metrics (NIQE, LIQE, MUSIQ, CLIP-IQA), indicating superior perceptual quality and naturalness.
  • Outperformed competitors (e.g., AirNet, PromptIR, DiffUIR, FoundIR, DiT4SR) in LPIPS (perceptual similarity) and FID (distribution distance) across Blur, Noise, Haze, Rain, and Low-light tasks.
  • For example, on the FoundIR Low-light task, MiM-DiT achieved an LPIPS of 0.2176 (best) and a CLIP-IQA of 0.6104 (best).

• Qualitative Results:

  • Visual comparisons show that MiM-DiT effectively removes blur while preserving sharp edges, eliminates haze without color distortion, and enhances low-light images with accurate colors and fine textures.
  • Competing methods often exhibited over-smoothing, color casts, or residual artifacts.

• Ablation Studies:

  • Intra-MoE: Removing the intra-structure MoE (replacing sub-experts with fixed experts) led to over-smoothing and lower perceptual scores, confirming the need for fine-grained specialization.
  • Inter-MoE Heterogeneity: Using a single attention type for all experts (homogeneous) significantly degraded performance compared to the heterogeneous design, proving that diverse inductive biases are essential.
  • Router Design: The combination of Dense Inter-MoE (for comprehensive feature integration) and Sparse Intra-MoE (for efficiency) was proven optimal. Purely sparse or purely dense configurations at both levels resulted in performance drops.

5. Significance

• Solving the "One-Size-Fits-All" Paradox: MiM-DiT addresses the fundamental conflict in all-in-one restoration by dynamically selecting the right "tool" (attention mechanism) for the job, rather than forcing a single architecture to handle all tasks.
• Bridging Generative and Restoration: It successfully bridges the gap between deterministic restoration (which lacks texture) and generative modeling (which can hallucinate details). By using MoE to guide a pretrained DiT, it achieves high-fidelity reconstruction with realistic textures.
• Scalability and Efficiency: The hierarchical MoE design allows the model to scale capacity efficiently (via sparse routing) while maintaining the ability to learn complex, diverse degradation patterns without requiring separate models for each task.

In summary, MiM-DiT represents a significant advancement in universal image restoration by leveraging a sophisticated, dual-level expert system to guide a powerful diffusion backbone, resulting in robust, high-quality restoration across a wide spectrum of real-world degradations.