Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

Le papier présente ProMoE, un cadre MoE pour les Transformers de diffusion qui surpasse les méthodes actuelles en introduisant un routage explicite à deux étapes guidé par des prototypes pour surmonter les défis de la spécialisation des experts dans les tokens visuels.

L'essentiel

Imaginez que vous essayez de peindre un tableau magnifique, mais au lieu d'avoir un seul artiste qui fait tout, vous engagez une équipe de 100 experts différents. C'est le principe du MoE (Mélange d'Experts), une technique très populaire dans l'intelligence artificielle moderne.

Dans le monde du texte (comme les chatbots), cette méthode fonctionne à merveille. Mais dans le monde de l'image, les choses se sont toujours mal passées. Pourquoi ? Parce que les images sont très différentes des mots.

Voici l'explication de la nouvelle méthode ProMoE, présentée dans ce papier, avec des analogies simples.

1. Le Problème : Pourquoi les experts se perdent dans les images

Pensez à un roman. Chaque mot est unique et porte une idée précise. Si vous demandez à un expert de lire le mot "pomme", il sait exactement de quoi il parle. C'est dense et clair.

Maintenant, pensez à une photo. Une photo est composée de milliers de petits carrés (des "patchs").

• Redondance spatiale : Si vous regardez un ciel bleu, 100 petits carrés sont presque identiques. C'est comme si vous demandiez à 100 experts de décrire la même chose 100 fois. Ils finissent par tous dire la même chose et ne développent pas de talents spécifiques.
• Hétérogénéité fonctionnelle : Dans la création d'images par IA, on utilise souvent une technique où l'ordinateur doit deviner à la fois "ce qu'il voit" (conditionnel) et "ce qu'il imagine sans indice" (inconditionnel). C'est comme demander à un chef de cuisine de préparer un plat spécifique et de cuisiner à l'aveugle en même temps. Si vous ne séparez pas les tâches, le chef se trompe.

Les anciennes méthodes traitaient tous les petits carrés de l'image de la même façon, ce qui empêchait les experts de se spécialiser.

2. La Solution : ProMoE, le Chef d'Orchestre Intelligent

Les auteurs proposent ProMoE, un système qui agit comme un chef d'orchestre très organisé. Au lieu de laisser les experts choisir au hasard, il utilise une guide de routage explicite en deux étapes.

Étape 1 : La Séparation des Tâches (Routage Conditionnel)

Imaginez une grande salle de réunion.

• Le chef d'orchestre (le routeur) regarde les participants.
• Il dit : "Toi, toi et toi, vous êtes les experts du silence (inconditionnels). Vous travaillez sur les parties de l'image qui ne dépendent d'aucune description."
• Il dit : "Vous, les autres, vous êtes les experts du contexte (conditionnels). Vous travaillez sur les parties qui dépendent de la description (par exemple, 'un chat rouge')."

C'est comme séparer les équipes de construction : une équipe s'occupe des fondations (inconditionnel) et l'autre de la décoration intérieure (conditionnel). Cela évite la confusion.

Étape 2 : L'Attribution par "Modèles" (Routage Prototypique)

Maintenant, pour les experts du contexte, comment savoir qui fait quoi ?
Au lieu de laisser les experts se battre pour les tâches, ProMoE utilise des modèles apprenables (des prototypes).

• Imaginez que chaque expert a une carte de visite avec une image de ce qu'il aime faire (ex: un expert "animaux", un expert "architecture", un expert "nourriture").
• Quand un petit carré de l'image arrive, le système compare son contenu à ces cartes de visite.
• Si le carré représente une patte de chien, il va vers l'expert "animaux". S'il représente une fenêtre, il va vers l'expert "architecture".

C'est beaucoup plus intelligent que de simplement compter les mots. Le système comprend la sémantique (le sens) de l'image.

3. L'Entraînement : Le "Cours de Conduite" (Contrastive Loss)

Pour que ces experts deviennent vraiment bons, il faut les entraîner. Les auteurs ont ajouté une règle spéciale appelée Perte de Contraste de Routage.

C'est comme un jeu éducatif :

• Si deux petits carrés de l'image sont très similaires (deux pattes de chien), le système les force à aller chez le même expert. Cela crée de la cohérence.
• Si deux carrés sont très différents (un chien et une voiture), le système les force à aller chez des experts différents. Cela crée de la diversité.

Cela empêche les experts de devenir paresseux et de tous faire la même chose.

4. Les Résultats : Pourquoi c'est génial

Grâce à cette organisation, ProMoE obtient des résultats incroyables :

• Moins de calcul, plus de qualité : Ils utilisent moins de "cerveaux" actifs (paramètres) que les modèles denses classiques, mais produisent des images plus belles.
• Rapidité : Ils apprennent plus vite.
• Polyvalence : Cela fonctionne aussi bien pour créer des images à partir de texte que pour d'autres tâches complexes.

En résumé

Si les anciennes méthodes d'IA pour les images étaient comme un groupe de 100 personnes essayant de peindre un tableau en criant toutes en même temps sans coordination, ProMoE est comme un chef d'orchestre qui :

1. Sépare les musiciens en deux groupes (fondations vs détails).
2. Donne à chaque musicien une partition précise basée sur ce qu'il joue le mieux.
3. Les entraîne à jouer ensemble de manière harmonieuse.

Résultat : Une symphonie visuelle parfaite, même avec moins de musiciens !

Résumé technique

1. Problématique

Bien que le paradigme Mixture-of-Experts (MoE) ait révolutionné l'évolutivité des grands modèles de langage (LLM) en augmentant la capacité du modèle tout en maintenant une efficacité computationnelle, son application aux Transformers de Diffusion (DiT) pour la génération d'images a donné des résultats décevants.

Les auteurs identifient deux différences fondamentales entre les tokens linguistiques et visuels qui expliquent cet échec :

• Redondance spatiale élevée : Contrairement aux tokens textuels discrets et sémantiquement denses, les tokens visuels (patchs d'images) sont continus, couplés spatialement et fortement redondants. Cela conduit les experts à apprendre des caractéristiques homogènes plutôt que spécialisées.
• Hétérogénéité fonctionnelle : Les modèles de diffusion utilisent souvent la guidance sans classifieur (CFG), générant deux types de tokens fonctionnellement distincts : des tokens conditionnels (avec une étiquette ou un texte) et des tokens inconditionnels (sans condition). Les approches MoE naïves traitent ces tokens de manière uniforme, ignorant leurs rôles différents, ce qui empêche une spécialisation efficace des experts.

Le résultat est un manque de diversité inter-experts et de cohérence intra-experts, limitant les gains de performance par rapport aux modèles denses.

2. Méthodologie : ProMoE

Pour surmonter ces limitations, les auteurs proposent ProMoE, un cadre MoE intégrant un routeur en deux étapes avec une guidance explicite. L'architecture repose sur trois piliers :

A. Routeur en deux étapes

Le routeur ne se contente pas de calculer des scores d'affinité implicites ; il utilise deux signaux de guidage distincts :

1. Routage Conditionnel (Étape 1) : Basé sur le rôle fonctionnel du token.

   • Le routeur partitionne les tokens d'image en deux ensembles : inconditionnels (issus de conditions nulles) et conditionnels.
   • Les tokens inconditionnels sont acheminés de manière déterministe vers des experts inconditionnels dédiés.
   • Les tokens conditionnels sont envoyés à une seconde étape pour un routage fin.
   • Objectif : Assurer une ségrégation fonctionnelle stricte.

2. Routage Prototypique (Étape 2) : Basé sur le contenu sémantique.

   • Les tokens conditionnels sont assignés aux experts via une similarité cosinus entre les embeddings des tokens et un ensemble de prototypes apprenables ($P$), chaque prototype étant associé à un expert spécifique.
   • Contrairement aux méthodes standard utilisant des couches linéaires, cette approche mesure directement la similarité sémantique dans l'espace latent.
   • Une fonction d'activation identité est utilisée (au lieu de Softmax) pour préserver les rangs relatifs et assurer une stabilité de l'entraînement.

B. Guidance Sémantique Explicite et Perte de Contraste de Routage

Les auteurs démontrent que le routage basé uniquement sur l'apprentissage implicite est insuffisant. Ils introduisent une Perte de Contraste de Routage (Routing Contrastive Loss - RCL) pour renforcer la guidance sémantique :

• Fonctionnement : La RCL encourage les tokens sémantiquement similaires à être routés vers le même expert (cohérence intra-expert) tout en repoussant les tokens dissimilaires vers des experts différents (diversité inter-expert).
• Mécanisme : Pour chaque prototype, les tokens assignés forment un ensemble positif, tandis que les autres forment des ensembles négatifs. La perte pousse le prototype vers le centre de son ensemble positif et l'éloigne des centres des ensembles négatifs.
• Avantage : Cette perte agit également comme un régularisateur d'équilibrage de charge (load-balancing) basé sur la sémantique, supérieur aux méthodes traditionnelles qui ne comptent que le nombre de tokens.

3. Contributions Clés

1. Analyse des différences Token Langage/Visuel : Identification de la redondance spatiale et de l'hétérogénéité fonctionnelle comme obstacles majeurs au MoE dans la vision.
2. Architecture ProMoE : Conception d'un routeur hybride combinant un routage conditionnel (fonctionnel) et un routage prototypique (sémantique).
3. Perte de Contraste de Routage (RCL) : Introduction d'une perte d'apprentissage qui améliore explicitement la spécialisation des experts sans nécessiter d'étiquettes manuelles coûteuses.
4. Validation Expérimentale : Démonstration que ProMoE surpasse les modèles denses et les méthodes MoE actuelles (SOTA) avec moins de paramètres activés.

4. Résultats Expérimentaux

Les expériences ont été menées sur le benchmark ImageNet (256x256) avec deux objectifs d'entraînement : Rectified Flow (RF) et DDPM.

• Performance Supérieure : ProMoE bat systématiquement les modèles denses (Dense-DiT) et les méthodes MoE SOTA (DiT-MoE, EC-DiT, DiffMoE).
  • Exemple (Rectified Flow, CFG=1.5) : ProMoE-L (458M paramètres activés) atteint un FID de 2.79, surpassant Dense-DiT-XL (675M paramètres activés) qui obtient 3.23.
  • ProMoE-XL (1.568B paramètres totaux) atteint un FID de 2.59, surpassant des modèles denses beaucoup plus grands.
• Efficacité des Paramètres : ProMoE obtient de meilleurs résultats avec moins de paramètres activés et souvent moins de paramètres totaux que les concurrents.
• Évolutivité (Scaling) : Les gains de performance sont monotones à mesure que la taille du modèle (B, L, XL) et le nombre d'experts augmentent.
• Convergence : ProMoE converge plus rapidement que les modèles denses et les autres MoE, comme le montrent les courbes de perte d'entraînement.
• Généralisation : Sur le benchmark GenEval (génération texte-à-image), ProMoE surpasse également les modèles denses et les MoE basés sur le choix de tokens, confirmant sa capacité à gérer des tâches complexes.

5. Signification et Impact

Ce travail est significatif car il résout le problème de l'inefficacité du MoE dans les modèles de diffusion visuelle en introduisant une guidance explicite dans le processus de routage.

• Il démontre que l'adaptation réussie du MoE à la vision nécessite de tenir compte de la nature spécifique des tokens visuels (redondance et dualité conditionnelle/inconditionnelle).
• La méthode proposée offre une voie viable pour passer à l'échelle (scaling) des modèles de diffusion de très grande taille tout en maintenant une efficacité computationnelle, comblant ainsi l'écart de performance observé entre les LLM et les DiT.
• L'approche est robuste, ne dépend pas de labels manuels pour le routage (grâce à la RCL), et s'applique à différents objectifs d'entraînement (DDPM et Rectified Flow).

En résumé, ProMoE établit un nouvel état de l'art pour les modèles de diffusion basés sur MoE, prouvant que "le routage compte" (routing matters) et que des mécanismes de guidage explicites sont essentiels pour la spécialisation des experts dans le domaine de la vision par ordinateur.