Routing doet ertoe in MoE: Schaalvergroting van Diffusion Transformers met Expliciete Routing-begeleiding

Samenvatting

Mixture-of-Experts (MoE) is naar voren gekomen als een krachtig paradigma voor het schalen van modelcapaciteit met behoud van computationele efficiëntie. Ondanks zijn opmerkelijke succes in grote taalmodel(len) (LLM's), hebben bestaande pogingen om MoE toe te passen op Diffusion Transformers (DiT's) beperkte winst opgeleverd. Wij schrijven dit verschil toe aan fundamentele verschillen tussen taal- en visuele tokens. Taaltokens zijn semantisch dicht met uitgesproken variatie tussen tokens, terwijl visuele tokens ruimtelijke redundantie en functionele heterogeniteit vertonen, wat expertspecialisatie in visuele MoE belemmert. Daartoe presenteren wij ProMoE, een MoE-raamwerk met een router in twee stappen met expliciete routeringsbegeleiding die expertspecialisatie bevordert. Specifiek moedigt deze begeleiding de router aan om beeldtokens te verdelen in conditionele en unconditionele sets via conditionele routering op basis van hun functionele rollen, en de toewijzingen van conditionele beeldtokens te verfijnen door prototypische routering met leerbare prototypes gebaseerd op semantische inhoud. Bovendien biedt de op gelijkenis gebaseerde experten-toewijzing in de latente ruimte, mogelijk gemaakt door prototypische routering, een natuurlijk mechanisme voor het incorporeren van expliciete semantische begeleiding, en wij valideren dat dergelijke begeleiding cruciaal is voor visuele MoE. Hierop voortbordurend stellen wij een routeringscontrastief verlies voor dat het prototypische routeringsproces expliciet versterkt, waardoor intra-expert coherentie en inter-expert diversiteit worden bevorderd. Uitgebreide experimenten op de ImageNet-benchmark tonen aan dat ProMoE state-of-the-art methoden overtreft onder zowel Rectified Flow- als DDPM-trainingsdoelstellingen. Code en modellen zullen openbaar beschikbaar worden gesteld.

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.

Routing doet ertoe in MoE: Schaalvergroting van Diffusion Transformers met Expliciete Routing-begeleiding

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

Samenvatting

Support