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

Ce papier propose MiM-DiT, un cadre unifié de restauration d'images qui combine une architecture à double niveau de mélange d'experts (MoE) avec des transformateurs de diffusion préentraînés pour gérer efficacement une variété de dégradations, des types majeurs aux variations fines.

L'essentiel

🌟 Le Problème : Un seul couteau pour toutes les tâches ?

Imaginez que vous êtes un chef cuisinier. On vous donne un ingrédient abîmé.

• Parfois, c'est une pomme écrasée (flou).
• Parfois, c'est une pomme couverte de boue (bruit).
• Parfois, c'est une pomme dans le brouillard (brume).
• Parfois, c'est une pomme dans le noir complet (faible luminosité).

Le problème, c'est que les restaurateurs d'images actuels essaient d'utiliser un seul et même outil pour tout réparer. C'est comme essayer de réparer une montre avec un marteau, ou de couper un steak avec une cuillère en bois. Le résultat est souvent décevant : l'image devient floue, les détails disparaissent, ou les couleurs sont bizarres.

🚀 La Solution : L'usine de réparation "MiM-DiT"

Les auteurs de cet article ont créé une nouvelle méthode appelée MiM-DiT. Pour comprendre comment ça marche, imaginez une usine de réparation ultra-sophistiquée qui fonctionne avec deux niveaux de spécialistes.

1. Le Premier Niveau : Le "Grand Chef" (Inter-MoE)

Imaginez que votre image abîmée arrive à l'entrée de l'usine. Au lieu de la donner à un seul ouvrier, elle est envoyée à quatre grands chefs spécialisés, chacun ayant un style de travail différent :

• Le Chef "Vue d'ensemble" (Attention spatiale) : Il regarde les grandes lignes et les formes lointaines.
• Le Chef "Couleurs" (Attention canal) : Il se concentre sur les teintes et les canaux de couleur.
• Le Chef "Détails locaux" (Attention Swin) : Il regarde les petits coins et les textures proches.
• Le Chef "Lumière" (Attention SE) : Il ajuste l'éclairage global.

Au lieu de choisir un seul chef, le Grand Chef (le routeur dense) écoute les quatre en même temps. Il dit : "Pour cette image floue, j'ai besoin de 40% du Chef Vue d'ensemble et 60% du Chef Détails locaux." Il mélange intelligemment leurs conseils pour créer une première ébauche de réparation.

2. Le Deuxième Niveau : Les "Sous-spécialistes" (Intra-MoE)

C'est là que la magie opère. Même si le Chef "Lumière" est bon, il y a des nuances : un brouillard léger n'est pas pareil qu'un brouillard très épais.
À l'intérieur de chaque équipe de chefs, il y a une équipe de sous-spécialistes.

• Le Chef "Lumière" a 10 assistants. L'un est expert en "brouillard matinal", l'autre en "brouillard de nuit".
• Un chef d'équipe rapide (le routeur sparse) regarde l'image et choisit uniquement l'assistant le plus pertinent pour ce cas précis.

C'est comme si, pour réparer une voiture, vous ne faisiez pas appel à tout le garage, mais uniquement au mécanicien qui connaît exactement le modèle de votre moteur.

3. Le Moteur Magique : Le "Diffusion Transformer" (DiT)

Toute cette usine est alimentée par un moteur très puissant appelé Diffusion Transformer. C'est un robot qui a déjà vu des millions de photos parfaites. Il sait à quoi une image "normale" doit ressembler.
Notre système MiM-DiT ne fait pas que réparer l'image ; il guide ce robot pour qu'il réinvente les détails manquants (comme les textures de la peau ou les feuilles d'un arbre) de manière réaliste, au lieu de simplement lisser l'image.

🎨 L'Analogie de l'Orchestre

Pour résumer avec une autre image :

• Les anciennes méthodes : C'est un orchestre où tous les musiciens jouent la même partition, tout le temps, peu importe la musique. Résultat : un son plat et ennuyeux.
• La méthode MiM-DiT : C'est un chef d'orchestre génial.
  1. Il choisit quels instruments (violons, cuivres, percussions) doivent jouer pour ce morceau (c'est le niveau Inter-MoE).
  2. Il choisit ensuite quel musicien précis dans chaque section doit jouer la note la plus importante (c'est le niveau Intra-MoE).
  3. Le tout est dirigé par une partition magique (le modèle de diffusion) qui garantit que la musique finale est magnifique.

🏆 Les Résultats

Grâce à cette approche en "boîte dans la boîte" (MoE dans MoE), le système est capable de :

• Enlever le flou sans rendre l'image molle.
• Dissiper le brouillard sans changer les couleurs.
• Éclaircir une photo sombre sans créer de bruit bizarre.

Les tests montrent que cette méthode est meilleure que toutes les autres actuellement disponibles, car elle est assez intelligente pour adapter sa stratégie à chaque type de problème, tout en gardant la qualité artistique et réaliste des images.

En bref : Au lieu d'essayer d'être bon partout avec une seule méthode, MiM-DiT assemble une équipe de spécialistes qui s'adaptent dynamiquement pour réparer chaque image comme un chef-d'œuvre.

Résumé technique

1. Problématique

La restauration d'images "tout-en-un" (All-in-One) vise à reconstruire des images de haute qualité à partir d'entrées dégradées de manière variée (flou, brouillard, bruit, faible luminosité, pluie, etc.). Ce domaine fait face à un défi fondamental : les différents types de dégradations imposent des exigences structurelles contradictoires.

• Certaines dégradations nécessitent l'exploitation de statistiques globales de l'image.
• D'autres exigent la reconstruction de structures locales nettes.
• Les approches existantes souffrent soit d'un sur-ajustement (nécessitant un entraînement spécifique par tâche), soit d'un compromis générique dû au partage de paramètres, conduisant à des résultats lissés ou à une perte de détails texturaux.
• De plus, l'intégration directe des modèles de diffusion (excellents pour la génération) avec des architectures Mixture-of-Experts (MoE) existantes reste sous-optimale car elle utilise souvent des experts homogènes qui ignorent les biais inductifs spécifiques à chaque tâche.

2. Méthodologie : MiM-DiT

Les auteurs proposent MiM-DiT, un cadre unifié qui intègre une architecture hiérarchique MoE dans MoE (Mixture-of-Experts) au sein d'un modèle de diffusion pré-entraîné basé sur un Transformateur de Diffusion (DiT, spécifiquement SD 3.5).

L'architecture repose sur deux niveaux de routage complémentaires :

A. Niveau Inter-MoE (Adaptation Macroscopique)

Ce niveau gère les grandes catégories de dégradations en fusionnant des groupes d'experts hétérogènes.

• Structure : Quatre groupes d'experts sont définis, chacun basé sur un mécanisme d'attention distinct :
  1. Self-attention spatiale (pour les dépendances à longue portée).
  2. Self-attention canal (pour la recalibration des caractéristiques par canal).
  3. Attention Swin (équilibre contexte local/global).
  4. Attention SE (modélisation de l'éclairage global, utile pour le débruitage et la faible luminosité).
• Routage : Un routeur dense (gating network) fusionne adaptativement les sorties de ces quatre groupes en fonction des caractéristiques de la dégradation d'entrée. Cela permet au modèle de combiner les biais inductifs complémentaires (ex: corriger la géométrie tout en ajustant l'éclairage) plutôt que d'en choisir un seul.

B. Niveau Intra-MoE (Adaptation Microscopique)

À l'intérieur de chaque groupe d'experts, ce niveau gère les variations fines au sein d'une même catégorie de dégradation (ex: brouillard léger vs brouillard épais).

• Structure : Chaque groupe contient plusieurs sous-experts partageant la même architecture d'attention mais avec des paramètres entraînés indépendamment.
• Routage : Un routeur sparse (mécanisme Top-k) sélectionne dynamiquement les sous-experts les plus pertinents pour l'entrée spécifique. Cela assure une spécialisation fine et une efficacité computationnelle.

C. Intégration avec le DiT (Diffusion Transformer)

Le module MiM agit comme un mécanisme de conditionnement puissant pour le backbone DiT pré-entraîné.

• L'image de basse qualité (LQ) est d'abord traitée par le module MiM pour extraire des caractéristiques de restauration.
• Ces caractéristiques sont injectées dans le processus de diffusion via une couche Zero-Linear (initialisée à zéro). Cette conception garantit que le signal de conditionnement a une influence minimale au début de l'entraînement (stabilité) et que le modèle apprend progressivement à l'utiliser pour guider la génération d'images de haute qualité.

3. Contributions Clés

1. Cadre Unifié MoE-Diffusion : Première intégration réussie d'une architecture MoE hiérarchique dans un modèle de diffusion pré-entraîné pour la restauration d'images tout-en-un, préservant la capacité générative tout en permettant une adaptation dynamique.
2. Architecture Hiérarchique "MoE dans MoE" : Introduction d'un mécanisme de routage à double niveau (Inter-MoE pour la sélection de l'architecture globale, Intra-MoE pour la spécialisation fine), permettant une adaptation granulaire aux motifs de dégradations inter-classe et intra-classe.
3. Fusion Dense et Routage Sparse : Combinaison stratégique d'une fusion dense au niveau supérieur (pour capturer la diversité structurelle) et d'un routage sparse au niveau inférieur (pour l'efficacité et la spécialisation), résolvant le compromis entre généralisation et précision.

4. Résultats Expérimentaux

Les auteurs ont évalué MiM-DiT sur plusieurs benchmarks publics, notamment FoundIR, 4KRD, RealRain-1K, HazeRD et UHD-LL.

• Performance Quantitative : La méthode obtient des résultats State-of-the-Art (SOTA) sur la majorité des métriques sans référence (NIQE, LIQE, MUSIQ, CLIP-IQA) et sur plusieurs métriques de référence (LPIPS, FID) pour des tâches variées (défloutage, débruitage, débrumage, amélioration faible luminosité).
  • Exemple : Sur FoundIR, MiM-DiT obtient le meilleur score LPIPS pour le bruit, le brouillard et la faible luminosité, et les meilleurs scores perceptuels (LIQE, MUSIQ) sur presque toutes les catégories.
• Performance Qualitative : Les comparaisons visuelles montrent que MiM-DiT préserve mieux les textures fines, les bords nets et les couleurs naturelles par rapport aux méthodes concurrentes (comme DiffIR, PromptIR, FoundIR). Elle évite le lissage excessif (over-smoothing) et les artefacts de couleur souvent observés dans les approches précédentes, en particulier dans des conditions extrêmes de faible luminosité.
• Études d'Abalation :
  • La suppression du niveau Intra-MoE entraîne une perte de détails et un lissage des textures.
  • L'utilisation d'experts homogènes (un seul type d'attention) dans le niveau Inter-MoE dégrade significativement les performances, prouvant la nécessité de l'hétérogénéité structurelle.
  • Le design hybride (Dense Inter + Sparse Intra) est supérieur aux configurations purement denses ou purement esparses.

5. Signification et Impact

Ce travail représente une avancée significative dans le domaine de la vision par ordinateur pour plusieurs raisons :

• Résolution du conflit de modélisation : Il résout le problème fondamental de la restauration tout-en-un en permettant au modèle de basculer dynamiquement entre des stratégies de traitement structurellement différentes selon la dégradation, sans sacrifier la cohérence globale.
• Synergie Diffusion-MoE : Il démontre que l'intégration de la puissance générative des modèles de diffusion (capables de recréer des détails plausibles) avec la spécialisation dynamique du MoE (capable de cibler la cause de la dégradation) produit des résultats supérieurs aux approches isolées.
• Robustesse Réelle : La capacité du modèle à généraliser sur des données réelles non vues lors de l'entraînement (Out-of-Distribution) en fait une solution pratique pour des applications réelles où les types de dégradations sont imprévisibles et mélangés.

En résumé, MiM-DiT établit un nouveau standard pour la restauration d'images universelle en combinant l'adaptabilité structurelle du MoE hiérarchique avec la fidélité générative des Transformers de diffusion.