L'importance du routage dans les MoE : Mise à l'échelle des Transformers de Diffusion avec un Guidage Explicite du Routage

papers.abstract

Le Mixture-of-Experts (MoE) est apparu comme un paradigme puissant pour augmenter la capacité des modèles tout en préservant l'efficacité computationnelle. Malgré son succès notable dans les grands modèles de langage (LLM), les tentatives existantes pour appliquer le MoE aux Transformers de Diffusion (DiT) n'ont obtenu que des gains limités. Nous attribuons cet écart à des différences fondamentales entre les tokens linguistiques et visuels. Les tokens linguistiques sont sémantiquement denses avec une variation inter-tokens prononcée, tandis que les tokens visuels présentent une redondance spatiale et une hétérogénéité fonctionnelle, entravant la spécialisation des experts dans le MoE visuel. Pour pallier cela, nous présentons ProMoE, un framework MoE doté d'un routeur à deux étapes avec un guidage de routage explicite qui favorise la spécialisation des experts. Concrètement, ce guidage encourage le routeur à partitionner les tokens d'image en ensembles conditionnels et non conditionnels via un routage conditionnel selon leurs rôles fonctionnels, et à affiner l'affectation des tokens d'image conditionnels grâce à un routage prototypique avec des prototypes apprenables basés sur le contenu sémantique. De plus, l'allocation des experts basée sur la similarité dans l'espace latent, permise par le routage prototypique, offre un mécanisme naturel pour incorporer un guidage sémantique explicite, et nous validons qu'un tel guidage est crucial pour le MoE visuel. Sur cette base, nous proposons une fonction de perte de contraste de routage qui améliore explicitement le processus de routage prototypique, en favorisant la cohérence intra-expert et la diversité inter-expert. Des expériences approfondies sur le benchmark ImageNet démontrent que ProMoE surpasse les méthodes de l'état de l'art sous les objectifs d'entraînement Rectified Flow et DDPM. Le code et les modèles seront rendus publics.

English

Mixture-of-Experts (MoE) has emerged as a powerful paradigm for scaling model capacity while preserving computational efficiency. Despite its notable success in large language models (LLMs), existing attempts to apply MoE to Diffusion Transformers (DiTs) have yielded limited gains. We attribute this gap to fundamental differences between language and visual tokens. Language tokens are semantically dense with pronounced inter-token variation, while visual tokens exhibit spatial redundancy and functional heterogeneity, hindering expert specialization in vision MoE. To this end, we present ProMoE, an MoE framework featuring a two-step router with explicit routing guidance that promotes expert specialization. Specifically, this guidance encourages the router to partition image tokens into conditional and unconditional sets via conditional routing according to their functional roles, and refine the assignments of conditional image tokens through prototypical routing with learnable prototypes based on semantic content. Moreover, the similarity-based expert allocation in latent space enabled by prototypical routing offers a natural mechanism for incorporating explicit semantic guidance, and we validate that such guidance is crucial for vision MoE. Building on this, we propose a routing contrastive loss that explicitly enhances the prototypical routing process, promoting intra-expert coherence and inter-expert diversity. Extensive experiments on ImageNet benchmark demonstrate that ProMoE surpasses state-of-the-art methods under both Rectified Flow and DDPM training objectives. Code and models will be made publicly available.

L'importance du routage dans les MoE : Mise à l'échelle des Transformers de Diffusion avec un Guidage Explicite du Routage

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

papers.abstract

Support