Routing Matters in MoE: Skalierung von Diffusion Transformern mit expliziter Routing-Anleitung

papers.abstract

Mixture-of-Experts (MoE) hat sich als leistungsstarkes Paradigma etabliert, um die Modellkapazität zu skalieren und gleichzeitig die Recheneffizienz zu erhalten. Trotz bemerkenswerter Erfolge bei großen Sprachmodellen (LLMs) haben bestehende Versuche, MoE auf Diffusion Transformer (DiTs) anzuwenden, nur begrenzte Gewinne erbracht. Wir führen diese Lücke auf grundlegende Unterschiede zwischen Sprach- und visuellen Tokens zurück. Sprach-Tokens sind semantisch dicht mit ausgeprägten Variationen zwischen den Tokens, während visuelle Tokens räumliche Redundanz und funktionale Heterogenität aufweisen, was die Expertenspezialisierung in visuellen MoE-Systemen behindert. Daher stellen wir ProMoE vor, ein MoE-Framework mit einem Zwei-Stufen-Router und expliziter Routing-Anleitung, die die Expertenspezialisierung fördert. Konkret ermutigt diese Anleitung den Router, Bild-Tokens gemäß ihrer funktionalen Rolle durch konditionelles Routing in konditionelle und unkonditionelle Sets aufzuteilen und die Zuordnungen der konditionellen Bild-Tokens durch prototypisches Routing mit lernbaren Prototypen auf Basis semantischen Inhalts zu verfeinern. Darüber hinaus bietet die durch prototypisches Routing ermöglichte ähnlichkeitsbasierte Expertenzuordnung im latenten Raum einen natürlichen Mechanismus zur Integration expliziter semantischer Guidance, und wir validieren, dass eine solche Guidance für visuelles MoE entscheidend ist. Aufbauend darauf schlagen wir einen Routing-Kontrastverlust vor, der den prototypischen Routing-Prozess explizit verbessert und Intra-Experten-Kohärenz sowie Inter-Experten-Diversität fördert. Umfangreiche Experimente auf dem ImageNet-Benchmark zeigen, dass ProMoE state-of-the-art Methoden unter sowohl Rectified Flow- als auch DDPM-Trainingszielen übertrifft. Code und Modelle werden öffentlich verfügbar gemacht.

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 Matters in MoE: Skalierung von Diffusion Transformern mit expliziter Routing-Anleitung

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

papers.abstract

Support